fix sqlserver named parameter tokenization (#40)

* fix sqlserver named parameter * fix test
DataDog · Oct 24, 2024 · 8c03f53 · 8c03f53
1 parent 9117b42
commit 8c03f53
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/obfuscate_and_normalize_test.go b/obfuscate_and_normalize_test.go
@@ -411,6 +411,20 @@ multiline comment */
 				Size:       11,
 			},
 		},
+		{
+			input:    `DELETE FROM [discount]  WHERE [description]=@1`,
+			expected: `DELETE FROM discount WHERE description = @1`,
+			statementMetadata: StatementMetadata{
+				Tables:     []string{"discount"},
+				Comments:   []string{},
+				Commands:   []string{"DELETE"},
+				Procedures: []string{},
+				Size:       14,
+			},
+			lexerOpts: []lexerOption{
+				WithDBMS(DBMSSQLServer),
+			},
+		},
 	}
 
 	obfuscator := NewObfuscator(

diff --git a/sqllexer.go b/sqllexer.go
@@ -370,7 +370,7 @@ func (s *Lexer) scanWhitespace() Token {
 func (s *Lexer) scanOperator(lastCh rune) Token {
 	s.start = s.cursor
 	ch := s.next()
-	for isOperator(ch) && !(lastCh == '=' && ch == '?') {
+	for isOperator(ch) && !(lastCh == '=' && (ch == '?' || ch == '@')) {
 		// hack: we don't want to treat "=?" as an single operator
 		lastCh = ch
 		ch = s.next()