From ca5d403a3fd82ed9f767f48c68500df0f4022c42 Mon Sep 17 00:00:00 2001
From: Yehoshua Sandler <accounts@tzed.io>
Date: Tue, 15 Apr 2025 16:41:27 -0500
Subject: [PATCH] feat: capture ASC and DESC

---
 q/select.go      | 75 ++++++++++++++++++++++++++++++++++--------------
 q/select_test.go | 17 +++++++++++
 2 files changed, 70 insertions(+), 22 deletions(-)

diff --git a/q/select.go b/q/select.go
index 6309504..07d7244 100644
--- a/q/select.go
+++ b/q/select.go
@@ -1,7 +1,7 @@
 package q
 
 import (
-	"fmt"
+	//	"fmt"
 	"strings"
 
 	"github.com/DataDog/go-sqllexer"
@@ -19,7 +19,7 @@ type Select struct {
 
 type OrderBy struct {
 	Key       string
-	IsDescend bool
+	IsDescend bool // SQL queries with no ASC|DESC on their ORDER BY are ASC by default, hence why this bool for the opposite
 }
 
 func (q *Select) GetFullSql() string {
@@ -61,7 +61,7 @@ func mutateSelectFromKeyword(query *Select, keyword string) {
 }
 
 // TODO: make this an array of tokens instead
-func unshiftBuffer(buf *[10]string, value string) {
+func unshiftBuffer(buf *[10]sqllexer.Token, value sqllexer.Token) {
 	for i := 9; i >= 1; i-- {
 		buf[i] = buf[i-1]
 	}
@@ -78,9 +78,11 @@ func ParseSelectStatement(sql string) Select {
 	passedTable := false
 	passedWHERE := false
 	passedConditionals := false
-	passedOrderBy := false
+	passedOrderByKeywords := false
+	passesOrderByColumns := false
+	//checkForOrderDirection := false
 
-	lookBehindBuffer := [10]string{} // TODO: make this an array of tokens instead
+	lookBehindBuffer := [10]sqllexer.Token{} // TODO: make this an array of tokens instead
 	var workingConditional = Conditional{}
 
 	var columns []string
@@ -88,9 +90,6 @@ func ParseSelectStatement(sql string) Select {
 	lexer := sqllexer.New(sql)
 	for {
 		token := lexer.Scan()
-		if IsTokenEndOfStatement(token) {
-			break
-		}
 
 		if !passedSELECT && strings.ToUpper(token.Value) != "SELECT" {
 			break
@@ -125,11 +124,9 @@ func ParseSelectStatement(sql string) Select {
 			}
 		}
 
-		if !passedFROM && strings.ToUpper(token.Value) == "FROM" {
+		if !passedFROM && strings.ToUpper(token.Value) == "FROM" { // TODO: make sure to check for other keywords that are allowed
 			passedFROM = true
 			continue
-		} else if !passedFROM {
-			// continue // TODO: make sure to check for other keywords that are allowed
 		}
 
 		if !passedTable && token.Type == sqllexer.IDENT {
@@ -165,21 +162,55 @@ func ParseSelectStatement(sql string) Select {
 		}
 
 		// Checking For ORDER BY
-		if !passedOrderBy && token.Type == sqllexer.KEYWORD {
-			unshiftBuffer(&lookBehindBuffer, token.Value)
+		if !passedOrderByKeywords && token.Type == sqllexer.KEYWORD {
+			unshiftBuffer(&lookBehindBuffer, *token)
 
-			if strings.ToUpper(lookBehindBuffer[1]) == "ORDER" && strings.ToUpper(lookBehindBuffer[0]) == "BY" {
-				passedOrderBy = true
+			if strings.ToUpper(lookBehindBuffer[1].Value) == "ORDER" && strings.ToUpper(lookBehindBuffer[0].Value) == "BY" {
+				passedOrderByKeywords = true
 			}
-
-			fmt.Printf("VALUE: %s %s \n", lookBehindBuffer[1], lookBehindBuffer[0])
 		}
 
-		if passedOrderBy && token.Type == sqllexer.IDENT {
-			orderBys = append(orderBys, OrderBy{
-				Key: token.Value, // TODO: need to be able to get the ASC or DESC keyword too
-			})
-			query.OrderBys = orderBys
+		if passedOrderByKeywords && !passesOrderByColumns {
+
+			if token.Type == sqllexer.IDENT || token.Type == sqllexer.KEYWORD {
+				unshiftBuffer(&lookBehindBuffer, *token)
+				continue
+			}
+
+			if token.Type == sqllexer.PUNCTUATION || token.Type == sqllexer.EOF {
+
+				var orderByColumnName string
+				var directionKeyword string
+				var isDescend bool = false
+
+				if lookBehindBuffer[0].Type == sqllexer.KEYWORD {
+					orderByColumnName = lookBehindBuffer[1].Value
+					directionKeyword = lookBehindBuffer[0].Value
+				} else if lookBehindBuffer[0].Type == sqllexer.IDENT {
+					orderByColumnName = lookBehindBuffer[0].Value
+				}
+
+				if strings.ToUpper(directionKeyword) == "DESC" {
+					isDescend = true
+				}
+
+				orderBys = append(orderBys, OrderBy{
+					Key:       orderByColumnName,
+					IsDescend: isDescend,
+				})
+
+				if IsTokenEndOfStatement(token) {
+					query.OrderBys = orderBys
+					break
+				}
+			}
+
+		}
+
+		query.OrderBys = orderBys
+
+		if IsTokenEndOfStatement(token) {
+			break
 		}
 
 	}
diff --git a/q/select_test.go b/q/select_test.go
index c999465..3545acd 100644
--- a/q/select_test.go
+++ b/q/select_test.go
@@ -61,6 +61,23 @@ func TestParseSelectStatement(t *testing.T) {
 				},
 			},
 		},
+		{
+			input: "SELECT * FROM ZipCodes ORDER BY Code ASC, StateName DESC",
+			expected: Select{
+				Table:      "ZipCodes",
+				IsWildcard: true,
+				OrderBys: []OrderBy{
+					{
+						Key:       "Code",
+						IsDescend: false,
+					},
+					{
+						Key:       "StateName",
+						IsDescend: true,
+					},
+				},
+			},
+		},
 	}
 
 	for _, sql := range testSqlStatements {