Golang key.ParseShardingSpec函数代码示例

// This function implements the restriction of handling one keyrange
// and one shard since streaming doesn't support merge sorting the results.
// The input/output api is generic though.
func (vtg *VTGate) mapKrToShardsForStreaming(streamQuery *proto.StreamQueryKeyRange) ([]string, error) {
	var keyRange key.KeyRange
	var err error
	if streamQuery.KeyRange == "" {
		keyRange = key.KeyRange{Start: "", End: ""}
	} else {
		krArray, err := key.ParseShardingSpec(streamQuery.KeyRange)
		if err != nil {
			return nil, err
		}
		keyRange = krArray[0]
	}
	shards, err := resolveKeyRangeToShards(vtg.scatterConn.toposerv,
		vtg.scatterConn.cell,
		streamQuery.Keyspace,
		streamQuery.TabletType,
		keyRange)
	if err != nil {
		return nil, err
	}

	if len(shards) != 1 {
		return nil, fmt.Errorf("KeyRange cannot map to more than one shard")
	}

	return shards, nil
}

func getAllShards(shardSpec string) (key.KeyRangeArray, error) {
	shardedKrArray, err := key.ParseShardingSpec(shardSpec)
	if err != nil {
		return nil, err
	}
	return shardedKrArray, nil
}

func multisnapshotCmd(mysqld *mysqlctl.Mysqld, subFlags *flag.FlagSet, args []string) {
	concurrency := subFlags.Int("concurrency", 8, "how many compression jobs to run simultaneously")
	spec := subFlags.String("spec", "-", "shard specification")
	tablesString := subFlags.String("tables", "", "dump only this comma separated list of tables")
	skipSlaveRestart := subFlags.Bool("skip-slave-restart", false, "after the snapshot is done, do not restart slave replication")
	maximumFilesize := subFlags.Uint64("maximum-file-size", 128*1024*1024, "the maximum size for an uncompressed data file")
	subFlags.Parse(args)
	if subFlags.NArg() != 2 {
		relog.Fatal("action multisnapshot requires <db name> <key name>")
	}

	shards, err := key.ParseShardingSpec(*spec)
	if err != nil {
		relog.Fatal("multisnapshot failed: %v", err)
	}
	var tables []string
	if *tablesString != "" {
		tables = strings.Split(*tablesString, ",")
	}
	filenames, err := mysqld.CreateMultiSnapshot(shards, subFlags.Arg(0), subFlags.Arg(1), tabletAddr, false, *concurrency, tables, *skipSlaveRestart, *maximumFilesize, nil)
	if err != nil {
		relog.Fatal("multisnapshot failed: %v", err)
	} else {
		relog.Info("manifest locations: %v", filenames)
	}
}

func TestVTGateSplitQuery(t *testing.T) {
	keyspace := "TestVTGateSplitQuery"
	keyranges, _ := key.ParseShardingSpec(DefaultShardSpec)
	createSandbox(keyspace)
	hcVTGateTest.Reset()
	port := int32(1001)
	for _, kr := range keyranges {
		hcVTGateTest.AddTestTablet("aa", "1.1.1.1", port, keyspace, key.KeyRangeString(kr), topodatapb.TabletType_RDONLY, true, 1, nil)
		port++
	}
	sql := "select col1, col2 from table"
	splitCount := 24
	splits, err := rpcVTGate.SplitQuery(context.Background(),
		keyspace,
		sql,
		nil,
		"",
		int64(splitCount))
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	_, err = getAllShards(DefaultShardSpec)
	// Total number of splits should be number of shards * splitsPerShard
	if splitCount != len(splits) {
		t.Errorf("wrong number of splits, want \n%+v, got \n%+v", splitCount, len(splits))
	}
	actualSqlsByKeyRange := map[string][]string{}
	for _, split := range splits {
		if split.Size != sandboxconn.SandboxSQRowCount {
			t.Errorf("wrong split size, want \n%+v, got \n%+v", sandboxconn.SandboxSQRowCount, split.Size)
		}
		if split.KeyRangePart.Keyspace != keyspace {
			t.Errorf("wrong keyspace, want \n%+v, got \n%+v", keyspace, split.KeyRangePart.Keyspace)
		}
		if len(split.KeyRangePart.KeyRanges) != 1 {
			t.Errorf("wrong number of keyranges, want \n%+v, got \n%+v", 1, len(split.KeyRangePart.KeyRanges))
		}
		kr := key.KeyRangeString(split.KeyRangePart.KeyRanges[0])
		actualSqlsByKeyRange[kr] = append(actualSqlsByKeyRange[kr], split.Query.Sql)
	}
	// Sort the sqls for each KeyRange so that we can compare them without
	// regard to the order in which they were returned by the vtgate.
	for _, sqlsForKeyRange := range actualSqlsByKeyRange {
		sort.Strings(sqlsForKeyRange)
	}
	expectedSqlsByKeyRange := map[string][]string{}
	for _, kr := range keyranges {
		expectedSqlsByKeyRange[key.KeyRangeString(kr)] = []string{
			"select col1, col2 from table /*split 0 */",
			"select col1, col2 from table /*split 1 */",
			"select col1, col2 from table /*split 2 */",
		}
	}
	if !reflect.DeepEqual(actualSqlsByKeyRange, expectedSqlsByKeyRange) {
		t.Errorf("splits contain the wrong sqls and/or keyranges, got: %v, want: %v", actualSqlsByKeyRange, expectedSqlsByKeyRange)
	}
}

func getAllShards() (key.KeyRangeArray, error) {
	if ShardedKrArray != nil {
		return ShardedKrArray, nil
	}
	shardedKrArray, err := key.ParseShardingSpec(ShardSpec)
	if err != nil {
		return nil, err
	}
	return shardedKrArray, nil
}

func TestVTGateSplitQuery(t *testing.T) {
	keyspace := "TestVTGateSplitQuery"
	keyranges, _ := key.ParseShardingSpec(DefaultShardSpec)
	s := createSandbox(keyspace)
	for _, kr := range keyranges {
		s.MapTestConn(fmt.Sprintf("%s-%s", kr.Start, kr.End), &sandboxConn{})
	}
	sql := "select col1, col2 from table"
	splitCount := 24
	req := proto.SplitQueryRequest{
		Keyspace: keyspace,
		Query: tproto.BoundQuery{
			Sql: sql,
		},
		SplitCount: splitCount,
	}
	result := new(proto.SplitQueryResult)
	err := rpcVTGate.SplitQuery(context.Background(), &req, result)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	_, err = getAllShards(DefaultShardSpec)
	// Total number of splits should be number of shards * splitsPerShard
	if splitCount != len(result.Splits) {
		t.Errorf("wrong number of splits, want \n%+v, got \n%+v", splitCount, len(result.Splits))
	}
	actualSqlsByKeyRange := map[kproto.KeyRange][]string{}
	for _, split := range result.Splits {
		if split.Size != sandboxSQRowCount {
			t.Errorf("wrong split size, want \n%+v, got \n%+v", sandboxSQRowCount, split.Size)
		}
		if split.Query.Keyspace != keyspace {
			t.Errorf("wrong split size, want \n%+v, got \n%+v", keyspace, split.Query.Keyspace)
		}
		if len(split.Query.KeyRanges) != 1 {
			t.Errorf("wrong number of keyranges, want \n%+v, got \n%+v", 1, len(split.Query.KeyRanges))
		}
		if split.Query.TabletType != topo.TYPE_RDONLY {
			t.Errorf("wrong tablet type, want \n%+v, got \n%+v", topo.TYPE_RDONLY, split.Query.TabletType)
		}
		kr := split.Query.KeyRanges[0]
		actualSqlsByKeyRange[kr] = append(actualSqlsByKeyRange[kr], split.Query.Sql)
	}
	expectedSqlsByKeyRange := map[kproto.KeyRange][]string{}
	for _, kr := range keyranges {
		expectedSqlsByKeyRange[kr] = []string{
			"select col1, col2 from table /*split 0 */",
			"select col1, col2 from table /*split 1 */",
			"select col1, col2 from table /*split 2 */",
		}
	}
	if !reflect.DeepEqual(actualSqlsByKeyRange, expectedSqlsByKeyRange) {
		t.Errorf("splits contain the wrong sqls and/or keyranges, got: %v, want: %v", actualSqlsByKeyRange, expectedSqlsByKeyRange)
	}
}

func TestVTGateSplitQuery(t *testing.T) {
	keyspace := "TestVTGateSplitQuery"
	keyranges, _ := key.ParseShardingSpec(DefaultShardSpec)
	s := createSandbox(keyspace)
	for _, kr := range keyranges {
		s.MapTestConn(key.KeyRangeString(kr), &sandboxConn{})
	}
	sql := "select col1, col2 from table"
	splitCount := 24
	splits, err := rpcVTGate.SplitQuery(context.Background(),
		keyspace,
		sql,
		nil,
		"",
		splitCount)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	_, err = getAllShards(DefaultShardSpec)
	// Total number of splits should be number of shards * splitsPerShard
	if splitCount != len(splits) {
		t.Errorf("wrong number of splits, want \n%+v, got \n%+v", splitCount, len(splits))
	}
	actualSqlsByKeyRange := map[string][]string{}
	for _, split := range splits {
		if split.Size != sandboxSQRowCount {
			t.Errorf("wrong split size, want \n%+v, got \n%+v", sandboxSQRowCount, split.Size)
		}
		if split.KeyRangePart.Keyspace != keyspace {
			t.Errorf("wrong split size, want \n%+v, got \n%+v", keyspace, split.KeyRangePart.Keyspace)
		}
		if len(split.KeyRangePart.KeyRanges) != 1 {
			t.Errorf("wrong number of keyranges, want \n%+v, got \n%+v", 1, len(split.KeyRangePart.KeyRanges))
		}
		kr := key.KeyRangeString(split.KeyRangePart.KeyRanges[0])
		actualSqlsByKeyRange[kr] = append(actualSqlsByKeyRange[kr], split.Query.Sql)
	}
	expectedSqlsByKeyRange := map[string][]string{}
	for _, kr := range keyranges {
		expectedSqlsByKeyRange[key.KeyRangeString(kr)] = []string{
			"select col1, col2 from table /*split 0 */",
			"select col1, col2 from table /*split 1 */",
			"select col1, col2 from table /*split 2 */",
		}
	}
	if !reflect.DeepEqual(actualSqlsByKeyRange, expectedSqlsByKeyRange) {
		t.Errorf("splits contain the wrong sqls and/or keyranges, got: %v, want: %v", actualSqlsByKeyRange, expectedSqlsByKeyRange)
	}
}

func TestMapKeyRangesToShards(t *testing.T) {
	ts := new(sandboxTopo)
	var testCases = []struct {
		keyspace string
		keyRange string
		shards   []string
	}{
		{keyspace: KsTestSharded, keyRange: "20-40", shards: []string{"20-40"}},
		// check for partial keyrange, spanning one shard
		{keyspace: KsTestSharded, keyRange: "10-18", shards: []string{"-20"}},
		// check for keyrange intersecting with multiple shards
		{keyspace: KsTestSharded, keyRange: "10-40", shards: []string{"-20", "20-40"}},
		// check for keyrange intersecting with multiple shards
		{keyspace: KsTestSharded, keyRange: "1c-2a", shards: []string{"-20", "20-40"}},
		// check for keyrange where kr.End is Max Key ""
		{keyspace: KsTestSharded, keyRange: "80-", shards: []string{"80-a0", "a0-c0", "c0-e0", "e0-"}},
		// test for sharded, non-partial keyrange spanning the entire space.
		{keyspace: KsTestSharded, keyRange: "", shards: []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"}},
		// test for unsharded, non-partial keyrange spanning the entire space.
		{keyspace: KsTestUnsharded, keyRange: "", shards: []string{"0"}},
	}

	for _, testCase := range testCases {
		var keyRange *topodatapb.KeyRange
		var err error
		if testCase.keyRange == "" {
			keyRange = &topodatapb.KeyRange{}
		} else {
			krArray, err := key.ParseShardingSpec(testCase.keyRange)
			if err != nil {
				t.Errorf("Got error while parsing sharding spec %v", err)
			}
			keyRange = krArray[0]
		}
		krs := []*topodatapb.KeyRange{keyRange}
		_, gotShards, err := mapKeyRangesToShards(context.Background(), ts, "", testCase.keyspace, topodatapb.TabletType_MASTER, krs)
		if err != nil {
			t.Errorf("want nil, got %v", err)
		}
		sort.Strings(gotShards)
		if !reflect.DeepEqual(testCase.shards, gotShards) {
			t.Errorf("want \n%#v, got \n%#v", testCase.shards, gotShards)
		}
	}
}

func TestKeyRangeToShardMap(t *testing.T) {
	ts := new(sandboxTopo)
	var testCases = []struct {
		keyspace string
		keyRange string
		shards   []string
	}{
		{keyspace: TEST_SHARDED, keyRange: "20-40", shards: []string{"20-40"}},
		// check for partial keyrange, spanning one shard
		{keyspace: TEST_SHARDED, keyRange: "10-18", shards: []string{"-20"}},
		// check for keyrange intersecting with multiple shards
		{keyspace: TEST_SHARDED, keyRange: "10-40", shards: []string{"-20", "20-40"}},
		// check for keyrange intersecting with multiple shards
		{keyspace: TEST_SHARDED, keyRange: "1C-2A", shards: []string{"-20", "20-40"}},
		// check for keyrange where kr.End is Max Key ""
		{keyspace: TEST_SHARDED, keyRange: "80-", shards: []string{"80-A0", "A0-C0", "C0-E0", "E0-"}},
		// test for sharded, non-partial keyrange spanning the entire space.
		{keyspace: TEST_SHARDED, keyRange: "", shards: []string{"-20", "20-40", "40-60", "60-80", "80-A0", "A0-C0", "C0-E0", "E0-"}},
		// test for unsharded, non-partial keyrange spanning the entire space.
		{keyspace: TEST_UNSHARDED, keyRange: "", shards: []string{"0"}},
	}

	for _, testCase := range testCases {
		var keyRange key.KeyRange
		var err error
		if testCase.keyRange == "" {
			keyRange = key.KeyRange{Start: "", End: ""}
		} else {
			krArray, err := key.ParseShardingSpec(testCase.keyRange)
			if err != nil {
				t.Errorf("Got error while parsing sharding spec %v", err)
			}
			keyRange = krArray[0]
		}
		allShards, err := getKeyspaceShards(ts, "", testCase.keyspace, topo.TYPE_MASTER)
		gotShards, err := resolveKeyRangeToShards(allShards, keyRange)
		if err != nil {
			t.Errorf("want nil, got %v", err)
		}
		if !reflect.DeepEqual(testCase.shards, gotShards) {
			t.Errorf("want \n%#v, got \n%#v", testCase.shards, gotShards)
		}
	}
}

func TestMapExactShards(t *testing.T) {
	ts := new(sandboxTopo)
	var testCases = []struct {
		keyspace string
		keyRange string
		shards   []string
		err      string
	}{
		{keyspace: KsTestSharded, keyRange: "20-40", shards: []string{"20-40"}},
		// check for partial keyrange, spanning one shard
		{keyspace: KsTestSharded, keyRange: "10-18", shards: nil, err: "keyrange {Start: 10, End: 18} does not exactly match shards"},
		// check for keyrange intersecting with multiple shards
		{keyspace: KsTestSharded, keyRange: "10-40", shards: nil, err: "keyrange {Start: 10, End: 40} does not exactly match shards"},
		// check for keyrange intersecting with multiple shards
		{keyspace: KsTestSharded, keyRange: "1c-2a", shards: nil, err: "keyrange {Start: 1c, End: 2a} does not exactly match shards"},
		// check for keyrange where kr.End is Max Key ""
		{keyspace: KsTestSharded, keyRange: "80-", shards: []string{"80-a0", "a0-c0", "c0-e0", "e0-"}},
		// test for sharded, non-partial keyrange spanning the entire space.
		{keyspace: KsTestSharded, keyRange: "", shards: []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"}},
	}

	for _, testCase := range testCases {
		var keyRange key.KeyRange
		var err error
		if testCase.keyRange == "" {
			keyRange = key.KeyRange{Start: "", End: ""}
		} else {
			krArray, err := key.ParseShardingSpec(testCase.keyRange)
			if err != nil {
				t.Errorf("Got error while parsing sharding spec %v", err)
			}
			keyRange = krArray[0]
		}
		_, gotShards, err := mapExactShards(context.Background(), ts, "", testCase.keyspace, pb.TabletType_MASTER, keyRange)
		if err != nil && err.Error() != testCase.err {
			t.Errorf("gotShards: %v, want %s", err, testCase.err)
		}
		if !reflect.DeepEqual(testCase.shards, gotShards) {
			t.Errorf("want \n%#v, got \n%#v", testCase.shards, gotShards)
		}
	}
}

func multisnapshotCmd(mysqld *mysqlctl.Mysqld, subFlags *flag.FlagSet, args []string) {
	concurrency := subFlags.Int("concurrency", 8, "how many compression jobs to run simultaneously")
	spec := subFlags.String("spec", "-", "shard specification")
	tablesString := subFlags.String("tables", "", "dump only this comma separated list of regexp for tables")
	excludeTablesString := subFlags.String("exclude_tables", "", "do not dump this comma separated list of regexp for tables")
	skipSlaveRestart := subFlags.Bool("skip_slave_restart", false, "after the snapshot is done, do not restart slave replication")
	maximumFilesize := subFlags.Uint64("maximum_file_size", 128*1024*1024, "the maximum size for an uncompressed data file")
	keyType := subFlags.String("key_type", "uint64", "type of the key column")
	subFlags.Parse(args)
	if subFlags.NArg() != 2 {
		log.Fatalf("action multisnapshot requires <db name> <key name>")
	}

	shards, err := key.ParseShardingSpec(*spec)
	if err != nil {
		log.Fatalf("multisnapshot failed: %v", err)
	}
	var tables []string
	if *tablesString != "" {
		tables = strings.Split(*tablesString, ",")
	}
	var excludedTables []string
	if *excludeTablesString != "" {
		excludedTables = strings.Split(*excludeTablesString, ",")
	}

	kit := key.KeyspaceIdType(*keyType)
	if !key.IsKeyspaceIdTypeInList(kit, key.AllKeyspaceIdTypes) {
		log.Fatalf("invalid key_type")
	}

	filenames, err := mysqld.CreateMultiSnapshot(shards, subFlags.Arg(0), subFlags.Arg(1), kit, tabletAddr, false, *concurrency, tables, excludedTables, *skipSlaveRestart, *maximumFilesize, nil)
	if err != nil {
		log.Fatalf("multisnapshot failed: %v", err)
	} else {
		log.Infof("manifest locations: %v", filenames)
	}
}

// TODO(erez): Rename after migration to SplitQuery V2 is done.
func TestVTGateSplitQueryV2Sharded(t *testing.T) {
	keyspace := "TestVTGateSplitQuery"
	keyranges, err := key.ParseShardingSpec(DefaultShardSpec)
	if err != nil {
		t.Fatalf("got: %v, want: nil", err)
	}
	createSandbox(keyspace)
	hcVTGateTest.Reset()
	port := int32(1001)
	for _, kr := range keyranges {
		hcVTGateTest.AddTestTablet("aa", "1.1.1.1", port, keyspace, key.KeyRangeString(kr), topodatapb.TabletType_RDONLY, true, 1, nil)
		port++
	}
	sql := "select col1, col2 from table"
	bindVars := map[string]interface{}{"bv1": nil}
	splitColumns := []string{"sc1", "sc2"}
	algorithm := querypb.SplitQueryRequest_FULL_SCAN
	type testCaseType struct {
		splitCount          int64
		numRowsPerQueryPart int64
	}
	testCases := []testCaseType{
		{splitCount: 100, numRowsPerQueryPart: 0},
		{splitCount: 0, numRowsPerQueryPart: 123},
	}
	for _, testCase := range testCases {
		splits, err := rpcVTGate.SplitQueryV2(
			context.Background(),
			keyspace,
			sql,
			bindVars,
			splitColumns,
			testCase.splitCount,
			testCase.numRowsPerQueryPart,
			algorithm)
		if err != nil {
			t.Errorf("got %v, want: nil. testCase: %+v", err, testCase)
		}
		// Total number of splits should be number of shards as our sandbox returns a single split
		// for its fake implementation of SplitQuery.
		if len(keyranges) != len(splits) {
			t.Errorf("wrong number of splits, got %+v, want %+v. testCase:\n%+v",
				len(splits), len(keyranges), testCase)
		}
		actualSqlsByKeyRange := map[string][]string{}
		for _, split := range splits {
			if split.KeyRangePart.Keyspace != keyspace {
				t.Errorf("wrong keyspace, got \n%+v, want \n%+v. testCase:\n%+v",
					keyspace, split.KeyRangePart.Keyspace, testCase)
			}
			if len(split.KeyRangePart.KeyRanges) != 1 {
				t.Errorf("wrong number of keyranges, got \n%+v, want \n%+v. testCase:\n%+v",
					1, len(split.KeyRangePart.KeyRanges), testCase)
			}
			kr := key.KeyRangeString(split.KeyRangePart.KeyRanges[0])
			actualSqlsByKeyRange[kr] = append(actualSqlsByKeyRange[kr], split.Query.Sql)
		}
		expectedSqlsByKeyRange := map[string][]string{}
		for _, kr := range keyranges {
			perShardSplitCount := int64(math.Ceil(float64(testCase.splitCount) / float64(len(keyranges))))
			shard := key.KeyRangeString(kr)
			expectedSqlsByKeyRange[shard] = []string{
				fmt.Sprintf(
					"query:%v, splitColumns:%v, splitCount:%v,"+
						" numRowsPerQueryPart:%v, algorithm:%v, shard:%v",
					querytypes.BoundQuery{Sql: sql, BindVariables: bindVars},
					splitColumns,
					perShardSplitCount,
					testCase.numRowsPerQueryPart,
					algorithm,
					shard,
				),
			}
		}
		if !reflect.DeepEqual(actualSqlsByKeyRange, expectedSqlsByKeyRange) {
			t.Errorf(
				"splits contain the wrong sqls and/or keyranges, "+
					"got:\n%+v\n, want:\n%+v\n. testCase:\n%+v",
				actualSqlsByKeyRange, expectedSqlsByKeyRange, testCase)
		}
	}
}