From 153c0c51f20572dc878bbe60c6a1905a3416ea0a Mon Sep 17 00:00:00 2001
From: Leon Zhang <martianzhang@hotmail.com>
Date: Mon, 29 Oct 2018 22:39:22 +0800
Subject: [PATCH] fix #61 :bug:

---
 ast/pretty_test.go             |  3 +++
 ast/testdata/TestPretty.golden | 16 ++++++++++++++--
 ast/token.go                   |  3 ++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/ast/pretty_test.go b/ast/pretty_test.go
index 877c83b..5c44c7d 100644
--- a/ast/pretty_test.go
+++ b/ast/pretty_test.go
@@ -122,6 +122,9 @@ var TestSqlsPretty = []string{
 	// MODIFY & CHANGE
 	"ALTER TABLE t1 MODIFY col1 BIGINT UNSIGNED DEFAULT 1 COMMENT 'my column';",
 	"ALTER TABLE t1 CHANGE b a INT NOT NULL;",
+
+	// COMMENT
+	"/*!40000 select 1*/;",
 }
 
 func TestPretty(t *testing.T) {
diff --git a/ast/testdata/TestPretty.golden b/ast/testdata/TestPretty.golden
index 5eaf8c5..57d66d9 100644
--- a/ast/testdata/TestPretty.golden
+++ b/ast/testdata/TestPretty.golden
@@ -176,6 +176,8 @@ FROM
 SELECT /*!40001 SQL_NO_CACHE */ * FROM `film`
 
 SELECT  
+  
+  /*!40001 SQL_NO_CACHE 
   */ * 
 FROM  
   `film`
@@ -363,7 +365,9 @@ SELECT
 FROM  
   table  
 WHERE  
-  FIELD  = 'value' */
+  FIELD  = 'value' 
+  /*arbitrary/31
+  */
 SELECT * FROM table WHERE field = 'value' /*arbitrary31*/ 
 
 SELECT  
@@ -371,7 +375,9 @@ SELECT
 FROM  
   table  
 WHERE  
-  FIELD  = 'value' */
+  FIELD  = 'value' 
+  /*arbitrary31
+  */
 SELECT *    FROM t WHERE 1=1 AND id=1
 
 SELECT  
@@ -503,6 +509,8 @@ HAVING
     1, 0)
   )= 0
 /* -- S++ SU ABORTABLE -- spd_user: rspadim */SELECT SQL_SMALL_RESULT SQL_CACHE DISTINCT centro_atividade FROM est_dia WHERE unidade_id=1001 AND item_id=67 AND item_id_red=573
+
+/* -- S++ SU ABORTABLE -- spd_user: rspadim 
 */ 
 SELECT  
   SQL_SMALL_RESULT  SQL_CACHE  DISTINCT  centro_atividade  
@@ -630,6 +638,10 @@ ALTER TABLE t1 CHANGE b a INT NOT NULL;
 
 ALTER TABLE  
   t1  CHANGE  b  a  INT  NOT  NULL;
+/*!40000 select 1*/;
+
+/*!40000 select 1
+*/;
 SELECT * FROM film WHERE length = 86;
 
 SELECT  
diff --git a/ast/token.go b/ast/token.go
index 43aa7c4..0b273e8 100644
--- a/ast/token.go
+++ b/ast/token.go
@@ -676,7 +676,7 @@ func getNextToken(buf string, previous Token) Token {
 	}
 
 	// Comment (#, --, /**/)
-	if buf[0] == '#' || (len(buf) > 1 && (buf[0] == '-' && buf[1] == '-')) || (buf[0] == '/' && buf[1] == '*') {
+	if buf[0] == '#' || (len(buf) > 1 && (buf[:2] == "--" || buf[:2] == "/*")) {
 		var last int
 		if buf[0] == '-' || buf[0] == '#' {
 			// Comment until end of line
@@ -685,6 +685,7 @@ func getNextToken(buf string, previous Token) Token {
 		} else {
 			// Comment until closing comment tag
 			last = strings.Index(buf[2:], "*/") + 2
+			typ = TokenTypeBlockComment
 		}
 		if last == 0 {
 			last = len(buf)
-- 
GitLab