Support for hex & shard in vindex query

go/vt/key/key.go

@@ -289,10 +289,16 @@ func KeyRangeIncludes(big, small *topodatapb.KeyRange) bool {
 // specification. a-b-c-d will be parsed as a-b, b-c, c-d. The empty
 // string may serve both as the start and end of the keyspace: -a-b-
 // will be parsed as start-a, a-b, b-end.
+// "0" is treated as "-", to allow us to not have to special-case
+// client code.
 func ParseShardingSpec(spec string) ([]*topodatapb.KeyRange, error) {
 	parts := strings.Split(spec, "-")
 	if len(parts) == 1 {
-		return nil, fmt.Errorf("malformed spec: doesn't define a range: %q", spec)
+		if spec == "0" {
+			parts = []string{"", ""}
+		} else {
+			return nil, fmt.Errorf("malformed spec: doesn't define a range: %q", spec)
+		}
 	}
 	old := parts[0]
 	ranges := make([]*topodatapb.KeyRange, len(parts)-1)

go/vt/key/key_test.go

@@ -272,6 +272,7 @@ func TestParseShardingSpec(t *testing.T) {
 	x80 := []byte{0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
 	goodTable := map[string][]*topodatapb.KeyRange{
 		"-": {{}},
+		"0": {{}},
 		"-4000000000000000-8000000000000000-": {
 			{End: x40},
 			{Start: x40, End: x80},

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -158,6 +158,10 @@ func (t noopVCursor) ExceedsMaxMemoryRows(numRows int) bool {
 	return !testIgnoreMaxMemoryRows && numRows > testMaxMemoryRows
 }
 
+func (t noopVCursor) GetKeyspace() string {
+	return ""
+}
+
 func (t noopVCursor) SetContextTimeout(timeout time.Duration) context.CancelFunc {
 	return func() {}
 }
@@ -294,6 +298,10 @@ func (f *loggingVCursor) ErrorGroupCancellableContext() (*errgroup.Group, func()
 	panic("implement me")
 }
 
+func (f *loggingVCursor) GetKeyspace() string {
+	return ""
+}
+
 func (f *loggingVCursor) RecordWarning(warning *querypb.QueryWarning) {
 	f.warnings = append(f.warnings, warning)
 }

go/vt/vtgate/engine/primitive.go

@@ -53,6 +53,7 @@ type (
 		// Context returns the context of the current request.
 		Context() context.Context
 
+		GetKeyspace() string
 		// MaxMemoryRows returns the maxMemoryRows flag value.
 		MaxMemoryRows() int
 

go/vt/vtgate/engine/vindex_func.go

@@ -18,6 +18,7 @@ package engine
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
@@ -131,10 +132,25 @@ func (vf *VindexFunc) mapVindex(vcursor VCursor, bindVars map[string]*querypb.Bi
 		}
 	case key.DestinationKeyspaceID:
 		if len(d) > 0 {
-			result.Rows = [][]sqltypes.Value{
-				vf.buildRow(vkey, d, nil),
+			if vcursor != nil {
+				resolvedShards, _, err := vcursor.ResolveDestinations(vcursor.GetKeyspace(), nil, []key.Destination{d})
+				if err != nil {
+					return nil, err
+				}
+				kr, err := key.ParseShardingSpec(resolvedShards[0].Target.Shard)
+				if err != nil {
+					return nil, err
+				}
+				result.Rows = [][]sqltypes.Value{
+					vf.buildRow(vkey, d, kr[0]),
+				}
+				result.RowsAffected = 1
+			} else {
+				result.Rows = [][]sqltypes.Value{
+					vf.buildRow(vkey, d, nil),
+				}
+				result.RowsAffected = 1
 			}
-			result.RowsAffected = 1
 		}
 	case key.DestinationKeyspaceIDs:
 		for _, ksid := range d {
@@ -173,6 +189,18 @@ func (vf *VindexFunc) buildRow(id sqltypes.Value, ksid []byte, kr *topodatapb.Ke
 			} else {
 				row = append(row, sqltypes.NULL)
 			}
+		case 4:
+			if ksid != nil {
+				row = append(row, sqltypes.NewVarBinary(fmt.Sprintf("%x", ksid)))
+			} else {
+				row = append(row, sqltypes.NULL)
+			}
+		case 5:
+			if ksid != nil {
+				row = append(row, sqltypes.NewVarBinary(key.KeyRangeString(kr)))
+			} else {
+				row = append(row, sqltypes.NULL)
+			}
 		default:
 			panic("BUG: unexpected column number")
 		}

go/vt/vtgate/engine/vindex_func_test.go

@@ -98,7 +98,7 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want := &sqltypes.Result{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}
 	if !reflect.DeepEqual(got, want) {
 		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
@@ -111,9 +111,13 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want = sqltypes.MakeTestResult(
-		sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
-		"1|foo",
+		sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
+		"1|foo|||666f6f",
 	)
+	for _, row := range want.Rows {
+		row[2] = sqltypes.NULL
+		row[3] = sqltypes.NULL
+	}
 	if !reflect.DeepEqual(got, want) {
 		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
 	}
@@ -125,12 +129,13 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want = &sqltypes.Result{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		Rows: [][]sqltypes.Value{{
 			sqltypes.NewVarBinary("1"),
 			sqltypes.NULL,
 			sqltypes.MakeTrusted(sqltypes.VarBinary, []byte{0x40}),
 			sqltypes.MakeTrusted(sqltypes.VarBinary, []byte{0x60}),
+			sqltypes.NULL,
 		}},
 		RowsAffected: 1,
 	}
@@ -145,7 +150,7 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want = &sqltypes.Result{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}
 	if !reflect.DeepEqual(got, want) {
 		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
@@ -158,9 +163,9 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want = sqltypes.MakeTestResult(
-		sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
-		"1|foo||",
-		"1|bar||",
+		sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
+		"1|foo|||666f6f",
+		"1|bar|||626172",
 	)
 	// Massage the rows because MakeTestResult doesn't do NULL values.
 	for _, row := range want.Rows {
@@ -178,12 +183,13 @@ func TestVindexFuncMap(t *testing.T) {
 		t.Fatal(err)
 	}
 	want = &sqltypes.Result{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		Rows: [][]sqltypes.Value{{
 			sqltypes.NewVarBinary("1"),
 			sqltypes.NULL,
 			sqltypes.MakeTrusted(sqltypes.VarBinary, []byte{0x40}),
 			sqltypes.MakeTrusted(sqltypes.VarBinary, []byte{0x60}),
+			sqltypes.NULL,
 		}},
 		RowsAffected: 1,
 	}
@@ -195,12 +201,12 @@ func TestVindexFuncMap(t *testing.T) {
 func TestVindexFuncStreamExecute(t *testing.T) {
 	vf := testVindexFunc(&nvindex{matchid: true})
 	want := []*sqltypes.Result{{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}, {
 		Rows: [][]sqltypes.Value{{
-			sqltypes.NewVarBinary("1"), sqltypes.NewVarBinary("foo"), sqltypes.NULL, sqltypes.NULL,
+			sqltypes.NewVarBinary("1"), sqltypes.NewVarBinary("foo"), sqltypes.NULL, sqltypes.NULL, sqltypes.NewVarBinary("666f6f"),
 		}, {
-			sqltypes.NewVarBinary("1"), sqltypes.NewVarBinary("bar"), sqltypes.NULL, sqltypes.NULL,
+			sqltypes.NewVarBinary("1"), sqltypes.NewVarBinary("bar"), sqltypes.NULL, sqltypes.NULL, sqltypes.NewVarBinary("626172"),
 		}},
 	}}
 	i := 0
@@ -223,7 +229,7 @@ func TestVindexFuncGetFields(t *testing.T) {
 		t.Fatal(err)
 	}
 	want := &sqltypes.Result{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}
 	if !reflect.DeepEqual(got, want) {
 		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
@@ -250,8 +256,8 @@ func TestFieldOrder(t *testing.T) {
 
 func testVindexFunc(v vindexes.SingleColumn) *VindexFunc {
 	return &VindexFunc{
-		Fields: sqltypes.MakeTestFields("id|keyspace_id|range_start|range_end", "varbinary|varbinary|varbinary|varbinary"),
-		Cols:   []int{0, 1, 2, 3},
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
+		Cols:   []int{0, 1, 2, 3, 4},
 		Opcode: VindexMap,
 		Vindex: v,
 		Value:  int64PlanValue(1),

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

@@ -1,22 +1,26 @@
 # vindex func read all cols
-"select id, keyspace_id, range_start, range_end from user_index where id = :id"
+"select id, keyspace_id, range_start, range_end, hex_keyspace_id, shard from user_index where id = :id"
 {
   "QueryType": "SELECT",
-  "Original": "select id, keyspace_id, range_start, range_end from user_index where id = :id",
+  "Original": "select id, keyspace_id, range_start, range_end, hex_keyspace_id, shard from user_index where id = :id",
   "Instructions": {
     "OperatorType": "VindexFunc",
     "Variant": "VindexMap",
     "Columns": [
       0,
       1,
       2,
-      3
+      3,
+      4,
+      5
     ],
     "Fields": {
+      "hex_keyspace_id": "VARBINARY",
       "id": "VARBINARY",
       "keyspace_id": "VARBINARY",
       "range_end": "VARBINARY",
-      "range_start": "VARBINARY"
+      "range_start": "VARBINARY",
+      "shard": "VARBINARY"
     },
     "Value": ":id",
     "Vindex": "user_index"
@@ -35,13 +39,17 @@
       0,
       1,
       2,
-      3
+      3,
+      4,
+      5
     ],
     "Fields": {
+      "hex_keyspace_id": "VARBINARY",
       "id": "VARBINARY",
       "keyspace_id": "VARBINARY",
       "range_end": "VARBINARY",
-      "range_start": "VARBINARY"
+      "range_start": "VARBINARY",
+      "shard": "VARBINARY"
     },
     "Value": ":id",
     "Vindex": "user_index"

go/vt/vtgate/planbuilder/vindex_func.go

@@ -59,6 +59,8 @@ func newVindexFunc(alias sqlparser.TableName, vindex vindexes.SingleColumn) (*vi
 	t.addColumn(sqlparser.NewColIdent("keyspace_id"), &column{origin: vf})
 	t.addColumn(sqlparser.NewColIdent("range_start"), &column{origin: vf})
 	t.addColumn(sqlparser.NewColIdent("range_end"), &column{origin: vf})
+	t.addColumn(sqlparser.NewColIdent("hex_keyspace_id"), &column{origin: vf})
+	t.addColumn(sqlparser.NewColIdent("shard"), &column{origin: vf})
 	t.isAuthoritative = true
 
 	st := newSymtab()

go/vt/vtgate/vcursor_impl.go

@@ -99,6 +99,10 @@ type vcursorImpl struct {
 	vm                    VSchemaOperator
 }
 
+func (vc *vcursorImpl) GetKeyspace() string {
+	return vc.keyspace
+}
+
 func (vc *vcursorImpl) ExecuteVSchema(keyspace string, vschemaDDL *sqlparser.DDL) error {
 	srvVschema := vc.vm.GetCurrentSrvVschema()
 	if srvVschema == nil {