diff --git a/parser/lexer.go b/parser/lexer.go index 70ff48c7fa368..8ce6c01379752 100644 --- a/parser/lexer.go +++ b/parser/lexer.go @@ -417,8 +417,26 @@ func startWithAt(s *Scanner) (tok int, pos Pos, lit string) { pos = s.r.pos() s.r.inc() ch1 := s.r.peek() - if isIdentFirstChar(ch1) { - s.r.incAsLongAs(isIdentChar) + if ch1 == '\'' || ch1 == '"' { + nTok, nPos, nLit := startString(s) + if nTok == stringLit { + tok = singleAtIdentifier + pos = nPos + lit = nLit + } else { + tok = int('@') + } + } else if ch1 == '`' { + nTok, nPos, nLit := scanQuotedIdent(s) + if nTok == quotedIdentifier { + tok = singleAtIdentifier + pos = nPos + lit = nLit + } else { + tok = int('@') + } + } else if isUserVarChar(ch1) { + s.r.incAsLongAs(isUserVarChar) tok, lit = singleAtIdentifier, s.r.data(&pos) } else if ch1 == '@' { s.r.inc() @@ -435,7 +453,7 @@ func startWithAt(s *Scanner) (tok int, pos Pos, lit string) { s.r.incAsLongAs(isIdentChar) tok, lit = doubleAtIdentifier, s.r.data(&pos) } else { - tok = int('@') + tok, lit = singleAtIdentifier, s.r.data(&pos) } return } diff --git a/parser/lexer_test.go b/parser/lexer_test.go index a2fe701a41d1f..a720d4acf814f 100644 --- a/parser/lexer_test.go +++ b/parser/lexer_test.go @@ -56,7 +56,6 @@ type testCaseItem struct { func (s *testLexerSuite) TestSingleCharOther(c *C) { defer testleak.AfterTest(c)() table := []testCaseItem{ - {"@", int('@')}, {"AT", identifier}, {"?", paramMarker}, {"PLACEHOLDER", identifier}, @@ -69,8 +68,17 @@ func (s *testLexerSuite) TestSingleCharOther(c *C) { func (s *testLexerSuite) TestAtLeadingIdentifier(c *C) { defer testleak.AfterTest(c)() table := []testCaseItem{ + {"@", singleAtIdentifier}, + {"@''", singleAtIdentifier}, + {"@1", singleAtIdentifier}, + {"@.1_", singleAtIdentifier}, + {"@-1.", singleAtIdentifier}, + {"@~", singleAtIdentifier}, + {"@$", singleAtIdentifier}, {"@a_3cbbc", singleAtIdentifier}, - {"@-3cbbc", int('@')}, + {"@`a_3cbbc`", singleAtIdentifier}, + {"@-3cbbc", singleAtIdentifier}, + {"@!3cbbc", singleAtIdentifier}, {"@@global.test", doubleAtIdentifier}, {"@@session.test", doubleAtIdentifier}, {"@@local.test", doubleAtIdentifier}, diff --git a/parser/misc.go b/parser/misc.go index 6f042e867e24b..e0183805bccc0 100644 --- a/parser/misc.go +++ b/parser/misc.go @@ -36,8 +36,8 @@ func isIdentExtend(ch rune) bool { return ch >= 0x80 && ch <= '\uffff' } -func isIdentFirstChar(ch rune) bool { - return isLetter(ch) || ch == '_' +func isUserVarChar(ch rune) bool { + return isLetter(ch) || isDigit(ch) || ch == '_' || ch == '$' || ch == '.' || isIdentExtend(ch) } type trieNode struct { diff --git a/parser/parser_test.go b/parser/parser_test.go index 2d07df6a91dde..ab8b842a604c9 100644 --- a/parser/parser_test.go +++ b/parser/parser_test.go @@ -550,8 +550,17 @@ func (s *testParserSuite) TestDBAStmt(c *C) { {"load stats '/tmp/stats.json'", true}, // set // user defined + {"SET @ = 1", true}, + {"SET @' ' = 1", true}, + {"SET @! = 1", false}, + {"SET @1 = 1", true}, {"SET @a = 1", true}, {"SET @b := 1", true}, + {"SET @.c = 1", true}, + {"SET @_d = 1", true}, + {"SET @_e._$. = 1", true}, + {"SET @~f = 1", false}, + {"SET @`g,` = 1", true}, // session system variables {"SET SESSION autocommit = 1", true}, {"SET @@session.autocommit = 1", true},