Generate workflow step for normal statement run (#2824)

* generate workflow step for normal statement run

* clean up

* build step image before run workflow test

* fix is_query

Author: typhoonzero

.github/workflows/main.yml

@@ -81,7 +81,10 @@ jobs:
         set -e
         bash scripts/test/prepare.sh
         source build/env/bin/activate
+        # build sqlflow binaries under build/
+        bash docker/dev/build.sh
         docker pull sqlflow/sqlflow:step
+        docker build --cache-from sqlflow/sqlflow:step -t sqlflow/sqlflow:step --build-arg FIND_FASTED_MIRROR="false" -f docker/step/Dockerfile .
         bash scripts/test/workflow.sh
         # bash scripts/travis/upload_codecov.sh
   push-images:

go/cmd/sqlflowserver/e2e_workflow_test.go

@@ -362,7 +362,9 @@ func TestEnd2EndFluidWorkflow(t *testing.T) {
 func CaseWorkflowTrainXgboost(t *testing.T) {
 	a := assert.New(t)
 
-	sqlProgram := `SELECT * FROM iris.train
+	sqlProgram := `SELECT * FROM iris.train LIMIT 100;
+
+SELECT * FROM iris.train
 TO TRAIN xgboost.gbtree
 WITH objective="multi:softmax",num_class=3
 LABEL class

go/codegen/experimental/codegen.go

@@ -29,6 +29,9 @@ func generateStepCode(stmt ir.SQLFlowStmt, stepIndex int, session *pb.Session) (
 			return XGBoostGenerateTrain(trainStmt, stepIndex, session)
 		}
 		return "", fmt.Errorf("not implemented estimator type %s", trainStmt.Estimator)
+	case *ir.NormalStmt:
+		stmt := stmt.(*ir.NormalStmt)
+		return GenerateNormalStmtStep(string(*stmt), session, stepIndex)
 	default:
 		return "", fmt.Errorf("not implemented stmt execution type %v", stmt)
 	}

go/codegen/experimental/codegen_normal_stmt.go

@@ -0,0 +1,58 @@
+// Copyright 2020 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package experimental
+
+import (
+	"bytes"
+	"text/template"
+
+	pb "sqlflow.org/sqlflow/go/proto"
+)
+
+var normalStmtStepTmpl = `
+def step_entry_{{.StepIndex}}():
+    import runtime
+    import runtime.dbapi
+    conn = runtime.dbapi.connect("{{.DataSource}}")
+    stmt = """{{.Stmt}}"""
+    if conn.is_query(stmt):
+        rs = conn.query(stmt)
+        # write rs to stdout using protobuf table writer
+    else:
+        success = conn.execute(stmt)
+        if not success:
+            raise Exception("execute statment error: " % stmt)
+`
+
+var normalStmtStepTemplate = template.Must(template.New("NormalStmtStep").Parse(normalStmtStepTmpl))
+
+type normalStmtFiller struct {
+	StepIndex  int
+	DataSource string
+	Stmt       string
+}
+
+// GenerateNormalStmtStep generate step Python code to run a normal SQL statement.
+func GenerateNormalStmtStep(stmt string, session *pb.Session, stepIndex int) (string, error) {
+	filler := &normalStmtFiller{
+		StepIndex:  stepIndex,
+		DataSource: session.DbConnStr,
+		Stmt:       stmt,
+	}
+	var program bytes.Buffer
+	if err := normalStmtStepTemplate.Execute(&program, filler); err != nil {
+		return "", err
+	}
+	return program.String(), nil
+}

python/runtime/dbapi/connection.py

@@ -143,6 +143,20 @@ def query(self, statement):
         """
         return self._get_result_set(statement)
 
+    def is_query(self, statement):
+        """Return true if the statement is a query SQL statement."""
+        s = statement.strip()
+        s = s.upper()
+
+        if s.startswith("SELECT") and s.find("INTO") == -1:
+            return True
+        if s.startswith("SHOW") and s.find("CREATE") >= 0 or s.find(
+                "DATABASES") >= 0 or s.find("TABLES") >= 0:
+            return True
+        if s.startswith("DESC") or s.startswith("EXPLAIN"):
+            return True
+        return False
+
     def execute(self, statement):
         """Execute given statement and return True on success