Fixing the parse error for @ symbol

Author: sivaalamp

parse/queryLex.go

@@ -114,17 +114,57 @@ func (l *Lexer) skipUntil(val rune) {
 	}
 }
 
-func (l *Lexer) consumeIdent() {
-	for !l.done() && !l.isEndIdent(l.next()) {
+func isNamedParamStart(r rune) bool {
+	return r == '_' || unicode.IsLetter(r)
+}
+
+func isNamedParamPart(r rune) bool {
+	return r == '_' || unicode.IsLetter(r) || unicode.IsDigit(r)
+}
+
+func isNamedParamBoundary(r rune) bool {
+	return !(r == '_' || unicode.IsLetter(r) || unicode.IsDigit(r) || r == '"')
+}
+
+func (l *Lexer) currentRuneStart() int {
+	if l.width == 0 {
+		return l.pos
+	}
+	return l.pos - l.width
+}
+
+func (l *Lexer) shouldLexNamedParam() bool {
+	at := l.currentRuneStart()
+	if at < 0 || at >= len(l.input) || l.input[at] != '@' {
+		return false
 	}
+
+	nextPos := at + len("@")
+	if nextPos >= len(l.input) {
+		return false
+	}
+
+	next, _ := utf8.DecodeRuneInString(l.input[nextPos:])
+	if !isNamedParamStart(next) {
+		return false
+	}
+
+	if at == 0 {
+		return true
+	}
+
+	prev, _ := utf8.DecodeLastRuneInString(l.input[:at])
+	return isNamedParamBoundary(prev)
 }
 
-func (l *Lexer) isEndIdent(r rune) bool {
-	shouldEnd := unicode.IsSpace(r) || strings.ContainsRune(",)", r)
-	if shouldEnd {
-		l.backup()
+func (l *Lexer) consumeNamedParam() {
+	for !l.done() {
+		r := l.next()
+		if !isNamedParamPart(r) {
+			l.backup()
+			return
+		}
 	}
-	return shouldEnd
 }
 
 func (l *Lexer) next() rune {
@@ -174,7 +214,7 @@ func lexQuery(l *Lexer) stateFunc {
 			return lexString
 		}
 
-		if r == '@' {
+		if r == '@' && l.shouldLexNamedParam() {
 			return lexNamedParam
 		}
 
@@ -220,8 +260,8 @@ func lexNamedParam(l *Lexer) stateFunc {
 	l.writeChunk()
 	l.next()
 	l.start = l.pos
-	// advance through the name
-	l.consumeIdent()
+	// advance through a valid named-parameter token
+	l.consumeNamedParam()
 	l.onNamed(strings.ToUpper(l.input[l.start:l.pos]))
 	l.start = l.pos
 	l.output.WriteRune('?')

parse/queryLex_test.go

@@ -155,6 +155,36 @@ func TestLexNamed(t *testing.T) {
 			a = ?
 `,
 		},
+		{
+			name:           "at sign inside invalid unquoted object name is not a named parameter",
+			query:          `CREATE VIEW vw_user@2024 AS SELECT user_id FROM "user"`,
+			expectedNamed:  []string{},
+			expectedOutput: `CREATE VIEW vw_user@2024 AS SELECT user_id FROM "user"`,
+		},
+		{
+			name:           "at sign inside quoted identifier is not a named parameter",
+			query:          `CREATE TABLE "user@2024" (user_id INT)`,
+			expectedNamed:  []string{},
+			expectedOutput: `CREATE TABLE "user@2024" (user_id INT)`,
+		},
+		{
+			name:           "at sign followed by digit is not a named parameter",
+			query:          `SELECT @2024`,
+			expectedNamed:  []string{},
+			expectedOutput: `SELECT @2024`,
+		},
+		{
+			name:           "named param stops before arithmetic operator",
+			query:          `select * from table where a = @name+1`,
+			expectedNamed:  []string{"NAME"},
+			expectedOutput: `select * from table where a = ?+1`,
+		},
+		{
+			name:           "named param stops before type cast",
+			query:          `select * from table where a = @name::int`,
+			expectedNamed:  []string{"NAME"},
+			expectedOutput: `select * from table where a = ?::int`,
+		},
 	}
 	var encounteredNames nameCapture
 	for _, tc := range testCases {

stmt_test.go

@@ -257,6 +257,16 @@ func TestNumInput(t *testing.T) {
 			query:    "select * from test where a = 'test?'",
 			expected: 0,
 		},
+		{
+			name:     "at sign in object name is not input",
+			query:    `CREATE VIEW vw_user@2024 AS SELECT user_id FROM "user"`,
+			expected: 0,
+		},
+		{
+			name:     "at sign in quoted identifier is not input",
+			query:    `CREATE TABLE "user@2024" (user_id INT)`,
+			expected: 0,
+		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {