// getQuerySplitToKeyRangePartFunc returns a function to use with scatterConn.SplitQueryV2
// that converts the given QuerySplit to a SplitQueryResponse_Part message whose KeyRangePart field
// is set.
func getQuerySplitToKeyRangePartFunc(
keyspace string,
shardReferenceByName map[string]*topodatapb.ShardReference) func(
querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
return func(
querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
// TODO(erez): Assert that shardReferenceByName contains an entry for 'shard'.
// Keyrange can be nil for the shard (e.g. for single-sharded keyspaces during resharding).
// In this case we append an empty keyrange that represents the entire keyspace.
keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
if shardReferenceByName[shard].KeyRange != nil {
keyranges = []*topodatapb.KeyRange{shardReferenceByName[shard].KeyRange}
}
bindVars, err := querytypes.BindVariablesToProto3(querySplit.BindVariables)
if err != nil {
return nil, err
}
return &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: querySplit.Sql,
BindVariables: bindVars,
},
KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
Keyspace: keyspace,
KeyRanges: keyranges,
},
Size: querySplit.RowCount,
}, nil
}
}
// StreamExecute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *sqltypes.Result, tabletconn.ErrFunc, error) {
bv, err := querytypes.BindVariablesToProto3(bindVars)
if err != nil {
return nil, nil, err
}
bindVars, err = querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, nil, err
}
result := make(chan *sqltypes.Result, 10)
var finalErr error
go func() {
finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, &querypb.Target{
Keyspace: itc.tablet.keyspace,
Shard: itc.tablet.shard,
TabletType: itc.tablet.tabletType,
}, query, bindVars, 0, func(reply *sqltypes.Result) error {
// We need to deep-copy the reply before returning,
// because the underlying buffers are reused.
result <- reply.Copy()
return nil
})
// the client will only access finalErr after the
// channel is closed, and then it's already set.
close(result)
}()
return result, func() error {
return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
}, nil
}
// StreamExecute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, options *querypb.ExecuteOptions) (sqltypes.ResultStream, error) {
bv, err := querytypes.BindVariablesToProto3(bindVars)
if err != nil {
return nil, err
}
bindVars, err = querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, err
}
result := make(chan *sqltypes.Result, 10)
var finalErr error
go func() {
finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, target, query, bindVars, options, func(reply *sqltypes.Result) error {
// We need to deep-copy the reply before returning,
// because the underlying buffers are reused.
result <- reply.Copy()
return nil
})
finalErr = tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(finalErr))
// the client will only access finalErr after the
// channel is closed, and then it's already set.
close(result)
}()
return &streamExecuteAdapter{result, &finalErr}, nil
}
// SplitQueryKeyRange scatters a SplitQuery request to all shards. For a set of
// splits received from a shard, it construct a KeyRange queries by
// appending that shard's keyrange to the splits. Aggregates all splits across
// all shards in no specific order and returns.
func (stc *ScatterConn) SplitQueryKeyRange(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, keyRangeByShard map[string]*topodatapb.KeyRange, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
tabletType := topodatapb.TabletType_RDONLY
actionFunc := func(shard string, transactionID int64, results chan<- interface{}) error {
// Get all splits from this shard
queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
if err != nil {
return err
}
// Append the keyrange for this shard to all the splits received,
// if keyrange is nil for the shard (e.g. for single-sharded keyspaces during resharding),
// append empty keyrange to represent the entire keyspace.
keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
if keyRangeByShard[shard] != nil {
keyranges = []*topodatapb.KeyRange{keyRangeByShard[shard]}
}
splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
for i, query := range queries {
q, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return err
}
splits[i] = &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: query.Sql,
BindVariables: q,
},
KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
Keyspace: keyspace,
KeyRanges: keyranges,
},
Size: query.RowCount,
}
}
// Push all the splits from this shard to results channel
results <- splits
return nil
}
shards := []string{}
for shard := range keyRangeByShard {
shards = append(shards, shard)
}
allSplits, allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
splits := []*vtgatepb.SplitQueryResponse_Part{}
for s := range allSplits {
splits = append(splits, s.([]*vtgatepb.SplitQueryResponse_Part)...)
}
if allErrors.HasErrors() {
err := allErrors.AggrError(stc.aggregateErrors)
return nil, err
}
return splits, nil
}
// SplitQueryCustomSharding scatters a SplitQuery request to all
// shards. For a set of splits received from a shard, it construct a
// KeyRange queries by appending that shard's name to the
// splits. Aggregates all splits across all shards in no specific
// order and returns.
func (stc *ScatterConn) SplitQueryCustomSharding(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, shards []string, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
tabletType := topodatapb.TabletType_RDONLY
// mu protects allSplits
var mu sync.Mutex
var allSplits []*vtgatepb.SplitQueryResponse_Part
actionFunc := func(target *querypb.Target) error {
// Get all splits from this shard
query := querytypes.BoundQuery{
Sql: sql,
BindVariables: bindVariables,
}
queries, err := stc.gateway.SplitQuery(ctx, target, query, splitColumn, splitCount)
if err != nil {
return err
}
// Use the shards list for all the splits received
shards := []string{target.Shard}
splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
for i, query := range queries {
q, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return err
}
splits[i] = &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: query.Sql,
BindVariables: q,
},
ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
Keyspace: keyspace,
Shards: shards,
},
Size: query.RowCount,
}
}
// aggregate splits
mu.Lock()
defer mu.Unlock()
allSplits = append(allSplits, splits...)
return nil
}
allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, actionFunc)
if allErrors.HasErrors() {
return nil, allErrors.AggrError(stc.aggregateErrors)
}
// See the comment for the analogues line in SplitQueryKeyRange for
// the motivation for shuffling.
shuffleQueryParts(allSplits)
return allSplits, nil
}
// Execute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, transactionID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
bv, err := querytypes.BindVariablesToProto3(bindVars)
if err != nil {
return nil, err
}
bindVars, err = querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, err
}
reply, err := itc.tablet.qsc.QueryService().Execute(ctx, target, query, bindVars, transactionID, options)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(err))
}
return reply, nil
}
// SplitQueryCustomSharding scatters a SplitQuery request to all
// shards. For a set of splits received from a shard, it construct a
// KeyRange queries by appending that shard's name to the
// splits. Aggregates all splits across all shards in no specific
// order and returns.
func (stc *ScatterConn) SplitQueryCustomSharding(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, shards []string, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
tabletType := topodatapb.TabletType_RDONLY
actionFunc := func(shard string, transactionID int64, results chan<- interface{}) error {
// Get all splits from this shard
queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
if err != nil {
return err
}
// Use the shards list for all the splits received
shards := []string{shard}
splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
for i, query := range queries {
q, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return err
}
splits[i] = &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: query.Sql,
BindVariables: q,
},
ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
Keyspace: keyspace,
Shards: shards,
},
Size: query.RowCount,
}
}
// Push all the splits from this shard to results channel
results <- splits
return nil
}
allSplits, allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
splits := []*vtgatepb.SplitQueryResponse_Part{}
for s := range allSplits {
splits = append(splits, s.([]*vtgatepb.SplitQueryResponse_Part)...)
}
if allErrors.HasErrors() {
err := allErrors.AggrError(stc.aggregateErrors)
return nil, err
}
return splits, nil
}
// Execute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*sqltypes.Result, error) {
bv, err := querytypes.BindVariablesToProto3(bindVars)
if err != nil {
return nil, err
}
bindVars, err = querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, err
}
reply, err := itc.tablet.qsc.QueryService().Execute(ctx, &querypb.Target{
Keyspace: itc.tablet.keyspace,
Shard: itc.tablet.shard,
TabletType: itc.tablet.tabletType,
}, query, bindVars, 0, transactionID)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
}
return reply, nil
}
func (c *echoClient) SplitQuery(ctx context.Context, keyspace string, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64) ([]*vtgatepb.SplitQueryResponse_Part, error) {
if strings.HasPrefix(sql, EchoPrefix) {
bv, err := querytypes.BindVariablesToProto3(bindVariables)
if err != nil {
return nil, err
}
return []*vtgatepb.SplitQueryResponse_Part{
{
Query: &querypb.BoundQuery{
Sql: fmt.Sprintf("%v:%v:%v", sql, splitColumn, splitCount),
BindVariables: bv,
},
KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
Keyspace: keyspace,
},
},
}, nil
}
return c.fallback.SplitQuery(ctx, sql, keyspace, bindVariables, splitColumn, splitCount)
}
// ExecuteBatch is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, target *querypb.Target, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64, options *querypb.ExecuteOptions) ([]sqltypes.Result, error) {
q := make([]querytypes.BoundQuery, len(queries))
for i, query := range queries {
bv, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return nil, err
}
bindVars, err := querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, err
}
q[i].Sql = query.Sql
q[i].BindVariables = bindVars
}
results, err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, target, q, asTransaction, transactionID, options)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(err))
}
return results, nil
}
// getQuerySplitToShardPartFunc returns a function to use with scatterConn.SplitQueryV2
// that converts the given QuerySplit to a SplitQueryResponse_Part message whose ShardPart field
// is set.
func getQuerySplitToShardPartFunc(keyspace string) func(
querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
return func(
querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
bindVars, err := querytypes.BindVariablesToProto3(querySplit.BindVariables)
if err != nil {
return nil, err
}
return &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: querySplit.Sql,
BindVariables: bindVars,
},
ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
Keyspace: keyspace,
Shards: []string{shard},
},
Size: querySplit.RowCount,
}, nil
}
}
// ExecuteBatch is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64) ([]sqltypes.Result, error) {
q := make([]querytypes.BoundQuery, len(queries))
for i, query := range queries {
bv, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return nil, err
}
bindVars, err := querytypes.Proto3ToBindVariables(bv)
if err != nil {
return nil, err
}
q[i].Sql = query.Sql
q[i].BindVariables = bindVars
}
results, err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, &querypb.Target{
Keyspace: itc.tablet.keyspace,
Shard: itc.tablet.shard,
TabletType: itc.tablet.tabletType,
}, q, 0, asTransaction, transactionID)
if err != nil {
return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
}
return results, nil
}
// SplitQueryKeyRange scatters a SplitQuery request to all shards. For a set of
// splits received from a shard, it construct a KeyRange queries by
// appending that shard's keyrange to the splits. Aggregates all splits across
// all shards in no specific order and returns.
func (stc *ScatterConn) SplitQueryKeyRange(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, keyRangeByShard map[string]*topodatapb.KeyRange, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
tabletType := topodatapb.TabletType_RDONLY
// mu protects allSplits
var mu sync.Mutex
var allSplits []*vtgatepb.SplitQueryResponse_Part
actionFunc := func(shard string, transactionID int64) error {
// Get all splits from this shard
queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
if err != nil {
return err
}
// Append the keyrange for this shard to all the splits received,
// if keyrange is nil for the shard (e.g. for single-sharded keyspaces during resharding),
// append empty keyrange to represent the entire keyspace.
keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
if keyRangeByShard[shard] != nil {
keyranges = []*topodatapb.KeyRange{keyRangeByShard[shard]}
}
splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
for i, query := range queries {
q, err := querytypes.BindVariablesToProto3(query.BindVariables)
if err != nil {
return err
}
splits[i] = &vtgatepb.SplitQueryResponse_Part{
Query: &querypb.BoundQuery{
Sql: query.Sql,
BindVariables: q,
},
KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
Keyspace: keyspace,
KeyRanges: keyranges,
},
Size: query.RowCount,
}
}
// aggregate splits
mu.Lock()
defer mu.Unlock()
allSplits = append(allSplits, splits...)
return nil
}
shards := []string{}
for shard := range keyRangeByShard {
shards = append(shards, shard)
}
allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
if allErrors.HasErrors() {
return nil, allErrors.AggrError(stc.aggregateErrors)
}
// We shuffle the query-parts here. External frameworks like MapReduce may
// "deal" these jobs to workers in the order they are in the list. Without
// shuffling workers can be very unevenly distributed among
// the shards they query. E.g. all workers will first query the first shard,
// then most of them to the second shard, etc, which results with uneven
// load balancing among shards.
shuffleQueryParts(allSplits)
return allSplits, nil
}
