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