use lexer to split multiple sql (#811)

* use lexer to split multiple sql

* normalize sqlflowserver_test

Author: typhoonzero

server/sqlflowserver.go

@@ -21,7 +21,6 @@ package server
 
 import (
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/golang/protobuf/proto"
@@ -56,26 +55,17 @@ func (s *Server) Run(req *pb.Request, stream pb.SQLFlow_RunServer) error {
 		}
 		defer db.Close()
 	}
-
-	// FIXME(typhoonzero): split by ; can not deal with situations like
-	// "SELECT * from mytable where col <> ';';", should be fixed.
-	sqlStatements := strings.Split(req.Sql, ";")
-	trimedStatements := []string{}
-	for _, singleSQL := range sqlStatements {
-		sqlToRun := strings.TrimSpace(singleSQL)
-		if sqlToRun == "" {
-			continue
-		}
-		trimedStatements = append(trimedStatements, sqlToRun)
+	sqlStatements, err := sf.SplitMultipleSQL(req.Sql)
+	if err != nil {
+		return err
 	}
-	for _, singleSQL := range trimedStatements {
-		sqlToRun := fmt.Sprintf("%s;", singleSQL)
+	for _, singleSQL := range sqlStatements {
 		var pr *sf.PipeReader
 		startTime := time.Now().UnixNano()
 		if s.enableSession == true {
-			pr = s.run(sqlToRun, db, s.modelDir, req.Session)
+			pr = s.run(singleSQL, db, s.modelDir, req.Session)
 		} else {
-			pr = s.run(sqlToRun, db, s.modelDir, nil)
+			pr = s.run(singleSQL, db, s.modelDir, nil)
 		}
 
 		defer pr.Close()
@@ -102,7 +92,7 @@ func (s *Server) Run(req *pb.Request, stream pb.SQLFlow_RunServer) error {
 			}
 		}
 		// Send EndOfExecution message if have multiple requests.
-		if len(trimedStatements) > 1 {
+		if len(sqlStatements) > 1 {
 			eoe := &pb.EndOfExecution{}
 			eoe.Sql = singleSQL
 			eoe.SpentTimeSeconds = time.Now().UnixNano() - startTime

server/sqlflowserver_test.go

@@ -20,7 +20,6 @@ import (
 	"log"
 	"net"
 	"os"
-	"strings"
 	"testing"
 	"time"
 
@@ -36,11 +35,12 @@ import (
 )
 
 const (
-	testErrorSQL             = "ERROR ..."
-	testQuerySQL             = "SELECT ..."
-	testExecuteSQL           = "INSERT ..."
-	testExtendedSQL          = "SELECT ... TRAIN ..."
-	testExtendedSQLWithSpace = "SELECT ... TRAIN ...; \n\t"
+	testErrorSQL               = "ERROR ..."
+	testQuerySQL               = "SELECT * FROM some_table;"
+	testExecuteSQL             = "INSERT INTO some_table VALUES (1,2,3,4);"
+	testExtendedSQL            = "SELECT * FROM some_table TRAIN SomeModel;"
+	testExtendedSQLNoSemicolon = "SELECT * FROM some_table TRAIN SomeModel"
+	testExtendedSQLWithSpace   = "SELECT * FROM some_table TRAIN SomeModel; \n\t"
 )
 
 var testServerAddress string
@@ -49,8 +49,6 @@ func mockRun(sql string, db *sf.DB, modelDir string, session *pb.Session) *sf.Pi
 	rd, wr := sf.Pipe()
 	go func() {
 		defer wr.Close()
-		// the server may automatically add a trailing ";", remove it
-		sql = strings.Trim(sql, ";")
 		switch sql {
 		case testErrorSQL:
 			wr.Write(fmt.Errorf("run error: %v", testErrorSQL))
@@ -65,7 +63,7 @@ func mockRun(sql string, db *sf.DB, modelDir string, session *pb.Session) *sf.Pi
 			wr.Write([]interface{}{time.Now(), nil})
 		case testExecuteSQL:
 			wr.Write("success; 0 rows affected")
-		case testExtendedSQL:
+		case testExtendedSQL, testExtendedSQLNoSemicolon, testExtendedSQLWithSpace:
 			wr.Write("log 0")
 			wr.Write("log 1")
 		default:
@@ -124,11 +122,10 @@ func TestSQL(t *testing.T) {
 	stream, err := c.Run(ctx, &pb.Request{Sql: testErrorSQL})
 	a.NoError(err)
 	_, err = stream.Recv()
-	a.Equal(status.Error(codes.Unknown, fmt.Sprintf("run error: %v", testErrorSQL)), err)
+	a.Equal(status.Error(codes.Unknown, "Lex: Unknown problem ..."), err)
 
-	testMultipleSQL := fmt.Sprintf("%s; %s", testQuerySQL, testExtendedSQL)
-
-	for _, s := range []string{testQuerySQL, testExecuteSQL, testExtendedSQL, testExtendedSQLWithSpace, testMultipleSQL} {
+	testMultipleSQL := fmt.Sprintf("%s %s", testQuerySQL, testExtendedSQL)
+	for _, s := range []string{testQuerySQL, testExecuteSQL, testExtendedSQL, testExtendedSQLWithSpace, testExtendedSQLNoSemicolon, testMultipleSQL} {
 		stream, err := c.Run(ctx, &pb.Request{Sql: s})
 		a.NoError(err)
 		for {

sql/executor.go

@@ -93,6 +93,35 @@ func splitExtendedSQL(slct string) ([]string, error) {
 	return []string{slct}, nil
 }
 
+// SplitMultipleSQL returns a list of SQL statements if the input statements contains mutiple
+// SQL statements separated by ;
+func SplitMultipleSQL(statements string) ([]string, error) {
+	l := newLexer(statements)
+	var n sqlSymType
+	var sqlList []string
+	splitPos := 0
+	for {
+		t := l.Lex(&n)
+		if t < 0 {
+			return []string{}, fmt.Errorf("Lex: Unknown problem %s", statements[0-t:])
+		}
+		if t == 0 {
+			if len(sqlList) == 0 {
+				// NOTE: this line support executing SQL statement without a trailing ";"
+				sqlList = append(sqlList, statements)
+			}
+			break
+		}
+		if t == ';' {
+			splited := statements[splitPos:l.pos]
+			splited = strings.TrimSpace(splited)
+			sqlList = append(sqlList, splited)
+			splitPos = l.pos
+		}
+	}
+	return sqlList, nil
+}
+
 // TODO(weiguo): isQuery is a hacky way to decide which API to call:
 // https://golang.org/pkg/database/sql/#DB.Exec .
 // We will need to extend our parser to be a full SQL parser in the future.

sql/executor_test.go

@@ -71,6 +71,16 @@ func TestSplitExtendedSQL(t *testing.T) {
 	a.Equal(`train a with b;`, s[0])
 }
 
+func TestSplitMulipleSQL(t *testing.T) {
+	a := assert.New(t)
+	splited, err := SplitMultipleSQL(`CREATE TABLE copy_table_1 AS SELECT a,b,c FROM table_1 WHERE c<>";";
+SELECT * FROM copy_table_1;SELECT * FROM copy_table_1 TRAIN DNNClassifier WITH n_classes=2 INTO test_model;`)
+	a.NoError(err)
+	a.Equal("CREATE TABLE copy_table_1 AS SELECT a,b,c FROM table_1 WHERE c<>\";\";", splited[0])
+	a.Equal("SELECT * FROM copy_table_1;", splited[1])
+	a.Equal("SELECT * FROM copy_table_1 TRAIN DNNClassifier WITH n_classes=2 INTO test_model;", splited[2])
+}
+
 func TestExecuteXGBoost(t *testing.T) {
 	a := assert.New(t)
 	modelDir := ""