func (conn *vtgateConn) StreamExecuteKeyspaceIds(ctx context.Context, query string, keyspace string, keyspaceIds [][]byte, bindVars map[string]interface{}, tabletType pbt.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
req := &pb.StreamExecuteKeyspaceIdsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Query: tproto.BoundQueryToProto3(query, bindVars),
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 *mproto.QueryResult, 10)
var finalError error
go func() {
for {
ser, err := stream.Recv()
if err != nil {
if err != io.EOF {
finalError = vterrors.FromGRPCError(err)
}
close(sr)
return
}
sr <- mproto.Proto3ToQueryResult(ser.Result)
}
}()
return sr, func() error {
return finalError
}, nil
}
func (conn *vtgateConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, tabletType topodatapb.TabletType) (<-chan *sqltypes.Result, vtgateconn.ErrFunc, error) {
q, err := tproto.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, nil, err
}
req := &pb.StreamExecuteRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Query: q,
TabletType: tabletType,
}
stream, err := conn.c.StreamExecute(ctx, req)
if err != nil {
return nil, nil, vterrors.FromGRPCError(err)
}
sr := make(chan *sqltypes.Result, 10)
var finalError error
go func() {
for {
ser, err := stream.Recv()
if err != nil {
if err != io.EOF {
finalError = vterrors.FromGRPCError(err)
}
close(sr)
return
}
sr <- sqltypes.Proto3ToResult(ser.Result)
}
}()
return sr, func() error {
return finalError
}, 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) Execute(ctx context.Context, query string, bindVars map[string]interface{}, keyspace string, 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.ExecuteRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Query: q,
Keyspace: keyspace,
TabletType: tabletType,
}
response, err := conn.c.Execute(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) ExecuteBatchKeyspaceIds(ctx context.Context, queries []proto.BoundKeyspaceIdQuery, tabletType pbt.TabletType, asTransaction bool, session interface{}) ([]mproto.QueryResult, interface{}, error) {
var s *pb.Session
if session != nil {
s = session.(*pb.Session)
}
qs, err := proto.BoundKeyspaceIdQueriesToProto(queries)
if err != nil {
return nil, session, err
}
request := &pb.ExecuteBatchKeyspaceIdsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Queries: qs,
TabletType: tabletType,
AsTransaction: asTransaction,
}
response, err := conn.c.ExecuteBatchKeyspaceIds(ctx, request)
if err != nil {
return nil, session, vterrors.FromGRPCError(err)
}
if response.Error != nil {
return nil, response.Session, vterrors.FromVtRPCError(response.Error)
}
return mproto.Proto3ToQueryResults(response.Results), response.Session, nil
}
func (conn *vtgateConn) ExecuteEntityIds(ctx context.Context, query string, keyspace string, entityColumnName string, entityKeyspaceIDs []*pb.ExecuteEntityIdsRequest_EntityId, bindVars map[string]interface{}, tabletType topodatapb.TabletType, notInTransaction bool, session interface{}) (*sqltypes.Result, interface{}, error) {
var s *pb.Session
if session != nil {
s = session.(*pb.Session)
}
q, err := tproto.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, session, err
}
request := &pb.ExecuteEntityIdsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Query: q,
Keyspace: keyspace,
EntityColumnName: entityColumnName,
EntityKeyspaceIds: entityKeyspaceIDs,
TabletType: tabletType,
NotInTransaction: notInTransaction,
}
response, err := conn.c.ExecuteEntityIds(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) ExecuteShards(ctx context.Context, query string, keyspace string, shards []string, bindVars map[string]interface{}, tabletType pbt.TabletType, notInTransaction bool, session interface{}) (*mproto.QueryResult, interface{}, error) {
var s *pb.Session
if session != nil {
s = session.(*pb.Session)
}
q, err := tproto.BoundQueryToProto3(query, bindVars)
if err != nil {
return nil, session, err
}
request := &pb.ExecuteShardsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Query: q,
Keyspace: keyspace,
Shards: shards,
TabletType: tabletType,
NotInTransaction: notInTransaction,
}
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 mproto.Proto3ToQueryResult(response.Result), response.Session, nil
}
func (e *eventStreamAdapter) Recv() (*logutilpb.Event, error) {
le, err := e.stream.Recv()
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return le.Event, nil
}
func (conn *vtgateConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, tabletType pbt.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
req := &pb.StreamExecuteRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Query: tproto.BoundQueryToProto3(query, bindVars),
TabletType: tabletType,
}
stream, err := conn.c.StreamExecute(ctx, req)
if err != nil {
return nil, nil, vterrors.FromGRPCError(err)
}
sr := make(chan *mproto.QueryResult, 10)
// TODO(aaijazi): I'm pretty sure ther error handling going on down here is overkill now...
var finalError error
go func() {
for {
ser, err := stream.Recv()
if err != nil {
if err != io.EOF {
finalError = err
}
close(sr)
return
}
if ser.Error != nil {
finalError = vterrors.FromVtRPCError(ser.Error)
close(sr)
return
}
sr <- mproto.Proto3ToQueryResult(ser.Result)
}
}()
return sr, func() error {
return finalError
}, nil
}
// MaxRates is part of the throttlerclient.Client interface and returns the
// current max rate for each throttler of the process.
func (c *client) MaxRates(ctx context.Context) (map[string]int64, error) {
response, err := c.gRPCClient.MaxRates(ctx, &throttlerdata.MaxRatesRequest{})
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Rates, 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
}
func (conn *vtgateConn) Rollback(ctx context.Context, session interface{}) error {
request := &pb.RollbackRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: session.(*pb.Session),
}
_, err := conn.c.Rollback(ctx, request)
return vterrors.FromGRPCError(err)
}
// ResetConfiguration is part of the throttlerclient.Client interface.
func (c *client) ResetConfiguration(ctx context.Context, throttlerName string) ([]string, error) {
response, err := c.gRPCClient.ResetConfiguration(ctx, &throttlerdata.ResetConfigurationRequest{
ThrottlerName: throttlerName,
})
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Names, nil
}
func (conn *vtgateConn) GetSrvKeyspace(ctx context.Context, keyspace string) (*pbt.SrvKeyspace, error) {
request := &pb.GetSrvKeyspaceRequest{
Keyspace: keyspace,
}
response, err := conn.c.GetSrvKeyspace(ctx, request)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.SrvKeyspace, nil
}
func (a *updateStreamAdapter) Recv() (*querypb.StreamEvent, int64, error) {
r, err := a.stream.Recv()
if err != nil {
if err != io.EOF {
err = vterrors.FromGRPCError(err)
}
return nil, 0, err
}
return r.Event, r.ResumeTimestamp, nil
}
func (conn *vtgateConn) Begin(ctx context.Context) (interface{}, error) {
request := &pb.BeginRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
}
response, err := conn.c.Begin(ctx, request)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Session, nil
}
func (conn *vtgateConn) Commit(ctx context.Context, session interface{}) error {
request := &pb.CommitRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: session.(*pb.Session),
}
_, err := conn.c.Commit(ctx, request)
if err != nil {
return vterrors.FromGRPCError(err)
}
return nil
}
// SetMaxRate is part of the throttlerclient.Client interface and sets the rate
// on all throttlers of the server.
func (c *client) SetMaxRate(ctx context.Context, rate int64) ([]string, error) {
request := &throttlerdata.SetMaxRateRequest{
Rate: rate,
}
response, err := c.gRPCClient.SetMaxRate(ctx, request)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Names, nil
}
// UpdateConfiguration is part of the throttlerclient.Client interface.
func (c *client) UpdateConfiguration(ctx context.Context, throttlerName string, configuration *throttlerdata.Configuration, copyZeroValues bool) ([]string, error) {
response, err := c.gRPCClient.UpdateConfiguration(ctx, &throttlerdata.UpdateConfigurationRequest{
ThrottlerName: throttlerName,
Configuration: configuration,
CopyZeroValues: copyZeroValues,
})
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return response.Names, nil
}
// ExecuteVtworkerCommand is part of the VtworkerClient interface.
func (client *gRPCVtworkerClient) ExecuteVtworkerCommand(ctx context.Context, args []string) (logutil.EventStream, error) {
query := &vtworkerdatapb.ExecuteVtworkerCommandRequest{
Args: args,
}
stream, err := client.c.ExecuteVtworkerCommand(ctx, query)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return &eventStreamAdapter{stream}, nil
}
func (a *streamExecuteAdapter) Recv() (*sqltypes.Result, error) {
qr, err := a.recv()
if err != nil {
if err != io.EOF {
err = vterrors.FromGRPCError(err)
}
return nil, err
}
if a.fields == nil {
a.fields = qr.Fields
}
return sqltypes.CustomProto3ToResult(a.fields, qr), nil
}
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
}
func (conn *vtgateConn) SplitQuery(ctx context.Context, keyspace string, query string, bindVars map[string]interface{}, splitColumn string, splitCount int) ([]proto.SplitQueryPart, error) {
request := &pb.SplitQueryRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Keyspace: keyspace,
Query: tproto.BoundQueryToProto3(query, bindVars),
SplitColumn: splitColumn,
SplitCount: int64(splitCount),
}
response, err := conn.c.SplitQuery(ctx, request)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return proto.ProtoToSplitQueryParts(response), nil
}
func (conn *vtgateConn) UpdateStream(ctx context.Context, keyspace string, shard string, keyRange *topodatapb.KeyRange, tabletType topodatapb.TabletType, timestamp int64, event *querypb.EventToken) (vtgateconn.UpdateStreamReader, error) {
req := &vtgatepb.UpdateStreamRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Keyspace: keyspace,
Shard: shard,
KeyRange: keyRange,
TabletType: tabletType,
Timestamp: timestamp,
Event: event,
}
stream, err := conn.c.UpdateStream(ctx, req)
if err != nil {
return nil, vterrors.FromGRPCError(err)
}
return &updateStreamAdapter{
stream: stream,
}, nil
}
func (conn *vtgateConn) ExecuteBatchShards(ctx context.Context, queries []*vtgatepb.BoundShardQuery, tabletType topodatapb.TabletType, asTransaction bool, session interface{}) ([]sqltypes.Result, interface{}, error) {
var s *vtgatepb.Session
if session != nil {
s = session.(*vtgatepb.Session)
}
request := &vtgatepb.ExecuteBatchShardsRequest{
CallerId: callerid.EffectiveCallerIDFromContext(ctx),
Session: s,
Queries: queries,
TabletType: tabletType,
AsTransaction: asTransaction,
}
response, err := conn.c.ExecuteBatchShards(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.Proto3ToResults(response.Results), response.Session, nil
}