From 471e0d359b5a22debaf378faf16af835d26a15b5 Mon Sep 17 00:00:00 2001
From: Leon Zhang <martianzhang@hotmail.com>
Date: Wed, 20 Feb 2019 14:37:00 +0800
Subject: [PATCH] fix -report-type qurery-type bug

  GRANT SELECT ON *.* to user@'host'
  REVOKE INSERT ON *.* FROM 'jeffrey'@'localhost';
---
 ast/testdata/TestQueryType.golden |  3 +++
 ast/token.go                      | 15 +++++++++------
 ast/token_test.go                 |  6 ++++++
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/ast/testdata/TestQueryType.golden b/ast/testdata/TestQueryType.golden
index 7b9b540..8f01697 100644
--- a/ast/testdata/TestQueryType.golden
+++ b/ast/testdata/TestQueryType.golden
@@ -1,5 +1,8 @@
 SELECT
 SELECT
+GRANT
+REVOKE
+SELECT
 SELECT
 SELECT
 SELECT
diff --git a/ast/token.go b/ast/token.go
index 9021661..70afb52 100644
--- a/ast/token.go
+++ b/ast/token.go
@@ -446,6 +446,7 @@ var mySQLKeywords = map[string]string{
 	"geometry":           "GEOMETRY",
 	"geometrycollection": "GEOMETRYCOLLECTION",
 	"global":             "GLOBAL",
+	"grant":              "GRANT",
 	"group":              "GROUP",
 	"group_concat":       "GROUP_CONCAT",
 	"having":             "HAVING",
@@ -512,6 +513,7 @@ var mySQLKeywords = map[string]string{
 	"reorganize":         "REORGANIZE",
 	"repair":             "REPAIR",
 	"replace":            "REPLACE",
+	"revoke":             "REVOKE",
 	"right":              "RIGHT",
 	"rlike":              "REGEXP",
 	"rollback":           "ROLLBACK",
@@ -989,13 +991,14 @@ func NewLines(buf []byte) int {
 
 // QueryType get query type such as SELECT/INSERT/DELETE/CREATE/ALTER
 func QueryType(sql string) string {
-	var typ string
-	tokens := Tokenizer(sql)
+	tokens := Tokenize(sql)
 	for _, token := range tokens {
-		if val, ok := mySQLKeywords[token.Val]; ok {
-			typ = val
-			break
+		// use strings.Fields for 'ALTER TABLE' token split
+		for _, tk := range strings.Fields(strings.TrimSpace(token.Val)) {
+			if val, ok := mySQLKeywords[strings.ToLower(tk)]; ok {
+				return val
+			}
 		}
 	}
-	return typ
+	return ""
 }
diff --git a/ast/token_test.go b/ast/token_test.go
index d72e718..c8f7ec9 100644
--- a/ast/token_test.go
+++ b/ast/token_test.go
@@ -241,8 +241,14 @@ func TestNewLines(t *testing.T) {
 func TestQueryType(t *testing.T) {
 	common.Log.Debug("Entering function: %s", common.GetFunctionName())
 	var testSQLs = []string{
+		`/*comment*/ select 1`,
 		`(select 1)`,
+		`grant select on *.* to user@'localhost'`,
+		`REVOKE INSERT ON *.* FROM 'jeffrey'@'localhost';`,
 	}
+	// fmt.Println(testSQLs[len(testSQLs)-1])
+	// fmt.Println(QueryType(testSQLs[len(testSQLs)-1]))
+	// return
 	err := common.GoldenDiff(func() {
 		for _, buf := range append(testSQLs, common.TestSQLs...) {
 			fmt.Println(QueryType(buf))
-- 
GitLab