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) } } }
// CheckServingGraph makes sure the serving graph functions work properly. func CheckServingGraph(ctx context.Context, t *testing.T, ts topo.Impl) { cell := getLocalCell(ctx, t, ts) // test individual cell/keyspace/shard/type entries if _, err := ts.GetSrvTabletTypesPerShard(ctx, cell, "test_keyspace", "-10"); err != topo.ErrNoNode { t.Errorf("GetSrvTabletTypesPerShard(invalid): %v", err) } if _, _, err := ts.GetEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER); err != topo.ErrNoNode { t.Errorf("GetEndPoints(invalid): %v", err) } endPoints := &topodatapb.EndPoints{ Entries: []*topodatapb.EndPoint{ &topodatapb.EndPoint{ Uid: 1, Host: "host1", PortMap: map[string]int32{ "vt": 1234, "mysql": 1235, "grpc": 1236, }, }, }, } if err := ts.CreateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints); err != nil { t.Fatalf("CreateEndPoints(master): %v", err) } // Try to create again. if err := ts.CreateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints); err != topo.ErrNodeExists { t.Fatalf("CreateEndPoints(master): err = %v, want topo.ErrNodeExists", err) } // Get version. _, version, err := ts.GetEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER) if err != nil { t.Fatalf("GetEndPoints(master): %v", err) } // Make a change. tmp := endPoints.Entries[0].Uid endPoints.Entries[0].Uid = tmp + 1 if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, -1); err != nil { t.Fatalf("UpdateEndPoints(master): %v", err) } endPoints.Entries[0].Uid = tmp // Try to delete with the wrong version. if err := ts.DeleteEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, version); err != topo.ErrBadVersion { t.Fatalf("DeleteEndPoints: err = %v, want topo.ErrBadVersion", err) } // Delete with the correct version. _, version, err = ts.GetEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER) if err != nil { t.Fatalf("GetEndPoints(master): %v", err) } if err := ts.DeleteEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, version); err != nil { t.Fatalf("DeleteEndPoints: %v", err) } // Recreate it with an unconditional update. if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, -1); err != nil { t.Fatalf("UpdateEndPoints(master): %v", err) } if types, err := ts.GetSrvTabletTypesPerShard(ctx, cell, "test_keyspace", "-10"); err != nil || len(types) != 1 || types[0] != topodatapb.TabletType_MASTER { t.Errorf("GetSrvTabletTypesPerShard(1): %v %v", err, types) } // Delete it unconditionally. if err := ts.DeleteEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, -1); err != nil { t.Fatalf("DeleteEndPoints: %v", err) } // Delete the SrvShard. if err := ts.DeleteSrvShard(ctx, cell, "test_keyspace", "-10"); err != nil { t.Fatalf("DeleteSrvShard: %v", err) } if _, err := ts.GetSrvShard(ctx, cell, "test_keyspace", "-10"); err != topo.ErrNoNode { t.Errorf("GetSrvShard(deleted) got %v, want ErrNoNode", err) } // Re-add endpoints. if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, -1); err != nil { t.Fatalf("UpdateEndPoints(master): %v", err) } addrs, version, err := ts.GetEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER) if err != nil { t.Errorf("GetEndPoints: %v", err) } if len(addrs.Entries) != 1 || addrs.Entries[0].Uid != 1 { t.Errorf("GetEndPoints(1): %v", addrs) } if pm := addrs.Entries[0].PortMap; pm["vt"] != 1234 || pm["mysql"] != 1235 || pm["grpc"] != 1236 { t.Errorf("GetSrcTabletType(1).PortMap: want %v, got %v", endPoints.Entries[0].PortMap, pm) } // Update with the wrong version. if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, version+1); err != topo.ErrBadVersion { t.Fatalf("UpdateEndPoints(master): err = %v, want topo.ErrBadVersion", err) } // Update with the right version. if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, version); err != nil { t.Fatalf("UpdateEndPoints(master): %v", err) } // Update existing EndPoints unconditionally. if err := ts.UpdateEndPoints(ctx, cell, "test_keyspace", "-10", topodatapb.TabletType_MASTER, endPoints, -1); err != nil { t.Fatalf("UpdateEndPoints(master): %v", err) } 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) } }