forked from dolthub/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
queryservice.go
118 lines (100 loc) · 4.3 KB
/
queryservice.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package queryservice contains the interface for the service definition
// of the Query Service.
package queryservice
import (
"github.com/youtube/vitess/go/sqltypes"
"golang.org/x/net/context"
querypb "github.com/youtube/vitess/go/vt/proto/query"
"github.com/youtube/vitess/go/vt/tabletserver/querytypes"
)
// QueryService is the interface implemented by the tablet's query service.
type QueryService interface {
// Transaction management
// Begin returns the transaction id to use for further operations
Begin(ctx context.Context, target *querypb.Target) (int64, error)
// Commit commits the current transaction
Commit(ctx context.Context, target *querypb.Target, transactionID int64) error
// Rollback aborts the current transaction
Rollback(ctx context.Context, target *querypb.Target, transactionID int64) error
// Query execution
Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]interface{}, transactionID int64) (*sqltypes.Result, error)
StreamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]interface{}, sendReply func(*sqltypes.Result) error) error
ExecuteBatch(ctx context.Context, target *querypb.Target, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64) ([]sqltypes.Result, error)
// SplitQuery is a map reduce helper function
// TODO(erez): Remove this and rename the following func to SplitQuery
// once we migrate to SplitQuery V2.
SplitQuery(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64) ([]querytypes.QuerySplit, error)
// SplitQueryV2 is a MapReduce helper function.
// This is version of SplitQuery supports multiple algorithms and multiple split columns.
// See the documentation of SplitQueryRequest in 'proto/vtgate.proto' for more information.
SplitQueryV2(
ctx context.Context,
target *querypb.Target,
sql string,
bindVariables map[string]interface{},
splitColumns []string,
splitCount int64,
numRowsPerQueryPart int64,
algorithm querypb.SplitQueryRequest_Algorithm,
) ([]querytypes.QuerySplit, error)
// StreamHealthRegister registers a listener for StreamHealth
StreamHealthRegister(chan<- *querypb.StreamHealthResponse) (int, error)
// StreamHealthUnregister unregisters a listener for StreamHealth
StreamHealthUnregister(int) error
// Helper for RPC panic handling: call this in a defer statement
// at the beginning of each RPC handling method.
HandlePanic(*error)
}
// CallCorrectSplitQuery calls the correct SplitQuery.
// This trivial logic is encapsulated in a function here so it can be easily tested.
// TODO(erez): Remove once the migration to SplitQueryV2 is done.
func CallCorrectSplitQuery(
queryService QueryService,
useSplitQueryV2 bool,
ctx context.Context,
target *querypb.Target,
sql string,
bindVariables map[string]interface{},
splitColumns []string,
splitCount int64,
numRowsPerQueryPart int64,
algorithm querypb.SplitQueryRequest_Algorithm,
) ([]querytypes.QuerySplit, error) {
if useSplitQueryV2 {
return queryService.SplitQueryV2(
ctx,
target,
sql,
bindVariables,
splitColumns,
splitCount,
numRowsPerQueryPart,
algorithm)
}
return queryService.SplitQuery(
ctx,
target,
sql,
bindVariables,
splitColumnsToSplitColumn(splitColumns),
splitCount)
}
// SplitColumnsToSplitColumn returns the first SplitColumn in the given slice or an empty
// string if the slice is empty.
//
// This method is used to get the traditional behavior when accessing the SplitColumn field in an
// older SplitQuery-V1 querypb.SplitQueryRequest represented in the newer SplitQuery-V2
// querypb.SplitQueryRequest message. In the new V2 message the SplitColumn field has been converted
// into a repeated string field.
// TODO(erez): Remove this function when migration to SplitQueryV2 is done.
func splitColumnsToSplitColumn(splitColumns []string) string {
if len(splitColumns) == 0 {
return ""
}
return splitColumns[0]
}
// Command to generate a mock for this interface with mockgen.
//go:generate mockgen -source $GOFILE -destination queryservice_testing/mock_queryservice.go -package queryservice_testing