func (conn *vtgateConn) ExecuteEntityIds(ctx context.Context, query string, keyspace string, entityColumnName string, entityKeyspaceIDs []*pbg.ExecuteEntityIdsRequest_EntityId, bindVars map[string]interface{}, tabletType pb.TabletType, notInTransaction bool, session interface{}) (*mproto.QueryResult, interface{}, error) {
var s *proto.Session
if session != nil {
s = session.(*proto.Session)
}
request := proto.EntityIdsQuery{
CallerID: getEffectiveCallerID(ctx),
Sql: query,
BindVariables: bindVars,
Keyspace: keyspace,
EntityColumnName: entityColumnName,
EntityKeyspaceIDs: proto.ProtoToEntityIds(entityKeyspaceIDs),
TabletType: topo.ProtoToTabletType(tabletType),
Session: s,
NotInTransaction: notInTransaction,
}
var result proto.QueryResult
if err := conn.rpcConn.Call(ctx, "VTGate.ExecuteEntityIds", request, &result); err != nil {
return nil, session, err
}
if err := vterrors.FromRPCError(result.Err); err != nil {
return nil, result.Session, err
}
return result.Result, result.Session, nil
}
func getKeyspaceShards(ctx context.Context, topoServ SrvTopoServer, cell, keyspace string, tabletType pb.TabletType) (string, *topo.SrvKeyspace, []topo.ShardReference, error) {
srvKeyspace, err := topoServ.GetSrvKeyspace(ctx, cell, keyspace)
if err != nil {
return "", nil, nil, vterrors.NewVitessError(
vtrpc.ErrorCode_INTERNAL_ERROR, err,
"keyspace %v fetch error: %v", keyspace, err,
)
}
// check if the keyspace has been redirected for this tabletType.
tt := topo.ProtoToTabletType(tabletType)
if servedFrom, ok := srvKeyspace.ServedFrom[tt]; ok {
keyspace = servedFrom
srvKeyspace, err = topoServ.GetSrvKeyspace(ctx, cell, keyspace)
if err != nil {
return "", nil, nil, vterrors.NewVitessError(
vtrpc.ErrorCode_INTERNAL_ERROR, err,
"keyspace %v fetch error: %v", keyspace, err,
)
}
}
partition, ok := srvKeyspace.Partitions[tt]
if !ok {
return "", nil, nil, vterrors.NewVitessError(
vtrpc.ErrorCode_INTERNAL_ERROR, err,
"No partition found for tabletType %v in keyspace %v", strings.ToLower(tabletType.String()), keyspace,
)
}
return keyspace, srvKeyspace, partition.ShardReferences, nil
}
// StreamExecute is part of the VTGateService interface
func (f *fakeVTGateService) StreamExecute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, sendReply func(*proto.QueryResult) error) error {
execCase, ok := execMap[sql]
if !ok {
return fmt.Errorf("no match for: %s", sql)
}
query := &proto.Query{
Sql: sql,
BindVariables: bindVariables,
TabletType: topo.ProtoToTabletType(tabletType),
}
if !reflect.DeepEqual(query, execCase.execQuery) {
return fmt.Errorf("request mismatch: got %+v, want %+v", query, execCase.execQuery)
}
if execCase.reply.Result != nil {
result := proto.QueryResult{Result: &mproto.QueryResult{}}
result.Result.Fields = execCase.reply.Result.Fields
if err := sendReply(&result); err != nil {
return err
}
for _, row := range execCase.reply.Result.Rows {
result := proto.QueryResult{Result: &mproto.QueryResult{}}
result.Result.Rows = [][]sqltypes.Value{row}
if err := sendReply(&result); err != nil {
return err
}
}
}
if execCase.reply.Error != "" {
return errors.New(execCase.reply.Error)
}
return nil
}
func (s *server) ChangeType(ctx context.Context, request *pb.ChangeTypeRequest) (*pb.ChangeTypeResponse, error) {
ctx = callinfo.GRPCCallInfo(ctx)
response := &pb.ChangeTypeResponse{}
return response, s.agent.RPCWrapLockAction(ctx, actionnode.TabletActionChangeType, request, response, true, func() error {
return s.agent.ChangeType(ctx, topo.ProtoToTabletType(request.TabletType))
})
}
// 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"))
query := &proto.KeyspaceIdBatchQuery{
Session: proto.ProtoToSession(request.Session),
Queries: proto.ProtoToBoundKeyspaceIdQueries(request.Queries),
TabletType: topo.ProtoToTabletType(request.TabletType),
AsTransaction: request.AsTransaction,
}
reply := new(proto.QueryResultList)
executeErr := vtg.server.ExecuteBatchKeyspaceIds(ctx, query, reply)
response = &pb.ExecuteBatchKeyspaceIdsResponse{
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
}
// 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"))
query := &proto.KeyspaceIdQuery{
Sql: string(request.Query.Sql),
BindVariables: tproto.Proto3ToBindVariables(request.Query.BindVariables),
Keyspace: request.Keyspace,
KeyspaceIds: key.ProtoToKeyspaceIds(request.KeyspaceIds),
TabletType: topo.ProtoToTabletType(request.TabletType),
Session: proto.ProtoToSession(request.Session),
NotInTransaction: request.NotInTransaction,
}
reply := new(proto.QueryResult)
executeErr := vtg.server.ExecuteKeyspaceIds(ctx, query, 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
}
// ExecuteShards is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) ExecuteShards(ctx context.Context, request *pb.ExecuteShardsRequest) (response *pb.ExecuteShardsResponse, err error) {
defer vtg.server.HandlePanic(&err)
query := &proto.QueryShard{
Sql: string(request.Query.Sql),
BindVariables: tproto.Proto3ToBindVariables(request.Query.BindVariables),
Keyspace: request.Keyspace,
Shards: request.Shards,
TabletType: topo.ProtoToTabletType(request.TabletType),
Session: proto.ProtoToSession(request.Session),
NotInTransaction: request.NotInTransaction,
}
reply := new(proto.QueryResult)
executeErr := vtg.server.ExecuteShard(ctx, query, reply)
response = &pb.ExecuteShardsResponse{
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
}
// Execute please see vtgateconn.Impl.Execute
func (conn *FakeVTGateConn) Execute(ctx context.Context, sql string, bindVars map[string]interface{}, tabletType pb.TabletType, notInTransaction bool, session interface{}) (*mproto.QueryResult, interface{}, error) {
var s *proto.Session
if session != nil {
s = session.(*proto.Session)
}
query := &proto.Query{
Sql: sql,
BindVariables: bindVars,
TabletType: topo.ProtoToTabletType(tabletType),
Session: s,
NotInTransaction: notInTransaction,
}
response, ok := conn.execMap[query.Sql]
if !ok {
return nil, nil, fmt.Errorf("no match for: %s", query.Sql)
}
if !reflect.DeepEqual(query, response.execQuery) {
return nil, nil, fmt.Errorf(
"Execute: %+v, want %+v", query, response.execQuery)
}
var reply mproto.QueryResult
reply = *response.reply
if s != nil {
s = newSession(true, "test_keyspace", []string{}, pb.TabletType_MASTER)
}
return &reply, s, nil
}
func (conn *vtgateConn) ExecuteShards(ctx context.Context, query string, keyspace string, shards []string, bindVars map[string]interface{}, tabletType pb.TabletType, notInTransaction bool, session interface{}) (*mproto.QueryResult, interface{}, error) {
var s *proto.Session
if session != nil {
s = session.(*proto.Session)
}
request := proto.QueryShard{
CallerID: getEffectiveCallerID(ctx),
Sql: query,
BindVariables: bindVars,
Keyspace: keyspace,
Shards: shards,
TabletType: topo.ProtoToTabletType(tabletType),
Session: s,
NotInTransaction: notInTransaction,
}
var result proto.QueryResult
if err := conn.rpcConn.Call(ctx, "VTGate.ExecuteShard", request, &result); err != nil {
return nil, session, err
}
if result.Error != "" {
return nil, result.Session, errors.New(result.Error)
}
if err := vterrors.FromRPCError(result.Err); err != nil {
return nil, result.Session, err
}
return result.Result, result.Session, nil
}
func (s *server) RunHealthCheck(ctx context.Context, request *pb.RunHealthCheckRequest) (*pb.RunHealthCheckResponse, error) {
ctx = callinfo.GRPCCallInfo(ctx)
response := &pb.RunHealthCheckResponse{}
return response, s.agent.RPCWrap(ctx, actionnode.TabletActionRunHealthCheck, request, response, func() error {
s.agent.RunHealthCheck(ctx, topo.ProtoToTabletType(request.TabletType))
return nil
})
}
func (conn *vtgateConn) StreamExecute2(ctx context.Context, query string, bindVars map[string]interface{}, tabletType pb.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
req := &proto.Query{
CallerID: getEffectiveCallerID(ctx),
Sql: query,
BindVariables: bindVars,
TabletType: topo.ProtoToTabletType(tabletType),
Session: nil,
}
sr := make(chan *proto.QueryResult, 10)
c := conn.rpcConn.StreamGo("VTGate.StreamExecute2", req, sr)
return sendStreamResults(c, sr)
}
// Find returns the transactionId, if any, for a session
func (session *SafeSession) Find(keyspace, shard string, tabletType pb.TabletType) int64 {
if session == nil {
return 0
}
tt := topo.ProtoToTabletType(tabletType)
session.mu.Lock()
defer session.mu.Unlock()
for _, shardSession := range session.ShardSessions {
if keyspace == shardSession.Keyspace && tt == shardSession.TabletType && shard == shardSession.Shard {
return shardSession.TransactionId
}
}
return 0
}
// StreamExecute is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) StreamExecute(request *pb.StreamExecuteRequest, stream pbs.Vitess_StreamExecuteServer) (err error) {
defer vtg.server.HandlePanic(&err)
query := &proto.Query{
Sql: string(request.Query.Sql),
BindVariables: tproto.Proto3ToBindVariables(request.Query.BindVariables),
TabletType: topo.ProtoToTabletType(request.TabletType),
}
return vtg.server.StreamExecute(stream.Context(), query, func(value *proto.QueryResult) error {
return stream.Send(&pb.StreamExecuteResponse{
Result: mproto.QueryResultToProto3(value.Result),
})
})
}
// StreamExecute is the RPC version of vtgateservice.VTGateService method
func (vtg *VTGate) StreamExecute(request *pb.StreamExecuteRequest, stream pbs.Vitess_StreamExecuteServer) (err error) {
defer vtg.server.HandlePanic(&err)
ctx := callerid.NewContext(callinfo.GRPCCallInfo(stream.Context()),
request.CallerId,
callerid.NewImmediateCallerID("grpc client"))
query := &proto.Query{
Sql: string(request.Query.Sql),
BindVariables: tproto.Proto3ToBindVariables(request.Query.BindVariables),
TabletType: topo.ProtoToTabletType(request.TabletType),
}
return vtg.server.StreamExecute(ctx, query, func(value *proto.QueryResult) error {
return stream.Send(&pb.StreamExecuteResponse{
Result: mproto.QueryResultToProto3(value.Result),
})
})
}
// Execute is part of the VTGateService interface
func (f *fakeVTGateService) Execute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
execCase, ok := execMap[sql]
if !ok {
return fmt.Errorf("no match for: %s", sql)
}
query := &proto.Query{
Sql: sql,
BindVariables: bindVariables,
TabletType: topo.ProtoToTabletType(tabletType),
Session: session,
NotInTransaction: notInTransaction,
}
if !reflect.DeepEqual(query, execCase.execQuery) {
return fmt.Errorf("request mismatch: got %+v, want %+v", query, execCase.execQuery)
}
*reply = *execCase.reply
return nil
}
// ProtoToSession transforms a proto3 Session into native types
func ProtoToSession(s *pb.Session) *Session {
if s == nil {
return nil
}
result := &Session{
InTransaction: s.InTransaction,
}
result.ShardSessions = make([]*ShardSession, len(s.ShardSessions))
for i, ss := range s.ShardSessions {
result.ShardSessions[i] = &ShardSession{
Keyspace: ss.Target.Keyspace,
Shard: ss.Target.Shard,
TabletType: topo.ProtoToTabletType(ss.Target.TabletType),
TransactionId: ss.TransactionId,
}
}
return result
}
func newSession(
inTransaction bool,
keyspace string,
shards []string,
tabletType pb.TabletType) *proto.Session {
shardSessions := make([]*proto.ShardSession, len(shards))
for _, shard := range shards {
shardSessions = append(shardSessions, &proto.ShardSession{
Keyspace: keyspace,
Shard: shard,
TabletType: topo.ProtoToTabletType(tabletType),
TransactionId: rand.Int63(),
})
}
return &proto.Session{
InTransaction: inTransaction,
ShardSessions: shardSessions,
}
}
func (conn *vtgateConn) ExecuteBatchKeyspaceIds(ctx context.Context, queries []proto.BoundKeyspaceIdQuery, tabletType pb.TabletType, asTransaction bool, session interface{}) ([]mproto.QueryResult, interface{}, error) {
var s *proto.Session
if session != nil {
s = session.(*proto.Session)
}
request := proto.KeyspaceIdBatchQuery{
CallerID: getEffectiveCallerID(ctx),
Queries: queries,
TabletType: topo.ProtoToTabletType(tabletType),
AsTransaction: asTransaction,
Session: s,
}
var result proto.QueryResultList
if err := conn.rpcConn.Call(ctx, "VTGate.ExecuteBatchKeyspaceIds", request, &result); err != nil {
return nil, session, err
}
if err := vterrors.FromRPCError(result.Err); err != nil {
return nil, result.Session, err
}
return result.List, result.Session, nil
}
// RefreshTablesByShard calls RefreshState on all the tables of a
// given type in a shard. It would work for the master, but the
// discovery wouldn't be very efficient.
func (wr *Wrangler) RefreshTablesByShard(ctx context.Context, si *topo.ShardInfo, tabletType pb.TabletType, cells []string) error {
wr.Logger().Infof("RefreshTablesByShard called on shard %v/%v", si.Keyspace(), si.ShardName())
tabletMap, err := topo.GetTabletMapForShardByCell(ctx, wr.ts, si.Keyspace(), si.ShardName(), cells)
switch err {
case nil:
// keep going
case topo.ErrPartialResult:
wr.Logger().Warningf("RefreshTablesByShard: got partial result for shard %v/%v, may not refresh all tablets everywhere", si.Keyspace(), si.ShardName())
default:
return err
}
// ignore errors in this phase
wg := sync.WaitGroup{}
for _, ti := range tabletMap {
if ti.Type != topo.ProtoToTabletType(tabletType) {
continue
}
wg.Add(1)
go func(ti *topo.TabletInfo) {
wr.Logger().Infof("Calling RefreshState on tablet %v", ti.Alias)
// Setting an upper bound timeout to fail faster in case of an error.
// Using 60 seconds because RefreshState should not take more than 30 seconds.
// (RefreshState will restart the tablet's QueryService and most time will be spent on the shutdown, i.e. waiting up to 30 seconds on transactions (see Config.TransactionTimeout)).
ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
if err := wr.tmc.RefreshState(ctx, ti); err != nil {
wr.Logger().Warningf("RefreshTablesByShard: failed to refresh %v: %v", ti.Alias, err)
}
cancel()
wg.Done()
}(ti)
}
wg.Wait()
return 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)
query := &proto.BatchQueryShard{
Session: proto.ProtoToSession(request.Session),
Queries: proto.ProtoToBoundShardQueries(request.Queries),
TabletType: topo.ProtoToTabletType(request.TabletType),
AsTransaction: request.AsTransaction,
}
reply := new(proto.QueryResultList)
executeErr := vtg.server.ExecuteBatchShard(ctx, query, 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
}
// StreamExecute please see vtgateconn.Impl.StreamExecute
func (conn *FakeVTGateConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, tabletType pb.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
response, ok := conn.execMap[query]
if !ok {
return nil, nil, fmt.Errorf("no match for: %s", query)
}
queryProto := &proto.Query{
Sql: query,
BindVariables: bindVars,
TabletType: topo.ProtoToTabletType(tabletType),
Session: nil,
}
if !reflect.DeepEqual(queryProto, response.execQuery) {
return nil, nil, fmt.Errorf("StreamExecute: %+v, want %+v", query, response.execQuery)
}
if response.err != nil {
return nil, nil, response.err
}
var resultChan chan *mproto.QueryResult
defer close(resultChan)
if response.reply != nil {
// create a result channel big enough to buffer all of
// the responses so we don't need to fork a go routine.
resultChan := make(chan *mproto.QueryResult, len(response.reply.Rows)+1)
result := &mproto.QueryResult{}
result.Fields = response.reply.Fields
resultChan <- result
for _, row := range response.reply.Rows {
result := &mproto.QueryResult{}
result.Rows = [][]sqltypes.Value{row}
resultChan <- result
}
} else {
resultChan = make(chan *mproto.QueryResult)
}
return resultChan, func() error { return nil }, nil
}
func (stc *ScatterConn) updateSession(
ctx context.Context,
keyspace, shard string,
tabletType pb.TabletType,
session *SafeSession,
notInTransaction bool,
) (transactionID int64, err error) {
if !session.InTransaction() {
return 0, nil
}
// No need to protect ourselves from the race condition between
// Find and Append. The higher level functions ensure that no
// duplicate (keyspace, shard, tabletType) tuples can execute
// this at the same time.
transactionID = session.Find(keyspace, shard, tabletType)
if transactionID != 0 {
return transactionID, nil
}
// We are in a transaction at higher level,
// but client requires not to start a transaction for this query.
// If a transaction was started on this conn, we will use it (as above).
if notInTransaction {
return 0, nil
}
transactionID, err = stc.gateway.Begin(ctx, keyspace, shard, tabletType)
if err != nil {
return 0, err
}
session.Append(&proto.ShardSession{
Keyspace: keyspace,
TabletType: topo.ProtoToTabletType(tabletType),
Shard: shard,
TransactionId: transactionID,
})
return transactionID, nil
}
// DbServingGraph returns the ServingGraph for the given cell.
func DbServingGraph(ctx context.Context, ts topo.Server, cell string) (servingGraph *ServingGraph) {
servingGraph = &ServingGraph{
Cell: cell,
Keyspaces: make(map[string]*KeyspaceNodes),
}
rec := concurrency.AllErrorRecorder{}
keyspaces, err := ts.GetSrvKeyspaceNames(ctx, cell)
if err != nil {
servingGraph.Errors = append(servingGraph.Errors, fmt.Sprintf("GetSrvKeyspaceNames failed: %v", err))
return
}
wg := sync.WaitGroup{}
servingTypes := []topo.TabletType{topo.TYPE_MASTER, topo.TYPE_REPLICA, topo.TYPE_RDONLY}
for _, keyspace := range keyspaces {
kn := newKeyspaceNodes()
servingGraph.Keyspaces[keyspace] = kn
wg.Add(1)
go func(keyspace string, kn *KeyspaceNodes) {
defer wg.Done()
ks, err := ts.GetSrvKeyspace(ctx, cell, keyspace)
if err != nil {
rec.RecordError(fmt.Errorf("GetSrvKeyspace(%v, %v) failed: %v", cell, keyspace, err))
return
}
if len(ks.ServedFrom) > 0 {
kn.ServedFrom = make(map[topo.TabletType]string)
for _, sf := range ks.ServedFrom {
kn.ServedFrom[topo.ProtoToTabletType(sf.TabletType)] = sf.Keyspace
}
}
displayedShards := make(map[string]bool)
for _, partitionTabletType := range servingTypes {
kp := topoproto.SrvKeyspaceGetPartition(ks, topo.TabletTypeToProto(partitionTabletType))
if kp == nil {
continue
}
for _, srvShard := range kp.ShardReferences {
shard := srvShard.Name
if displayedShards[shard] {
continue
}
displayedShards[shard] = true
sn := &ShardNodes{
Name: shard,
}
kn.ShardNodes = append(kn.ShardNodes, sn)
wg.Add(1)
go func(shard string, sn *ShardNodes) {
defer wg.Done()
tabletTypes, err := ts.GetSrvTabletTypesPerShard(ctx, cell, keyspace, shard)
if err != nil {
rec.RecordError(fmt.Errorf("GetSrvTabletTypesPerShard(%v, %v, %v) failed: %v", cell, keyspace, shard, err))
return
}
for _, tabletType := range tabletTypes {
endPoints, _, err := ts.GetEndPoints(ctx, cell, keyspace, shard, tabletType)
if err != nil {
rec.RecordError(fmt.Errorf("GetEndPoints(%v, %v, %v, %v) failed: %v", cell, keyspace, shard, tabletType, err))
continue
}
for _, endPoint := range endPoints.Entries {
var tabletNode *TabletNodesByType
for _, t := range sn.TabletNodes {
if t.TabletType == tabletType {
tabletNode = t
break
}
}
if tabletNode == nil {
tabletNode = &TabletNodesByType{
TabletType: tabletType,
}
sn.TabletNodes = append(sn.TabletNodes, tabletNode)
}
tabletNode.Nodes = append(tabletNode.Nodes, newTabletNodeFromEndPoint(endPoint, cell))
}
}
}(shard, sn)
}
}
}(keyspace, kn)
}
wg.Wait()
servingGraph.Errors = rec.ErrorStrings()
return
}
