Merge pull request #15 from rzhaovi/sql

SQL parsing fixes

script/parser.go

@@ -4,6 +4,7 @@ import (
 	"github.com/viant/toolbox"
 	"io"
 	"io/ioutil"
+	"regexp"
 	"strings"
 )
 
@@ -24,6 +25,7 @@ const (
 	commandEnd
 	pgDelimiter
 	plSQLBlock
+	inlineComment
 )
 
 //ParseWithReader splits SQL blob into separate commands
@@ -39,8 +41,10 @@ func ParseWithReader(reader io.Reader) []string {
 func appendMatched(terminator, pending *string, result *[]string) func(text string) {
 	return func(text string) {
 		SQL := strings.TrimSpace(*pending + text)
-		quotesCount := strings.Count(SQL, `'`)
-		if quotesCount%2 == 1 { //missing closing quote
+		SQL = regexp.MustCompile(`--.*\n`).ReplaceAllString(SQL, "")
+		SQL = regexp.MustCompile(`--.*\r`).ReplaceAllString(SQL, "")
+		quotesCount := strings.Count(SQL, `'`) - strings.Count(SQL, `\'`)
+		if quotesCount % 2 == 1 { //missing closing quote
 			*pending = SQL + *terminator
 			return
 		}
@@ -68,14 +72,13 @@ func parse(expression string, terminator string, delimiterMode bool) []string {
 		quoteTerminator:   toolbox.NewTerminatorMatcher(`'`),
 		delimiterKeyword:  toolbox.NewKeywordsMatcher(false, "delimiter"),
 		pgDelimiter:       toolbox.NewTerminatorMatcher("$$"),
-		plSQLBlock:        toolbox.NewBodyMatcher("BEGIN", "END;"),
+		plSQLBlock:        toolbox.NewBlockMatcher(false, "begin", "end;", []string{"CASE"}, []string{"END IF"}),
+		inlineComment:     toolbox.NewBodyMatcher("--", "\n"),
 		beginKeyword:      toolbox.NewTerminatorMatcher("BEGIN"),
 		createKeyword:     toolbox.NewKeywordsMatcher(false, "create"),
 		orKeyword:         toolbox.NewKeywordsMatcher(false, "or"),
 		replaceKeyword:    toolbox.NewKeywordsMatcher(false, "replace"),
-
 		lineBreakTerminator: toolbox.NewTerminatorMatcher("\n"),
-
 		functionKeyword: toolbox.NewKeywordsMatcher(false, "function"),
 		whitespaces:     toolbox.CharactersMatcher{" \n\t"},
 		lineBreak:       toolbox.CharactersMatcher{"\n"},
@@ -85,14 +88,15 @@ func parse(expression string, terminator string, delimiterMode bool) []string {
 
 	pending := ""
 	appendMatched := appendMatched(&terminator, &pending, &result)
-	done := false
 
 outer:
-	for tokenizer.Index < len(expression) && !done {
-		match := tokenizer.Nexts(whitespaces, createKeyword, delimiterKeyword, plSQLBlock, commandTerminator, commandEnd, eofToken)
+	for tokenizer.Index < len(expression) {
+		match := tokenizer.Nexts(whitespaces, inlineComment, createKeyword, delimiterKeyword, plSQLBlock, commandTerminator, commandEnd, eofToken)
 		switch match.Token {
 		case whitespaces:
 			pending += match.Matched
+		case inlineComment:
+			appendMatched("")
 		case delimiterKeyword:
 			if match := tokenizer.Nexts(lineBreakTerminator, eofToken); match.Token == lineBreakTerminator {
 				delimiter := strings.TrimSpace(string(match.Matched[:len(match.Matched)]))
@@ -117,7 +121,7 @@ outer:
 
 			if match := tokenizer.Nexts(whitespaces, eofToken); match.Token == whitespaces {
 				pending += match.Matched
-				candidates := []int{orKeyword, whitespaces, replaceKeyword, whitespaces, functionKeyword, beginKeyword, orKeyword}
+				candidates := []int{orKeyword, whitespaces, replaceKeyword, whitespaces, functionKeyword,  beginKeyword, orKeyword}
 				match := tokenizer.Nexts(candidates...)
 
 				for len(candidates) > 3 {
@@ -133,14 +137,22 @@ outer:
 
 				case functionKeyword:
 					pending += match.Matched
-					if match = tokenizer.Nexts(pgDelimiter, eofToken); match.Token == pgDelimiter {
-						pending += match.Matched + "$$"
-						tokenizer.Index += 2
+					if delimiterMode {
+						if match = tokenizer.Nexts(delimiterKeyword, eofToken); match.Token==delimiterKeyword {
+							pending += match.Matched + "$$"
+							tokenizer.Index += 2
+						}
+					} else {
 						if match = tokenizer.Nexts(pgDelimiter, eofToken); match.Token == pgDelimiter {
 							pending += match.Matched + "$$"
 							tokenizer.Index += 2
+							if match = tokenizer.Nexts(pgDelimiter, eofToken); match.Token == pgDelimiter {
+								pending += match.Matched + "$$"
+								tokenizer.Index += 2
+							}
 						}
 					}
+
 				case beginKeyword:
 					if strings.Contains(match.Matched, ";") || strings.Contains(match.Matched, "$$") {
 						tokenizer.Index -= len(match.Matched)
@@ -153,6 +165,10 @@ outer:
 
 				}
 			}
+		case pgDelimiter:
+			pending += match.Matched
+			appendMatched("")
+			tokenizer.Index += len("$$")
 		case eofToken:
 			pending += match.Matched
 		case plSQLBlock:

script/parser_test.go

@@ -201,6 +201,96 @@ INSERT INTO DUMMY(ID, NAME) VALUES(2, 'xyz');
 				`INSERT INTO DUMMY(ID, NAME) VALUES(2, 'xyz')`,
 			},
 		},
+		{
+			description: "Create function, begin end blocks",
+			SQL: `DELIMITER $$
+START TRANSACTION $$
+
+CREATE FUNCTION get_version()
+	RETURNS INT
+	READS SQL DATA
+BEGIN
+	DECLARE result INT DEFAULT 1;
+	SELECT coalesce(max(version_number), 1) INTO result FROM version;
+	RETURN result;
+END $$
+
+DROP PROCEDURE IF EXISTS set_version $$
+
+CREATE PROCEDURE set_version(version INT)
+BEGIN
+	DELETE FROM version;
+	INSERT INTO version VALUES(version);
+END $$
+
+COMMIT $$
+DELIMITER ;
+`,
+			SQLs: []string{
+				"START TRANSACTION",
+				"CREATE FUNCTION get_version()\n\tRETURNS INT\n\tREADS SQL DATA\nBEGIN\n\tDECLARE result INT DEFAULT 1;\n\tSELECT coalesce(max(version_number), 1) INTO result FROM version;\n\tRETURN result;\nEND",
+				"DROP PROCEDURE IF EXISTS set_version",
+				"CREATE PROCEDURE set_version(version INT)\nBEGIN\n\tDELETE FROM version;\n\tINSERT INTO version VALUES(version);\nEND",
+				"COMMIT",
+			},
+		},
+		{
+			description: "Comments and IF condition inside BEGIN-END block",
+			SQL:`DELIMITER $$
+DROP PROCEDURE IF EXISTS test_func $$
+CREATE PROCEDURE test_func()
+BEGIN
+  IF get_version() = 17
+  THEN
+    -- Set this to false by default
+    alter table TABLE_NAME add column IS_BLACK BOOLEAN DEFAULT FALSE;
+    CALL set_version(18);
+END IF;
+END $$
+DELIMITER ;
+`,
+			SQLs : []string{
+				`DROP PROCEDURE IF EXISTS test_func`,
+`CREATE PROCEDURE test_func()
+BEGIN
+  IF get_version() = 17
+  THEN
+        alter table TABLE_NAME add column IS_BLACK BOOLEAN DEFAULT FALSE;
+    CALL set_version(18);
+END IF;
+END`,
+			},
+
+		},
+		{
+			description: "'END;' string within BEGIN block before intended END;",
+			SQL:`DELIMITER $$
+DROP PROCEDURE IF EXISTS test_proc $$
+CREATE PROCEDURE test_proc()
+BEGIN
+  IF get_version() = 14
+  THEN
+    DROP VIEW IF EXISTS TABLE_NAME;
+    DROP VIEW IF EXISTS MONEY_SPEND;
+    DROP VIEW IF EXISTS OTHER_TABLE;
+    CALL set_version(15);
+  END IF;
+END $$
+DELIMITER ;`,
+			SQLs  : []string{
+				"DROP PROCEDURE IF EXISTS test_proc",
+				`CREATE PROCEDURE test_proc()
+BEGIN
+  IF get_version() = 14
+  THEN
+    DROP VIEW IF EXISTS TABLE_NAME;
+    DROP VIEW IF EXISTS MONEY_SPEND;
+    DROP VIEW IF EXISTS OTHER_TABLE;
+    CALL set_version(15);
+  END IF;
+END`,
+			},
+		},
 	}
 
 	for _, useCase := range useCases {