// ExecuteEntityIds executes a non-streaming query based on given KeyspaceId map.
// It retries query if new keyspace/shards are re-resolved after a retryable error.
func (res *Resolver) ExecuteEntityIds(
context context.Context,
query *proto.EntityIdsQuery,
) (*mproto.QueryResult, error) {
newKeyspace, shardIDMap, err := mapEntityIdsToShards(
res.scatterConn.toposerv,
res.scatterConn.cell,
query.Keyspace,
query.EntityKeyspaceIDs,
query.TabletType)
if err != nil {
return nil, err
}
query.Keyspace = newKeyspace
shards, sqls, bindVars := buildEntityIds(shardIDMap, query.Sql, query.EntityColumnName, query.BindVariables)
for {
qr, err := res.scatterConn.ExecuteEntityIds(
context,
shards,
sqls,
bindVars,
query.Keyspace,
query.TabletType,
NewSafeSession(query.Session))
if connError, ok := err.(*ShardConnError); ok && connError.Code == tabletconn.ERR_RETRY {
resharding := false
newKeyspace, newShardIDMap, err := mapEntityIdsToShards(
res.scatterConn.toposerv,
res.scatterConn.cell,
query.Keyspace,
query.EntityKeyspaceIDs,
query.TabletType)
if err != nil {
return nil, err
}
// check keyspace change for vertical resharding
if newKeyspace != query.Keyspace {
query.Keyspace = newKeyspace
resharding = true
}
// check shards change for horizontal resharding
newShards, newSqls, newBindVars := buildEntityIds(newShardIDMap, query.Sql, query.EntityColumnName, query.BindVariables)
if !StrsEquals(newShards, shards) {
shards = newShards
sqls = newSqls
bindVars = newBindVars
resharding = true
}
// retry if resharding happened
if resharding {
continue
}
}
if err != nil {
return nil, err
}
return qr, err
}
}
func TestVTGateExecuteEntityIds(t *testing.T) {
s := createSandbox("TestVTGateExecuteEntityIds")
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.EntityIdsQuery{
Sql: "query",
Keyspace: "TestVTGateExecuteEntityIds",
EntityColumnName: "kid",
EntityKeyspaceIdMap: map[string]key.KeyspaceId{"id1": kid10},
TabletType: topo.TYPE_MASTER,
}
// Test for successful execution
qr := new(proto.QueryResult)
err = RpcVTGate.ExecuteEntityIds(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.ExecuteEntityIds(nil, &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteEntityIds",
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.EntityKeyspaceIdMap["id2"] = kid30
RpcVTGate.ExecuteEntityIds(nil, &q, qr)
if qr.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.RowsAffected)
}
}
func TestVTGateExecuteEntityIds(t *testing.T) {
s := createSandbox("TestVTGateExecuteEntityIds")
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.EntityIdsQuery{
Sql: "query",
Keyspace: "TestVTGateExecuteEntityIds",
EntityColumnName: "kid",
EntityKeyspaceIDs: []proto.EntityId{
proto.EntityId{
ExternalID: "id1",
KeyspaceID: kid10,
},
},
TabletType: topo.TYPE_MASTER,
}
// Test for successful execution
qr := new(proto.QueryResult)
err = rpcVTGate.ExecuteEntityIds(context.Background(), &q, qr)
if err != nil {
t.Errorf("want nil, got %v", err)
}
wantqr := new(proto.QueryResult)
wantqr.Result = singleRowResult
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 execCount := sbc1.ExecCount.Get(); execCount != 1 {
t.Errorf("want 1, got %v\n", execCount)
}
// Test for successful execution in transaction
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.ExecuteEntityIds(context.Background(), &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteEntityIds",
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(context.Background(), q.Session)
if commitCount := sbc1.CommitCount.Get(); commitCount != 1 {
t.Errorf("want 1, got %d", commitCount)
}
// Test for multiple shards
kid30, err := key.HexKeyspaceId("30").Unhex()
if err != nil {
t.Errorf("want nil, got %+v", err)
}
q.EntityKeyspaceIDs = append(q.EntityKeyspaceIDs, proto.EntityId{ExternalID: "id2", KeyspaceID: kid30})
rpcVTGate.ExecuteEntityIds(context.Background(), &q, qr)
if qr.Result.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.Result.RowsAffected)
}
}