vreplication: keyspace_id() support

go/vt/vttablet/tabletmanager/vreplication/replicator_plan_test.go

@@ -461,6 +461,56 @@ func TestBuildPlayerPlan(t *testing.T) {
 				},
 			},
 		},
+	}, {
+		// keyspace_id
+		input: &binlogdatapb.Filter{
+			Rules: []*binlogdatapb.Rule{{
+				Match:  "t1",
+				Filter: "select c1, c2, keyspace_id() ksid from t1",
+			}},
+		},
+		plan: &TestReplicatorPlan{
+			VStreamFilter: &binlogdatapb.Filter{
+				Rules: []*binlogdatapb.Rule{{
+					Match:  "t1",
+					Filter: "select c1, c2, keyspace_id() from t1",
+				}},
+			},
+			TargetTables: []string{"t1"},
+			TablePlans: map[string]*TestTablePlan{
+				"t1": {
+					TargetName:   "t1",
+					SendRule:     "t1",
+					PKReferences: []string{"c1"},
+					InsertFront:  "insert into t1(c1,c2,ksid)",
+					InsertValues: "(:a_c1,:a_c2,:a_keyspace_id)",
+					Insert:       "insert into t1(c1,c2,ksid) values (:a_c1,:a_c2,:a_keyspace_id)",
+					Update:       "update t1 set c2=:a_c2, ksid=:a_keyspace_id where c1=:b_c1",
+					Delete:       "delete from t1 where c1=:b_c1",
+				},
+			},
+		},
+		planpk: &TestReplicatorPlan{
+			VStreamFilter: &binlogdatapb.Filter{
+				Rules: []*binlogdatapb.Rule{{
+					Match:  "t1",
+					Filter: "select c1, c2, keyspace_id(), pk1, pk2 from t1",
+				}},
+			},
+			TargetTables: []string{"t1"},
+			TablePlans: map[string]*TestTablePlan{
+				"t1": {
+					TargetName:   "t1",
+					SendRule:     "t1",
+					PKReferences: []string{"c1", "pk1", "pk2"},
+					InsertFront:  "insert into t1(c1,c2,ksid)",
+					InsertValues: "(:a_c1,:a_c2,:a_keyspace_id)",
+					Insert:       "insert into t1(c1,c2,ksid) select :a_c1, :a_c2, :a_keyspace_id from dual where (:a_pk1,:a_pk2) <= (1,'aaa')",
+					Update:       "update t1 set c2=:a_c2, ksid=:a_keyspace_id where c1=:b_c1 and (:b_pk1,:b_pk2) <= (1,'aaa')",
+					Delete:       "delete from t1 where c1=:b_c1 and (:b_pk1,:b_pk2) <= (1,'aaa')",
+				},
+			},
+		},
 	}, {
 		// syntax error
 		input: &binlogdatapb.Filter{

go/vt/vttablet/tabletmanager/vreplication/table_plan_builder.go

@@ -49,7 +49,9 @@ type colExpr struct {
 	// operation==opCount: nothing is set.
 	// operation==opSum: for 'sum(a)', expr is set to 'a'.
 	operation operation
-	expr      sqlparser.Expr
+	// expr stores the expected field name from vstreamer and dictates
+	// the generated bindvar names, like a_col or b_col.
+	expr sqlparser.Expr
 	// references contains all the column names referenced in the expression.
 	references map[string]bool
 
@@ -332,6 +334,14 @@ func (tpb *tablePlanBuilder) analyzeExpr(selExpr sqlparser.SelectExpr) (*colExpr
 			tpb.addCol(innerCol.Name)
 			cexpr.references[innerCol.Name.Lowered()] = true
 			return cexpr, nil
+		case "keyspace_id":
+			if len(expr.Exprs) != 0 {
+				return nil, fmt.Errorf("unexpected: %v", sqlparser.String(expr))
+			}
+			tpb.sendSelect.SelectExprs = append(tpb.sendSelect.SelectExprs, &sqlparser.AliasedExpr{Expr: aliased.Expr})
+			// The vstreamer responds with "keyspace_id" as the field name for this request.
+			cexpr.expr = &sqlparser.ColName{Name: sqlparser.NewColIdent("keyspace_id")}
+			return cexpr, nil
 		}
 	}
 	err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {

go/vt/vttablet/tabletmanager/vreplication/vplayer_test.go

@@ -447,6 +447,82 @@ func TestPlayerKeywordNames(t *testing.T) {
 		}
 	}
 }
+
+var shardedVSchema = `{
+  "sharded": true,
+  "vindexes": {
+    "hash": {
+      "type": "hash"
+    }
+  },
+  "tables": {
+    "src1": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "hash"
+        }
+      ]
+    }
+  }
+}`
+
+func TestPlayerKeyspaceID(t *testing.T) {
+	defer deleteTablet(addTablet(100))
+
+	execStatements(t, []string{
+		"create table src1(id int, val varbinary(128), primary key(id))",
+		fmt.Sprintf("create table %s.dst1(id int, val varbinary(128), primary key(id))", vrepldb),
+	})
+	defer execStatements(t, []string{
+		"drop table src1",
+		fmt.Sprintf("drop table %s.dst1", vrepldb),
+	})
+	env.SchemaEngine.Reload(context.Background())
+
+	if err := env.SetVSchema(shardedVSchema); err != nil {
+		t.Fatal(err)
+	}
+	defer env.SetVSchema("{}")
+
+	filter := &binlogdatapb.Filter{
+		Rules: []*binlogdatapb.Rule{{
+			Match:  "dst1",
+			Filter: "select id, keyspace_id() as val from src1",
+		}},
+	}
+	cancel, _ := startVReplication(t, filter, binlogdatapb.OnDDLAction_IGNORE, "")
+	defer cancel()
+
+	testcases := []struct {
+		input  string
+		output []string
+		table  string
+		data   [][]string
+	}{{
+		// insert with insertNormal
+		input: "insert into src1 values(1, 'aaa')",
+		output: []string{
+			"begin",
+			"insert into dst1(id,val) values (1,'\x16k@\xb4J\xbaK\xd6')",
+			"/update _vt.vreplication set pos=",
+			"commit",
+		},
+		table: "dst1",
+		data: [][]string{
+			{"1", "\x16k@\xb4J\xbaK\xd6"},
+		},
+	}}
+
+	for _, tcases := range testcases {
+		execStatements(t, []string{tcases.input})
+		expectDBClientQueries(t, tcases.output)
+		if tcases.table != "" {
+			expectData(t, tcases.table, tcases.data)
+		}
+	}
+}
+
 func TestUnicode(t *testing.T) {
 	defer deleteTablet(addTablet(100))
 

go/vt/vttablet/tabletserver/vstreamer/planbuilder.go

@@ -45,6 +45,7 @@ type Plan struct {
 // ColExpr represents a column expression.
 type ColExpr struct {
 	ColNum int
+	Vindex vindexes.Vindex
 	Alias  sqlparser.ColIdent
 	Type   querypb.Type
 }
@@ -75,28 +76,43 @@ func (plan *Plan) filter(values []sqltypes.Value) (bool, []sqltypes.Value, error
 		if colExpr.ColNum >= len(values) {
 			return false, nil, fmt.Errorf("index out of range, colExpr.ColNum: %d, len(values): %d", colExpr.ColNum, len(values))
 		}
-		result[i] = values[colExpr.ColNum]
+		val := values[colExpr.ColNum]
+		if colExpr.Vindex != nil {
+			ksid, err := getKeyspaceID(val, colExpr.Vindex)
+			if err != nil {
+				return false, nil, err
+			}
+			val = sqltypes.MakeTrusted(sqltypes.VarBinary, []byte(ksid))
+		}
+		result[i] = val
 	}
 	if plan.Vindex == nil {
 		return true, result, nil
 	}
 
-	// Filter by Vindex.
-	destinations, err := plan.Vindex.Map(nil, []sqltypes.Value{result[plan.VindexColumn]})
+	ksid, err := getKeyspaceID(result[plan.VindexColumn], plan.Vindex)
 	if err != nil {
 		return false, nil, err
 	}
+	if !key.KeyRangeContains(plan.KeyRange, ksid) {
+		return false, nil, nil
+	}
+	return true, result, nil
+}
+
+func getKeyspaceID(value sqltypes.Value, vindex vindexes.Vindex) (key.DestinationKeyspaceID, error) {
+	destinations, err := vindex.Map(nil, []sqltypes.Value{value})
+	if err != nil {
+		return nil, err
+	}
 	if len(destinations) != 1 {
-		return false, nil, fmt.Errorf("mapping row to keyspace id returned an invalid array of destinations: %v", key.DestinationsString(destinations))
+		return nil, fmt.Errorf("mapping row to keyspace id returned an invalid array of destinations: %v", key.DestinationsString(destinations))
 	}
 	ksid, ok := destinations[0].(key.DestinationKeyspaceID)
 	if !ok || len(ksid) == 0 {
-		return false, nil, fmt.Errorf("could not map %v to a keyspace id, got destination %v", result[plan.VindexColumn], destinations[0])
+		return nil, fmt.Errorf("could not map %v to a keyspace id, got destination %v", value, destinations[0])
 	}
-	if !key.KeyRangeContains(plan.KeyRange, ksid) {
-		return false, nil, nil
-	}
-	return true, result, nil
+	return ksid, nil
 }
 
 func mustSendDDL(query mysql.Query, dbname string, filter *binlogdatapb.Filter) bool {
@@ -233,7 +249,7 @@ func buildTablePlan(ti *Table, kschema *vindexes.KeyspaceSchema, query string) (
 	plan := &Plan{
 		Table: ti,
 	}
-	if err := plan.analyzeExprs(sel.SelectExprs); err != nil {
+	if err := plan.analyzeExprs(kschema, sel.SelectExprs); err != nil {
 		return nil, err
 	}
 
@@ -277,10 +293,10 @@ func analyzeSelect(query string) (sel *sqlparser.Select, fromTable sqlparser.Tab
 	return sel, fromTable, nil
 }
 
-func (plan *Plan) analyzeExprs(selExprs sqlparser.SelectExprs) error {
+func (plan *Plan) analyzeExprs(kschema *vindexes.KeyspaceSchema, selExprs sqlparser.SelectExprs) error {
 	if _, ok := selExprs[0].(*sqlparser.StarExpr); !ok {
 		for _, expr := range selExprs {
-			cExpr, err := plan.analyzeExpr(expr)
+			cExpr, err := plan.analyzeExpr(kschema, expr)
 			if err != nil {
 				return err
 			}
@@ -300,27 +316,48 @@ func (plan *Plan) analyzeExprs(selExprs sqlparser.SelectExprs) error {
 	return nil
 }
 
-func (plan *Plan) analyzeExpr(selExpr sqlparser.SelectExpr) (cExpr ColExpr, err error) {
+func (plan *Plan) analyzeExpr(kschema *vindexes.KeyspaceSchema, selExpr sqlparser.SelectExpr) (cExpr ColExpr, err error) {
 	aliased, ok := selExpr.(*sqlparser.AliasedExpr)
 	if !ok {
 		return ColExpr{}, fmt.Errorf("unsupported: %v", sqlparser.String(selExpr))
 	}
-	as := aliased.As
-	if as.IsEmpty() {
-		as = sqlparser.NewColIdent(sqlparser.String(aliased.Expr))
-	}
-	colname, ok := aliased.Expr.(*sqlparser.ColName)
-	if !ok {
+	switch inner := aliased.Expr.(type) {
+	case *sqlparser.ColName:
+		if !inner.Qualifier.IsEmpty() {
+			return ColExpr{}, fmt.Errorf("unsupported qualifier for column: %v", sqlparser.String(inner))
+		}
+		colnum, err := findColumn(plan.Table, inner.Name)
+		if err != nil {
+			return ColExpr{}, err
+		}
+		as := aliased.As
+		if as.IsEmpty() {
+			as = sqlparser.NewColIdent(sqlparser.String(aliased.Expr))
+		}
+		return ColExpr{ColNum: colnum, Alias: as, Type: plan.Table.Columns[colnum].Type}, nil
+	case *sqlparser.FuncExpr:
+		if inner.Name.Lowered() != "keyspace_id" {
+			return ColExpr{}, fmt.Errorf("unsupported function: %v", sqlparser.String(inner))
+		}
+		if len(inner.Exprs) != 0 {
+			return ColExpr{}, fmt.Errorf("unexpected: %v", sqlparser.String(inner))
+		}
+		table := kschema.Tables[plan.Table.Name]
+		if table == nil {
+			return ColExpr{}, fmt.Errorf("no vschema definition for table %s", plan.Table.Name)
+		}
+		// Get Primary Vindex.
+		if len(table.ColumnVindexes) == 0 {
+			return ColExpr{}, fmt.Errorf("table %s has no primary vindex", plan.Table.Name)
+		}
+		colnum, err := findColumn(plan.Table, table.ColumnVindexes[0].Columns[0])
+		if err != nil {
+			return ColExpr{}, err
+		}
+		return ColExpr{ColNum: colnum, Vindex: table.ColumnVindexes[0].Vindex, Alias: sqlparser.NewColIdent("keyspace_id"), Type: sqltypes.VarBinary}, nil
+	default:
 		return ColExpr{}, fmt.Errorf("unsupported: %v", sqlparser.String(aliased.Expr))
 	}
-	if !colname.Qualifier.IsEmpty() {
-		return ColExpr{}, fmt.Errorf("unsupported qualifier for column: %v", sqlparser.String(colname))
-	}
-	colnum, err := findColumn(plan.Table, colname.Name)
-	if err != nil {
-		return ColExpr{}, err
-	}
-	return ColExpr{ColNum: colnum, Alias: as, Type: plan.Table.Columns[colnum].Type}, nil
 }
 
 func (plan *Plan) analyzeInKeyRange(kschema *vindexes.KeyspaceSchema, exprs sqlparser.SelectExprs) error {

go/vt/vttablet/tabletserver/vstreamer/planbuilder_test.go

@@ -383,7 +383,7 @@ func TestPlanbuilder(t *testing.T) {
 	}, {
 		inTable: t1,
 		inRule:  &binlogdatapb.Rule{Match: "t1", Filter: "select id, val, max(val) from t1"},
-		outErr:  `unsupported: max(val)`,
+		outErr:  `unsupported function: max(val)`,
 	}, {
 		inTable: t1,
 		inRule:  &binlogdatapb.Rule{Match: "t1", Filter: "select id+1, val from t1"},