lexer: support more syntaxes regarding 'SET Syntax' (#7020)

pingcap · Jul 10, 2018 · fd37061 · fd37061
1 parent 20a985d
commit fd37061
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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},