Generate workflow step for normal statement run

.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

@@ -144,6 +144,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