tokenize sqlserver @@ system functions

obfuscate_and_normalize_test.go

@@ -281,6 +281,22 @@ multiline comment */
 				WithDBMS(DBMSSQLServer),
 			},
 		},
+		{
+			input: `SET NOCOUNT ON
+			IF @@OPTIONS & 512 > 0
+			RAISERROR ('Current user has SET NOCOUNT turned on.', 1, 1)`,
+			expected: `SET NOCOUNT ON IF @@OPTIONS & ? > ? RAISERROR ( ? )`,
+			statementMetadata: StatementMetadata{
+				Tables:     []string{},
+				Comments:   []string{},
+				Commands:   []string{},
+				Procedures: []string{},
+				Size:       0,
+			},
+			lexerOpts: []lexerOption{
+				WithDBMS(DBMSSQLServer),
+			},
+		},
 	}
 
 	obfuscator := NewObfuscator(

obfuscator_test.go

@@ -403,6 +403,19 @@ func TestObfuscator(t *testing.T) {
 			expected: "SELECT * FROM #users where id = @id and name = @1",
 			dbms:     DBMSSQLServer,
 		},
+		{
+			input: `SELECT @@OPTIONS AS OriginalOptionsValue;
+			SET CONCAT_NULL_YIELDS_NULL OFF;
+			SELECT 'abc' + NULL AS ResultWhen_OFF, @@OPTIONS AS OptionsValueWhen_OFF;
+			SET CONCAT_NULL_YIELDS_NULL ON;
+			SELECT 'abc' + NULL AS ResultWhen_ON, @@OPTIONS AS OptionsValueWhen_ON;`,
+			expected: `SELECT @@OPTIONS AS OriginalOptionsValue;
+			SET CONCAT_NULL_YIELDS_NULL OFF;
+			SELECT ? + NULL AS ResultWhen_OFF, @@OPTIONS AS OptionsValueWhen_OFF;
+			SET CONCAT_NULL_YIELDS_NULL ON;
+			SELECT ? + NULL AS ResultWhen_ON, @@OPTIONS AS OptionsValueWhen_ON;`,
+			dbms: DBMSSQLServer,
+		},
 		{
 			input:    "SELECT * FROM users where id = :id and name = :1",
 			expected: "SELECT * FROM users where id = :id and name = :1",

sqllexer.go

@@ -23,6 +23,7 @@ const (
 	POSITIONAL_PARAMETER   // numbered parameter
 	BIND_PARAMETER         // bind parameter
 	FUNCTION               // function
+	SYSTEM_VARIABLE        // system variable
 	UNKNOWN                // unknown token
 )
 
@@ -137,11 +138,6 @@ func (s *Lexer) Scan() Token {
 			return s.scanBindParameter()
 		}
 		return s.scanOperator(ch)
-	case ch == '@':
-		if s.config.DBMS == DBMSSQLServer && isAlphaNumeric(s.lookAhead(1)) {
-			return s.scanBindParameter()
-		}
-		return s.scanOperator(ch)
 	case ch == '`':
 		if s.config.DBMS == DBMSMySQL {
 			return s.scanDoubleQuotedIdentifier('`')
@@ -155,6 +151,13 @@ func (s *Lexer) Scan() Token {
 			return s.scanSingleLineComment()
 		}
 		fallthrough
+	case ch == '@':
+		if isAlphaNumeric(s.lookAhead(1)) {
+			return s.scanBindParameter()
+		} else if s.lookAhead(1) == '@' {
+			return s.scanSystemVariable()
+		}
+		fallthrough
 	case isOperator(ch):
 		return s.scanOperator(ch)
 	case isPunctuation(ch):
@@ -459,6 +462,18 @@ func (s *Lexer) scanBindParameter() Token {
 	return Token{BIND_PARAMETER, s.src[s.start:s.cursor]}
 }
 
+func (s *Lexer) scanSystemVariable() Token {
+	s.start = s.cursor
+	ch := s.nextBy(2) // consume @@
+	for {
+		if !isAlphaNumeric(ch) {
+			break
+		}
+		ch = s.next()
+	}
+	return Token{SYSTEM_VARIABLE, s.src[s.start:s.cursor]}
+}
+
 func (s *Lexer) scanUnknown() Token {
 	// When we see an unknown token, we advance the cursor until we see something that looks like a token boundary.
 	s.start = s.cursor

sqllexer_test.go

@@ -528,7 +528,19 @@ func TestLexer(t *testing.T) {
 				{WS, " "},
 				{BIND_PARAMETER, "@1"},
 			},
-			lexerOpts: []lexerOption{WithDBMS(DBMSSQLServer)},
+		},
+		{
+			name:  "select with system variable",
+			input: "SELECT @@VERSION AS SqlServerVersion",
+			expected: []Token{
+				{IDENT, "SELECT"},
+				{WS, " "},
+				{SYSTEM_VARIABLE, "@@VERSION"},
+				{WS, " "},
+				{IDENT, "AS"},
+				{WS, " "},
+				{IDENT, "SqlServerVersion"},
+			},
 		},
 		{
 			name:  "SQL Server quoted identifier",