// ExecuteKeyRanges executes a non-streaming query based on the specified keyranges.
func (vtg *VTGate) ExecuteKeyRanges(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, keyRanges []*pb.KeyRange, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
startTime := time.Now()
statsKey := []string{"ExecuteKeyRanges", keyspace, strings.ToLower(tabletType.String())}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
sql = sqlannotation.AddFilteredReplicationUnfriendlyIfDML(sql)
qr, err := vtg.resolver.ExecuteKeyRanges(ctx, sql, bindVariables, keyspace, keyRanges, tabletType, session, notInTransaction)
if err == nil {
reply.Result = qr
vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
} else {
query := map[string]interface{}{
"Sql": sql,
"BindVariables": bindVariables,
"Keyspace": keyspace,
"KeyRanges": keyRanges,
"TabletType": strings.ToLower(tabletType.String()),
"Session": session,
"NotInTransaction": notInTransaction,
}
reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteKeyRanges).Error()
reply.Err = vterrors.RPCErrFromVtError(err)
}
reply.Session = 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 context.Context, query *proto.KeyRangeQuery, sendReply func(*proto.QueryResult) error) (err error) {
defer handlePanic(&err)
startTime := time.Now()
statsKey := []string{"StreamExecuteKeyRanges", query.Keyspace, string(query.TabletType)}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return ErrTooManyInFlight
}
err = vtg.resolver.StreamExecuteKeyRanges(
context,
query,
func(mreply *mproto.QueryResult) error {
reply := new(proto.QueryResult)
reply.Result = mreply
// Note we don't populate reply.Session here,
// as it may change incrementaly as responses are sent.
return sendReply(reply)
})
if err != nil {
normalErrors.Add(statsKey, 1)
vtg.logStreamExecuteKeyRanges.Errorf("%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
}
func TestVTGateStreamExecuteShards(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteShards")
sbc := &sandboxConn{}
s.MapTestConn("0", sbc)
// Test for successful execution
var qrs []*proto.QueryResult
err := rpcVTGate.StreamExecuteShards(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteShards",
[]string{"0"},
topodatapb.TabletType_MASTER,
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)
row.Result = singleRowResult
want := []*proto.QueryResult{row}
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 context.Context, query *proto.EntityIdsQuery, reply *proto.QueryResult) (err error) {
defer handlePanic(&err)
startTime := time.Now()
statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return ErrTooManyInFlight
}
qr, err := vtg.resolver.ExecuteEntityIds(context, query)
if err == nil {
reply.Result = qr
} else {
reply.Error = err.Error()
if strings.Contains(reply.Error, errDupKey) {
infoErrors.Add("DupKey", 1)
} else {
normalErrors.Add(statsKey, 1)
vtg.logExecuteEntityIds.Errorf("%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(ctx context.Context, query *proto.QueryShard, reply *proto.QueryResult) error {
startTime := time.Now()
statsKey := []string{"ExecuteShard", query.Keyspace, string(query.TabletType)}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
qr, err := vtg.resolver.Execute(
ctx,
query.Sql,
query.BindVariables,
query.Keyspace,
query.TabletType,
query.Session,
func(keyspace string) (string, []string, error) {
return query.Keyspace, query.Shards, nil
},
query.NotInTransaction,
)
if err == nil {
reply.Result = qr
vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
} else {
reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteShard)
}
reply.Session = query.Session
return nil
}
// ExecuteShard executes a non-streaming query on the specified shards.
func (vtg *VTGate) ExecuteShard(context context.Context, 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 {
reply.Result = qr
} else {
reply.Error = err.Error()
if strings.Contains(reply.Error, errDupKey) {
vtg.infoErrors.Add("DupKey", 1)
} else {
vtg.errors.Add(statsKey, 1)
vtg.logExecuteShard.Errorf("%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(ctx context.Context, 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)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
var rowCount int64
err := vtg.resolver.StreamExecuteKeyRanges(
ctx,
query,
func(mreply *mproto.QueryResult) error {
reply := new(proto.QueryResult)
reply.Result = mreply
rowCount += int64(len(mreply.Rows))
vtg.rowsReturned.Add(statsKey, int64(len(mreply.Rows)))
// Note we don't populate reply.Session here,
// as it may change incrementaly as responses are sent.
return sendReply(reply)
})
if err != nil {
normalErrors.Add(statsKey, 1)
logError(err, query, vtg.logStreamExecuteKeyRanges)
}
// Now we can send the final Sessoin info.
if query.Session != nil {
sendReply(&proto.QueryResult{Session: query.Session})
}
return formatError(err)
}
// StreamExecuteShard executes a streaming query on the specified shards.
func (vtg *VTGate) StreamExecuteShard(context context.Context, 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)
reply.Result = mreply
// 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.logStreamExecuteShard.Errorf("%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
}
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, entityColumnName string, entityKeyspaceIDs []*pbg.ExecuteEntityIdsRequest_EntityId, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
startTime := time.Now()
statsKey := []string{"ExecuteEntityIds", keyspace, strings.ToLower(tabletType.String())}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
qr, err := vtg.resolver.ExecuteEntityIds(ctx, sql, bindVariables, keyspace, entityColumnName, entityKeyspaceIDs, tabletType, session, notInTransaction)
if err == nil {
reply.Result = qr
vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
} else {
query := map[string]interface{}{
"Sql": sql,
"BindVariables": bindVariables,
"Keyspace": keyspace,
"EntityColumnName": entityColumnName,
"EntityKeyspaceIDs": entityKeyspaceIDs,
"TabletType": strings.ToLower(tabletType.String()),
"Session": session,
"NotInTransaction": notInTransaction,
}
reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteEntityIds).Error()
reply.Err = rpcErrFromVtGateError(err)
}
reply.Session = session
return nil
}
func TestVTGateExecuteShard(t *testing.T) {
sandbox := createSandbox("TestVTGateExecuteShard")
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
q := proto.QueryShard{
Sql: "query",
Keyspace: "TestVTGateExecuteShard",
Shards: []string{"0"},
TabletType: topo.TYPE_REPLICA,
}
qr := new(proto.QueryResult)
err := rpcVTGate.ExecuteShard(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)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.ExecuteShard(context.Background(), &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteShard",
Shard: "0",
TabletType: topo.TYPE_REPLICA,
TransactionId: 1,
}},
}
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 := sbc.CommitCount.Get(); commitCount != 1 {
t.Errorf("want 1, got %d", commitCount)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
rpcVTGate.ExecuteShard(context.Background(), &q, qr)
rpcVTGate.Rollback(context.Background(), q.Session)
/*
// Flaky: This test should be run manually.
runtime.Gosched()
if sbc.RollbackCount != 1 {
t.Errorf("want 1, got %d", sbc.RollbackCount)
}
*/
}
// Execute executes a non-streaming query by routing based on the values in the query.
func (vtg *VTGate) Execute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
startTime := time.Now()
statsKey := []string{"Execute", "Any", strings.ToLower(tabletType.String())}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
qr, err := vtg.router.Execute(ctx, sql, bindVariables, tabletType, session, notInTransaction)
if err == nil {
reply.Result = qr
vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
} else {
query := map[string]interface{}{
"Sql": sql,
"BindVariables": bindVariables,
"TabletType": strings.ToLower(tabletType.String()),
"Session": session,
"NotInTransaction": notInTransaction,
}
reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecute)
}
reply.Session = 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(&context.DummyContext{}, &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)
row.Result = singleRowResult
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(&context.DummyContext{}, q.Session)
err = RpcVTGate.StreamExecuteShard(&context.DummyContext{}, &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)
}
}
func TestVTGateExecute(t *testing.T) {
sandbox := createSandbox(KsTestUnsharded)
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
q := proto.Query{
Sql: "select * from t1",
TabletType: topo.TYPE_MASTER,
}
qr := new(proto.QueryResult)
err := rpcVTGate.Execute(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)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.Execute(context.Background(), &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: KsTestUnsharded,
Shard: "0",
TabletType: topo.TYPE_MASTER,
TransactionId: 1,
}},
}
if !reflect.DeepEqual(wantSession, q.Session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
}
rpcVTGate.Commit(context.Background(), q.Session)
if sbc.CommitCount != 1 {
t.Errorf("want 1, got %d", sbc.CommitCount)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
rpcVTGate.Execute(context.Background(), &q, qr)
rpcVTGate.Rollback(context.Background(), q.Session)
}
func (c *echoClient) Execute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
if strings.HasPrefix(sql, EchoPrefix) {
reply.Result = echoQueryResult(map[string]interface{}{
"callerId": callerid.EffectiveCallerIDFromContext(ctx),
"query": sql,
"bindVars": bindVariables,
"tabletType": tabletType,
"session": session,
"notInTransaction": notInTransaction,
})
reply.Session = session
return nil
}
return c.fallbackClient.Execute(ctx, sql, bindVariables, tabletType, session, notInTransaction, reply)
}
func TestVTGateStreamExecuteKeyRanges(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteKeyRanges")
sbc := &sandboxConn{}
s.MapTestConn("-20", sbc)
sbc1 := &sandboxConn{}
s.MapTestConn("20-40", sbc1)
kr, err := key.ParseKeyRangeParts("", "20")
// Test for successful execution
var qrs []*proto.QueryResult
err = rpcVTGate.StreamExecuteKeyRanges(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteKeyRanges",
[]key.KeyRange{kr},
pb.TabletType_MASTER,
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)
row.Result = singleRowResult
want := []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
// Test for successful execution - multiple shards
kr, err = key.ParseKeyRangeParts("10", "40")
err = rpcVTGate.StreamExecuteKeyRanges(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteKeyRanges",
[]key.KeyRange{kr},
pb.TabletType_MASTER,
func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
if err != nil {
t.Errorf("want nil, got %v", err)
}
}
func (c *echoClient) ExecuteEntityIds(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, entityColumnName string, entityKeyspaceIDs []*pbg.ExecuteEntityIdsRequest_EntityId, tabletType pb.TabletType, session *pbg.Session, notInTransaction bool, reply *proto.QueryResult) error {
if strings.HasPrefix(sql, EchoPrefix) {
reply.Result = echoQueryResult(map[string]interface{}{
"callerId": callerid.EffectiveCallerIDFromContext(ctx),
"query": sql,
"bindVars": bindVariables,
"keyspace": keyspace,
"entityColumnName": entityColumnName,
"entityIds": entityKeyspaceIDs,
"tabletType": tabletType,
"session": session,
"notInTransaction": notInTransaction,
})
reply.Session = session
return nil
}
return c.fallbackClient.ExecuteEntityIds(ctx, sql, bindVariables, keyspace, entityColumnName, entityKeyspaceIDs, tabletType, session, notInTransaction, reply)
}
// StreamExecuteShards executes a streaming query on the specified shards.
func (vtg *VTGate) StreamExecuteShards(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, shards []string, tabletType pb.TabletType, sendReply func(*proto.QueryResult) error) error {
startTime := time.Now()
statsKey := []string{"StreamExecuteShards", keyspace, strings.ToLower(tabletType.String())}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
var rowCount int64
err := vtg.resolver.StreamExecute(
ctx,
sql,
bindVariables,
keyspace,
tabletType,
func(keyspace string) (string, []string, error) {
return keyspace, shards, nil
},
func(mreply *mproto.QueryResult) error {
reply := new(proto.QueryResult)
reply.Result = mreply
rowCount += int64(len(mreply.Rows))
vtg.rowsReturned.Add(statsKey, int64(len(mreply.Rows)))
// Note we don't populate reply.Session here,
// as it may change incrementaly as responses are sent.
return sendReply(reply)
})
if err != nil {
normalErrors.Add(statsKey, 1)
query := map[string]interface{}{
"Sql": sql,
"BindVariables": bindVariables,
"Keyspace": keyspace,
"Shards": shards,
"TabletType": strings.ToLower(tabletType.String()),
}
logError(err, query, vtg.logStreamExecuteShards)
}
return formatError(err)
}
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(context context.Context, 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 {
reply.Result = qr
} else {
reply.Error = err.Error()
if strings.Contains(reply.Error, errDupKey) {
vtg.infoErrors.Add("DupKey", 1)
} else {
vtg.errors.Add(statsKey, 1)
vtg.logExecuteEntityIds.Errorf("%v, query: %+v", err, query)
}
}
reply.Session = query.Session
return nil
}
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(ctx context.Context, query *proto.EntityIdsQuery, reply *proto.QueryResult) error {
startTime := time.Now()
statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)}
defer vtg.timings.Record(statsKey, startTime)
x := vtg.inFlight.Add(1)
defer vtg.inFlight.Add(-1)
if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
return errTooManyInFlight
}
qr, err := vtg.resolver.ExecuteEntityIds(ctx, query)
if err == nil {
reply.Result = qr
vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
} else {
reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteEntityIds)
}
reply.Session = query.Session
return nil
}
func TestVTGateStreamExecute(t *testing.T) {
sandbox := createSandbox(KsTestUnsharded)
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
var qrs []*proto.QueryResult
err := rpcVTGate.StreamExecute(context.Background(),
"select * from t1",
nil,
topodatapb.TabletType_MASTER,
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)
row.Result = singleRowResult
want := []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
}
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(&context.DummyContext{}, &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 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(&context.DummyContext{}, q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
RpcVTGate.ExecuteKeyspaceIds(&context.DummyContext{}, &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(&context.DummyContext{}, 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(&context.DummyContext{}, &q, qr)
if qr.Result.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.Result.RowsAffected)
}
}
func TestVTGateExecuteKeyRanges(t *testing.T) {
s := createSandbox("TestVTGateExecuteKeyRanges")
sbc1 := &sandboxConn{}
sbc2 := &sandboxConn{}
s.MapTestConn("-20", sbc1)
s.MapTestConn("20-40", sbc2)
// Test for successful execution
qr := new(proto.QueryResult)
err := rpcVTGate.ExecuteKeyRanges(context.Background(),
"query",
nil,
"TestVTGateExecuteKeyRanges",
[]*pb.KeyRange{&pb.KeyRange{End: []byte{0x20}}},
pb.TabletType_MASTER,
nil,
false,
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
session := new(proto.Session)
rpcVTGate.Begin(context.Background(), session)
if !session.InTransaction {
t.Errorf("want true, got false")
}
err = rpcVTGate.ExecuteKeyRanges(context.Background(),
"query",
nil,
"TestVTGateExecuteKeyRanges",
[]*pb.KeyRange{&pb.KeyRange{End: []byte{0x20}}},
pb.TabletType_MASTER,
session,
false,
qr)
if err != nil {
t.Errorf("want nil, got %v", err)
}
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteKeyRanges",
Shard: "-20",
TransactionId: 1,
TabletType: topo.TYPE_MASTER,
}},
}
if !reflect.DeepEqual(wantSession, session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, session)
}
rpcVTGate.Commit(context.Background(), session)
if commitCount := sbc1.CommitCount.Get(); commitCount != 1 {
t.Errorf("want 1, got %v", commitCount)
}
// Test for multiple shards
rpcVTGate.ExecuteKeyRanges(context.Background(), "query",
nil,
"TestVTGateExecuteKeyRanges",
[]*pb.KeyRange{&pb.KeyRange{Start: []byte{0x10}, End: []byte{0x30}}},
pb.TabletType_MASTER,
nil,
false, qr)
if qr.Result.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.Result.RowsAffected)
}
}
func TestVTGateExecuteKeyRanges(t *testing.T) {
s := createSandbox("TestVTGateExecuteKeyRanges")
sbc1 := &sandboxConn{}
sbc2 := &sandboxConn{}
s.MapTestConn("-20", sbc1)
s.MapTestConn("20-40", sbc2)
kr, err := key.ParseKeyRangeParts("", "20")
q := proto.KeyRangeQuery{
Sql: "query",
Keyspace: "TestVTGateExecuteKeyRanges",
KeyRanges: []key.KeyRange{kr},
TabletType: topo.TYPE_MASTER,
}
// Test for successful execution
qr := new(proto.QueryResult)
err = rpcVTGate.ExecuteKeyRanges(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 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(context.Background(), q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
err = rpcVTGate.ExecuteKeyRanges(context.Background(), &q, qr)
if err != nil {
t.Errorf("want nil, got %v", err)
}
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteKeyRanges",
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 sbc1.CommitCount.Get() != 1 {
t.Errorf("want 1, got %v", sbc1.CommitCount.Get())
}
// Test for multiple shards
kr, err = key.ParseKeyRangeParts("10", "30")
q.KeyRanges = []key.KeyRange{kr}
rpcVTGate.ExecuteKeyRanges(context.Background(), &q, qr)
if qr.Result.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.Result.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 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(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 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.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)
}
}
func TestVTGateStreamExecuteKeyspaceIds(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteKeyspaceIds")
sbc := &sandboxConn{}
s.MapTestConn("-20", sbc)
sbc1 := &sandboxConn{}
s.MapTestConn("20-40", sbc1)
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(context.Background(), &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)
row.Result = singleRowResult
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(context.Background(), sq.Session)
err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(), &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\ngot\n%#v", want, qrs)
}
rpcVTGate.Commit(context.Background(), 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(context.Background(), &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)
row.Result = singleRowResult
want = []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
// Test for successful execution - multiple keyspaceids in 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(context.Background(), &sq, func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
if err != nil {
t.Errorf("want nil, got %v", err)
}
}
func TestVTGateStreamExecuteKeyspaceIds(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteKeyspaceIds")
sbc := &sandboxConn{}
s.MapTestConn("-20", sbc)
sbc1 := &sandboxConn{}
s.MapTestConn("20-40", sbc1)
// Test for successful execution
var qrs []*proto.QueryResult
err := rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteKeyspaceIds",
[][]byte{[]byte{0x10}},
topodatapb.TabletType_MASTER,
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)
row.Result = singleRowResult
want := []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
// Test for successful execution - multiple keyspaceids in single shard
qrs = nil
err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteKeyspaceIds",
[][]byte{[]byte{0x10}, []byte{0x15}},
topodatapb.TabletType_MASTER,
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)
row.Result = singleRowResult
want = []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
// Test for successful execution - multiple keyspaceids in multiple shards
err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
"query",
nil,
"TestVTGateStreamExecuteKeyspaceIds",
[][]byte{[]byte{0x10}, []byte{0x30}},
topodatapb.TabletType_MASTER,
func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
if err != nil {
t.Errorf("want nil, got %v", err)
}
}
func TestVTGateExecute(t *testing.T) {
sandbox := createSandbox(KsTestUnsharded)
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
qr := new(proto.QueryResult)
err := rpcVTGate.Execute(context.Background(),
"select * from t1",
nil,
topodatapb.TabletType_MASTER,
nil,
false,
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)
}
session := new(pbg.Session)
rpcVTGate.Begin(context.Background(), session)
if !session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.Execute(context.Background(),
"select * from t1",
nil,
topodatapb.TabletType_MASTER,
session,
false,
qr)
wantSession := &pbg.Session{
InTransaction: true,
ShardSessions: []*pbg.Session_ShardSession{{
Target: &querypb.Target{
Keyspace: KsTestUnsharded,
Shard: "0",
TabletType: topodatapb.TabletType_MASTER,
},
TransactionId: 1,
}},
}
if !reflect.DeepEqual(wantSession, session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, session)
}
rpcVTGate.Commit(context.Background(), session)
if commitCount := sbc.CommitCount.Get(); commitCount != 1 {
t.Errorf("want 1, got %d", commitCount)
}
session = new(pbg.Session)
rpcVTGate.Begin(context.Background(), session)
rpcVTGate.Execute(context.Background(),
"select * from t1",
nil,
topodatapb.TabletType_MASTER,
session,
false,
qr)
rpcVTGate.Rollback(context.Background(), session)
}
func TestVTGateExecuteEntityIds(t *testing.T) {
s := createSandbox("TestVTGateExecuteEntityIds")
sbc1 := &sandboxConn{}
sbc2 := &sandboxConn{}
s.MapTestConn("-20", sbc1)
s.MapTestConn("20-40", sbc2)
// Test for successful execution
qr := new(proto.QueryResult)
err := rpcVTGate.ExecuteEntityIds(context.Background(),
"query",
nil,
"TestVTGateExecuteEntityIds",
"kid",
[]*pbg.ExecuteEntityIdsRequest_EntityId{
&pbg.ExecuteEntityIdsRequest_EntityId{
XidType: sqltypes.VarBinary,
XidValue: []byte("id1"),
KeyspaceId: []byte{0x10},
},
},
topodatapb.TabletType_MASTER,
nil,
false,
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
session := new(pbg.Session)
rpcVTGate.Begin(context.Background(), session)
if !session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.ExecuteEntityIds(context.Background(),
"query",
nil,
"TestVTGateExecuteEntityIds",
"kid",
[]*pbg.ExecuteEntityIdsRequest_EntityId{
&pbg.ExecuteEntityIdsRequest_EntityId{
XidType: sqltypes.VarBinary,
XidValue: []byte("id1"),
KeyspaceId: []byte{0x10},
},
},
topodatapb.TabletType_MASTER,
session,
false,
qr)
wantSession := &pbg.Session{
InTransaction: true,
ShardSessions: []*pbg.Session_ShardSession{{
Target: &querypb.Target{
Keyspace: "TestVTGateExecuteEntityIds",
Shard: "-20",
TabletType: topodatapb.TabletType_MASTER,
},
TransactionId: 1,
}},
}
if !reflect.DeepEqual(wantSession, session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, session)
}
rpcVTGate.Commit(context.Background(), session)
if commitCount := sbc1.CommitCount.Get(); commitCount != 1 {
t.Errorf("want 1, got %d", commitCount)
}
// Test for multiple shards
rpcVTGate.ExecuteEntityIds(context.Background(), "query",
nil,
"TestVTGateExecuteEntityIds",
"kid",
[]*pbg.ExecuteEntityIdsRequest_EntityId{
&pbg.ExecuteEntityIdsRequest_EntityId{
XidType: sqltypes.VarBinary,
XidValue: []byte("id1"),
KeyspaceId: []byte{0x10},
},
&pbg.ExecuteEntityIdsRequest_EntityId{
XidType: sqltypes.VarBinary,
XidValue: []byte("id2"),
KeyspaceId: []byte{0x30},
},
},
topodatapb.TabletType_MASTER,
nil,
false,
qr)
if qr.Result.RowsAffected != 2 {
t.Errorf("want 2, got %v", qr.Result.RowsAffected)
}
}
func TestVTGateStreamExecuteKeyRanges(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteKeyRanges")
sbc := &sandboxConn{}
s.MapTestConn("-20", sbc)
sbc1 := &sandboxConn{}
s.MapTestConn("20-40", sbc1)
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(context.Background(), &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)
row.Result = singleRowResult
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(context.Background(), sq.Session)
err = rpcVTGate.StreamExecuteKeyRanges(context.Background(), &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 successful execution - multiple shards
kr, err = key.ParseKeyRangeParts("10", "40")
sq.KeyRanges = []key.KeyRange{kr}
err = rpcVTGate.StreamExecuteKeyRanges(context.Background(), &sq, func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
if err != nil {
t.Errorf("want nil, got %v", err)
}
}