Repository: vitessio/vitess
Pull request #6692 (2020-09-10 09:58:44)

Use statement support in olap mode
Support use statement in streaming mode.


Patch:
diff --git a/go/vt/vtgate/engine/update_target.go b/go/vt/vtgate/engine/update_target.go
index f637d4166e..a4451f808d 100644
--- a/go/vt/vtgate/engine/update_target.go
+++ b/go/vt/vtgate/engine/update_target.go
@@ -44,22 +44,22 @@ func (updTarget *UpdateTarget) description() PrimitiveDescription {
 }
 
 // RouteType implements the Primitive interface
-func (updTarget UpdateTarget) RouteType() string {
+func (updTarget *UpdateTarget) RouteType() string {
 	return "UpdateTarget"
 }
 
 // GetKeyspaceName implements the Primitive interface
-func (updTarget UpdateTarget) GetKeyspaceName() string {
+func (updTarget *UpdateTarget) GetKeyspaceName() string {
 	return updTarget.Target
 }
 
 // GetTableName implements the Primitive interface
-func (updTarget UpdateTarget) GetTableName() string {
+func (updTarget *UpdateTarget) GetTableName() string {
 	return ""
 }
 
 // Execute implements the Primitive interface
-func (updTarget UpdateTarget) Execute(vcursor VCursor, bindVars map[string]*query.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+func (updTarget *UpdateTarget) Execute(vcursor VCursor, bindVars map[string]*query.BindVariable, wantfields bool) (*sqltypes.Result, error) {
 	err := vcursor.Session().SetTarget(updTarget.Target)
 	if err != nil {
 		return nil, err
@@ -68,11 +68,15 @@ func (updTarget UpdateTarget) Execute(vcursor VCursor, bindVars map[string]*quer
 }
 
 // StreamExecute implements the Primitive interface
-func (updTarget UpdateTarget) StreamExecute(vcursor VCursor, bindVars map[string]*query.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
-	return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "use cannot be used for streaming")
+func (updTarget *UpdateTarget) StreamExecute(vcursor VCursor, bindVars map[string]*query.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	result, err := updTarget.Execute(vcursor, bindVars, wantfields)
+	if err != nil {
+		return err
+	}
+	return callback(result)
 }
 
 // GetFields implements the Primitive interface
-func (updTarget UpdateTarget) GetFields(vcursor VCursor, bindVars map[string]*query.BindVariable) (*sqltypes.Result, error) {
+func (updTarget *UpdateTarget) GetFields(vcursor VCursor, bindVars map[string]*query.BindVariable) (*sqltypes.Result, error) {
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "use cannot be used for get fields")
 }


Test patch:
diff --git a/go/test/endtoend/vtgate/misc_test.go b/go/test/endtoend/vtgate/misc_test.go
index 53664b06f0..c0b2beed76 100644
--- a/go/test/endtoend/vtgate/misc_test.go
+++ b/go/test/endtoend/vtgate/misc_test.go
@@ -345,6 +345,22 @@ func TestSwitchBetweenOlapAndOltp(t *testing.T) {
 	exec(t, conn, "set workload='oltp'")
 }
 
+func TestUseStmtInOLAP(t *testing.T) {
+	defer cluster.PanicHandler(t)
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	queries := []string{"set workload='olap'", "use `ks:80-`"}
+	for i, q := range queries {
+		t.Run(fmt.Sprintf("%d-%s", i, q), func(t *testing.T) {
+			exec(t, conn, q)
+			require.NoError(t, err)
+		})
+	}
+}
+
 func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
 	t.Helper()
 	qr := exec(t, conn, query)
diff --git a/go/vt/vtgate/engine/update_target_test.go b/go/vt/vtgate/engine/update_target_test.go
index d571f1409f..2b9479a20a 100644
--- a/go/vt/vtgate/engine/update_target_test.go
+++ b/go/vt/vtgate/engine/update_target_test.go
@@ -54,17 +54,15 @@ func TestUpdateTargetTable(t *testing.T) {
 			_, err := updateTarget.Execute(vc, map[string]*querypb.BindVariable{}, false)
 			require.NoError(t, err)
 			vc.ExpectLog(t, tc.expectedQueryLog)
+
+			vc = &loggingVCursor{}
+			_, err = wrapStreamExecute(updateTarget, vc, map[string]*querypb.BindVariable{}, false)
+			require.NoError(t, err)
+			vc.ExpectLog(t, tc.expectedQueryLog)
 		})
 	}
 }
 
-func TestUpdateTargetStreamExecute(t *testing.T) {
-	updateTarget := &UpdateTarget{}
-	vc := &noopVCursor{}
-	err := updateTarget.StreamExecute(vc, map[string]*querypb.BindVariable{}, false, nil)
-	require.EqualError(t, err, "use cannot be used for streaming")
-}
-
 func TestUpdateTargetGetFields(t *testing.T) {
 	updateTarget := &UpdateTarget{}
 	vc := &noopVCursor{}