Golang key.KeyRangeString函數代碼示例

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 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)
	}
}

// StreamKeyRange is part of the binlogplayer.Client interface
func (fbc *fakeBinlogClient) StreamKeyRange(ctx context.Context, position string, keyRange *topodatapb.KeyRange, charset *binlogdatapb.Charset) (binlogplayer.BinlogTransactionStream, error) {
	actualKeyRange := key.KeyRangeString(keyRange)
	if actualKeyRange != fbc.expectedKeyRange {
		return nil, fmt.Errorf("Got wrong keyrange %v, expected %v", actualKeyRange, fbc.expectedKeyRange)
	}
	return &testStreamEventAdapter{c: fbc.keyRangeChannel, ctx: ctx}, nil
}

// initializeKeyRangeRule will create and set the key range rules
func (agent *ActionAgent) initializeKeyRangeRule(ctx context.Context, keyspace string, keyRange *topodatapb.KeyRange) error {
	// check we have a partial key range
	if !key.KeyRangeIsPartial(keyRange) {
		log.Infof("Tablet covers the full KeyRange, not adding KeyRange query rule")
		return nil
	}

	// read the keyspace to get the sharding column name
	keyspaceInfo, err := agent.TopoServer.GetKeyspace(ctx, keyspace)
	if err != nil {
		return fmt.Errorf("cannot read keyspace %v to get sharding key: %v", keyspace, err)
	}
	if keyspaceInfo.ShardingColumnName == "" {
		log.Infof("Keyspace %v has an empty ShardingColumnName, not adding KeyRange query rule", keyspace)
		return nil
	}

	// create the rules
	log.Infof("Restricting to keyrange: %v", key.KeyRangeString(keyRange))
	keyrangeRules := tabletserver.NewQueryRules()
	dmlPlans := []struct {
		planID   planbuilder.PlanType
		onAbsent bool
	}{
		{planbuilder.PlanInsertPK, true},
		{planbuilder.PlanInsertSubquery, true},
		{planbuilder.PlanPassDML, false},
		{planbuilder.PlanDMLPK, false},
		{planbuilder.PlanDMLSubquery, false},
		{planbuilder.PlanUpsertPK, false},
	}
	for _, plan := range dmlPlans {
		qr := tabletserver.NewQueryRule(
			fmt.Sprintf("enforce %v range for %v", keyspaceInfo.ShardingColumnName, plan.planID),
			fmt.Sprintf("%v_not_in_range_%v", keyspaceInfo.ShardingColumnName, plan.planID),
			tabletserver.QRFail,
		)
		qr.AddPlanCond(plan.planID)
		err := qr.AddBindVarCond(keyspaceInfo.ShardingColumnName, plan.onAbsent, true, tabletserver.QRNotIn, keyRange)
		if err != nil {
			return fmt.Errorf("Unable to add key range rule: %v", err)
		}
		keyrangeRules.Add(qr)
	}

	// and load them
	agent.QueryServiceControl.RegisterQueryRuleSource(keyrangeQueryRules)
	if err := agent.QueryServiceControl.SetQueryRules(keyrangeQueryRules, keyrangeRules); err != nil {
		return fmt.Errorf("failed to load query rule set %s: %s", keyrangeQueryRules, err)
	}
	return nil
}

func createShardedSrvKeyspace(shardSpec, servedFromKeyspace string) (*topodatapb.SrvKeyspace, error) {
	shardKrArray, err := getAllShards(shardSpec)
	if err != nil {
		return nil, err
	}
	shards := make([]*topodatapb.ShardReference, 0, len(shardKrArray))
	for i := 0; i < len(shardKrArray); i++ {
		shard := &topodatapb.ShardReference{
			Name:     key.KeyRangeString(shardKrArray[i]),
			KeyRange: shardKrArray[i],
		}
		shards = append(shards, shard)
	}
	shardedSrvKeyspace := &topodatapb.SrvKeyspace{
		ShardingColumnName: "user_id", // exact value is ignored
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
			{
				ServedType:      topodatapb.TabletType_MASTER,
				ShardReferences: shards,
			},
			{
				ServedType:      topodatapb.TabletType_REPLICA,
				ShardReferences: shards,
			},
			{
				ServedType:      topodatapb.TabletType_RDONLY,
				ShardReferences: shards,
			},
		},
	}
	if servedFromKeyspace != "" {
		shardedSrvKeyspace.ServedFrom = []*topodatapb.SrvKeyspace_ServedFrom{
			{
				TabletType: topodatapb.TabletType_RDONLY,
				Keyspace:   servedFromKeyspace,
			},
			{
				TabletType: topodatapb.TabletType_MASTER,
				Keyspace:   servedFromKeyspace,
			},
		}
	}
	return shardedSrvKeyspace, nil
}

// mapExactShards maps a keyrange to shards only if there's a complete
// match. If there's any partial match the function returns no match.
func mapExactShards(ctx context.Context, topoServ SrvTopoServer, cell, keyspace string, tabletType pb.TabletType, kr *pb.KeyRange) (newkeyspace string, shards []string, err error) {
	keyspace, _, allShards, err := getKeyspaceShards(ctx, topoServ, cell, keyspace, tabletType)
	if err != nil {
		return "", nil, err
	}
	shardnum := 0
	for shardnum < len(allShards) {
		if bytes.Compare(kr.Start, []byte(allShards[shardnum].KeyRange.Start)) == 0 {
			break
		}
		shardnum++
	}
	for shardnum < len(allShards) {
		shards = append(shards, allShards[shardnum].Name)
		if bytes.Compare(kr.End, []byte(allShards[shardnum].KeyRange.End)) == 0 {
			return keyspace, shards, nil
		}
		shardnum++
	}
	return keyspace, nil, fmt.Errorf("keyrange %v does not exactly match shards", key.KeyRangeString(kr))
}

// StreamKeyRange is part of the binlogplayer.Client interface
func (fbc *fakeBinlogClient) StreamKeyRange(ctx context.Context, position string, keyRange *topodatapb.KeyRange, charset *binlogdatapb.Charset) (chan *binlogdatapb.BinlogTransaction, binlogplayer.ErrFunc, error) {
	actualKeyRange := key.KeyRangeString(keyRange)
	if actualKeyRange != fbc.expectedKeyRange {
		return nil, nil, fmt.Errorf("Got wrong keyrange %v, expected %v", actualKeyRange, fbc.expectedKeyRange)
	}

	c := make(chan *binlogdatapb.BinlogTransaction)
	var finalErr error
	go func() {
		for {
			select {
			case bt := <-fbc.keyRangeChannel:
				c <- bt
			case <-ctx.Done():
				finalErr = ctx.Err()
				close(c)
				return
			}
		}
	}()
	return c, func() error {
		return finalErr
	}, nil
}

// 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)
		}
	}
}

//.........這裏部分代碼省略.........

	if err := ts.DeleteEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_REPLICA, -1); err != topo.ErrNoNode {
		t.Errorf("DeleteEndPoints(unknown): %v", err)
	}
	if err := ts.DeleteEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, -1); err != nil {
		t.Errorf("DeleteEndPoints(master): %v", err)
	}

	// test cell/keyspace/shard entries (SrvShard)
	srvShard := &topodatapb.SrvShard{
		Name:       "-10",
		KeyRange:   newKeyRange("-10"),
		MasterCell: "test",
	}
	if err := ts.UpdateSrvShard(ctx, cell, "test_keyspace", "-10", srvShard); err != nil {
		t.Fatalf("UpdateSrvShard(1): %v", err)
	}
	if _, err := ts.GetSrvShard(ctx, cell, "test_keyspace", "666"); err != topo.ErrNoNode {
		t.Errorf("GetSrvShard(invalid): %v", err)
	}
	if s, err := ts.GetSrvShard(ctx, cell, "test_keyspace", "-10"); err != nil ||
		s.Name != "-10" ||
		!key.KeyRangeEqual(s.KeyRange, newKeyRange("-10")) ||
		s.MasterCell != "test" {
		t.Errorf("GetSrvShard(valid): %v", err)
	}

	// test cell/keyspace entries (SrvKeyspace)
	srvKeyspace := topodatapb.SrvKeyspace{
		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
			&topodatapb.SrvKeyspace_KeyspacePartition{
				ServedType: topodatapb.TabletType_MASTER,
				ShardReferences: []*topodatapb.ShardReference{
					&topodatapb.ShardReference{
						Name: "-80",
						KeyRange: &topodatapb.KeyRange{
							End: []byte{0x80},
						},
					},
				},
			},
		},
		ShardingColumnName: "video_id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
		ServedFrom: []*topodatapb.SrvKeyspace_ServedFrom{
			&topodatapb.SrvKeyspace_ServedFrom{
				TabletType: topodatapb.TabletType_REPLICA,
				Keyspace:   "other_keyspace",
			},
		},
	}
	if err := ts.UpdateSrvKeyspace(ctx, cell, "test_keyspace", &srvKeyspace); err != nil {
		t.Errorf("UpdateSrvKeyspace(1): %v", err)
	}
	if _, err := ts.GetSrvKeyspace(ctx, cell, "test_keyspace666"); err != topo.ErrNoNode {
		t.Errorf("GetSrvKeyspace(invalid): %v", err)
	}
	if k, err := ts.GetSrvKeyspace(ctx, cell, "test_keyspace"); err != nil ||
		len(k.Partitions) != 1 ||
		k.Partitions[0].ServedType != topodatapb.TabletType_MASTER ||
		len(k.Partitions[0].ShardReferences) != 1 ||
		k.Partitions[0].ShardReferences[0].Name != "-80" ||
		key.KeyRangeString(k.Partitions[0].ShardReferences[0].KeyRange) != "-80" ||
		k.ShardingColumnName != "video_id" ||
		k.ShardingColumnType != topodatapb.KeyspaceIdType_UINT64 ||
		len(k.ServedFrom) != 1 ||
		k.ServedFrom[0].TabletType != topodatapb.TabletType_REPLICA ||
		k.ServedFrom[0].Keyspace != "other_keyspace" {
		t.Errorf("GetSrvKeyspace(valid): %v %v", err, k)
	}
	if k, err := ts.GetSrvKeyspaceNames(ctx, cell); err != nil || len(k) != 1 || k[0] != "test_keyspace" {
		t.Errorf("GetSrvKeyspaceNames(): %v", err)
	}

	// check that updating a SrvKeyspace out of the blue works
	if err := ts.UpdateSrvKeyspace(ctx, cell, "unknown_keyspace_so_far", &srvKeyspace); err != nil {
		t.Fatalf("UpdateSrvKeyspace(2): %v", err)
	}
	if k, err := ts.GetSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != nil ||
		len(k.Partitions) != 1 ||
		k.Partitions[0].ServedType != topodatapb.TabletType_MASTER ||
		len(k.Partitions[0].ShardReferences) != 1 ||
		k.Partitions[0].ShardReferences[0].Name != "-80" ||
		key.KeyRangeString(k.Partitions[0].ShardReferences[0].KeyRange) != "-80" ||
		k.ShardingColumnName != "video_id" ||
		k.ShardingColumnType != topodatapb.KeyspaceIdType_UINT64 ||
		len(k.ServedFrom) != 1 ||
		k.ServedFrom[0].TabletType != topodatapb.TabletType_REPLICA ||
		k.ServedFrom[0].Keyspace != "other_keyspace" {
		t.Errorf("GetSrvKeyspace(out of the blue): %v %v", err, *k)
	}

	// Delete the SrvKeyspace.
	if err := ts.DeleteSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != nil {
		t.Fatalf("DeleteSrvShard: %v", err)
	}
	if _, err := ts.GetSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != topo.ErrNoNode {
		t.Errorf("GetSrvKeyspace(deleted) got %v, want ErrNoNode", err)
	}
}