Additional options support for SELECT INTO and LOAD DATA

go/test/endtoend/vtgate/unsharded/main_test.go

@@ -0,0 +1,160 @@
+/*
+Copyright 2020 The Vitess Authors.
+
+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 unsharded
+
+import (
+ "context"
+ "flag"
+ "fmt"
+ "os"
+ "testing"
+
+ "github.com/google/go-cmp/cmp"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "vitess.io/vitess/go/mysql"
+ "vitess.io/vitess/go/sqltypes"
+ "vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+ clusterInstance *cluster.LocalProcessCluster
+ cell = "zone1"
+ hostname = "localhost"
+ KeyspaceName = "customer"
+ SchemaSQL = `
+CREATE TABLE t1 (
+ c1 BIGINT NOT NULL,
+ c2 BIGINT NOT NULL,
+ c3 BIGINT,
+ c4 varchar(100),
+ PRIMARY KEY (c1),
+ UNIQUE KEY (c2),
+ UNIQUE KEY (c3),
+ UNIQUE KEY (c4)
+) ENGINE=Innodb;
+`
+ VSchema = `
+{
+ "sharded": false,
+ "tables": {
+ "t1": {
+ "columns": [
+ {
+ "name": "c1",
+ "type": "INT64"
+ },
+ {
+ "name": "c2",
+ "type": "INT64"
+ },
+ {
+ "name": "c3",
+ "type": "INT64"
+ },
+ {
+ "name": "c4",
+ "type": "VARCHAR"
+ }
+ ],
+ }
+ }
+}
+`
+)
+
+func TestMain(m *testing.M) {
+ defer cluster.PanicHandler(nil)
+ flag.Parse()
+
+ exitCode := func() int {
+ clusterInstance = cluster.NewCluster(cell, hostname)
+ defer clusterInstance.Teardown()
+
+ // Start topo server
+ if err := clusterInstance.StartTopo(); err != nil {
+ return 1
+ }
+
+ // Start keyspace
+ Keyspace := &cluster.Keyspace{
+ Name: KeyspaceName,
+ SchemaSQL: SchemaSQL,
+ VSchema: VSchema,
+ }
+ if err := clusterInstance.StartUnshardedKeyspace(*Keyspace, 0, false); err != nil {
+ return 1
+ }
+
+ // Start vtgate
+ if err := clusterInstance.StartVtgate(); err != nil {
+ return 1
+ }
+
+ return m.Run()
+ }()
+ os.Exit(exitCode)
+}
+
+func TestSelectIntoAndLoadFrom(t *testing.T) {
+ defer cluster.PanicHandler(t)
+ ctx := context.Background()
+ vtParams := mysql.ConnParams{
+ Host: "localhost",
+ Port: clusterInstance.VtgateMySQLPort,
+ }
+ conn, err := mysql.Connect(ctx, &vtParams)
+ require.Nil(t, err)
+ defer conn.Close()
+
+ defer exec(t, conn, `delete from t1`)
+ exec(t, conn, `insert into t1(c1, c2, c3, c4) values (300,100,300,'abc')`)
+
+ exec(t, conn, `select * from t1 into outfile 'x.txt'`)
+ execAssertError(t, conn, `load data infile 'x.txt' into table t1`, "ERROR 1062 (23000): Duplicate entry '300' for key 'PRIMARY'")
+ exec(t, conn, `delete from t1`)
+ exec(t, conn, `load data infile 'x.txt' into table t1`)
+ assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
+ exec(t, conn, `select * from t1 into dumpfile 'x1.txt'`)
+ exec(t, conn, `select * from t1 into outfile 'x2.txt' Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`)
+ exec(t, conn, `load data infile 'x.txt' into replace table t1 Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`)
+ assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
+}
+
+func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
+ t.Helper()
+ qr, err := conn.ExecuteFetch(query, 1000, true)
+ require.NoError(t, err)
+ return qr
+}
+
+func execAssertError(t *testing.T, conn *mysql.Conn, query string, errorString string) {
+ t.Helper()
+ _, err := conn.ExecuteFetch(query, 1000, true)
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), errorString)
+}
+
+func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
+ t.Helper()
+ qr := exec(t, conn, query)
+ got := fmt.Sprintf("%v", qr.Rows)
+ diff := cmp.Diff(expected, got)
+ if diff != "" {
+ t.Errorf("Query: %s (-want +got):\n%s", query, diff)
+ }
+}

go/vt/sqlparser/ast.go

@@ -69,9 +69,23 @@ type (
  OrderBy OrderBy
  Limit *Limit
  Lock Lock
- IntoOutfileS3 string
+ Into  *SelectInto
  }
 
+ // SelectInto is a struct that represent the INTO part of a select query
+ SelectInto struct {
+ Type SelectIntoType
+ FileName string
+ Charset string
+ FormatOption string
+ ExportOption string
+ Manifest string
+ Overwrite string
+ }
+
+ // SelectIntoType is an enum for SelectInto.Type
+ SelectIntoType int8
+
  // Lock is an enum for the type of lock in the statement
  Lock int8
 
@@ -242,9 +256,8 @@ type (
  // DDLAction is an enum for DDL.Action
  DDLAction int8
 
- // Load is for s3 statement
+ // Load represents a LOAD statement
  Load struct {
- InfileS3 string
  }
 
  // ParenSelect is a parenthesized SELECT statement.
@@ -1014,14 +1027,11 @@ func (node *Select) Format(buf *TrackedBuffer) {
  addIf(node.StraightJoinHint, StraightJoinHint)
  addIf(node.SQLCalcFoundRows, SQLCalcFoundRowsStr)
 
- buf.astPrintf(node, "select %v%s%v from %v%v%v%v%v%v%s",
+ buf.astPrintf(node, "select %v%s%v from %v%v%v%v%v%v%s%v",
  node.Comments, options, node.SelectExprs,
  node.From, node.Where,
  node.GroupBy, node.Having, node.OrderBy,
- node.Limit, node.Lock.ToString())
- if node.IntoOutfileS3 != "" {
- buf.astPrintf(node, " into outfile s3 '%s'", node.IntoOutfileS3)
- }
+ node.Limit, node.Lock.ToString(), node.Into)
 }
 
 // Format formats the node.
@@ -2138,3 +2148,15 @@ func (node *ShowTableStatus) Format(buf *TrackedBuffer) {
  }
  buf.astPrintf(node, "%v", node.Filter)
 }
+
+// Format formats the node.
+func (node *SelectInto) Format(buf *TrackedBuffer) {
+ if node == nil {
+ return
+ }
+ buf.astPrintf(node, "%s'%s'", node.Type.ToString(), node.FileName)
+ if node.Charset != "" {
+ buf.astPrintf(node, " character set %s", node.Charset)
+ }
+ buf.astPrintf(node, "%s%s%s%s", node.FormatOption, node.ExportOption, node.Manifest, node.Overwrite)
+}

go/vt/sqlparser/ast_funcs.go

@@ -1122,6 +1122,19 @@ func (ty ExplainType) ToString() string {
  }
 }
 
+func (sel SelectIntoType) ToString() string {
+ switch sel {
+ case IntoOutfile:
+ return IntoOutfileStr
+ case IntoOutfileS3:
+ return IntoOutfileS3Str
+ case IntoDumpfile:
+ return IntoDumpfileStr
+ default:
+ return "Unknown Select Into Type"
+ }
+}
+
 // AtCount represents the '@' count in ColIdent
 type AtCount int
 

go/vt/sqlparser/constants.go

@@ -156,6 +156,11 @@ const (
  NaturalLanguageModeWithQueryExpansionStr = " in natural language mode with query expansion"
  QueryExpansionStr = " with query expansion"
 
+ // INTO OUTFILE
+ IntoOutfileStr = " into outfile "
+ IntoOutfileS3Str = " into outfile s3 "
+ IntoDumpfileStr = " into dumpfile "
+
  // Order.Direction
  AscScr = "asc"
  DescScr = "desc"
@@ -373,3 +378,10 @@ const (
  TraditionalType
  AnalyzeType
 )
+
+// Constant for Enum Type - SelectIntoType
+const (
+ IntoOutfile SelectIntoType = iota
+ IntoOutfileS3
+ IntoDumpfile
+)

go/vt/sqlparser/parse_test.go

@@ -26,6 +26,7 @@ import (
  "sync"
  "testing"
 
+ "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
 )
 
@@ -1051,12 +1052,6 @@ var (
  output: "alter database d",
  }, {
  input: "create table a",
- }, {
- input: "load data from s3 'x.txt'",
- output: "AST node missing for Load type",
- }, {
- input: "load data from s3 'x.txt' into table x",
- output: "AST node missing for Load type",
  }, {
  input: "create table a (\n\t`a` int\n)",
  output: "create table a (\n\ta int\n)",
@@ -2129,13 +2124,19 @@ func TestConvert(t *testing.T) {
  }
 }
 
-func TestIntoOutfileS3(t *testing.T) {
+func TestSelectInto(t *testing.T) {
  validSQL := []struct {
  input string
  output string
  }{{
  input: "select * from t order by name limit 100 into outfile s3 'out_file_name'",
  output: "select * from t order by name asc limit 100 into outfile s3 'out_file_name'",
+ }, {
+ input: "select * from t into dumpfile 'out_file_name'",
+ }, {
+ input: "select * from t into outfile 'out_file_name' character set binary fields terminated by 'term' optionally enclosed by 'c' escaped by 'e' lines starting by 'a' terminated by '\n'",
+ }, {
+ input: "select * from t into outfile s3 'out_file_name' character set binary format csv header fields terminated by 'term' optionally enclosed by 'c' escaped by 'e' lines starting by 'a' terminated by '\n' manifest on overwrite off",
  }, {
  input: "select * from (select * from t union select * from t2) as t3 where t3.name in (select col from t4) into outfile s3 'out_file_name'",
  }, {
@@ -2155,9 +2156,7 @@ func TestIntoOutfileS3(t *testing.T) {
  continue
  }
  out := String(tree)
- if out != tcase.output {
- t.Errorf("out: %s, want %s", out, tcase.output)
- }
+ assert.Equal(t, tcase.output, out)
  }
 
  invalidSQL := []struct {
@@ -2275,6 +2274,19 @@ func TestSubStr(t *testing.T) {
  }
 }
 
+func TestLoadData(t *testing.T) {
+ validSQL := []string{
+ "load data from s3 'x.txt'",
+ "load data from s3 manifest 'x.txt'",
+ "load data from s3 file 'x.txt'",
+ "load data infile 'x.txt' into table 'c'",
+ "load data from s3 'x.txt' into table x"}
+ for _, tcase := range validSQL {
+ _, err := Parse(tcase)
+ require.NoError(t, err)
+ }
+}
+
 func TestCreateTable(t *testing.T) {
  validSQL := []string{
  // test all the data types and options