// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map. func (vtg *VTGate) ExecuteEntityIds(context interface{}, query *proto.EntityIdsQuery, reply *proto.QueryResult) error { shardMap, err := mapEntityIdsToShards( vtg.scatterConn.toposerv, vtg.scatterConn.cell, query.Keyspace, query.EntityKeyspaceIdMap, query.TabletType) if err != nil { return err } shards, sqls, bindVars := buildEntityIds(shardMap, query.Sql, query.EntityColumnName, query.BindVariables) qr, err := vtg.scatterConn.ExecuteEntityIds( context, shards, sqls, bindVars, query.Keyspace, query.TabletType, NewSafeSession(query.Session)) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() log.Errorf("ExecuteEntityIds: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
// StreamExecuteShard executes a streaming query on the specified shards. func (vtg *VTGate) StreamExecuteShard(context interface{}, query *proto.QueryShard, sendReply func(*proto.QueryResult) error) error { err := vtg.scatterConn.StreamExecute( context, query.Sql, query.BindVariables, query.Keyspace, query.Shards, query.TabletType, NewSafeSession(query.Session), func(mreply *mproto.QueryResult) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply, reply) // Note we don't populate reply.Session here, // as it may change incrementaly as responses // are sent. return sendReply(reply) }) if err != nil { log.Errorf("StreamExecuteShard: %v, query: %+v", err, query) } // now we can send the final Session info. if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
// ExecuteShard executes a non-streaming query on the specified shards. func (vtg *VTGate) ExecuteShard(context interface{}, query *proto.QueryShard, reply *proto.QueryResult) error { startTime := time.Now() statsKey := []string{"ExecuteShard", query.Keyspace, string(query.TabletType)} defer vtg.timings.Record(statsKey, startTime) qr, err := vtg.resolver.Execute( context, query.Sql, query.BindVariables, query.Keyspace, query.TabletType, query.Session, func(keyspace string) (string, []string, error) { return query.Keyspace, query.Shards, nil }, ) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() vtg.errors.Add(statsKey, 1) // FIXME(alainjobart): throttle this log entry log.Errorf("ExecuteShard: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
// ExecuteKeyRanges executes a non-streaming query based on the specified keyranges. func (vtg *VTGate) ExecuteKeyRanges(context interface{}, query *proto.KeyRangeQuery, reply *proto.QueryResult) error { shards, err := mapKeyRangesToShards( vtg.scatterConn.toposerv, vtg.scatterConn.cell, query.Keyspace, query.TabletType, query.KeyRanges) if err != nil { return err } qr, err := vtg.scatterConn.Execute( context, query.Sql, query.BindVariables, query.Keyspace, shards, query.TabletType, NewSafeSession(query.Session)) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() log.Errorf("ExecuteKeyRange: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
// StreamExecuteKeyRanges executes a streaming query on the specified KeyRanges. // The KeyRanges are resolved to shards using the serving graph. // This function currently temporarily enforces the restriction of executing on // one shard since it cannot merge-sort the results to guarantee ordering of // response which is needed for checkpointing. // The api supports supplying multiple keyranges to make it future proof. func (vtg *VTGate) StreamExecuteKeyRanges(context interface{}, query *proto.KeyRangeQuery, sendReply func(*proto.QueryResult) error) error { startTime := time.Now() statsKey := []string{"StreamExecuteKeyRanges", query.Keyspace, string(query.TabletType)} defer vtg.timings.Record(statsKey, startTime) err := vtg.resolver.StreamExecuteKeyRanges( context, query, func(mreply *mproto.QueryResult) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply, reply) // Note we don't populate reply.Session here, // as it may change incrementaly as responses are sent. return sendReply(reply) }) if err != nil { vtg.errors.Add(statsKey, 1) vtg.logStreamExecuteKeyRanges.Errorf("%v, query: %+v", err, query) } // now we can send the final Session info. if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
// StreamExecuteShard executes a streaming query on the specified shards. func (vtg *VTGate) StreamExecuteShard(context interface{}, query *proto.QueryShard, sendReply func(*proto.QueryResult) error) error { startTime := time.Now() statsKey := []string{"StreamExecuteShard", query.Keyspace, string(query.TabletType)} defer vtg.timings.Record(statsKey, startTime) err := vtg.resolver.StreamExecute( context, query.Sql, query.BindVariables, query.Keyspace, query.TabletType, query.Session, func(keyspace string) (string, []string, error) { return query.Keyspace, query.Shards, nil }, func(mreply *mproto.QueryResult) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply, reply) // Note we don't populate reply.Session here, // as it may change incrementaly as responses // are sent. return sendReply(reply) }) if err != nil { vtg.errors.Add(statsKey, 1) // FIXME(alainjobart): throttle this log entry log.Errorf("StreamExecuteShard: %v, query: %+v", err, query) } // now we can send the final Session info. if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
func TestVTGateExecuteShard(t *testing.T) { sandbox := createSandbox("TestVTGateExecuteShard") sbc := &sandboxConn{} sandbox.MapTestConn("0", sbc) q := proto.QueryShard{ Sql: "query", Keyspace: "TestVTGateExecuteShard", Shards: []string{"0"}, } qr := new(proto.QueryResult) err := RpcVTGate.ExecuteShard(nil, &q, qr) if err != nil { t.Errorf("want nil, got %v", err) } wantqr := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, wantqr) if !reflect.DeepEqual(wantqr, qr) { t.Errorf("want \n%+v, got \n%+v", singleRowResult, qr) } if qr.Session != nil { t.Errorf("want nil, got %+v\n", qr.Session) } q.Session = new(proto.Session) RpcVTGate.Begin(nil, q.Session) if !q.Session.InTransaction { t.Errorf("want true, got false") } RpcVTGate.ExecuteShard(nil, &q, qr) wantSession := &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateExecuteShard", Shard: "0", TransactionId: 1, }}, } if !reflect.DeepEqual(wantSession, q.Session) { t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session) } RpcVTGate.Commit(nil, q.Session) if sbc.CommitCount != 1 { t.Errorf("want 1, got %d", sbc.CommitCount) } q.Session = new(proto.Session) RpcVTGate.Begin(nil, q.Session) RpcVTGate.ExecuteShard(nil, &q, qr) RpcVTGate.Rollback(nil, q.Session) /* // Flaky: This test should be run manually. runtime.Gosched() if sbc.RollbackCount != 1 { t.Errorf("want 1, got %d", sbc.RollbackCount) } */ }
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map. func (vtg *VTGate) ExecuteEntityIds(context interface{}, query *proto.EntityIdsQuery, reply *proto.QueryResult) error { qr, err := vtg.resolver.ExecuteEntityIds(context, query) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() log.Errorf("ExecuteEntityIds: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
func TestVTGateStreamExecuteShard(t *testing.T) { s := createSandbox("TestVTGateStreamExecuteShard") sbc := &sandboxConn{} s.MapTestConn("0", sbc) q := proto.QueryShard{ Sql: "query", Keyspace: "TestVTGateStreamExecuteShard", Shards: []string{"0"}, TabletType: topo.TYPE_MASTER, } // Test for successful execution var qrs []*proto.QueryResult err := RpcVTGate.StreamExecuteShard(nil, &q, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } q.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, q.Session) err = RpcVTGate.StreamExecuteShard(nil, &q, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) want = []*proto.QueryResult{ row, &proto.QueryResult{ Session: &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateStreamExecuteShard", Shard: "0", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } }
// ExecuteKeyspaceIds executes a non-streaming query based on the specified keyspace ids. func (vtg *VTGate) ExecuteKeyspaceIds(context interface{}, query *proto.KeyspaceIdQuery, reply *proto.QueryResult) error { startTime := time.Now() statsKey := []string{"ExecuteKeyspaceIds", query.Keyspace, string(query.TabletType)} defer vtg.timings.Record(statsKey, startTime) qr, err := vtg.resolver.ExecuteKeyspaceIds(context, query) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() vtg.errors.Add(statsKey, 1) vtg.logExecuteKeyspaceIds.Errorf("%v, query: %+v", err, query) } reply.Session = query.Session return nil }
func TestVTGateStreamExecuteShard(t *testing.T) { resetSandbox() sbc := &sandboxConn{} testConns[0] = sbc q := proto.QueryShard{ Sql: "query", Shards: []string{"0"}, } var qrs []*proto.QueryResult err := RpcVTGate.StreamExecuteShard(nil, &q, func(r interface{}) error { qrs = append(qrs, r.(*proto.QueryResult)) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%#v, got \n%#v", want, qrs) } q.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, nil, q.Session) err = RpcVTGate.StreamExecuteShard(nil, &q, func(r interface{}) error { qrs = append(qrs, r.(*proto.QueryResult)) return nil }) want = []*proto.QueryResult{ row, &proto.QueryResult{ Session: &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Shard: "0", TransactionId: 1, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%#v, got \n%#v", want, qrs) } }
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map. func (vtg *VTGate) ExecuteEntityIds(context interface{}, query *proto.EntityIdsQuery, reply *proto.QueryResult) error { startTime := time.Now() statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)} defer vtg.timings.Record(statsKey, startTime) qr, err := vtg.resolver.ExecuteEntityIds(context, query) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() vtg.errors.Add(statsKey, 1) // FIXME(alainjobart): throttle this log entry log.Errorf("ExecuteEntityIds: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
// ExecuteShard executes a non-streaming query on the specified shards. func (vtg *VTGate) ExecuteShard(context *rpcproto.Context, query *proto.QueryShard, reply *proto.QueryResult) error { qr, err := vtg.scatterConn.Execute( query.Sql, query.BindVariables, query.Keyspace, query.Shards, query.TabletType, NewSafeSession(query.Session)) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() log.Errorf("ExecuteShard: %v, query: %#v", err, query) } reply.Session = query.Session return nil }
// StreamExecuteKeyspaceIds executes a streaming query on the specified KeyspaceIds. // The KeyspaceIds are resolved to shards using the serving graph. // This function currently temporarily enforces the restriction of executing on // one shard since it cannot merge-sort the results to guarantee ordering of // response which is needed for checkpointing. // The api supports supplying multiple KeyspaceIds to make it future proof. func (vtg *VTGate) StreamExecuteKeyspaceIds(context interface{}, query *proto.KeyspaceIdQuery, sendReply func(*proto.QueryResult) error) error { err := vtg.resolver.StreamExecuteKeyspaceIds( context, query, func(mreply *mproto.QueryResult) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply, reply) // Note we don't populate reply.Session here, // as it may change incrementaly as responses are sent. return sendReply(reply) }) if err != nil { log.Errorf("StreamExecuteKeyspaceIds: %v, query: %+v", err, query) } // now we can send the final Sessoin info. if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
// ExecuteShard executes a non-streaming query on the specified shards. func (vtg *VTGate) ExecuteShard(context interface{}, query *proto.QueryShard, reply *proto.QueryResult) error { qr, err := vtg.resolver.Execute( context, query.Sql, query.BindVariables, query.Keyspace, query.TabletType, query.Session, func(keyspace string) ([]string, error) { return query.Shards, nil }, ) if err == nil { proto.PopulateQueryResult(qr, reply) } else { reply.Error = err.Error() log.Errorf("ExecuteShard: %v, query: %+v", err, query) } reply.Session = query.Session return nil }
// StreamExecuteShard executes a streaming query on the specified shards. func (vtg *VTGate) StreamExecuteShard(context *rpcproto.Context, query *proto.QueryShard, sendReply func(interface{}) error) error { err := vtg.scatterConn.StreamExecute( query.Sql, query.BindVariables, query.Keyspace, query.Shards, query.TabletType, NewSafeSession(query.Session), func(mreply interface{}) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply.(*mproto.QueryResult), reply) return sendReply(reply) }) if err != nil { log.Errorf("StreamExecuteShard: %v, query: %#v", err, query) } if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
// StreamExecuteKeyRanges executes a streaming query on the specified KeyRanges. // The KeyRanges are resolved to shards using the serving graph. // This function currently temporarily enforces the restriction of executing on // one shard since it cannot merge-sort the results to guarantee ordering of // response which is needed for checkpointing. The api supports supplying multiple keyranges // to make it future proof. func (vtg *VTGate) StreamExecuteKeyRanges(context interface{}, query *proto.KeyRangeQuery, sendReply func(*proto.QueryResult) error) error { shards, err := mapKeyRangesToShards( vtg.scatterConn.toposerv, vtg.scatterConn.cell, query.Keyspace, query.TabletType, query.KeyRanges) if err != nil { return err } if len(shards) != 1 { return fmt.Errorf("KeyRanges cannot map to more than one shard") } err = vtg.scatterConn.StreamExecute( context, query.Sql, query.BindVariables, query.Keyspace, shards, query.TabletType, NewSafeSession(query.Session), func(mreply *mproto.QueryResult) error { reply := new(proto.QueryResult) proto.PopulateQueryResult(mreply, reply) // Note we don't populate reply.Session here, // as it may change incrementaly as responses are sent. return sendReply(reply) }) if err != nil { log.Errorf("StreamExecuteKeyRange: %v, query: %+v", err, query) } // now we can send the final Session info. if query.Session != nil { sendReply(&proto.QueryResult{Session: query.Session}) } return err }
func TestVTGateExecuteKeyspaceIds(t *testing.T) { s := createSandbox("TestVTGateExecuteKeyspaceIds") sbc1 := &sandboxConn{} sbc2 := &sandboxConn{} s.MapTestConn("-20", sbc1) s.MapTestConn("20-40", sbc2) kid10, err := key.HexKeyspaceId("10").Unhex() if err != nil { t.Errorf("want nil, got %+v", err) } q := proto.KeyspaceIdQuery{ Sql: "query", Keyspace: "TestVTGateExecuteKeyspaceIds", KeyspaceIds: []key.KeyspaceId{kid10}, TabletType: topo.TYPE_MASTER, } // Test for successful execution qr := new(proto.QueryResult) err = RpcVTGate.ExecuteKeyspaceIds(nil, &q, qr) if err != nil { t.Errorf("want nil, got %v", err) } wantqr := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, wantqr) if !reflect.DeepEqual(wantqr, qr) { t.Errorf("want \n%+v, got \n%+v", singleRowResult, qr) } if qr.Session != nil { t.Errorf("want nil, got %+v\n", qr.Session) } if sbc1.ExecCount != 1 { t.Errorf("want 1, got %v\n", sbc1.ExecCount) } // Test for successful execution in transaction q.Session = new(proto.Session) RpcVTGate.Begin(nil, q.Session) if !q.Session.InTransaction { t.Errorf("want true, got false") } RpcVTGate.ExecuteKeyspaceIds(nil, &q, qr) wantSession := &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateExecuteKeyspaceIds", Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, } if !reflect.DeepEqual(wantSession, q.Session) { t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session) } RpcVTGate.Commit(nil, q.Session) if sbc1.CommitCount.Get() != 1 { t.Errorf("want 1, got %d", sbc1.CommitCount.Get()) } // Test for multiple shards kid30, err := key.HexKeyspaceId("30").Unhex() if err != nil { t.Errorf("want nil, got %+v", err) } q.KeyspaceIds = []key.KeyspaceId{kid10, kid30} RpcVTGate.ExecuteKeyspaceIds(nil, &q, qr) if qr.RowsAffected != 2 { t.Errorf("want 2, got %v", qr.RowsAffected) } }
func TestVTGateStreamExecuteKeyRange(t *testing.T) { resetSandbox() sbc := &sandboxConn{} mapTestConn("-20", sbc) sq := proto.StreamQueryKeyRange{ Sql: "query", KeyRange: "-20", TabletType: topo.TYPE_MASTER, } // Test for successful execution var qrs []*proto.QueryResult err := RpcVTGate.StreamExecuteKeyRange(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%#v, got \n%#v", want, qrs) } sq.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, sq.Session) err = RpcVTGate.StreamExecuteKeyRange(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) want = []*proto.QueryResult{ row, &proto.QueryResult{ Session: &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%#v, got \n%#v", want, qrs) } // Test for error condition - multiple shards sq.KeyRange = "10-40" err = RpcVTGate.StreamExecuteKeyRange(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } // Test for error condition - multiple shards, non-partial keyspace sq.KeyRange = "" err = RpcVTGate.StreamExecuteKeyRange(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } }
func TestVTGateStreamExecuteKeyRanges(t *testing.T) { s := createSandbox("TestVTGateStreamExecuteKeyRanges") sbc := &sandboxConn{} s.MapTestConn("-20", sbc) kr, err := key.ParseKeyRangeParts("", "20") sq := proto.KeyRangeQuery{ Sql: "query", Keyspace: "TestVTGateStreamExecuteKeyRanges", KeyRanges: []key.KeyRange{kr}, TabletType: topo.TYPE_MASTER, } // Test for successful execution var qrs []*proto.QueryResult err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } sq.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, sq.Session) err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) want = []*proto.QueryResult{ row, &proto.QueryResult{ Session: &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateStreamExecuteKeyRanges", Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } // Test for error condition - multiple shards kr, err = key.ParseKeyRangeParts("10", "40") sq.KeyRanges = []key.KeyRange{kr} err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } // Test for error condition - multiple shards, non-partial keyspace kr, err = key.ParseKeyRangeParts("", "") sq.KeyRanges = []key.KeyRange{kr} err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } }
func TestVTGateStreamExecuteKeyspaceIds(t *testing.T) { s := createSandbox("TestVTGateStreamExecuteKeyspaceIds") sbc := &sandboxConn{} s.MapTestConn("-20", sbc) kid10, err := key.HexKeyspaceId("10").Unhex() if err != nil { t.Errorf("want nil, got %+v", err) } sq := proto.KeyspaceIdQuery{ Sql: "query", Keyspace: "TestVTGateStreamExecuteKeyspaceIds", KeyspaceIds: []key.KeyspaceId{kid10}, TabletType: topo.TYPE_MASTER, } // Test for successful execution var qrs []*proto.QueryResult err = RpcVTGate.StreamExecuteKeyspaceIds(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } // Test for successful execution in transaction sq.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, sq.Session) err = RpcVTGate.StreamExecuteKeyspaceIds(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) want = []*proto.QueryResult{ row, &proto.QueryResult{ Session: &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateStreamExecuteKeyspaceIds", Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } RpcVTGate.Commit(nil, sq.Session) if sbc.CommitCount.Get() != 1 { t.Errorf("want 1, got %d", sbc.CommitCount.Get()) } // Test for successful execution - multiple keyspaceids in single shard sq.Session = nil qrs = nil kid15, err := key.HexKeyspaceId("15").Unhex() if err != nil { t.Errorf("want nil, got %+v", err) } sq.KeyspaceIds = []key.KeyspaceId{kid10, kid15} err = RpcVTGate.StreamExecuteKeyspaceIds(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err != nil { t.Errorf("want nil, got %v", err) } row = new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, row) want = []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } // Test for error condition - multiple shards kid30, err := key.HexKeyspaceId("30").Unhex() if err != nil { t.Errorf("want nil, got %+v", err) } sq.KeyspaceIds = []key.KeyspaceId{kid10, kid30} err = RpcVTGate.StreamExecuteKeyspaceIds(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got nil") } }