query routing: vschema and DML tests

Signed-off-by: Sugu Sougoumarane <ssougou@gmail.com>

go/vt/vtgate/planbuilder/delete.go

@@ -30,18 +30,13 @@ import (
 
 // buildDeletePlan builds the instructions for a DELETE statement.
 func buildDeletePlan(del *sqlparser.Delete, vschema ContextVSchema) (*engine.Delete, error) {
-	edel := &engine.Delete{
-		Query: generateQuery(del),
-	}
+	edel := &engine.Delete{}
 	pb := newPrimitiveBuilder(vschema, newJointab(sqlparser.GetBindvars(del)))
-	if err := pb.processTableExprs(del.TableExprs); err != nil {
+	ro, err := pb.processDMLTable(del.TableExprs)
+	if err != nil {
 		return nil, err
 	}
-	rb, ok := pb.bldr.(*route)
-	if !ok {
-		return nil, errors.New("unsupported: multi-table/vindex delete statement in sharded keyspace")
-	}
-	ro := rb.routeOptions[0]
+	edel.Query = generateQuery(del)
 	edel.Keyspace = ro.eroute.Keyspace
 	if !edel.Keyspace.Sharded {
 		// We only validate non-table subexpressions because the previous analysis has already validated them.
@@ -65,15 +60,14 @@ func buildDeletePlan(del *sqlparser.Delete, vschema ContextVSchema) (*engine.Del
 	}
 
 	edel.QueryTimeout = queryTimeout(directives)
-	if rb.routeOptions[0].eroute.TargetDestination != nil {
-		if rb.routeOptions[0].eroute.TargetTabletType != topodatapb.TabletType_MASTER {
+	if ro.eroute.TargetDestination != nil {
+		if ro.eroute.TargetTabletType != topodatapb.TabletType_MASTER {
 			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "unsupported: DELETE statement with a replica target")
 		}
 		edel.Opcode = engine.DeleteByDestination
-		edel.TargetDestination = rb.routeOptions[0].eroute.TargetDestination
+		edel.TargetDestination = ro.eroute.TargetDestination
 		return edel, nil
 	}
-	var err error
 	edel.Vindex, edel.Values, err = getDMLRouting(del.Where, edel.Table)
 	// We couldn't generate a route for a single shard
 	// Execute a delete sharded

go/vt/vtgate/planbuilder/from.go

@@ -17,6 +17,7 @@ limitations under the License.
 package planbuilder
 
 import (
+	"errors"
 	"fmt"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -27,6 +28,22 @@ import (
 
 // This file has functions to analyze the FROM clause.
 
+// processDMLTable analyzes the FROM clause for DMLs and returns a routeOption.
+func (pb *primitiveBuilder) processDMLTable(tableExprs sqlparser.TableExprs) (*routeOption, error) {
+	if err := pb.processTableExprs(tableExprs); err != nil {
+		return nil, err
+	}
+	rb, ok := pb.bldr.(*route)
+	if !ok {
+		return nil, errors.New("unsupported: multi-shard or vindex write statement")
+	}
+	ro := rb.routeOptions[0]
+	for _, sub := range ro.substitutions {
+		*sub.oldExpr = *sub.newExpr
+	}
+	return ro, nil
+}
+
 // processTableExprs analyzes the FROM clause. It produces a builder
 // with all the routes identified.
 func (pb *primitiveBuilder) processTableExprs(tableExprs sqlparser.TableExprs) error {
@@ -192,7 +209,7 @@ func (pb *primitiveBuilder) buildTablePrimitive(tableExpr *sqlparser.AliasedTabl
 			if tableName.Name != vst.Name {
 				// Table name does not match. Change and alias it to old name.
 				sub.newExpr = &sqlparser.AliasedTableExpr{
-					Expr: &sqlparser.TableName{Name: vst.Name},
+					Expr: sqlparser.TableName{Name: vst.Name},
 					As:   tableName.Name,
 				}
 			}
@@ -201,7 +218,7 @@ func (pb *primitiveBuilder) buildTablePrimitive(tableExpr *sqlparser.AliasedTabl
 			if tableName.Name != vst.Name {
 				// Table name does not match. Change it and reuse existing alias.
 				sub.newExpr = &sqlparser.AliasedTableExpr{
-					Expr: &sqlparser.TableName{Name: vst.Name},
+					Expr: sqlparser.TableName{Name: vst.Name},
 					As:   tableExpr.As,
 				}
 			}

go/vt/vtgate/planbuilder/insert.go

@@ -31,16 +31,13 @@ import (
 // buildInsertPlan builds the route for an INSERT statement.
 func buildInsertPlan(ins *sqlparser.Insert, vschema ContextVSchema) (*engine.Insert, error) {
 	pb := newPrimitiveBuilder(vschema, newJointab(sqlparser.GetBindvars(ins)))
-	aliased := &sqlparser.AliasedTableExpr{Expr: ins.Table}
-	if err := pb.processAliasedTable(aliased); err != nil {
+	exprs := sqlparser.TableExprs{&sqlparser.AliasedTableExpr{Expr: ins.Table}}
+	ro, err := pb.processDMLTable(exprs)
+	if err != nil {
 		return nil, err
 	}
-	rb, ok := pb.bldr.(*route)
-	if !ok {
-		// This can happen only for vindexes right now.
-		return nil, fmt.Errorf("inserting into a vindex not allowed: %s", sqlparser.String(ins.Table))
-	}
-	ro := rb.routeOptions[0]
+	// The table might have been routed to a different one.
+	ins.Table = exprs[0].(*sqlparser.AliasedTableExpr).Expr.(sqlparser.TableName)
 	if ro.eroute.TargetDestination != nil {
 		return nil, errors.New("unsupported: INSERT with a target destination")
 	}

go/vt/vtgate/planbuilder/testdata/dml_cases.txt

@@ -90,6 +90,25 @@
   }
 }
 
+# update of a routed table
+"update route1 set a=1 where id=1"
+{
+  "Original": "update route1 set a=1 where id=1",
+  "Instructions": {
+    "Opcode": "UpdateEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user as route1 set a = 1 where id = 1",
+    "Vindex": "user_index",
+    "Values": [
+      1
+    ],
+    "Table": "user"
+  }
+}
+
 # delete unsharded
 "delete from unsharded"
 {
@@ -339,6 +358,26 @@
   }
 }
 
+# delete from a routed table
+"delete from route1 where id = 1"
+{
+  "Original": "delete from route1 where id = 1",
+  "Instructions": {
+    "Opcode": "DeleteEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "delete from user as route1 where id = 1",
+    "Vindex": "user_index",
+    "Values": [
+      1
+    ],
+    "Table": "user",
+    "OwnedVindexQuery": "select Name, Costly from user where id = 1 for update"
+  }
+}
+
 # update by lookup
 "update music set val = 1 where id = 1"
 {
@@ -682,6 +721,52 @@
   }
 }
 
+# insert into a routed table
+"insert into route1(id) values (1)"
+{
+  "Original": "insert into route1(id) values (1)",
+  "Instructions": {
+    "Opcode": "InsertSharded",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "insert into user(id, Name, Costly) values (:_Id0, :_Name0, :_Costly0)",
+    "Values": [
+      [
+        [
+          ":__seq0"
+        ]
+      ],
+      [
+        [
+          null
+        ]
+      ],
+      [
+        [
+          null
+        ]
+      ]
+    ],
+    "Table": "user",
+    "Generate": {
+      "Keyspace": {
+        "Name": "main",
+        "Sharded": false
+      },
+      "Query": "select next :n values from seq",
+      "Values": [
+        1
+      ]
+    },
+    "Prefix": "insert into user(id, Name, Costly) values ",
+    "Mid": [
+      "(:_Id0, :_Name0, :_Costly0)"
+    ]
+  }
+}
+
 # insert no column list
 "insert into user values(1, 2, 3)"
 "no column list"
@@ -1044,7 +1129,7 @@
 
 # insert into a vindex not allowed
 "insert into user_index(id) values(1)"
-"inserting into a vindex not allowed: user_index"
+"unsupported: multi-shard or vindex write statement"
 
 # simple replace unsharded
 "replace into unsharded values(1, 2)"

go/vt/vtgate/planbuilder/testdata/schema_test.json

@@ -1,4 +1,10 @@
 {
+  "routing_rules": {
+    "rules": [{
+      "from_table": "route1",
+      "to_tables": ["user.user", "user.user_metadata"]
+    }]
+  },
   "keyspaces": {
     "user": {
       "sharded": true,

go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt

@@ -249,7 +249,7 @@
 
 # multi delete multi table
 "delete user from user join user_extra on user.id = user_extra.id where user.name = 'foo'"
-"unsupported: multi-table/vindex delete statement in sharded keyspace"
+"unsupported: multi-shard or vindex write statement"
 
 # scatter delete with owned lookup vindex
 "delete from user"
@@ -285,11 +285,11 @@
 
 # join in update tables
 "update user join user_extra on user.id = user_extra.id set user.name = 'foo'"
-"unsupported: multi-table/vindex update statement in sharded keyspace"
+"unsupported: multi-shard or vindex write statement"
 
 # multiple tables in update
 "update user as u, user_extra as ue set u.name = 'foo' where u.id = ue.id"
-"unsupported: multi-table/vindex update statement in sharded keyspace"
+"unsupported: multi-shard or vindex write statement"
 
 # unsharded insert with cross-shard join"
 "insert into unsharded select u.col from user u join user u1"

go/vt/vtgate/planbuilder/update.go

@@ -32,18 +32,14 @@ import (
 // buildUpdatePlan builds the instructions for an UPDATE statement.
 func buildUpdatePlan(upd *sqlparser.Update, vschema ContextVSchema) (*engine.Update, error) {
 	eupd := &engine.Update{
-		Query:               generateQuery(upd),
 		ChangedVindexValues: make(map[string][]sqltypes.PlanValue),
 	}
 	pb := newPrimitiveBuilder(vschema, newJointab(sqlparser.GetBindvars(upd)))
-	if err := pb.processTableExprs(upd.TableExprs); err != nil {
+	ro, err := pb.processDMLTable(upd.TableExprs)
+	if err != nil {
 		return nil, err
 	}
-	rb, ok := pb.bldr.(*route)
-	if !ok {
-		return nil, errors.New("unsupported: multi-table/vindex update statement in sharded keyspace")
-	}
-	ro := rb.routeOptions[0]
+	eupd.Query = generateQuery(upd)
 	eupd.Keyspace = ro.eroute.Keyspace
 	if !eupd.Keyspace.Sharded {
 		// We only validate non-table subexpressions because the previous analysis has already validated them.
@@ -71,7 +67,6 @@ func buildUpdatePlan(upd *sqlparser.Update, vschema ContextVSchema) (*engine.Upd
 	if eupd.Table == nil {
 		return nil, errors.New("internal error: table.vindexTable is mysteriously nil")
 	}
-	var err error
 
 	if ro.eroute.TargetDestination != nil {
 		if ro.eroute.TargetTabletType != topodatapb.TabletType_MASTER {

go/vt/vtgate/vindexes/vschema_test.go

@@ -1798,6 +1798,18 @@ func TestFindTable(t *testing.T) {
 
 func TestFindTablesOrVindex(t *testing.T) {
 	input := vschemapb.SrvVSchema{
+		RoutingRules: &vschemapb.RoutingRules{
+			Rules: []*vschemapb.RoutingRule{{
+				FromTable: "unqualified",
+				ToTables:  []string{"ksa.ta", "ksb.t1"},
+			}, {
+				FromTable: "newks.qualified",
+				ToTables:  []string{"ksa.ta"},
+			}, {
+				FromTable: "notarget",
+				ToTables:  []string{},
+			}},
+		},
 		Keyspaces: map[string]*vschemapb.Keyspace{
 			"ksa": {
 				Tables: map[string]*vschemapb.Table{
@@ -1841,6 +1853,8 @@ func TestFindTablesOrVindex(t *testing.T) {
 		},
 	}
 	vschema, _ := BuildVSchema(&input)
+	ta := vschema.Keyspaces["ksa"].Tables["ta"]
+	t1 := vschema.Keyspaces["ksb"].Tables["t1"]
 
 	_, _, err := vschema.FindTablesOrVindex("", "t1")
 	wantErr := "ambiguous table reference: t1"
@@ -1858,12 +1872,6 @@ func TestFindTablesOrVindex(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ta := &Table{
-		Name: sqlparser.NewTableIdent("ta"),
-		Keyspace: &Keyspace{
-			Name: "ksa",
-		},
-	}
 	if !reflect.DeepEqual(got, []*Table{ta}) {
 		t.Errorf("FindTablesOrVindex(\"t1a\"): %+v, want %+v", got, ta)
 	}
@@ -1895,6 +1903,28 @@ func TestFindTablesOrVindex(t *testing.T) {
 	if err == nil || err.Error() != wantErr {
 		t.Errorf("FindTablesOrVindex(\"\"): %v, want %s", err, wantErr)
 	}
+
+	got, _, err = vschema.FindTablesOrVindex("", "unqualified")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := []*Table{ta, t1}; !reflect.DeepEqual(got, want) {
+		t.Errorf("FindTablesOrVindex(unqualified): %+v, want %+v", got, want)
+	}
+
+	got, _, err = vschema.FindTablesOrVindex("newks", "qualified")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := []*Table{ta}; !reflect.DeepEqual(got, want) {
+		t.Errorf("FindTablesOrVindex(unqualified): %+v, want %+v", got, want)
+	}
+
+	_, _, err = vschema.FindTablesOrVindex("", "notarget")
+	wantErr = "table notarget has been disabled"
+	if err == nil || err.Error() != wantErr {
+		t.Errorf("FindTablesOrVindex(\"\"): %v, want %s", err, wantErr)
+	}
 }
 
 func TestBuildKeyspaceSchema(t *testing.T) {