vtexplain2

go/cmd/vtexplain/vtexplain.go

@@ -29,18 +29,20 @@ import (
 )
 
 var (
-	sqlFlag         = flag.String("sql", "", "A list of semicolon-delimited SQL commands to analyze")
-	sqlFileFlag     = flag.String("sql-file", "", "Identifies the file that contains the SQL commands to analyze")
-	schemaFlag      = flag.String("schema", "", "The SQL table schema")
-	schemaFileFlag  = flag.String("schema-file", "", "Identifies the file that contains the SQL table schema")
-	vschemaFlag     = flag.String("vschema", "", "Identifies the VTGate routing schema")
-	vschemaFileFlag = flag.String("vschema-file", "", "Identifies the VTGate routing schema file")
-	numShards       = flag.Int("shards", 2, "Number of shards per keyspace")
-	executionMode   = flag.String("execution-mode", "multi", "The execution mode to simulate -- must be set to multi, legacy-autocommit, or twopc")
-	replicationMode = flag.String("replication-mode", "ROW", "The replication mode to simulate -- must be set to either ROW or STATEMENT")
-	normalize       = flag.Bool("normalize", false, "Whether to enable vtgate normalization")
-	outputMode      = flag.String("output-mode", "text", "Output in human-friendly text or json")
-	dbName          = flag.String("dbname", "", "Optional database target to override normal routing")
+	sqlFlag            = flag.String("sql", "", "A list of semicolon-delimited SQL commands to analyze")
+	sqlFileFlag        = flag.String("sql-file", "", "Identifies the file that contains the SQL commands to analyze")
+	schemaFlag         = flag.String("schema", "", "The SQL table schema")
+	schemaFileFlag     = flag.String("schema-file", "", "Identifies the file that contains the SQL table schema")
+	vschemaFlag        = flag.String("vschema", "", "Identifies the VTGate routing schema")
+	vschemaFileFlag    = flag.String("vschema-file", "", "Identifies the VTGate routing schema file")
+	ksShardMapFlag     = flag.String("ks-shard-map", "", "JSON map of keyspace name -> shard name -> ShardReference object. The inner map is the same as the output of FindAllShardsInKeyspace")
+	ksShardMapFileFlag = flag.String("ks-shard-map-file", "", "File containing json blob of keyspace name -> shard name -> ShardReference object")
+	numShards          = flag.Int("shards", 2, "Number of shards per keyspace. Passing -ks-shard-map/-ks-shard-map-file causes this flag to be ignored.")
+	executionMode      = flag.String("execution-mode", "multi", "The execution mode to simulate -- must be set to multi, legacy-autocommit, or twopc")
+	replicationMode    = flag.String("replication-mode", "ROW", "The replication mode to simulate -- must be set to either ROW or STATEMENT")
+	normalize          = flag.Bool("normalize", false, "Whether to enable vtgate normalization")
+	outputMode         = flag.String("output-mode", "text", "Output in human-friendly text or json")
+	dbName             = flag.String("dbname", "", "Optional database target to override normal routing")
 
 	// vtexplainFlags lists all the flags that should show in usage
 	vtexplainFlags = []string{
@@ -54,6 +56,8 @@ var (
 		"sql-file",
 		"vschema",
 		"vschema-file",
+		"ks-shard-map",
+		"ks-shard-map-file",
 		"dbname",
 		"queryserver-config-passthrough-dmls",
 	}
@@ -104,7 +108,7 @@ func init() {
 
 // getFileParam returns a string containing either flag is not "",
 // or the content of the file named flagFile
-func getFileParam(flag, flagFile, name string) (string, error) {
+func getFileParam(flag, flagFile, name string, required bool) (string, error) {
 	if flag != "" {
 		if flagFile != "" {
 			return "", fmt.Errorf("action requires only one of %v or %v-file", name, name)
@@ -113,7 +117,11 @@ func getFileParam(flag, flagFile, name string) (string, error) {
 	}
 
 	if flagFile == "" {
-		return "", fmt.Errorf("action requires one of %v or %v-file", name, name)
+		if required {
+			return "", fmt.Errorf("action requires one of %v or %v-file", name, name)
+		}
+
+		return "", nil
 	}
 	data, err := ioutil.ReadFile(flagFile)
 	if err != nil {
@@ -137,17 +145,22 @@ func main() {
 }
 
 func parseAndRun() error {
-	sql, err := getFileParam(*sqlFlag, *sqlFileFlag, "sql")
+	sql, err := getFileParam(*sqlFlag, *sqlFileFlag, "sql", true)
+	if err != nil {
+		return err
+	}
+
+	schema, err := getFileParam(*schemaFlag, *schemaFileFlag, "schema", true)
 	if err != nil {
 		return err
 	}
 
-	schema, err := getFileParam(*schemaFlag, *schemaFileFlag, "schema")
+	vschema, err := getFileParam(*vschemaFlag, *vschemaFileFlag, "vschema", true)
 	if err != nil {
 		return err
 	}
 
-	vschema, err := getFileParam(*vschemaFlag, *vschemaFileFlag, "vschema")
+	ksShardMap, err := getFileParam(*ksShardMapFlag, *ksShardMapFileFlag, "ks-shard-map", false)
 	if err != nil {
 		return err
 	}
@@ -164,7 +177,7 @@ func parseAndRun() error {
 	log.V(100).Infof("schema %s\n", schema)
 	log.V(100).Infof("vschema %s\n", vschema)
 
-	err = vtexplain.Init(vschema, schema, opts)
+	err = vtexplain.Init(vschema, schema, ksShardMap, opts)
 	if err != nil {
 		return err
 	}

go/vt/vtexplain/testdata/multi-output/select-sharded-8-output.txt

@@ -0,0 +1,18 @@
+----------------------------------------------------------------------
+select * from user
+
+1 ks_sharded/-20: select * from user limit 10001
+1 ks_sharded/20-40: select * from user limit 10001
+1 ks_sharded/40-60: select * from user limit 10001
+1 ks_sharded/60-80: select * from user limit 10001
+1 ks_sharded/80-a0: select * from user limit 10001
+1 ks_sharded/a0-c0: select * from user limit 10001
+1 ks_sharded/c0-e0: select * from user limit 10001
+1 ks_sharded/e0-: select * from user limit 10001
+
+----------------------------------------------------------------------
+select * from user where id in (1, 2)
+
+1 ks_sharded/-20: select * from user where id in (1, 2) limit 10001
+
+----------------------------------------------------------------------

go/vt/vtexplain/testdata/multi-output/uneven-keyspace-output.txt

@@ -0,0 +1,16 @@
+----------------------------------------------------------------------
+select * from user
+
+1 ks_sharded/-80: select * from user limit 10001
+1 ks_sharded/80-90: select * from user limit 10001
+1 ks_sharded/90-a0: select * from user limit 10001
+1 ks_sharded/a0-e8: select * from user limit 10001
+1 ks_sharded/e8-: select * from user limit 10001
+
+----------------------------------------------------------------------
+select * from user where id in (10, 17, 42, 100000)
+
+1 ks_sharded/-80: select * from user where id in (10, 17, 42) limit 10001
+1 ks_sharded/80-90: select * from user where id in (100000) limit 10001
+
+----------------------------------------------------------------------

go/vt/vtexplain/testdata/select-sharded-8-queries.sql

@@ -0,0 +1,2 @@
+select * from user;
+select * from user where id in (1, 2);

go/vt/vtexplain/testdata/uneven-keyspace-queries.sql

@@ -0,0 +1,2 @@
+select * from user;
+select * from user where id in (10, 17, 42, 100000);

go/vt/vtexplain/vtexplain.go

@@ -144,7 +144,7 @@ type Explain struct {
 }
 
 // Init sets up the fake execution environment
-func Init(vSchemaStr, sqlSchema string, opts *Options) error {
+func Init(vSchemaStr, sqlSchema, ksShardMapStr string, opts *Options) error {
 	// Verify options
 	if opts.ReplicationMode != "ROW" && opts.ReplicationMode != "STATEMENT" {
 		return fmt.Errorf("invalid replication mode \"%s\"", opts.ReplicationMode)
@@ -160,7 +160,7 @@ func Init(vSchemaStr, sqlSchema string, opts *Options) error {
 		return fmt.Errorf("initTabletEnvironment: %v", err)
 	}
 
-	err = initVtgateExecutor(vSchemaStr, opts)
+	err = initVtgateExecutor(vSchemaStr, ksShardMapStr, opts)
 	if err != nil {
 		return fmt.Errorf("initVtgateExecutor: %v", err)
 	}

go/vt/vtexplain/vtexplain_flaky_test.go

@@ -26,6 +26,9 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
+	"vitess.io/vitess/go/vt/key"
+	"vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/topo"
 )
 
 var testOutputTempDir string
@@ -39,15 +42,27 @@ func defaultTestOpts() *Options {
 	}
 }
 
-func initTest(mode string, opts *Options, t *testing.T) {
+type testopts struct {
+	shardmap map[string]map[string]*topo.ShardInfo
+}
+
+func initTest(mode string, opts *Options, topts *testopts, t *testing.T) {
 	schema, err := ioutil.ReadFile("testdata/test-schema.sql")
 	require.NoError(t, err)
 
 	vSchema, err := ioutil.ReadFile("testdata/test-vschema.json")
 	require.NoError(t, err)
 
+	shardmap := ""
+	if topts.shardmap != nil {
+		shardmapBytes, err := json.Marshal(topts.shardmap)
+		require.NoError(t, err)
+
+		shardmap = string(shardmapBytes)
+	}
+
 	opts.ExecutionMode = mode
-	err = Init(string(vSchema), string(schema), opts)
+	err = Init(string(vSchema), string(schema), shardmap, opts)
 	require.NoError(t, err, "vtexplain Init error\n%s", string(schema))
 }
 
@@ -61,13 +76,13 @@ func testExplain(testcase string, opts *Options, t *testing.T) {
 	}
 
 	for _, mode := range modes {
-		runTestCase(testcase, mode, opts, t)
+		runTestCase(testcase, mode, opts, &testopts{}, t)
 	}
 }
 
-func runTestCase(testcase, mode string, opts *Options, t *testing.T) {
+func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) {
 	t.Run(testcase, func(t *testing.T) {
-		initTest(mode, opts, t)
+		initTest(mode, opts, topts, t)
 
 		sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
 		sql, err := ioutil.ReadFile(sqlFile)
@@ -130,7 +145,7 @@ func TestExplain(t *testing.T) {
 }
 
 func TestErrors(t *testing.T) {
-	initTest(ModeMulti, defaultTestOpts(), t)
+	initTest(ModeMulti, defaultTestOpts(), &testopts{}, t)
 
 	tests := []struct {
 		SQL string
@@ -229,3 +244,60 @@ func TestJSONOutput(t *testing.T) {
 		t.Errorf(diff)
 	}
 }
+
+func testShardInfo(ks, start, end string, t *testing.T) *topo.ShardInfo {
+	kr, err := key.ParseKeyRangeParts(start, end)
+	require.NoError(t, err)
+
+	return topo.NewShardInfo(
+		ks,
+		fmt.Sprintf("%s-%s", start, end),
+		&topodata.Shard{KeyRange: kr},
+		&vtexplainTestTopoVersion{},
+	)
+}
+
+func TestUsingKeyspaceShardMap(t *testing.T) {
+	tests := []struct {
+		testcase      string
+		ShardRangeMap map[string]map[string]*topo.ShardInfo
+	}{
+		{
+			testcase: "select-sharded-8",
+			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
+				"ks_sharded": {
+					"-20":   testShardInfo("ks_sharded", "", "20", t),
+					"20-40": testShardInfo("ks_sharded", "20", "40", t),
+					"40-60": testShardInfo("ks_sharded", "40", "60", t),
+					"60-80": testShardInfo("ks_sharded", "60", "80", t),
+					"80-a0": testShardInfo("ks_sharded", "80", "a0", t),
+					"a0-c0": testShardInfo("ks_sharded", "a0", "c0", t),
+					"c0-e0": testShardInfo("ks_sharded", "c0", "e0", t),
+					"e0-":   testShardInfo("ks_sharded", "e0", "", t),
+				},
+			},
+		},
+		{
+			testcase: "uneven-keyspace",
+			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
+				// Have mercy on the poor soul that has this keyspace sharding.
+				// But, hey, vtexplain still works so they have that going for them.
+				"ks_sharded": {
+					"-80":   testShardInfo("ks_sharded", "", "80", t),
+					"80-90": testShardInfo("ks_sharded", "80", "90", t),
+					"90-a0": testShardInfo("ks_sharded", "90", "a0", t),
+					"a0-e8": testShardInfo("ks_sharded", "a0", "e8", t),
+					"e8-":   testShardInfo("ks_sharded", "e8", "", t),
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		runTestCase(test.testcase, ModeMulti, defaultTestOpts(), &testopts{test.ShardRangeMap}, t)
+	}
+}
+
+type vtexplainTestTopoVersion struct{}
+
+func (vtexplain *vtexplainTestTopoVersion) String() string { return "vtexplain-test-topo" }

go/vt/vtexplain/vtexplain_topo.go

@@ -22,7 +22,6 @@ import (
 
 	"golang.org/x/net/context"
 
-	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/topo"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -38,6 +37,8 @@ type ExplainTopo struct {
 	// Map of ks/shard to test tablet connection
 	TabletConns map[string]*explainTablet
 
+	KeyspaceShards map[string]map[string]*topodatapb.ShardReference
+
 	// Synchronization lock
 	Lock sync.Mutex
 
@@ -83,68 +84,31 @@ func (et *ExplainTopo) GetSrvKeyspace(ctx context.Context, cell, keyspace string
 		return nil, fmt.Errorf("no vschema for keyspace %s", keyspace)
 	}
 
-	var srvKeyspace *topodatapb.SrvKeyspace
-	if vschema.Sharded {
-		shards := make([]*topodatapb.ShardReference, 0, et.NumShards)
-		for i := 0; i < et.NumShards; i++ {
-			kr, err := key.EvenShardsKeyRange(i, et.NumShards)
-			if err != nil {
-				return nil, err
-			}
-
-			shard := &topodatapb.ShardReference{
-				Name:     key.KeyRangeString(kr),
-				KeyRange: kr,
-			}
-			shards = append(shards, shard)
-		}
-
-		srvKeyspace = &topodatapb.SrvKeyspace{
-			ShardingColumnName: "", // exact value is ignored
-			ShardingColumnType: 0,
-			Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
-				{
-					ServedType:      topodatapb.TabletType_MASTER,
-					ShardReferences: shards,
-				},
-				{
-					ServedType:      topodatapb.TabletType_REPLICA,
-					ShardReferences: shards,
-				},
-				{
-					ServedType:      topodatapb.TabletType_RDONLY,
-					ShardReferences: shards,
-				},
+	shards := make([]*topodatapb.ShardReference, 0, len(et.KeyspaceShards[keyspace]))
+	for _, shard := range et.KeyspaceShards[keyspace] {
+		shards = append(shards, shard)
+	}
+
+	srvKeyspace := &topodatapb.SrvKeyspace{
+		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
+			{
+				ServedType:      topodatapb.TabletType_MASTER,
+				ShardReferences: shards,
 			},
-		}
-
-	} else {
-		// unsharded
-		kr, err := key.EvenShardsKeyRange(0, 1)
-		if err != nil {
-			return nil, err
-		}
-
-		shard := &topodatapb.ShardReference{
-			Name: key.KeyRangeString(kr),
-		}
-
-		srvKeyspace = &topodatapb.SrvKeyspace{
-			Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
-				{
-					ServedType:      topodatapb.TabletType_MASTER,
-					ShardReferences: []*topodatapb.ShardReference{shard},
-				},
-				{
-					ServedType:      topodatapb.TabletType_REPLICA,
-					ShardReferences: []*topodatapb.ShardReference{shard},
-				},
-				{
-					ServedType:      topodatapb.TabletType_RDONLY,
-					ShardReferences: []*topodatapb.ShardReference{shard},
-				},
+			{
+				ServedType:      topodatapb.TabletType_REPLICA,
+				ShardReferences: shards,
 			},
-		}
+			{
+				ServedType:      topodatapb.TabletType_RDONLY,
+				ShardReferences: shards,
+			},
+		},
+	}
+
+	if vschema.Sharded {
+		srvKeyspace.ShardingColumnName = "" // exact value is ignored
+		srvKeyspace.ShardingColumnType = 0
 	}
 
 	return srvKeyspace, nil