add split sql method

zr-hebo · bramp · commit c4b3a90a35b4 · 2018-06-05T07:41:09.000-07:00
Signed-off-by: zr-hebo &lt;zrhebo@gmail.com&gt;
diff --git a/ast.go b/ast.go
@@ -114,6 +114,37 @@ func SplitStatement(blob string) (string, string, error) {
 	return blob, "", nil
 }
 
+// SplitStatementToPieces split raw sql statement that may have multi sql pieces to sql pieces
+// returns the sql pieces blob contains; or error if sql cannot be parsed
+func SplitStatementToPieces(blob string) (pieces []string, err error) {
+	pieces = make([]string, 0, 16)
+	tokenizer := NewStringTokenizer(blob)
+
+	tkn := 0
+	var stmt string
+	stmtBegin := 0
+	for {
+		tkn, _ = tokenizer.Scan()
+		if tkn == ';' {
+			stmt = blob[stmtBegin : tokenizer.Position-2]
+			pieces = append(pieces, stmt)
+			stmtBegin = tokenizer.Position - 1
+
+		} else if tkn == 0 || tkn == eofChar {
+			blobTail := tokenizer.Position - 2
+
+			if stmtBegin < blobTail {
+				stmt = blob[stmtBegin : blobTail+1]
+				pieces = append(pieces, stmt)
+			}
+			break
+		}
+	}
+
+	err = tokenizer.LastError
+	return
+}
+
 // SQLNode defines the interface for all nodes
 // generated by the parser.
 type SQLNode interface {
diff --git a/ast_test.go b/ast_test.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"reflect"
+	"strings"
 	"testing"
 	"unsafe"
 
@@ -441,3 +442,48 @@ func TestColumns_FindColumn(t *testing.T) {
 		}
 	}
 }
+
+func TestSplitStatementToPieces(t *testing.T) {
+	testcases := []struct {
+		input  string
+		output string
+	}{{
+		input: "select * from table",
+	}, {
+		input:  "select * from table1; select * from table2;",
+		output: "select * from table1; select * from table2",
+	}, {
+		input:  "select * from /* comment ; */ table;",
+		output: "select * from /* comment ; */ table",
+	}, {
+		input:  "select * from table where semi = ';';",
+		output: "select * from table where semi = ';'",
+	}, {
+		input:  "select * from table1;--comment;\nselect * from table2;",
+		output: "select * from table1;--comment;\nselect * from table2",
+	}, {
+		input: "CREATE TABLE `total_data` (`id` int(11) NOT NULL AUTO_INCREMENT COMMENT 'id', " +
+			"`region` varchar(32) NOT NULL COMMENT 'region name, like zh; th; kepler'," +
+			"`data_size` bigint NOT NULL DEFAULT '0' COMMENT 'data size;'," +
+			"`createtime` datetime NOT NULL DEFAULT NOW() COMMENT 'create time;'," +
+			"`comment` varchar(100) NOT NULL DEFAULT '' COMMENT 'comment'," +
+			"PRIMARY KEY (`id`))",
+	}}
+
+	for _, tcase := range testcases {
+		if tcase.output == "" {
+			tcase.output = tcase.input
+		}
+
+		stmtPieces, err := SplitStatementToPieces(tcase.input)
+		if err != nil {
+			t.Errorf("input: %s, err: %v", tcase.input, err)
+			continue
+		}
+
+		out := strings.Join(stmtPieces, ";")
+		if out != tcase.output {
+			t.Errorf("out: %s, want %s", out, tcase.output)
+		}
+	}
+}
