func (conn *vtgateConn) ExecuteShards(ctx context.Context, query string, keyspace string, shards []string, bindVars map[string]interface{}, tabletType topodatapb.TabletType, session interface{}) (*sqltypes.Result, interface{}, error) {
var s *vtgatepb.Session
if session != nil {
s = session.(*vtgatepb.Session)
}
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, session, err
}
request := &vtgatepb.ExecuteShardsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Query: q,
Keyspace: keyspace,
Shards: shards,
TabletType: tabletType,
}
response, err := conn.c.ExecuteShards(ctx, request)
if err != nil {
return nil, session, vterrors.FromGRPCError(err)
}
if response.Error != nil {
return nil, response.Session, vterrors.FromVtRPCError(response.Error)
}
return sqltypes.Proto3ToResult(response.Result), response.Session, nil
}
func (conn *vtgateConn) StreamExecuteKeyspaceIds(ctx context.Context, query string, keyspace string, keyspaceIds [][]byte, bindVars map[string]interface{}, tabletType topodatapb.TabletType) (sqltypes.ResultStream, error) {
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, err
}
req := &vtgatepb.StreamExecuteKeyspaceIdsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Query: q,
Keyspace: keyspace,
KeyspaceIds: keyspaceIds,
TabletType: tabletType,
}
stream, err := conn.c.StreamExecuteKeyspaceIds(ctx, req)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return &streamExecuteAdapter{
recv: func() (*querypb.QueryResult, error) {
ser, err := stream.Recv()
if err != nil {
return nil, err
}
return ser.Result, nil
},
}, nil
}
func (conn *vtgateConn) SplitQueryV2(
ctx context.Context,
keyspace string,
query string,
bindVars map[string]interface{},
splitColumns []string,
splitCount int64,
numRowsPerQueryPart int64,
algorithm querypb.SplitQueryRequest_Algorithm) ([]*vtgatepb.SplitQueryResponse_Part, error) {
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, err
}
request := &vtgatepb.SplitQueryRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Keyspace: keyspace,
Query: q,
SplitColumn: splitColumns,
SplitCount: splitCount,
NumRowsPerQueryPart: numRowsPerQueryPart,
Algorithm: algorithm,
UseSplitQueryV2: true,
}
response, err := conn.c.SplitQuery(ctx, request)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Splits, nil
}
// Execute sends the query to VTTablet.
func (conn *gRPCQueryClient) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*sqltypes.Result, error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, tabletconn.ConnClosed
}
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, err
}
req := &querypb.ExecuteRequest{
Target: conn.target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
TransactionId: transactionID,
SessionId: conn.sessionID,
}
er, err := conn.c.Execute(ctx, req)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
return sqltypes.Proto3ToResult(er.Result), nil
}
// ExecuteBatch sends a batch query to VTTablet.
func (conn *gRPCQueryClient) ExecuteBatch(ctx context.Context, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64) ([]sqltypes.Result, error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, tabletconn.ConnClosed
}
req := &querypb.ExecuteBatchRequest{
Target: conn.target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Queries: make([]*querypb.BoundQuery, len(queries)),
AsTransaction: asTransaction,
TransactionId: transactionID,
SessionId: conn.sessionID,
}
for i, q := range queries {
qq, err := querytypes.BoundQueryToProto3(q.Sql, q.BindVariables)
if err != nil {
return nil, err
}
req.Queries[i] = qq
}
ebr, err := conn.c.ExecuteBatch(ctx, req)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
return sqltypes.Proto3ToResults(ebr.Results), nil
}
// SplitQuery is the stub for TabletServer.SplitQuery RPC
func (conn *gRPCQueryClient) SplitQuery(ctx context.Context, query querytypes.BoundQuery, splitColumn string, splitCount int64) (queries []querytypes.QuerySplit, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
err = tabletconn.ConnClosed
return
}
q, err := querytypes.BoundQueryToProto3(query.Sql, query.BindVariables)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
req := &querypb.SplitQueryRequest{
Target: conn.target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
SplitColumn: splitColumn,
SplitCount: splitCount,
SessionId: conn.sessionID,
}
sqr, err := conn.c.SplitQuery(ctx, req)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
split, err := querytypes.Proto3ToQuerySplits(sqr.Queries)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
return split, nil
}
// SplitQuery is the stub for TabletServer.SplitQuery RPC
// TODO(erez): Remove this method and rename SplitQueryV2 to SplitQuery once
// the migration to SplitQuery V2 is done.
func (conn *gRPCQueryClient) SplitQuery(ctx context.Context, target *querypb.Target, query querytypes.BoundQuery, splitColumn string, splitCount int64) (queries []querytypes.QuerySplit, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
err = tabletconn.ConnClosed
return
}
q, err := querytypes.BoundQueryToProto3(query.Sql, query.BindVariables)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
req := &querypb.SplitQueryRequest{
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
SplitColumn: []string{splitColumn},
SplitCount: splitCount,
NumRowsPerQueryPart: 0,
Algorithm: querypb.SplitQueryRequest_EQUAL_SPLITS,
UseSplitQueryV2: false,
}
sqr, err := conn.c.SplitQuery(ctx, req)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
split, err := querytypes.Proto3ToQuerySplits(sqr.Queries)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
return split, nil
}
// BeginExecute starts a transaction and runs an Execute.
func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, options *querypb.ExecuteOptions) (result *sqltypes.Result, transactionID int64, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, 0, tabletconn.ConnClosed
}
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, 0, err
}
if *combo {
// If combo is enabled, we combine both calls
req := &querypb.BeginExecuteRequest{
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
Options: options,
}
reply, err := conn.c.BeginExecute(ctx, req)
if err != nil {
return nil, 0, tabletconn.TabletErrorFromGRPC(err)
}
if reply.Error != nil {
return nil, reply.TransactionId, tabletconn.TabletErrorFromRPCError(reply.Error)
}
return sqltypes.Proto3ToResult(reply.Result), reply.TransactionId, nil
}
// Begin part.
breq := &querypb.BeginRequest{
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
}
br, err := conn.c.Begin(ctx, breq)
if err != nil {
return nil, 0, tabletconn.TabletErrorFromGRPC(err)
}
transactionID = br.TransactionId
// Execute part.
ereq := &querypb.ExecuteRequest{
Target: target,
EffectiveCallerId: breq.EffectiveCallerId,
ImmediateCallerId: breq.ImmediateCallerId,
Query: q,
TransactionId: transactionID,
Options: options,
}
er, err := conn.c.Execute(ctx, ereq)
if err != nil {
return nil, transactionID, tabletconn.TabletErrorFromGRPC(err)
}
return sqltypes.Proto3ToResult(er.Result), transactionID, nil
}
// StreamExecute starts a streaming query to VTTablet.
func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *sqltypes.Result, tabletconn.ErrFunc, error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, nil, tabletconn.ConnClosed
}
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, nil, err
}
req := &querypb.StreamExecuteRequest{
Target: conn.target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
SessionId: conn.sessionID,
}
stream, err := conn.c.StreamExecute(ctx, req)
if err != nil {
return nil, nil, tabletconn.TabletErrorFromGRPC(err)
}
sr := make(chan *sqltypes.Result, 10)
var finalError error
go func() {
var fields []*querypb.Field
for {
ser, err := stream.Recv()
if err != nil {
if err != io.EOF {
finalError = tabletconn.TabletErrorFromGRPC(err)
}
close(sr)
return
}
if fields == nil {
fields = ser.Result.Fields
}
sr <- sqltypes.CustomProto3ToResult(fields, ser.Result)
}
}()
return sr, func() error {
return finalError
}, nil
}
// BoundKeyspaceIdQueriesToProto transforms a list of BoundKeyspaceIdQuery to proto3
func BoundKeyspaceIdQueriesToProto(bsq []BoundKeyspaceIdQuery) ([]*vtgatepb.BoundKeyspaceIdQuery, error) {
if len(bsq) == 0 {
return nil, nil
}
result := make([]*vtgatepb.BoundKeyspaceIdQuery, len(bsq))
for i, q := range bsq {
qq, err := querytypes.BoundQueryToProto3(q.Sql, q.BindVariables)
if err != nil {
return nil, err
}
result[i] = &vtgatepb.BoundKeyspaceIdQuery{
Query: qq,
Keyspace: q.Keyspace,
KeyspaceIds: q.KeyspaceIds,
}
}
return result, nil
}
func testEchoSplitQuery(t *testing.T, conn *vtgateconn.VTGateConn) {
q, err := querytypes.BoundQueryToProto3(echoPrefix+query+":split_column:123", bindVars)
if err != nil {
t.Fatalf("BoundQueryToProto3 error: %v", err)
}
want := &vtgatepb.SplitQueryResponse_Part{
Query: q,
KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{Keyspace: keyspace},
}
got, err := conn.SplitQuery(context.Background(), keyspace, echoPrefix+query, bindVars, "split_column", 123)
if err != nil {
t.Fatalf("SplitQuery error: %v", err)
}
// For some reason, proto.Equal() is calling them unequal even though no diffs
// are found.
gotstr, wantstr := got[0].String(), want.String()
if gotstr != wantstr {
t.Errorf("SplitQuery() = %v, want %v", gotstr, wantstr)
}
}
func (conn *vtgateConn) StreamExecuteKeyspaceIds(ctx context.Context, query string, keyspace string, keyspaceIds [][]byte, bindVars map[string]interface{}, tabletType topodatapb.TabletType) (<-chan *sqltypes.Result, vtgateconn.ErrFunc, error) {
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, nil, err
}
req := &vtgatepb.StreamExecuteKeyspaceIdsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Query: q,
Keyspace: keyspace,
KeyspaceIds: keyspaceIds,
TabletType: tabletType,
}
stream, err := conn.c.StreamExecuteKeyspaceIds(ctx, req)
if err != nil {
return nil, nil, vterrors.FromGRPCError(err)
}
sr := make(chan *sqltypes.Result, 10)
var finalError error
go func() {
var fields []*querypb.Field
for {
ser, err := stream.Recv()
if err != nil {
if err != io.EOF {
finalError = vterrors.FromGRPCError(err)
}
close(sr)
return
}
if fields == nil {
fields = ser.Result.Fields
}
sr <- sqltypes.CustomProto3ToResult(fields, ser.Result)
}
}()
return sr, func() error {
return finalError
}, nil
}
// StreamExecute starts a streaming query to VTTablet.
func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}) (sqltypes.ResultStream, error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, tabletconn.ConnClosed
}
q, err := querytypes.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, err
}
req := &querypb.StreamExecuteRequest{
Target: conn.target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Query: q,
}
stream, err := conn.c.StreamExecute(ctx, req)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(err)
}
return &streamExecuteAdapter{stream: stream}, err
}
// BeginExecuteBatch starts a transaction and runs an ExecuteBatch.
func (conn *gRPCQueryClient) BeginExecuteBatch(ctx context.Context, target *querypb.Target, queries []querytypes.BoundQuery, asTransaction bool, options *querypb.ExecuteOptions) (results []sqltypes.Result, transactionID int64, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
return nil, 0, tabletconn.ConnClosed
}
if *combo {
// If combo is enabled, we combine both calls
req := &querypb.BeginExecuteBatchRequest{
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
Queries: make([]*querypb.BoundQuery, len(queries)),
AsTransaction: asTransaction,
Options: options,
}
for i, q := range queries {
qq, err := querytypes.BoundQueryToProto3(q.Sql, q.BindVariables)
if err != nil {
return nil, transactionID, err
}
req.Queries[i] = qq
}
reply, err := conn.c.BeginExecuteBatch(ctx, req)
if err != nil {
return nil, 0, tabletconn.TabletErrorFromGRPC(err)
}
if reply.Error != nil {
return nil, reply.TransactionId, tabletconn.TabletErrorFromRPCError(reply.Error)
}
return sqltypes.Proto3ToResults(reply.Results), reply.TransactionId, nil
}
breq := &querypb.BeginRequest{
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
}
br, err := conn.c.Begin(ctx, breq)
if err != nil {
return nil, 0, tabletconn.TabletErrorFromGRPC(err)
}
transactionID = br.TransactionId
ereq := &querypb.ExecuteBatchRequest{
Target: target,
EffectiveCallerId: breq.EffectiveCallerId,
ImmediateCallerId: breq.ImmediateCallerId,
Queries: make([]*querypb.BoundQuery, len(queries)),
AsTransaction: asTransaction,
TransactionId: transactionID,
Options: options,
}
for i, q := range queries {
qq, err := querytypes.BoundQueryToProto3(q.Sql, q.BindVariables)
if err != nil {
return nil, transactionID, err
}
ereq.Queries[i] = qq
}
ebr, err := conn.c.ExecuteBatch(ctx, ereq)
if err != nil {
return nil, transactionID, tabletconn.TabletErrorFromGRPC(err)
}
return sqltypes.Proto3ToResults(ebr.Results), transactionID, nil
}