From 046f31098deceaacc93d757ba72dac5f163e33fc Mon Sep 17 00:00:00 2001
From: Zhengda Lu <zhengda.lu@datadoghq.com>
Date: Tue, 10 Oct 2023 12:32:42 -0400
Subject: [PATCH] support sqlserver quoted identifier

---
 sqllexer.go      | 18 ++++++++++++++----
 sqllexer_test.go | 22 ++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/sqllexer.go b/sqllexer.go
index a6f6025..2374f80 100644
--- a/sqllexer.go
+++ b/sqllexer.go
@@ -100,7 +100,7 @@ func (s *Lexer) Scan() Token {
 	case isLetter(ch):
 		return s.scanIdentifier()
 	case isDoubleQuote(ch):
-		return s.scanDoubleQuotedIdentifier()
+		return s.scanDoubleQuotedIdentifier('"')
 	case isSingleQuote(ch):
 		return s.scanString()
 	case isSingleLineComment(ch, s.lookAhead(1)):
@@ -141,6 +141,9 @@ func (s *Lexer) Scan() Token {
 	case isOperator(ch):
 		return s.scanOperator()
 	case isPunctuation(ch):
+		if ch == '[' && s.config.DBMS == DBMSSQLServer {
+			return s.scanDoubleQuotedIdentifier('[')
+		}
 		return s.scanPunctuation()
 	case isEOF(ch):
 		return Token{EOF, ""}
@@ -296,15 +299,22 @@ func (s *Lexer) scanIdentifier() Token {
 	return Token{IDENT, s.src[s.start:s.cursor]}
 }
 
-func (s *Lexer) scanDoubleQuotedIdentifier() Token {
+func (s *Lexer) scanDoubleQuotedIdentifier(delimiter rune) Token {
+	closingDelimiter := delimiter
+	if delimiter == '[' {
+		closingDelimiter = ']'
+	}
+
 	s.start = s.cursor
 	ch := s.next() // consume the opening quote
 	for {
 		// encountered the closing quote
 		// BUT if it's followed by .", then we should keep going
 		// e.g. postgre "foo"."bar"
-		if ch == '"' {
-			if s.matchAt([]rune(`"."`)) {
+		// e.g. sqlserver [foo].[bar]
+		if ch == closingDelimiter {
+			specialCase := []rune{closingDelimiter, '.', delimiter}
+			if s.matchAt([]rune(specialCase)) {
 				ch = s.nextBy(3) // consume the "."
 				continue
 			}
diff --git a/sqllexer_test.go b/sqllexer_test.go
index 443f0f3..886aa1d 100644
--- a/sqllexer_test.go
+++ b/sqllexer_test.go
@@ -514,6 +514,28 @@ func TestLexer(t *testing.T) {
 			},
 			lexerOpts: []lexerOption{WithDBMS(DBMSSQLServer)},
 		},
+		{
+			name:  "SQL Server quoted identifier",
+			input: "SELECT [user] FROM [test].[table] WHERE [id] = 1",
+			expected: []Token{
+				{IDENT, "SELECT"},
+				{WS, " "},
+				{IDENT, "[user]"},
+				{WS, " "},
+				{IDENT, "FROM"},
+				{WS, " "},
+				{IDENT, "[test].[table]"},
+				{WS, " "},
+				{IDENT, "WHERE"},
+				{WS, " "},
+				{IDENT, "[id]"},
+				{WS, " "},
+				{OPERATOR, "="},
+				{WS, " "},
+				{NUMBER, "1"},
+			},
+			lexerOpts: []lexerOption{WithDBMS(DBMSSQLServer)},
+		},
 	}
 
 	for _, tt := range tests {