func testUpdateStream(t *testing.T, conn tabletconn.TabletConn, f *FakeQueryService) {
t.Log("testUpdateStream")
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
stream, err := conn.UpdateStream(ctx, TestTarget, UpdateStreamPosition, UpdateStreamTimestamp)
if err != nil {
t.Fatalf("UpdateStream failed: %v", err)
}
qr, err := stream.Recv()
if err != nil {
t.Fatalf("UpdateStream failed: cannot read result1: %v", err)
}
if !reflect.DeepEqual(*qr, UpdateStreamStreamEvent1) {
t.Errorf("Unexpected result1 from UpdateStream: got %v wanted %v", qr, UpdateStreamStreamEvent1)
}
qr, err = stream.Recv()
if err != nil {
t.Fatalf("UpdateStream failed: cannot read result2: %v", err)
}
if !reflect.DeepEqual(*qr, UpdateStreamStreamEvent2) {
t.Errorf("Unexpected result2 from UpdateStream: got %v wanted %v", qr, UpdateStreamStreamEvent2)
}
qr, err = stream.Recv()
if err != io.EOF {
t.Fatalf("UpdateStream errFunc failed: %v", err)
}
}
// SplitQuery is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) SplitQuery(ctx context.Context, request *vtgatepb.SplitQueryRequest) (response *vtgatepb.SplitQueryResponse, err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
bv, err := querytypes.Proto3ToBindVariables(request.Query.BindVariables)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
splits, vtgErr := vtgateservice.CallCorrectSplitQuery(
vtg.server,
request.UseSplitQueryV2,
ctx,
request.Keyspace,
string(request.Query.Sql),
bv,
request.SplitColumn,
request.SplitCount,
request.NumRowsPerQueryPart,
request.Algorithm)
if vtgErr != nil {
return nil, vterrors.ToGRPCError(vtgErr)
}
return &vtgatepb.SplitQueryResponse{
Splits: splits,
}, nil
}
func testUpdateStreamError(t *testing.T, conn tabletconn.TabletConn, f *FakeQueryService) {
t.Log("testUpdateStreamError")
f.HasError = true
testErrorHelper(t, f, "UpdateStream", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
stream, err := conn.UpdateStream(ctx, TestTarget, UpdateStreamPosition, UpdateStreamTimestamp)
if err != nil {
t.Fatalf("UpdateStream failed: %v", err)
}
qr, err := stream.Recv()
if err != nil {
t.Fatalf("UpdateStream failed: cannot read result1: %v", err)
}
if !reflect.DeepEqual(*qr, UpdateStreamStreamEvent1) {
t.Errorf("Unexpected result1 from UpdateStream: got %v wanted %v", qr, UpdateStreamStreamEvent1)
}
// signal to the server that the first result has been received
close(f.ErrorWait)
// After 1 result, we expect to get an error (no more results).
qr, err = stream.Recv()
if err == nil {
t.Fatalf("UpdateStream channel wasn't closed")
}
return err
})
f.HasError = false
}
// ExecuteBatch is part of the queryservice.QueryServer interface
func (q *query) ExecuteBatch(ctx context.Context, request *pb.ExecuteBatchRequest) (response *pb.ExecuteBatchResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
reply := new(proto.QueryResultList)
bql, err := proto.Proto3ToBoundQueryList(request.Queries)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
if err := q.server.ExecuteBatch(ctx, request.Target, &proto.QueryList{
Queries: bql,
SessionId: request.SessionId,
AsTransaction: request.AsTransaction,
TransactionId: request.TransactionId,
}, reply); err != nil {
return nil, tabletserver.ToGRPCError(err)
}
results, err := proto.QueryResultListToProto3(reply.List)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
return &pb.ExecuteBatchResponse{Results: results}, nil
}
// ExecuteEntityIds is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGateP3) ExecuteEntityIds(ctx context.Context, request *pb.ExecuteEntityIdsRequest, response *pb.ExecuteEntityIdsResponse) (err error) {
defer vtg.server.HandlePanic(&err)
ctx, cancel := context.WithDeadline(ctx, time.Now().Add(*rpcTimeout))
defer cancel()
ctx = callerid.NewContext(ctx,
request.CallerId,
callerid.NewImmediateCallerID("bsonp3 client"))
reply := new(proto.QueryResult)
executeErr := vtg.server.ExecuteEntityIds(ctx,
string(request.Query.Sql),
tproto.Proto3ToBindVariables(request.Query.BindVariables),
request.Keyspace,
request.EntityColumnName,
request.EntityKeyspaceIds,
request.TabletType,
proto.ProtoToSession(request.Session),
request.NotInTransaction,
reply)
if executeErr == nil {
response.Error = vtgate.RPCErrorToVtRPCError(reply.Err)
response.Result = mproto.QueryResultToProto3(reply.Result)
response.Session = proto.SessionToProto(reply.Session)
}
return vterrors.ToJSONError(executeErr)
}
// StreamExecute is part of the queryservice.QueryServer interface
func (q *query) StreamExecute(request *pb.StreamExecuteRequest, stream pbs.Query_StreamExecuteServer) (err error) {
defer q.server.HandlePanic(&err)
ctx := callerid.NewContext(callinfo.GRPCCallInfo(stream.Context()),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
bv, err := proto.Proto3ToBindVariables(request.Query.BindVariables)
if err != nil {
return tabletserver.ToGRPCError(err)
}
if err := q.server.StreamExecute(ctx, request.Target, &proto.Query{
Sql: request.Query.Sql,
BindVariables: bv,
SessionId: request.SessionId,
}, func(reply *mproto.QueryResult) error {
result, err := mproto.QueryResultToProto3(reply)
if err != nil {
return err
}
return stream.Send(&pb.StreamExecuteResponse{Result: result})
}); err != nil {
return tabletserver.ToGRPCError(err)
}
return nil
}
// StreamExecuteKeyRanges2 is the RPC version of
// vtgateservice.VTGateService method
func (vtg *VTGateP3) StreamExecuteKeyRanges2(ctx context.Context, request *pb.StreamExecuteKeyRangesRequest, sendReply func(interface{}) error) (err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(ctx,
request.CallerId,
callerid.NewImmediateCallerID("gorpc client"))
vtgErr := vtg.server.StreamExecuteKeyRanges(ctx,
string(request.Query.Sql),
tproto.Proto3ToBindVariables(request.Query.BindVariables),
request.Keyspace,
request.KeyRanges,
request.TabletType,
func(reply *proto.QueryResult) error {
return sendReply(&pb.StreamExecuteKeyRangesResponse{
Error: vtgate.VtGateErrorToVtRPCError(nil, reply.Error),
Result: mproto.QueryResultToProto3(reply.Result),
})
})
if vtgErr == nil {
return nil
}
if *vtgate.RPCErrorOnlyInReply {
// If there was an app error, send a QueryResult back with it.
// Sending back errors this way is not backwards compatible. If a (new) server sends an additional
// QueryResult with an error, and the (old) client doesn't know how to read it, it will cause
// problems where the client will get out of sync with the number of QueryResults sent.
// That's why this the error is only sent this way when the --rpc_errors_only_in_reply flag is set
// (signalling that all clients are able to handle new-style errors).
return sendReply(&pb.StreamExecuteKeyRangesResponse{
Error: vtgate.VtGateErrorToVtRPCError(vtgErr, ""),
})
}
return vtgErr
}
// ExecuteKeyspaceIds is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) ExecuteKeyspaceIds(ctx context.Context, request *pb.ExecuteKeyspaceIdsRequest) (response *pb.ExecuteKeyspaceIdsResponse, err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
reply := new(proto.QueryResult)
bv, err := tproto.Proto3ToBindVariables(request.Query.BindVariables)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
executeErr := vtg.server.ExecuteKeyspaceIds(ctx,
string(request.Query.Sql),
bv,
request.Keyspace,
request.KeyspaceIds,
request.TabletType,
request.Session,
request.NotInTransaction,
reply)
response = &pb.ExecuteKeyspaceIdsResponse{
Error: vtgate.RPCErrorToVtRPCError(reply.Err),
}
if executeErr != nil {
return nil, vterrors.ToGRPCError(executeErr)
}
response.Result = sqltypes.ResultToProto3(reply.Result)
response.Session = reply.Session
return response, nil
}
// SplitQuery is part of the queryservice.QueryServer interface
func (q *query) SplitQuery(ctx context.Context, request *pb.SplitQueryRequest) (response *pb.SplitQueryResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
reply := &proto.SplitQueryResult{}
bq, err := proto.Proto3ToBoundQuery(request.Query)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
if err := q.server.SplitQuery(ctx, request.Target, &proto.SplitQueryRequest{
Query: *bq,
SplitColumn: request.SplitColumn,
SplitCount: int(request.SplitCount),
SessionID: request.SessionId,
}, reply); err != nil {
return nil, tabletserver.ToGRPCError(err)
}
qs, err := proto.QuerySplitsToProto3(reply.Queries)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
return &pb.SplitQueryResponse{Queries: qs}, nil
}
// StreamExecuteKeyRanges2 is the RPC version of
// vtgateservice.VTGateService method
func (vtg *VTGate) StreamExecuteKeyRanges2(ctx context.Context, request *gorpcvtgatecommon.KeyRangeQuery, sendReply func(interface{}) error) (err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(ctx,
callerid.GoRPCEffectiveCallerID(request.CallerID),
callerid.NewImmediateCallerID("gorpc client"))
vtgErr := vtg.server.StreamExecuteKeyRanges(ctx,
request.Sql,
request.BindVariables,
request.Keyspace,
request.KeyRanges,
request.TabletType,
func(value *sqltypes.Result) error {
return sendReply(&gorpcvtgatecommon.QueryResult{
Result: value,
})
})
if vtgErr == nil {
return nil
}
// If there was an app error, send a QueryResult back with it.
return sendReply(&gorpcvtgatecommon.QueryResult{
Err: vterrors.RPCErrFromVtError(vtgErr),
})
}
func testStreamExecuteError(t *testing.T, conn tabletconn.TabletConn, fake *FakeQueryService) {
ctx := context.Background()
ctx = callerid.NewContext(ctx, testCallerID, testVTGateCallerID)
stream, errFunc, err := conn.StreamExecute(ctx, streamExecuteQuery, streamExecuteBindVars, streamExecuteTransactionID)
if err != nil {
t.Fatalf("StreamExecute failed: %v", err)
}
qr, ok := <-stream
if !ok {
t.Fatalf("StreamExecute failed: cannot read result1")
}
if len(qr.Rows) == 0 {
qr.Rows = nil
}
if !reflect.DeepEqual(*qr, streamExecuteQueryResult1) {
t.Errorf("Unexpected result1 from StreamExecute: got %v wanted %v", qr, streamExecuteQueryResult1)
}
// signal to the server that the first result has been received
close(fake.errorWait)
// After 1 result, we expect to get an error (no more results).
qr, ok = <-stream
if ok {
t.Fatalf("StreamExecute channel wasn't closed")
}
err = errFunc()
verifyError(t, err, "StreamExecute")
}
// Execute is part of the queryservice.QueryServer interface
func (q *query) Execute(ctx context.Context, request *pb.ExecuteRequest) (response *pb.ExecuteResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
reply := new(mproto.QueryResult)
bv, err := proto.Proto3ToBindVariables(request.Query.BindVariables)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
if err := q.server.Execute(ctx, request.Target, &proto.Query{
Sql: request.Query.Sql,
BindVariables: bv,
SessionId: request.SessionId,
TransactionId: request.TransactionId,
}, reply); err != nil {
return nil, tabletserver.ToGRPCError(err)
}
result, err := mproto.QueryResultToProto3(reply)
if err != nil {
return nil, tabletserver.ToGRPCError(err)
}
return &pb.ExecuteResponse{Result: result}, nil
}
// ExecuteBatchShards is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) ExecuteBatchShards(ctx context.Context, request *pb.ExecuteBatchShardsRequest) (response *pb.ExecuteBatchShardsResponse, err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
reply := new(proto.QueryResultList)
executeErr := vtg.server.ExecuteBatchShards(ctx,
proto.ProtoToBoundShardQueries(request.Queries),
request.TabletType,
request.AsTransaction,
proto.ProtoToSession(request.Session),
reply)
response = &pb.ExecuteBatchShardsResponse{
Error: vtgate.VtGateErrorToVtRPCError(executeErr, reply.Error),
}
if executeErr == nil {
response.Results = tproto.QueryResultListToProto3(reply.List)
response.Session = proto.SessionToProto(reply.Session)
return response, nil
}
if *vtgate.RPCErrorOnlyInReply {
return response, nil
}
return nil, executeErr
}
func testStreamExecute(t *testing.T, conn tabletconn.TabletConn, f *FakeQueryService) {
t.Log("testStreamExecute")
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
stream, err := conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars)
if err != nil {
t.Fatalf("StreamExecute failed: %v", err)
}
qr, err := stream.Recv()
if err != nil {
t.Fatalf("StreamExecute failed: cannot read result1: %v", err)
}
if len(qr.Rows) == 0 {
qr.Rows = nil
}
if !reflect.DeepEqual(*qr, StreamExecuteQueryResult1) {
t.Errorf("Unexpected result1 from StreamExecute: got %v wanted %v", qr, StreamExecuteQueryResult1)
}
qr, err = stream.Recv()
if err != nil {
t.Fatalf("StreamExecute failed: cannot read result2: %v", err)
}
if len(qr.Fields) == 0 {
qr.Fields = nil
}
if !reflect.DeepEqual(*qr, StreamExecuteQueryResult2) {
t.Errorf("Unexpected result2 from StreamExecute: got %v wanted %v", qr, StreamExecuteQueryResult2)
}
qr, err = stream.Recv()
if err != io.EOF {
t.Fatalf("StreamExecute errFunc failed: %v", err)
}
}
func testStreamExecuteError(t *testing.T, conn tabletconn.TabletConn, f *FakeQueryService) {
t.Log("testStreamExecuteError")
f.HasError = true
testErrorHelper(t, f, "StreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
stream, err := conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars)
if err != nil {
t.Fatalf("StreamExecute failed: %v", err)
}
qr, err := stream.Recv()
if err != nil {
t.Fatalf("StreamExecute failed: cannot read result1: %v", err)
}
if len(qr.Rows) == 0 {
qr.Rows = nil
}
if !reflect.DeepEqual(*qr, StreamExecuteQueryResult1) {
t.Errorf("Unexpected result1 from StreamExecute: got %v wanted %v", qr, StreamExecuteQueryResult1)
}
// signal to the server that the first result has been received
close(f.ErrorWait)
// After 1 result, we expect to get an error (no more results).
qr, err = stream.Recv()
if err == nil {
t.Fatalf("StreamExecute channel wasn't closed")
}
return err
})
f.HasError = false
}
// StreamExecute is part of the queryservice.QueryServer interface
func (q *query) StreamExecute(request *pb.StreamExecuteRequest, stream pbs.Query_StreamExecuteServer) (err error) {
defer q.server.HandlePanic(&err)
ctx := callerid.NewContext(callinfo.GRPCCallInfo(stream.Context()),
request.GetEffectiveCallerId(),
request.GetImmediateCallerId(),
)
seErr := q.server.StreamExecute(ctx, &proto.Query{
Sql: string(request.Query.Sql),
BindVariables: proto.Proto3ToBindVariables(request.Query.BindVariables),
SessionId: request.SessionId,
}, func(reply *mproto.QueryResult) error {
return stream.Send(&pb.StreamExecuteResponse{
Result: mproto.QueryResultToProto3(reply),
})
})
if seErr != nil {
response := &pb.StreamExecuteResponse{
Error: tabletserver.TabletErrorToRPCError(seErr),
}
if err := stream.Send(response); err != nil {
return err
}
}
return nil
}
// StreamExecuteKeyRanges2 is the RPC version of
// vtgateservice.VTGateService method
func (vtg *VTGate) StreamExecuteKeyRanges2(ctx context.Context, request *proto.KeyRangeQuery, sendReply func(interface{}) error) (err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(ctx,
callerid.GoRPCEffectiveCallerID(request.CallerID),
callerid.NewImmediateCallerID("gorpc client"))
vtgErr := vtg.server.StreamExecuteKeyRanges(ctx,
request.Sql,
request.BindVariables,
request.Keyspace,
key.KeyRangesToProto(request.KeyRanges),
topo.TabletTypeToProto(request.TabletType),
func(value *proto.QueryResult) error {
return sendReply(value)
})
if vtgErr == nil {
return nil
}
if *vtgate.RPCErrorOnlyInReply {
// If there was an app error, send a QueryResult back with it.
qr := new(proto.QueryResult)
vtgate.AddVtGateErrorToQueryResult(vtgErr, qr)
// Sending back errors this way is not backwards compatible. If a (new) server sends an additional
// QueryResult with an error, and the (old) client doesn't know how to read it, it will cause
// problems where the client will get out of sync with the number of QueryResults sent.
// That's why this the error is only sent this way when the --rpc_errors_only_in_reply flag is set
// (signalling that all clients are able to handle new-style errors).
return sendReply(qr)
}
return vtgErr
}
// SplitQuery is part of the queryservice.QueryServer interface
func (q *query) SplitQuery(ctx context.Context, request *querypb.SplitQueryRequest) (response *querypb.SplitQueryResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
bq, err := querytypes.Proto3ToBoundQuery(request.Query)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
splits := []querytypes.QuerySplit{}
splits, err = queryservice.CallCorrectSplitQuery(
q.server,
request.UseSplitQueryV2,
ctx,
request.Target,
bq.Sql,
bq.BindVariables,
request.SplitColumn,
request.SplitCount,
request.NumRowsPerQueryPart,
request.Algorithm)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
qs, err := querytypes.QuerySplitsToProto3(splits)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
return &querypb.SplitQueryResponse{Queries: qs}, nil
}
// BeginExecuteBatch is part of the queryservice.QueryServer interface
func (q *query) BeginExecuteBatch(ctx context.Context, request *querypb.BeginExecuteBatchRequest) (response *querypb.BeginExecuteBatchResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
bql, err := querytypes.Proto3ToBoundQueryList(request.Queries)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
results, transactionID, err := q.server.BeginExecuteBatch(ctx, request.Target, bql, request.AsTransaction, request.Options)
if err != nil {
// if we have a valid transactionID, return the error in-band
if transactionID != 0 {
return &querypb.BeginExecuteBatchResponse{
Error: vterrors.VtRPCErrorFromVtError(err),
TransactionId: transactionID,
}, nil
}
return nil, vterrors.ToGRPCError(err)
}
return &querypb.BeginExecuteBatchResponse{
Results: sqltypes.ResultsToProto3(results),
TransactionId: transactionID,
}, nil
}
// BeginExecute is part of the queryservice.QueryServer interface
func (q *query) BeginExecute(ctx context.Context, request *querypb.BeginExecuteRequest) (response *querypb.BeginExecuteResponse, err error) {
defer q.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.EffectiveCallerId,
request.ImmediateCallerId,
)
bv, err := querytypes.Proto3ToBindVariables(request.Query.BindVariables)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
result, transactionID, err := q.server.BeginExecute(ctx, request.Target, request.Query.Sql, bv)
if err != nil {
// if we have a valid transactionID, return the error in-band
if transactionID != 0 {
return &querypb.BeginExecuteResponse{
Error: vterrors.VtRPCErrorFromVtError(err),
TransactionId: transactionID,
}, nil
}
return nil, vterrors.ToGRPCError(err)
}
return &querypb.BeginExecuteResponse{
Result: sqltypes.ResultToProto3(result),
TransactionId: transactionID,
}, nil
}
// ExecuteKeyspaceIds is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGateP3) ExecuteKeyspaceIds(ctx context.Context, request *pb.ExecuteKeyspaceIdsRequest, response *pb.ExecuteKeyspaceIdsResponse) (err error) {
defer vtg.server.HandlePanic(&err)
ctx, cancel := context.WithDeadline(ctx, time.Now().Add(*rpcTimeout))
defer cancel()
ctx = callerid.NewContext(ctx,
request.CallerId,
callerid.NewImmediateCallerID("gorpc client"))
reply := &proto.QueryResult{}
vtgErr := vtg.server.ExecuteKeyspaceIds(ctx,
string(request.Query.Sql),
tproto.Proto3ToBindVariables(request.Query.BindVariables),
request.Keyspace,
request.KeyspaceIds,
request.TabletType,
proto.ProtoToSession(request.Session),
request.NotInTransaction,
reply)
response.Error = vtgate.VtGateErrorToVtRPCError(vtgErr, reply.Error)
response.Result = mproto.QueryResultToProto3(reply.Result)
response.Session = proto.SessionToProto(reply.Session)
if *vtgate.RPCErrorOnlyInReply {
return nil
}
return vtgErr
}
// StreamExecute2 should not be used by anything other than tests.
// It will eventually replace Rollback, but it breaks compatibility with older clients.
// Once all clients are upgraded, it can be replaced.
func (sq *SqlQuery) StreamExecute2(ctx context.Context, req *proto.StreamExecuteRequest, sendReply func(reply interface{}) error) (err error) {
defer sq.server.HandlePanic(&err)
if req == nil || req.Query == nil {
return nil
}
ctx = callerid.NewContext(ctx,
callerid.GoRPCEffectiveCallerID(req.EffectiveCallerID),
callerid.GoRPCImmediateCallerID(req.ImmediateCallerID),
)
tErr := sq.server.StreamExecute(callinfo.RPCWrapCallInfo(ctx), proto.TargetToProto3(req.Target), req.Query, func(reply *mproto.QueryResult) error {
return sendReply(reply)
})
if tErr == nil {
return nil
}
if *tabletserver.RPCErrorOnlyInReply {
// If there was an app error, send a QueryResult back with it.
qr := new(mproto.QueryResult)
tabletserver.AddTabletErrorToQueryResult(tErr, qr)
// Sending back errors this way is not backwards compatible. If a (new) server sends an additional
// QueryResult with an error, and the (old) client doesn't know how to read it, it will cause
// problems where the client will get out of sync with the number of QueryResults sent.
// That's why this the error is only sent this way when the --rpc_errors_only_in_reply flag is set
// (signalling that all clients are able to handle new-style errors).
return sendReply(qr)
}
return tErr
}
// ExecuteKeyspaceIds is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) ExecuteKeyspaceIds(ctx context.Context, request *pb.ExecuteKeyspaceIdsRequest) (response *pb.ExecuteKeyspaceIdsResponse, err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
reply := new(proto.QueryResult)
executeErr := vtg.server.ExecuteKeyspaceIds(ctx,
string(request.Query.Sql),
tproto.Proto3ToBindVariables(request.Query.BindVariables),
request.Keyspace,
key.ProtoToKeyspaceIds(request.KeyspaceIds),
request.TabletType,
proto.ProtoToSession(request.Session),
request.NotInTransaction,
reply)
response = &pb.ExecuteKeyspaceIdsResponse{
Error: vtgate.VtGateErrorToVtRPCError(executeErr, reply.Error),
}
if executeErr == nil {
response.Result = mproto.QueryResultToProto3(reply.Result)
response.Session = proto.SessionToProto(reply.Session)
return response, nil
}
if *vtgate.RPCErrorOnlyInReply {
return response, nil
}
return nil, executeErr
}
func testStreamExecute2(t *testing.T, conn tabletconn.TabletConn) {
ctx := context.Background()
ctx = callerid.NewContext(ctx, testCallerID, testVTGateCallerID)
stream, errFunc, err := conn.StreamExecute2(ctx, streamExecuteQuery, streamExecuteBindVars, streamExecuteTransactionID)
if err != nil {
t.Fatalf("StreamExecute2 failed: %v", err)
}
qr, ok := <-stream
if !ok {
t.Fatalf("StreamExecute2 failed: cannot read result1")
}
if len(qr.Rows) == 0 {
qr.Rows = nil
}
if !reflect.DeepEqual(*qr, streamExecuteQueryResult1) {
t.Errorf("Unexpected result1 from StreamExecute2: got %v wanted %v", qr, streamExecuteQueryResult1)
}
qr, ok = <-stream
if !ok {
t.Fatalf("StreamExecute2 failed: cannot read result2")
}
if len(qr.Fields) == 0 {
qr.Fields = nil
}
if !reflect.DeepEqual(*qr, streamExecuteQueryResult2) {
t.Errorf("Unexpected result2 from StreamExecute2: got %v wanted %v", qr, streamExecuteQueryResult2)
}
qr, ok = <-stream
if ok {
t.Fatalf("StreamExecute2 channel wasn't closed")
}
if err := errFunc(); err != nil {
t.Fatalf("StreamExecute2 errFunc failed: %v", err)
}
}
// ExecuteBatchKeyspaceIds is the RPC version of
// vtgateservice.VTGateService method
func (vtg *VTGate) ExecuteBatchKeyspaceIds(ctx context.Context, request *pb.ExecuteBatchKeyspaceIdsRequest) (response *pb.ExecuteBatchKeyspaceIdsResponse, err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(callinfo.GRPCCallInfo(ctx),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
reply := new(proto.QueryResultList)
bq, err := proto.ProtoToBoundKeyspaceIdQueries(request.Queries)
if err != nil {
return nil, vterrors.ToGRPCError(err)
}
executeErr := vtg.server.ExecuteBatchKeyspaceIds(ctx,
bq,
request.TabletType,
request.AsTransaction,
request.Session,
reply)
response = &pb.ExecuteBatchKeyspaceIdsResponse{
Error: vtgate.RPCErrorToVtRPCError(reply.Err),
}
if executeErr != nil {
return nil, vterrors.ToGRPCError(executeErr)
}
response.Results = tproto.QueryResultListToProto3(reply.List)
response.Session = reply.Session
return response, nil
}
func testRollback2(t *testing.T, conn tabletconn.TabletConn) {
ctx := context.Background()
ctx = callerid.NewContext(ctx, testCallerID, testVTGateCallerID)
err := conn.Rollback2(ctx, rollbackTransactionID)
if err != nil {
t.Fatalf("Rollback2 failed: %v", err)
}
}
func testCommit(t *testing.T, conn tabletconn.TabletConn) {
ctx := context.Background()
ctx = callerid.NewContext(ctx, testCallerID, testVTGateCallerID)
err := conn.Commit(ctx, commitTransactionID)
if err != nil {
t.Fatalf("Commit failed: %v", err)
}
}
func testRollback(t *testing.T, conn tabletconn.TabletConn, f *FakeQueryService) {
t.Log("testRollback")
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
err := conn.Rollback(ctx, TestTarget, RollbackTransactionID)
if err != nil {
t.Fatalf("Rollback failed: %v", err)
}
}
func TestQueryExecutorTableAclDryRun(t *testing.T) {
aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63())
tableacl.Register(aclName, &simpleacl.Factory{})
tableacl.SetDefaultACL(aclName)
db := setUpQueryExecutorTest()
query := "select * from test_table limit 1000"
want := &mproto.QueryResult{
Fields: getTestTableFields(),
RowsAffected: 0,
Rows: [][]sqltypes.Value{},
}
db.AddQuery(query, want)
db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
Fields: getTestTableFields(),
})
username := "******"
callerID := &querypb.VTGateCallerID{
Username: username,
}
ctx := callerid.NewContext(context.Background(), nil, callerID)
config := &tableaclpb.Config{
TableGroups: []*tableaclpb.TableGroupSpec{{
Name: "group02",
TableNamesOrPrefixes: []string{"test_table"},
Readers: []string{"u1"},
}},
}
if err := tableacl.InitFromProto(config); err != nil {
t.Fatalf("unable to load tableacl config, error: %v", err)
}
tableACLStatsKey := strings.Join([]string{
"test_table",
"group02",
planbuilder.PLAN_PASS_SELECT.String(),
username,
}, ".")
// enable Config.StrictTableAcl
sqlQuery := newTestSQLQuery(ctx, enableRowCache|enableSchemaOverrides|enableStrict|enableStrictTableAcl)
sqlQuery.qe.enableTableAclDryRun = true
qre := newTestQueryExecutor(ctx, sqlQuery, query, 0)
defer sqlQuery.disallowQueries()
checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
beforeCount := sqlQuery.qe.tableaclPseudoDenied.Counters.Counts()[tableACLStatsKey]
// query should fail because current user do not have read permissions
_, err := qre.Execute()
if err != nil {
t.Fatalf("qre.Execute() = %v, want: nil", err)
}
afterCount := sqlQuery.qe.tableaclPseudoDenied.Counters.Counts()[tableACLStatsKey]
if afterCount-beforeCount != 1 {
t.Fatalf("table acl pseudo denied count should increase by one. got: %d, want: %d", afterCount, beforeCount+1)
}
}
// StreamExecuteKeyspaceIds is the RPC version of
// vtgateservice.VTGateService method
func (vtg *VTGate) StreamExecuteKeyspaceIds(ctx context.Context, request *proto.KeyspaceIdQuery, sendReply func(interface{}) error) (err error) {
defer vtg.server.HandlePanic(&err)
ctx = callerid.NewContext(ctx,
callerid.GoRPCEffectiveCallerID(request.CallerID),
callerid.NewImmediateCallerID("gorpc client"))
return vtg.server.StreamExecuteKeyspaceIds(ctx, request, func(value *proto.QueryResult) error {
return sendReply(value)
})
}