From 147ee2d0817637f4e63a959681c580e2ada85cd8 Mon Sep 17 00:00:00 2001
From: Leon Zhang <martianzhang@hotmail.com>
Date: Wed, 21 Nov 2018 23:15:50 +0800
Subject: [PATCH] fix #116

  SplitStatement check if single comment line is in multi-line sql.

  pure comment line should break quickly, single comment in multi-line
  sql should wait for line delimiter
---
 ast/testdata/TestSplitStatement.golden | 15 +++++++++++++++
 ast/token.go                           |  7 ++++++-
 ast/token_test.go                      | 13 +++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/ast/testdata/TestSplitStatement.golden b/ast/testdata/TestSplitStatement.golden
index e4d6516..8d08708 100644
--- a/ast/testdata/TestSplitStatement.golden
+++ b/ast/testdata/TestSplitStatement.golden
@@ -16,6 +16,21 @@
 13 select
 *
 -- comment
+from tb
+where col = 1
+14 select
+* --
+from tb
+where col = 1
+15 select
+* #
+from tb
+where col = 1
+16 select
+*
+--
+from tb
+where col = 1
 0 select * from test\Ghello
 1 select 'hello\Gworld', col from test\Ghello
 2 -- select * from test\Ghello
diff --git a/ast/token.go b/ast/token.go
index 8502b63..c326127 100644
--- a/ast/token.go
+++ b/ast/token.go
@@ -928,7 +928,12 @@ func SplitStatement(buf []byte, delimiter []byte) (string, string, []byte) {
 		if singleLineComment {
 			if b == '\r' || b == '\n' {
 				sql = string(buf[:i])
-				break
+				if strings.HasPrefix(sql, "--") || strings.HasPrefix(sql, "#") {
+					// just comment, query start with '--', '#'
+					break
+				}
+				// comment in multi-line sql
+				continue
 			}
 		}
 
diff --git a/ast/token_test.go b/ast/token_test.go
index 8bb79dd..a62b3fa 100644
--- a/ast/token_test.go
+++ b/ast/token_test.go
@@ -131,6 +131,19 @@ func TestSplitStatement(t *testing.T) {
 *
 -- comment
 from tb
+where col = 1`),
+		[]byte(`select
+* --
+from tb
+where col = 1`),
+		[]byte(`select
+* #
+from tb
+where col = 1`),
+		[]byte(`select
+*
+--
+from tb
 where col = 1`),
 	}
 	buf2s := [][]byte{
-- 
GitLab