diff --git a/changelog/24.0/24.0.0/summary.md b/changelog/24.0/24.0.0/summary.md
index 303942748da..10d8da8b337 100644
--- a/changelog/24.0/24.0.0/summary.md
+++ b/changelog/24.0/24.0.0/summary.md
@@ -9,6 +9,7 @@
- **[Minor Changes](#minor-changes)**
- **[VTGate](#minor-changes-vtgate)**
- [New default for `--legacy-replication-lag-algorithm` flag](#vtgate-new-default-legacy-replication-lag-algorithm)
+ - [New "session" mode for `--vtgate-balancer-mode` flag](#vtgate-session-balancer-mode)
- **[Query Serving](#minor-changes-query-serving)**
- [JSON_EXTRACT now supports dynamic path arguments](#query-serving-json-extract-dynamic-args)
- **[VTTablet](#minor-changes-vttablet)**
@@ -42,6 +43,16 @@ Instead, a simpler algorithm purely based on low lag, high lag and minimum numbe
In v25 this flag will become deprecated and in the following release it will be removed. In the meantime, the legacy behaviour can be used by setting `--legacy-replication-lag-algorithm=true`. This deprecation is tracked in https://github.com/vitessio/vitess/issues/18914.
+#### New "session" mode for `--vtgate-balancer-mode` flag
+
+The VTGate flag `--vtgate-balancer-mode` now supports a new "session" mode in addition to the existing "cell", "prefer-cell", and "random" modes. Session mode routes each session consistently to the same tablet for the session's duration.
+
+To enable session mode, set the flag when starting VTGate:
+
+```
+--vtgate-balancer-mode=session
+```
+
### Query Serving
#### JSON_EXTRACT now supports dynamic path arguments
diff --git a/go/flags/endtoend/vtgate.txt b/go/flags/endtoend/vtgate.txt
index b1f4ef4b0b0..e6ae18ee32a 100644
--- a/go/flags/endtoend/vtgate.txt
+++ b/go/flags/endtoend/vtgate.txt
@@ -246,7 +246,7 @@ Flags:
-v, --version print binary version
--vmodule vModuleFlag comma-separated list of pattern=N settings for file-filtered logging
--vschema-ddl-authorized-users string List of users authorized to execute vschema ddl operations, or '%' to allow all users.
- --vtgate-balancer-mode string Tablet balancer mode (options: cell, prefer-cell, random). Defaults to 'cell' which shuffles tablets in the local cell.
+ --vtgate-balancer-mode string Tablet balancer mode (options: cell, prefer-cell, random, session). Defaults to 'cell' which shuffles tablets in the local cell.
--vtgate-config-terse-errors prevent bind vars from escaping in returned errors
--warming-reads-concurrency int Number of concurrent warming reads allowed (default 500)
--warming-reads-percent int Percentage of reads on the primary to forward to replicas. Useful for keeping buffer pools warm
diff --git a/go/test/endtoend/cluster/cluster_process.go b/go/test/endtoend/cluster/cluster_process.go
index 3d95b39915e..5f21c749ec0 100644
--- a/go/test/endtoend/cluster/cluster_process.go
+++ b/go/test/endtoend/cluster/cluster_process.go
@@ -45,12 +45,14 @@ import (
"vitess.io/vitess/go/vt/grpcclient"
"vitess.io/vitess/go/vt/log"
"vitess.io/vitess/go/vt/vterrors"
+ "vitess.io/vitess/go/vt/vtgate/executorcontext"
"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
"vitess.io/vitess/go/vt/vtgate/vtgateconn"
"vitess.io/vitess/go/vt/vttablet/tabletconn"
querypb "vitess.io/vitess/go/vt/proto/query"
topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
// Ensure dialers are registered (needed by ExecOnTablet and ExecOnVTGate).
@@ -807,7 +809,7 @@ func NewBareCluster(cell string, hostname string) *LocalProcessCluster {
// path/to/whatever exists
cluster.ReusingVTDATAROOT = true
} else {
- err = createDirectory(cluster.CurrentVTDATAROOT, 0700)
+ err = createDirectory(cluster.CurrentVTDATAROOT, 0o700)
if err != nil {
log.Fatal(err)
}
@@ -948,11 +950,12 @@ func (cluster *LocalProcessCluster) ExecOnTablet(ctx context.Context, vttablet *
txID, reservedID := 0, 0
- return conn.Execute(ctx, &querypb.Target{
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: opts})
+ return conn.Execute(ctx, session, &querypb.Target{
Keyspace: tablet.Keyspace,
Shard: tablet.Shard,
TabletType: tablet.Type,
- }, sql, bindvars, int64(txID), int64(reservedID), opts)
+ }, sql, bindvars, int64(txID), int64(reservedID))
}
// ExecOnVTGate executes a query on a local cluster VTGate with the provided
@@ -1165,7 +1168,8 @@ func (cluster *LocalProcessCluster) waitForMySQLProcessToExit(mysqlctlProcessLis
// StartVtbackup starts a vtbackup
func (cluster *LocalProcessCluster) StartVtbackup(newInitDBFile string, initialBackup bool,
- keyspace string, shard string, cell string, extraArgs ...string) error {
+ keyspace string, shard string, cell string, extraArgs ...string,
+) error {
log.Info("Starting vtbackup")
cluster.VtbackupProcess = *VtbackupProcessInstance(
cluster.GetAndReserveTabletUID(),
@@ -1195,7 +1199,6 @@ func (cluster *LocalProcessCluster) GetAndReservePort() int {
cluster.nextPortForProcess = cluster.nextPortForProcess + 1
log.Infof("Attempting to reserve port: %v", cluster.nextPortForProcess)
ln, err := net.Listen("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(cluster.nextPortForProcess)))
-
if err != nil {
log.Errorf("Can't listen on port %v: %s, trying next port", cluster.nextPortForProcess, err)
continue
@@ -1218,7 +1221,7 @@ const portFileTimeout = 1 * time.Hour
// If yes, then return that port, and save port + 200 in the same file
// here, assumptions is 200 ports might be consumed for all tests in a package
func getPort() int {
- portFile, err := os.OpenFile(path.Join(os.TempDir(), "endtoend.port"), os.O_CREATE|os.O_RDWR, 0644)
+ portFile, err := os.OpenFile(path.Join(os.TempDir(), "endtoend.port"), os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
panic(err)
}
diff --git a/go/test/endtoend/vtgate/tabletbalancer/session_test.go b/go/test/endtoend/vtgate/tabletbalancer/session_test.go
new file mode 100644
index 00000000000..d4ed13c0e53
--- /dev/null
+++ b/go/test/endtoend/vtgate/tabletbalancer/session_test.go
@@ -0,0 +1,223 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package tabletbalancer
+
+import (
+ "context"
+ "fmt"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+
+ "vitess.io/vitess/go/mysql"
+ "vitess.io/vitess/go/test/endtoend/cluster"
+ "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+)
+
+// TestSessionModeBalancer tests the "session" mode routes each session consistently to the same tablet.
+func TestSessionModeBalancer(t *testing.T) {
+ vtgateProcess, vtParams, _, _ := setupCluster(t)
+ defer vtgateProcess.TearDown()
+
+ // Create 2 session connections that route to different tablets
+ conns := createSessionConnections(t, &vtParams, 2)
+ for conn := range conns {
+ defer conn.Close()
+ }
+
+ verifyStickiness(t, conns, 20)
+}
+
+// TestSessionModeRemoveTablet tests that when a tablet is killed, connections switch to remaining tablets
+func TestSessionModeRemoveTablet(t *testing.T) {
+ vtgateProcess, vtParams, replicaTablets, aliases := setupCluster(t)
+ defer vtgateProcess.TearDown()
+
+ // Create 2 connections to different tablets
+ conns := createSessionConnections(t, &vtParams, 2)
+ for conn := range conns {
+ defer conn.Close()
+ }
+
+ // Find the first replica tablet that one of our connections is using
+ var tabletToKill *cluster.Vttablet
+ var affectedConn *mysql.Conn
+ var killedServerID int64
+
+ for _, tablet := range replicaTablets {
+ tabletServerID := aliases[tablet.Alias]
+
+ // Check if any connection is using this tablet
+ for conn, connServerID := range conns {
+ if connServerID != tabletServerID {
+ continue
+ }
+
+ // We found a connection that's using this tablet, let's kill this tablet
+ tabletToKill = tablet
+ affectedConn = conn
+ killedServerID = tabletServerID
+ break
+ }
+
+ // We found a tablet, no need to check other tablets
+ if tabletToKill != nil {
+ break
+ }
+ }
+
+ require.NotNil(t, tabletToKill, "Should find a tablet to kill")
+
+ // Kill the tablet immediately
+ err := tabletToKill.VttabletProcess.Kill()
+ require.Error(t, err)
+
+ // Wait for the connection to switch to a new tablet and update the map
+ require.Eventually(t, func() bool {
+ newServerID := getServerID(t, affectedConn)
+ if newServerID != killedServerID {
+ conns[affectedConn] = newServerID
+ return true
+ }
+
+ return false
+ }, 10*time.Millisecond, 1*time.Millisecond, "Connection should switch to a different tablet")
+
+ verifyStickiness(t, conns, 20)
+}
+
+// setupCluster sets up a cluster with a vtgate using the session balancer.
+func setupCluster(t *testing.T) (*cluster.VtgateProcess, mysql.ConnParams, []*cluster.Vttablet, map[string]int64) {
+ t.Helper()
+
+ // Start vtgate in cell1 with session mode
+ vtgateProcess := cluster.VtgateProcessInstance(
+ clusterInstance.GetAndReservePort(),
+ clusterInstance.GetAndReservePort(),
+ clusterInstance.GetAndReservePort(),
+ cell1,
+ fmt.Sprintf("%s,%s", cell1, cell2),
+ clusterInstance.Hostname,
+ replicaStr,
+ clusterInstance.TopoProcess.Port,
+ clusterInstance.TmpDirectory,
+ []string{
+ "--vtgate-balancer-mode", "session",
+ },
+ plancontext.PlannerVersion(0),
+ )
+ require.NoError(t, vtgateProcess.Setup())
+ require.True(t, vtgateProcess.WaitForStatus())
+
+ vtParams := mysql.ConnParams{
+ Host: clusterInstance.Hostname,
+ Port: vtgateProcess.MySQLServerPort,
+ }
+
+ allTablets := clusterInstance.Keyspaces[0].Shards[0].Vttablets
+ shardName := clusterInstance.Keyspaces[0].Shards[0].Name
+ replicaTablets := replicaTablets(allTablets)
+
+ conn, err := mysql.Connect(context.Background(), &vtParams)
+ require.NoError(t, err)
+ defer conn.Close()
+
+ // Wait for tablets to be discovered
+ err = vtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.primary", keyspaceName, shardName), 1, 30*time.Second)
+ require.NoError(t, err)
+
+ err = vtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.replica", keyspaceName, shardName), len(replicaTablets), 30*time.Second)
+ require.NoError(t, err)
+
+ aliases := mapTabletAliasToMySQLServerID(t, allTablets)
+
+ // Insert test data
+ testValue := fmt.Sprintf("session_test_%d", time.Now().UnixNano())
+ _, err = conn.ExecuteFetch(fmt.Sprintf("INSERT INTO balancer_test (value) VALUES ('%s')", testValue), 1, false)
+ require.NoError(t, err)
+ waitForReplication(t, replicaTablets, testValue)
+
+ return vtgateProcess, vtParams, replicaTablets, aliases
+}
+
+// getServerID returns the server ID that the connection is currently routing to.
+func getServerID(t *testing.T, conn *mysql.Conn) int64 {
+ t.Helper()
+
+ res, err := conn.ExecuteFetch("SELECT @@server_id", 1, false)
+ require.NoError(t, err)
+ require.Equal(t, 1, len(res.Rows), "expected one row from server_id query")
+
+ serverID, err := res.Rows[0][0].ToInt64()
+ require.NoError(t, err)
+
+ return serverID
+}
+
+// createSessionConnections creates `n` connections that route to different tablets.
+// Returns a map of mysql.Conn -> serverID.
+func createSessionConnections(t *testing.T, vtParams *mysql.ConnParams, numConnections int) map[*mysql.Conn]int64 {
+ t.Helper()
+
+ conns := make(map[*mysql.Conn]int64)
+ seenServerIDs := make(map[int64]bool)
+
+ // Try up to 50 times to get numConnections with different server IDs
+ for range 50 {
+ conn, err := mysql.Connect(context.Background(), vtParams)
+ require.NoError(t, err)
+
+ _, err = conn.ExecuteFetch("USE @replica", 1, false)
+ require.NoError(t, err)
+
+ // Get the server ID this connection routes to
+ serverID := getServerID(t, conn)
+
+ // If this is a new tablet, keep the connection
+ if !seenServerIDs[serverID] {
+ seenServerIDs[serverID] = true
+ conns[conn] = serverID
+
+ // If we have enough connections, return
+ if len(conns) == numConnections {
+ return conns
+ }
+
+ continue
+ }
+
+ // Already seen this tablet, close and try again
+ conn.Close()
+ }
+
+ t.Fatalf("could not create %d connections with different tablets after 50 attempts, only got %d", numConnections, len(conns))
+ return nil
+}
+
+// verifyStickiness validates whether the given connections remain connected to the same
+// server `n` times in a row.
+func verifyStickiness(t *testing.T, conns map[*mysql.Conn]int64, n uint) {
+ t.Helper()
+
+ for conn, expectedServerID := range conns {
+ for range n {
+ currentServerID := getServerID(t, conn)
+ require.Equal(t, expectedServerID, currentServerID, "Connection should stick to tablet %d, got %d", expectedServerID, currentServerID)
+ }
+ }
+}
diff --git a/go/vt/vtcombo/tablet_map.go b/go/vt/vtcombo/tablet_map.go
index 79b9cfb4e7f..45166d04665 100644
--- a/go/vt/vtcombo/tablet_map.go
+++ b/go/vt/vtcombo/tablet_map.go
@@ -460,14 +460,14 @@ var _ queryservice.QueryService = (*internalTabletConn)(nil)
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) Execute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
query string,
bindVars map[string]*querypb.BindVariable,
transactionID, reservedID int64,
- options *querypb.ExecuteOptions,
) (*sqltypes.Result, error) {
bindVars = sqltypes.CopyBindVariables(bindVars)
- reply, err := itc.tablet.qsc.QueryService().Execute(ctx, target, query, bindVars, transactionID, reservedID, options)
+ reply, err := itc.tablet.qsc.QueryService().Execute(ctx, session, target, query, bindVars, transactionID, reservedID)
if err != nil {
return nil, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
@@ -478,26 +478,26 @@ func (itc *internalTabletConn) Execute(
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) StreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
query string,
bindVars map[string]*querypb.BindVariable,
transactionID int64,
reservedID int64,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) error {
bindVars = sqltypes.CopyBindVariables(bindVars)
- err := itc.tablet.qsc.QueryService().StreamExecute(ctx, target, query, bindVars, transactionID, reservedID, options, callback)
+ err := itc.tablet.qsc.QueryService().StreamExecute(ctx, session, target, query, bindVars, transactionID, reservedID, callback)
return tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
// Begin is part of queryservice.QueryService
func (itc *internalTabletConn) Begin(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
- options *querypb.ExecuteOptions,
) (queryservice.TransactionState, error) {
- state, err := itc.tablet.qsc.QueryService().Begin(ctx, target, options)
+ state, err := itc.tablet.qsc.QueryService().Begin(ctx, session, target)
return state, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
@@ -575,31 +575,31 @@ func (itc *internalTabletConn) UnresolvedTransactions(ctx context.Context, targe
// BeginExecute is part of queryservice.QueryService
func (itc *internalTabletConn) BeginExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
query string,
bindVars map[string]*querypb.BindVariable,
reserveID int64,
- options *querypb.ExecuteOptions,
) (queryservice.TransactionState, *sqltypes.Result, error) {
bindVars = sqltypes.CopyBindVariables(bindVars)
- state, result, err := itc.tablet.qsc.QueryService().BeginExecute(ctx, target, preQueries, query, bindVars, reserveID, options)
+ state, result, err := itc.tablet.qsc.QueryService().BeginExecute(ctx, session, target, preQueries, query, bindVars, reserveID)
return state, result, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
// BeginStreamExecute is part of queryservice.QueryService
func (itc *internalTabletConn) BeginStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
query string,
bindVars map[string]*querypb.BindVariable,
reservedID int64,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (queryservice.TransactionState, error) {
bindVars = sqltypes.CopyBindVariables(bindVars)
- state, err := itc.tablet.qsc.QueryService().BeginStreamExecute(ctx, target, preQueries, query, bindVars, reservedID, options, callback)
+ state, err := itc.tablet.qsc.QueryService().BeginStreamExecute(ctx, session, target, preQueries, query, bindVars, reservedID, callback)
return state, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
@@ -622,62 +622,62 @@ func (itc *internalTabletConn) HandlePanic(err *error) {
// ReserveBeginExecute is part of the QueryService interface.
func (itc *internalTabletConn) ReserveBeginExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
postBeginQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
- options *querypb.ExecuteOptions,
) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
bindVariables = sqltypes.CopyBindVariables(bindVariables)
- state, result, err := itc.tablet.qsc.QueryService().ReserveBeginExecute(ctx, target, preQueries, postBeginQueries, sql, bindVariables, options)
+ state, result, err := itc.tablet.qsc.QueryService().ReserveBeginExecute(ctx, session, target, preQueries, postBeginQueries, sql, bindVariables)
return state, result, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
// ReserveBeginStreamExecute is part of the QueryService interface.
func (itc *internalTabletConn) ReserveBeginStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
postBeginQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (queryservice.ReservedTransactionState, error) {
bindVariables = sqltypes.CopyBindVariables(bindVariables)
- state, err := itc.tablet.qsc.QueryService().ReserveBeginStreamExecute(ctx, target, preQueries, postBeginQueries, sql, bindVariables, options, callback)
+ state, err := itc.tablet.qsc.QueryService().ReserveBeginStreamExecute(ctx, session, target, preQueries, postBeginQueries, sql, bindVariables, callback)
return state, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
// ReserveExecute is part of the QueryService interface.
func (itc *internalTabletConn) ReserveExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
transactionID int64,
- options *querypb.ExecuteOptions,
) (queryservice.ReservedState, *sqltypes.Result, error) {
bindVariables = sqltypes.CopyBindVariables(bindVariables)
- state, result, err := itc.tablet.qsc.QueryService().ReserveExecute(ctx, target, preQueries, sql, bindVariables, transactionID, options)
+ state, result, err := itc.tablet.qsc.QueryService().ReserveExecute(ctx, session, target, preQueries, sql, bindVariables, transactionID)
return state, result, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
// ReserveStreamExecute is part of the QueryService interface.
func (itc *internalTabletConn) ReserveStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
preQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
transactionID int64,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (queryservice.ReservedState, error) {
bindVariables = sqltypes.CopyBindVariables(bindVariables)
- state, err := itc.tablet.qsc.QueryService().ReserveStreamExecute(ctx, target, preQueries, sql, bindVariables, transactionID, options, callback)
+ state, err := itc.tablet.qsc.QueryService().ReserveStreamExecute(ctx, session, target, preQueries, sql, bindVariables, transactionID, callback)
return state, tabletconn.ErrorFromGRPC(vterrors.ToGRPC(err))
}
@@ -1103,6 +1103,7 @@ func (itmc *internalTabletManagerClient) ResetReplicationParameters(context.Cont
func (itmc *internalTabletManagerClient) ReplicaWasRestarted(context.Context, *topodatapb.Tablet, *topodatapb.TabletAlias) error {
return errors.New("not implemented in vtcombo")
}
+
func (itmc *internalTabletManagerClient) ResetSequences(ctx context.Context, tablet *topodatapb.Tablet, tables []string) error {
return errors.New("not implemented in vtcombo")
}
diff --git a/go/vt/vtexplain/vtexplain_vttablet.go b/go/vt/vtexplain/vtexplain_vttablet.go
index 36e90176d16..6dd2fc250a9 100644
--- a/go/vt/vtexplain/vtexplain_vttablet.go
+++ b/go/vt/vtexplain/vtexplain_vttablet.go
@@ -127,7 +127,7 @@ func (vte *VTExplain) newTablet(ctx context.Context, env *vtenv.Environment, opt
tablet.QueryService = queryservice.Wrap(
nil,
- func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, inTransaction bool, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {
+ func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, opts queryservice.WrapOpts, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {
return fmt.Errorf("explainTablet does not implement %s", name)
},
)
@@ -155,7 +155,7 @@ func (vte *VTExplain) newTablet(ctx context.Context, env *vtenv.Environment, opt
var _ queryservice.QueryService = (*explainTablet)(nil) // compile-time interface check
// Begin is part of the QueryService interface.
-func (t *explainTablet) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (t *explainTablet) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (queryservice.TransactionState, error) {
t.mu.Lock()
t.currentTime = t.vte.batchTime.Wait()
t.tabletQueries = append(t.tabletQueries, &TabletQuery{
@@ -165,7 +165,7 @@ func (t *explainTablet) Begin(ctx context.Context, target *querypb.Target, optio
t.mu.Unlock()
- return t.tsv.Begin(ctx, target, options)
+ return t.tsv.Begin(ctx, session, target)
}
// Commit is part of the QueryService interface.
@@ -190,7 +190,7 @@ func (t *explainTablet) Rollback(ctx context.Context, target *querypb.Target, tr
}
// Execute is part of the QueryService interface.
-func (t *explainTablet) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (t *explainTablet) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
t.mu.Lock()
t.currentTime = t.vte.batchTime.Wait()
@@ -204,7 +204,7 @@ func (t *explainTablet) Execute(ctx context.Context, target *querypb.Target, sql
})
t.mu.Unlock()
- return t.tsv.Execute(ctx, target, sql, bindVariables, transactionID, reservedID, options)
+ return t.tsv.Execute(ctx, session, target, sql, bindVariables, transactionID, reservedID)
}
// Prepare is part of the QueryService interface.
@@ -264,7 +264,7 @@ func (t *explainTablet) ReadTransaction(ctx context.Context, target *querypb.Tar
}
// BeginExecute is part of the QueryService interface.
-func (t *explainTablet) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
+func (t *explainTablet) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64) (queryservice.TransactionState, *sqltypes.Result, error) {
t.mu.Lock()
t.currentTime = t.vte.batchTime.Wait()
bindVariables = sqltypes.CopyBindVariables(bindVariables)
@@ -275,7 +275,7 @@ func (t *explainTablet) BeginExecute(ctx context.Context, target *querypb.Target
})
t.mu.Unlock()
- return t.tsv.BeginExecute(ctx, target, preQueries, sql, bindVariables, reservedID, options)
+ return t.tsv.BeginExecute(ctx, session, target, preQueries, sql, bindVariables, reservedID)
}
// Close is part of the QueryService interface.
@@ -408,7 +408,6 @@ func newTabletEnvironment(ddls []sqlparser.DDLStatement, opts *Options, collatio
Rows: [][]sqltypes.Value{},
},
"create table if not exists `_vt`.dt_participant(\n dtid varbinary(512),\n\tid bigint,\n\tkeyspace varchar(256),\n\tshard varchar(256),\n primary key(dtid, id)\n\t) engine=InnoDB": {
-
Fields: []*querypb.Field{{
Type: sqltypes.Uint64,
Charset: collations.CollationBinaryID,
diff --git a/go/vt/vtgate/balancer/balancer.go b/go/vt/vtgate/balancer/balancer.go
index a9bc5b4c808..7f5631bd959 100644
--- a/go/vt/vtgate/balancer/balancer.go
+++ b/go/vt/vtgate/balancer/balancer.go
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/
+// Package balancer contains a number of different tablet balancing algorithms and their implementations.
package balancer
import (
@@ -96,6 +97,7 @@ const (
ModeCell
ModePreferCell
ModeRandom
+ ModeSession
)
func ParseMode(ms string) Mode {
@@ -106,6 +108,8 @@ func ParseMode(ms string) Mode {
return ModePreferCell
case "random":
return ModeRandom
+ case "session":
+ return ModeSession
default:
return ModeInvalid
}
@@ -119,19 +123,21 @@ func (m Mode) String() string {
return "prefer-cell"
case ModeRandom:
return "random"
+ case ModeSession:
+ return "session"
default:
return "invalid"
}
}
func GetAvailableModeNames() []string {
- return []string{ModeCell.String(), ModePreferCell.String(), ModeRandom.String()}
+ return []string{ModeCell.String(), ModePreferCell.String(), ModeRandom.String(), ModeSession.String()}
}
type TabletBalancer interface {
// Pick is the main entry point to the balancer. Returns the best tablet out of the list
// for a given query to maintain the desired balanced allocation over multiple executions.
- Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth
+ Pick(target *querypb.Target, tablets []*discovery.TabletHealth, opts ...PickOption) *discovery.TabletHealth
// DebugHandler provides a summary of tablet balancer state
DebugHandler(w http.ResponseWriter, r *http.Request)
@@ -142,6 +148,7 @@ type TabletBalancer interface {
// - "prefer-cell": Flow-based balancer that maintains cell affinity while balancing load
// - See the RFC here: https://github.com/vitessio/vitess/issues/12241
// - "random": Random balancer that uniformly distributes load without cell affinity
+// - "session": Session balancer that pins a session to the same tablet for the duration of the session. If the tablet goes away, the session is automatically and transparently migrated to another tablet of the same type.
//
// Note: "cell" mode is handled by the gateway and does not create a balancer instance.
// operates as a round robin inside of the vtgate's cell
@@ -152,6 +159,8 @@ func NewTabletBalancer(mode Mode, localCell string, vtGateCells []string) (Table
return newFlowBalancer(localCell, vtGateCells), nil
case ModeRandom:
return newRandomBalancer(localCell, vtGateCells), nil
+ case ModeSession:
+ return newSessionBalancer(localCell), nil
case ModeCell:
return nil, errors.New("cell mode should be handled by the gateway, not the balancer factory")
default:
@@ -230,7 +239,7 @@ func (b *flowBalancer) DebugHandler(w http.ResponseWriter, _ *http.Request) {
// Given the total allocation for the set of tablets, choose the best target
// by a weighted random sample so that over time the system will achieve the
// desired balanced allocation.
-func (b *flowBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth {
+func (b *flowBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth, _ ...PickOption) *discovery.TabletHealth {
numTablets := len(tablets)
if numTablets == 0 {
return nil
@@ -342,9 +351,6 @@ func (b *flowBalancer) allocateFlows(allTablets []*discovery.TabletHealth) *targ
}
}
- // fmt.Printf("outflows %v over %v under %v\n", a.Outflows, overAllocated, underAllocated)
-
- //
// For each overallocated cell, proportionally shift flow from targets that are overallocated
// to targets that are underallocated.
//
@@ -366,9 +372,6 @@ func (b *flowBalancer) allocateFlows(allTablets []*discovery.TabletHealth) *targ
// to avoid truncating the integer values.
shiftFlow := overAllocatedFlow * currentFlow * underAllocatedFlow / a.Inflows[overAllocatedCell] / unbalancedFlow
- //fmt.Printf("shift %d %s %s -> %s (over %d current %d in %d under %d unbalanced %d) \n", shiftFlow, vtgateCell, overAllocatedCell, underAllocatedCell,
- // overAllocatedFlow, currentFlow, a.Inflows[overAllocatedCell], underAllocatedFlow, unbalancedFlow)
-
a.Outflows[vtgateCell][overAllocatedCell] -= shiftFlow
a.Inflows[overAllocatedCell] -= shiftFlow
diff --git a/go/vt/vtgate/balancer/options.go b/go/vt/vtgate/balancer/options.go
new file mode 100644
index 00000000000..9ca065b7285
--- /dev/null
+++ b/go/vt/vtgate/balancer/options.go
@@ -0,0 +1,62 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package balancer
+
+// pickOptions configure a Pick call. pickOptions are set by the PickOption
+// values passed to the Pick function.
+type pickOptions struct {
+ sessionUUID string
+}
+
+// PickOption configures how we perform the pick operation.
+type PickOption interface {
+ apply(*pickOptions)
+}
+
+// funcPickOption wraps a function that modifies pickOptions into an
+// implementation of the PickOption interface.
+type funcPickOption struct {
+ f func(*pickOptions)
+}
+
+func (fpo *funcPickOption) apply(po *pickOptions) {
+ fpo.f(po)
+}
+
+func newFuncPickOption(f func(*pickOptions)) *funcPickOption {
+ return &funcPickOption{
+ f: f,
+ }
+}
+
+// WithSessionUUID allows you to specify the session UUID for the session balancer.
+func WithSessionUUID(sessionUUID string) PickOption {
+ return newFuncPickOption(func(o *pickOptions) {
+ o.sessionUUID = sessionUUID
+ })
+}
+
+// getOptions applies the given options to a new pickOptions struct
+// and returns it.
+func getOptions(opts []PickOption) *pickOptions {
+ options := &pickOptions{}
+ for _, opt := range opts {
+ opt.apply(options)
+ }
+
+ return options
+}
diff --git a/go/vt/vtgate/balancer/random_balancer.go b/go/vt/vtgate/balancer/random_balancer.go
index f8ad723c910..52c28792af7 100644
--- a/go/vt/vtgate/balancer/random_balancer.go
+++ b/go/vt/vtgate/balancer/random_balancer.go
@@ -73,7 +73,7 @@ type randomBalancer struct {
// Pick returns a random tablet from the list with uniform probability (1/N).
// If vtGateCells is configured, only tablets in those cells are considered.
-func (b *randomBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth) *discovery.TabletHealth {
+func (b *randomBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth, _ ...PickOption) *discovery.TabletHealth {
if len(tablets) == 0 {
return nil
}
diff --git a/go/vt/vtgate/balancer/session.go b/go/vt/vtgate/balancer/session.go
new file mode 100644
index 00000000000..7f70bc6dcac
--- /dev/null
+++ b/go/vt/vtgate/balancer/session.go
@@ -0,0 +1,106 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package balancer
+
+import (
+ "fmt"
+ "net/http"
+
+ "github.com/cespare/xxhash/v2"
+
+ "vitess.io/vitess/go/vt/discovery"
+ querypb "vitess.io/vitess/go/vt/proto/query"
+ "vitess.io/vitess/go/vt/topo/topoproto"
+)
+
+// SessionBalancer implements the TabletBalancer interface. For a given session,
+// it will return the same tablet for its duration, with preference to tablets in
+// the local cell.
+type SessionBalancer struct {
+ // localCell is the cell the gateway is currently running in.
+ localCell string
+}
+
+// newSessionBalancer creates a new session balancer.
+func newSessionBalancer(localCell string) TabletBalancer {
+ return &SessionBalancer{localCell: localCell}
+}
+
+// Pick is the main entry point to the balancer.
+//
+// For a given session, it will return the same tablet for its duration, with preference to tablets
+// in the local cell.
+func (b *SessionBalancer) Pick(target *querypb.Target, tablets []*discovery.TabletHealth, opts ...PickOption) *discovery.TabletHealth {
+ options := getOptions(opts)
+ if options.sessionUUID == "" {
+ return nil
+ }
+
+ // Find the highest weight local and external tablets
+ var maxLocalWeight, maxExternalWeight uint64
+ var maxLocalTablet, maxExternalTablet *discovery.TabletHealth
+
+ for _, tablet := range tablets {
+ alias := tabletAlias(tablet)
+ weight := tabletWeight(alias, options.sessionUUID)
+
+ if b.isLocal(tablet) && ((maxLocalTablet == nil) || (weight > maxLocalWeight)) {
+ maxLocalWeight = weight
+ maxLocalTablet = tablet
+ }
+
+ // We can consider all tablets here since we'd only use this if there were no
+ // valid local tablets (meaning we'd have only considered external tablets anyway).
+ if (maxExternalTablet == nil) || (weight > maxExternalWeight) {
+ maxExternalWeight = weight
+ maxExternalTablet = tablet
+ }
+ }
+
+ // If we found a valid local tablet, use that
+ if maxLocalTablet != nil {
+ return maxLocalTablet
+ }
+
+ // Otherwise, use the max external tablet (if it exists)
+ return maxExternalTablet
+}
+
+// tabletWeight computes the weight of a tablet by hashing its alias and the session UUID together.
+func tabletWeight(alias string, sessionUUID string) uint64 {
+ h := xxhash.New()
+ _, _ = h.WriteString(alias)
+ _, _ = h.WriteString("#")
+ _, _ = h.WriteString(sessionUUID)
+ return h.Sum64()
+}
+
+// tabletAlias returns the tablet's alias as a string.
+func tabletAlias(tablet *discovery.TabletHealth) string {
+ return topoproto.TabletAliasString(tablet.Tablet.Alias)
+}
+
+// isLocal returns true if the tablet is in the local cell.
+func (b *SessionBalancer) isLocal(tablet *discovery.TabletHealth) bool {
+ return tablet.Tablet.Alias.Cell == b.localCell
+}
+
+// DebugHandler provides a summary of the session balancer state.
+func (b *SessionBalancer) DebugHandler(w http.ResponseWriter, r *http.Request) {
+ w.Header().Set("Content-Type", "text/plain")
+ fmt.Fprintf(w, "Local cell: %s", b.localCell)
+}
diff --git a/go/vt/vtgate/balancer/session_test.go b/go/vt/vtgate/balancer/session_test.go
new file mode 100644
index 00000000000..3b127221bcf
--- /dev/null
+++ b/go/vt/vtgate/balancer/session_test.go
@@ -0,0 +1,351 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package balancer
+
+import (
+ "fmt"
+ "slices"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+
+ "vitess.io/vitess/go/vt/discovery"
+ querypb "vitess.io/vitess/go/vt/proto/query"
+ topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ "vitess.io/vitess/go/vt/topo/topoproto"
+)
+
+func createSessionBalancer(t *testing.T) TabletBalancer {
+ t.Helper()
+
+ return newSessionBalancer("local")
+}
+
+func TestSessionPickNoTablets(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ result := b.Pick(target, nil, WithSessionUUID("a"))
+ require.Nil(t, result)
+}
+
+func TestSessionPickLocalOnly(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ tablets := []*discovery.TabletHealth{
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 100,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 101,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+ }
+
+ // Pick for a specific session UUID
+ picked1 := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.NotNil(t, picked1)
+
+ // Pick again with same session hash, should return same tablet
+ picked2 := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.Equal(t, picked1, picked2, fmt.Sprintf("expected %s, got %s", tabletAlias(picked1), tabletAlias(picked2)))
+
+ // Pick with different session hash, empirically know that it should return tablet2
+ picked3 := b.Pick(target, tablets, WithSessionUUID("b"))
+ require.NotNil(t, picked3)
+ require.NotEqual(t, picked2, picked3, fmt.Sprintf("expected different tablets, got %s for both", tabletAlias(picked3)))
+}
+
+func TestSessionPickPreferLocal(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ tablets := []*discovery.TabletHealth{
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 100,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 101,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "external",
+ Uid: 200,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "external",
+ },
+ Serving: true,
+ },
+ }
+
+ // Pick should prefer local cell
+ picked1 := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.NotNil(t, picked1)
+ require.Equal(t, "local", picked1.Target.Cell)
+
+ // Pick should pick the same tablet consistently
+ for range 20 {
+ picked := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.Equal(t, picked1, picked, fmt.Sprintf("expected %s, got %s", tabletAlias(picked1), tabletAlias(picked)))
+ }
+}
+
+func TestSessionPickNoLocal(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ tablets := []*discovery.TabletHealth{
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "external",
+ Uid: 200,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "external",
+ },
+ Serving: true,
+ },
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "external",
+ Uid: 201,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "external",
+ },
+ Serving: true,
+ },
+ }
+
+ // Pick should return external cell since there are no local cells
+ picked1 := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.NotNil(t, picked1)
+ require.Equal(t, "external", picked1.Target.Cell)
+
+ // Pick should pick the same tablet consistently
+ for range 20 {
+ picked := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.Equal(t, picked1, picked, fmt.Sprintf("expected %s, got %s", tabletAlias(picked1), tabletAlias(picked)))
+ }
+}
+
+func TestSessionPickNoOpts(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ tablets := []*discovery.TabletHealth{
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 100,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+ }
+
+ // Test with no opts (no session UUID)
+ result := b.Pick(target, tablets)
+ require.Nil(t, result)
+}
+
+func TestSessionPickInvalidTablets(t *testing.T) {
+ b := createSessionBalancer(t)
+
+ target := &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ }
+
+ tablets := []*discovery.TabletHealth{
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 100,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+
+ {
+ Tablet: &topodatapb.Tablet{
+ Alias: &topodatapb.TabletAlias{
+ Cell: "local",
+ Uid: 101,
+ },
+ Keyspace: "keyspace",
+ Shard: "0",
+ },
+ Target: &querypb.Target{
+ Keyspace: "keyspace",
+ Shard: "0",
+ TabletType: topodatapb.TabletType_REPLICA,
+ Cell: "local",
+ },
+ Serving: true,
+ },
+ }
+
+ // Get a tablet regularly
+ tablet := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.NotNil(t, tablet)
+
+ // Filter out the returned tablet as invalid
+ tablets = slices.DeleteFunc(tablets, func(t *discovery.TabletHealth) bool {
+ return topoproto.TabletAliasString(t.Tablet.Alias) == topoproto.TabletAliasString(tablet.Tablet.Alias)
+ })
+
+ // Pick should now return a different tablet
+ tablet2 := b.Pick(target, tablets, WithSessionUUID("a"))
+ require.NotNil(t, tablet2)
+ require.NotEqual(t, tablet, tablet2)
+
+ // Filter out the last tablet, Pick should return nothing
+ tablet3 := b.Pick(target, []*discovery.TabletHealth{}, WithSessionUUID("a"))
+ require.Nil(t, tablet3)
+}
diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
index 19f41fe7643..f12aeb0d0d8 100644
--- a/go/vt/vtgate/executor.go
+++ b/go/vt/vtgate/executor.go
@@ -155,12 +155,16 @@ type (
var executorOnce sync.Once
-const pathQueryPlans = "/debug/query_plans"
-const pathScatterStats = "/debug/scatter_stats"
-const pathVSchema = "/debug/vschema"
+const (
+ pathQueryPlans = "/debug/query_plans"
+ pathScatterStats = "/debug/scatter_stats"
+ pathVSchema = "/debug/vschema"
+)
-type PlanCacheKey = theine.HashKey256
-type PlanCache = theine.Store[PlanCacheKey, *engine.Plan]
+type (
+ PlanCacheKey = theine.HashKey256
+ PlanCache = theine.Store[PlanCacheKey, *engine.Plan]
+)
func DefaultPlanCache() *PlanCache {
// when being endtoend tested, disable the doorkeeper to ensure reproducible results
@@ -365,7 +369,6 @@ func (e *Executor) StreamExecute(
err := vc.StreamExecutePrimitive(ctx, plan.Instructions, bindVars, true, func(qr *sqltypes.Result) error {
return srr.storeResultStats(plan.QueryType, qr)
})
-
// Check if there was partial DML execution. If so, rollback the effect of the partially executed query.
if err != nil {
if safeSession.InTransaction() && e.rollbackOnFatalTxError(ctx, safeSession, err) {
@@ -1009,7 +1012,7 @@ func (e *Executor) ShowVitessReplicationStatus(ctx context.Context, filter *sqlp
replicaSQLRunningField = "Slave_SQL_Running"
secondsBehindSourceField = "Seconds_Behind_Master"
}
- results, err := e.txConn.tabletGateway.Execute(ctx, ts.Target, sql, nil, 0, 0, nil)
+ results, err := e.txConn.tabletGateway.Execute(ctx, nil, ts.Target, sql, nil, 0, 0)
if err != nil || results == nil {
log.Warningf("Could not get replication status from %s: %v", tabletHostPort, err)
} else if row := results.Named().Row(); row != nil {
@@ -1119,7 +1122,8 @@ func (e *Executor) fetchOrCreatePlan(
logStats *logstats.LogStats,
isExecutePath bool, // this means we are trying to execute the query - this is not a PREPARE call
) (
- plan *engine.Plan, vcursor *econtext.VCursorImpl, stmt sqlparser.Statement, err error) {
+ plan *engine.Plan, vcursor *econtext.VCursorImpl, stmt sqlparser.Statement, err error,
+) {
if e.VSchema() == nil {
return nil, nil, nil, vterrors.VT13001("vschema not initialized")
}
diff --git a/go/vt/vtgate/scatter_conn.go b/go/vt/vtgate/scatter_conn.go
index 77401f07f92..a2e46442c31 100644
--- a/go/vt/vtgate/scatter_conn.go
+++ b/go/vt/vtgate/scatter_conn.go
@@ -193,6 +193,7 @@ func (stc *ScatterConn) ExecuteMultiShard(
if opts == nil && fetchLastInsertID {
opts = &querypb.ExecuteOptions{FetchLastInsertId: fetchLastInsertID}
+ session = econtext.NewSafeSession(&vtgatepb.Session{Options: opts})
}
if autocommit {
@@ -223,21 +224,21 @@ func (stc *ScatterConn) ExecuteMultiShard(
switch info.actionNeeded {
case nothing:
- innerqr, err = qs.Execute(ctx, rs.Target, queries[i].Sql, queries[i].BindVariables, info.transactionID, info.reservedID, opts)
+ innerqr, err = qs.Execute(ctx, session, rs.Target, queries[i].Sql, queries[i].BindVariables, info.transactionID, info.reservedID)
if err != nil {
retryRequest(func() {
// we seem to have lost our connection. it was a reserved connection, let's try to recreate it
info.actionNeeded = reserve
info.ignoreOldSession = true
var state queryservice.ReservedState
- state, innerqr, err = qs.ReserveExecute(ctx, rs.Target, session.SetPreQueries(), queries[i].Sql, queries[i].BindVariables, 0 /*transactionId*/, opts)
+ state, innerqr, err = qs.ReserveExecute(ctx, session, rs.Target, session.SetPreQueries(), queries[i].Sql, queries[i].BindVariables, 0 /*transactionId*/)
reservedID = state.ReservedID
alias = state.TabletAlias
})
}
case begin:
var state queryservice.TransactionState
- state, innerqr, err = qs.BeginExecute(ctx, rs.Target, session.SavePoints(), queries[i].Sql, queries[i].BindVariables, reservedID, opts)
+ state, innerqr, err = qs.BeginExecute(ctx, session, rs.Target, session.SavePoints(), queries[i].Sql, queries[i].BindVariables, reservedID)
transactionID = state.TransactionID
alias = state.TabletAlias
if err != nil {
@@ -246,7 +247,7 @@ func (stc *ScatterConn) ExecuteMultiShard(
info.actionNeeded = reserveBegin
info.ignoreOldSession = true
var state queryservice.ReservedTransactionState
- state, innerqr, err = qs.ReserveBeginExecute(ctx, rs.Target, session.SetPreQueries(), session.SavePoints(), queries[i].Sql, queries[i].BindVariables, opts)
+ state, innerqr, err = qs.ReserveBeginExecute(ctx, session, rs.Target, session.SetPreQueries(), session.SavePoints(), queries[i].Sql, queries[i].BindVariables)
transactionID = state.TransactionID
reservedID = state.ReservedID
alias = state.TabletAlias
@@ -254,12 +255,12 @@ func (stc *ScatterConn) ExecuteMultiShard(
}
case reserve:
var state queryservice.ReservedState
- state, innerqr, err = qs.ReserveExecute(ctx, rs.Target, session.SetPreQueries(), queries[i].Sql, queries[i].BindVariables, transactionID, opts)
+ state, innerqr, err = qs.ReserveExecute(ctx, session, rs.Target, session.SetPreQueries(), queries[i].Sql, queries[i].BindVariables, transactionID)
reservedID = state.ReservedID
alias = state.TabletAlias
case reserveBegin:
var state queryservice.ReservedTransactionState
- state, innerqr, err = qs.ReserveBeginExecute(ctx, rs.Target, session.SetPreQueries(), session.SavePoints(), queries[i].Sql, queries[i].BindVariables, opts)
+ state, innerqr, err = qs.ReserveBeginExecute(ctx, session, rs.Target, session.SetPreQueries(), session.SavePoints(), queries[i].Sql, queries[i].BindVariables)
transactionID = state.TransactionID
reservedID = state.ReservedID
alias = state.TabletAlias
@@ -424,6 +425,7 @@ func (stc *ScatterConn) StreamExecuteMulti(
if opts == nil && fetchLastInsertID {
opts = &querypb.ExecuteOptions{FetchLastInsertId: fetchLastInsertID}
+ session = econtext.NewSafeSession(&vtgatepb.Session{Options: opts})
}
if autocommit {
@@ -454,20 +456,20 @@ func (stc *ScatterConn) StreamExecuteMulti(
switch info.actionNeeded {
case nothing:
- err = qs.StreamExecute(ctx, rs.Target, query, bindVars[i], transactionID, reservedID, opts, observedCallback)
+ err = qs.StreamExecute(ctx, session, rs.Target, query, bindVars[i], transactionID, reservedID, observedCallback)
if err != nil {
retryRequest(func() {
// we seem to have lost our connection. it was a reserved connection, let's try to recreate it
info.actionNeeded = reserve
var state queryservice.ReservedState
- state, err = qs.ReserveStreamExecute(ctx, rs.Target, session.SetPreQueries(), query, bindVars[i], 0 /*transactionId*/, opts, observedCallback)
+ state, err = qs.ReserveStreamExecute(ctx, session, rs.Target, session.SetPreQueries(), query, bindVars[i], 0 /*transactionId*/, observedCallback)
reservedID = state.ReservedID
alias = state.TabletAlias
})
}
case begin:
var state queryservice.TransactionState
- state, err = qs.BeginStreamExecute(ctx, rs.Target, session.SavePoints(), query, bindVars[i], reservedID, opts, observedCallback)
+ state, err = qs.BeginStreamExecute(ctx, session, rs.Target, session.SavePoints(), query, bindVars[i], reservedID, observedCallback)
transactionID = state.TransactionID
alias = state.TabletAlias
if err != nil {
@@ -475,7 +477,7 @@ func (stc *ScatterConn) StreamExecuteMulti(
// we seem to have lost our connection. it was a reserved connection, let's try to recreate it
info.actionNeeded = reserveBegin
var state queryservice.ReservedTransactionState
- state, err = qs.ReserveBeginStreamExecute(ctx, rs.Target, session.SetPreQueries(), session.SavePoints(), query, bindVars[i], opts, observedCallback)
+ state, err = qs.ReserveBeginStreamExecute(ctx, session, rs.Target, session.SetPreQueries(), session.SavePoints(), query, bindVars[i], observedCallback)
transactionID = state.TransactionID
reservedID = state.ReservedID
alias = state.TabletAlias
@@ -483,12 +485,12 @@ func (stc *ScatterConn) StreamExecuteMulti(
}
case reserve:
var state queryservice.ReservedState
- state, err = qs.ReserveStreamExecute(ctx, rs.Target, session.SetPreQueries(), query, bindVars[i], transactionID, opts, observedCallback)
+ state, err = qs.ReserveStreamExecute(ctx, session, rs.Target, session.SetPreQueries(), query, bindVars[i], transactionID, observedCallback)
reservedID = state.ReservedID
alias = state.TabletAlias
case reserveBegin:
var state queryservice.ReservedTransactionState
- state, err = qs.ReserveBeginStreamExecute(ctx, rs.Target, session.SetPreQueries(), session.SavePoints(), query, bindVars[i], opts, observedCallback)
+ state, err = qs.ReserveBeginStreamExecute(ctx, session, rs.Target, session.SetPreQueries(), session.SavePoints(), query, bindVars[i], observedCallback)
transactionID = state.TransactionID
reservedID = state.ReservedID
alias = state.TabletAlias
@@ -759,7 +761,6 @@ func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedSha
var (
qr *sqltypes.Result
err error
- opts *querypb.ExecuteOptions
alias *topodatapb.TabletAlias
)
allErrors := new(concurrency.AllErrorRecorder)
@@ -770,7 +771,6 @@ func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedSha
return nil, vterrors.VT13001("session cannot be nil")
}
- opts = session.Options
info, err := lockInfo(rs.Target, session, lockFuncType)
// Lock session is created on alphabetic sorted keyspace.
// This error will occur if the existing session target does not match the current target.
@@ -788,7 +788,7 @@ func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedSha
switch info.actionNeeded {
case nothing:
- qr, err = qs.Execute(ctx, rs.Target, query.Sql, query.BindVariables, 0 /* transactionID */, reservedID, opts)
+ qr, err = qs.Execute(ctx, session, rs.Target, query.Sql, query.BindVariables, 0 /* transactionID */, reservedID)
if err != nil && wasConnectionClosed(err) {
// TODO: try to acquire lock again.
session.ResetLock()
@@ -799,7 +799,7 @@ func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedSha
}
case reserve:
var state queryservice.ReservedState
- state, qr, err = qs.ReserveExecute(ctx, rs.Target, session.SetPreQueries(), query.Sql, query.BindVariables, 0 /* transactionID */, opts)
+ state, qr, err = qs.ReserveExecute(ctx, session, rs.Target, session.SetPreQueries(), query.Sql, query.BindVariables, 0 /* transactionID */)
reservedID = state.ReservedID
alias = state.TabletAlias
if err != nil && reservedID != 0 {
@@ -877,7 +877,7 @@ func actionInfo(ctx context.Context, target *querypb.Target, session *econtext.S
shouldReserve := session.InReservedConn() && (shardSession == nil || shardSession.ReservedId == 0)
shouldBegin := session.InTransaction() && (shardSession == nil || shardSession.TransactionId == 0) && !autocommit
- var act = nothing
+ act := nothing
switch {
case shouldBegin && shouldReserve:
act = reserveBegin
diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go
index be83726f0e4..bfe9508691c 100644
--- a/go/vt/vtgate/tabletgateway.go
+++ b/go/vt/vtgate/tabletgateway.go
@@ -109,6 +109,9 @@ type TabletGateway struct {
// balancer used for routing to tablets
balancer balancer.TabletBalancer
+
+ // balancerMode is the current tablet balancer mode.
+ balancerMode balancer.Mode
}
func createHealthCheck(ctx context.Context, retryDelay, timeout time.Duration, ts *topo.Server, cell, cellsToWatch string) discovery.HealthCheck {
@@ -182,38 +185,37 @@ func (gw *TabletGateway) setupBalancer() {
}
// Determine the effective mode: new flag takes precedence, then deprecated flag, then default
- var mode balancer.Mode
if balancerModeFlag != "" {
// Explicit new flag
- mode = balancer.ParseMode(balancerModeFlag)
+ gw.balancerMode = balancer.ParseMode(balancerModeFlag)
} else if balancerEnabled {
// Deprecated flag for backwards compatibility
log.Warning("Flag --enable-balancer is deprecated. Please use --vtgate-balancer-mode=prefer-cell instead.")
- mode = balancer.ModePreferCell
+ gw.balancerMode = balancer.ModePreferCell
} else {
// Default: no flags set
- mode = balancer.ModeCell
+ gw.balancerMode = balancer.ModeCell
}
// Cell mode uses the default shuffleTablets behavior, no balancer needed
- if mode == balancer.ModeCell {
+ if gw.balancerMode == balancer.ModeCell {
log.Info("Tablet balancer using 'cell' mode (shuffle tablets in local cell)")
return
}
// Validate mode-specific requirements
- if mode == balancer.ModePreferCell && len(balancerVtgateCells) == 0 {
+ if gw.balancerMode == balancer.ModePreferCell && len(balancerVtgateCells) == 0 {
log.Exitf("--balancer-vtgate-cells is required when using --vtgate-balancer-mode=prefer-cell")
}
// Create the balancer for prefer-cell or random modes
var err error
- gw.balancer, err = balancer.NewTabletBalancer(mode, gw.localCell, balancerVtgateCells)
+ gw.balancer, err = balancer.NewTabletBalancer(gw.balancerMode, gw.localCell, balancerVtgateCells)
if err != nil {
log.Exitf("Failed to create tablet balancer: %v", err)
}
- log.Infof("Tablet balancer enabled with mode: %s", mode)
+ log.Infof("Tablet balancer enabled with mode: %s", gw.balancerMode)
}
// QueryServiceByAlias satisfies the Gateway interface
@@ -322,9 +324,10 @@ func (gw *TabletGateway) DebugBalancerHandler(w http.ResponseWriter, r *http.Req
// withRetry also adds shard information to errors returned from the inner QueryService, so
// withShardError should not be combined with withRetry.
func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target, _ queryservice.QueryService,
- _ string, inTransaction bool, inner func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService) (bool, error)) error {
+ _ string, opts queryservice.WrapOpts, inner func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService) (bool, error),
+) error {
// for transactions, we connect to a specific tablet instead of letting gateway choose one
- if inTransaction && target.TabletType != topodatapb.TabletType_PRIMARY {
+ if opts.InTransaction && target.TabletType != topodatapb.TabletType_PRIMARY {
return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "tabletGateway's query service can only be used for non-transactional queries on replicas")
}
var tabletLastUsed *topodatapb.Tablet
@@ -350,7 +353,7 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target,
// Note: We only buffer once and only "!inTransaction" queries i.e.
// a) no transaction is necessary (e.g. critical reads) or
// b) no transaction was created yet.
- if gw.buffer != nil && !bufferedOnce && !inTransaction && target.TabletType == topodatapb.TabletType_PRIMARY {
+ if gw.buffer != nil && !bufferedOnce && !opts.InTransaction && target.TabletType == topodatapb.TabletType_PRIMARY {
// The next call blocks if we should buffer during a failover.
retryDone, bufferErr := gw.buffer.WaitForFailoverEnd(ctx, target.Keyspace, target.Shard, gw.kev, err)
@@ -401,36 +404,7 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target,
break
}
- var th *discovery.TabletHealth
-
- // Determine if we should use the balancer for this target
- useBalancer := gw.balancer != nil
- if useBalancer && len(balancerKeyspaces) > 0 {
- useBalancer = slices.Contains(balancerKeyspaces, target.Keyspace)
- }
-
- if useBalancer {
- // Filter out the tablets that we've tried before (if any)
- if len(invalidTablets) > 0 {
- tablets = slices.DeleteFunc(tablets, func(t *discovery.TabletHealth) bool {
- _, isInvalid := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]
- return isInvalid
- })
- }
-
- th = gw.balancer.Pick(target, tablets)
- } else {
- gw.shuffleTablets(gw.localCell, tablets)
-
- // skip tablets we tried before
- for _, t := range tablets {
- if _, ok := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]; !ok {
- th = t
- break
- }
- }
- }
-
+ th := gw.getBalancerTablet(target, tablets, invalidTablets, opts)
if th == nil {
// do not override error from last attempt.
if err == nil {
@@ -462,9 +436,59 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target,
return NewShardError(err, target)
}
+// getBalancerTablet selects a tablet for the given query target, using the configured balancer if enabled. Otherwise, it will
+// select a random tablet, with preference to the local cell.
+func (gw *TabletGateway) getBalancerTablet(target *querypb.Target, tablets []*discovery.TabletHealth, invalidTablets map[string]bool, opts queryservice.WrapOpts) *discovery.TabletHealth {
+ // Return early if no tablets are available
+ if len(tablets) == 0 {
+ return nil
+ }
+
+ // Filter out the tablets that we've tried before (if any)
+ if len(invalidTablets) > 0 {
+ tablets = slices.DeleteFunc(tablets, func(t *discovery.TabletHealth) bool {
+ _, isInvalid := invalidTablets[topoproto.TabletAliasString(t.Tablet.Alias)]
+ return isInvalid
+ })
+
+ // If all tablets are invalid, let's return early
+ if len(tablets) == 0 {
+ return nil
+ }
+ }
+
+ // Determine if we should use the balancer for this target
+ useBalancer := gw.balancer != nil
+ if useBalancer && len(balancerKeyspaces) > 0 {
+ useBalancer = slices.Contains(balancerKeyspaces, target.Keyspace)
+ }
+
+ // Get the tablet from the balancer if enabled
+ if useBalancer {
+ var pickOpts []balancer.PickOption
+
+ // Add the session UUID to the options if the session balancer is enabled and a session is present.
+ if gw.balancerMode == balancer.ModeSession && opts.Session != nil {
+ pickOpts = append(pickOpts, balancer.WithSessionUUID(opts.Session.GetSessionUUID()))
+ }
+
+ tablet := gw.balancer.Pick(target, tablets, pickOpts...)
+ if tablet != nil {
+ return tablet
+ }
+ }
+
+ // If the balancer isn't enabled, or it didn't return a tablet, shuffle the tablets
+ // and return the first one. (This will always contain at least one tablet due to the
+ // check above).
+ gw.shuffleTablets(gw.localCell, tablets)
+ return tablets[0]
+}
+
// withShardError adds shard information to errors returned from the inner QueryService.
func (gw *TabletGateway) withShardError(ctx context.Context, target *querypb.Target, conn queryservice.QueryService,
- _ string, _ bool, inner func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService) (bool, error)) error {
+ _ string, _ queryservice.WrapOpts, inner func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService) (bool, error),
+) error {
_, err := inner(ctx, target, conn)
return NewShardError(err, target)
}
diff --git a/go/vt/vtgate/tabletgateway_flaky_test.go b/go/vt/vtgate/tabletgateway_flaky_test.go
index e0b8c36b7f8..875ae9c18fb 100644
--- a/go/vt/vtgate/tabletgateway_flaky_test.go
+++ b/go/vt/vtgate/tabletgateway_flaky_test.go
@@ -94,7 +94,7 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) {
sbc.SetResults([]*sqltypes.Result{sqlResult1})
// run a query that we indeed get the result added to the sandbox connection back
- res, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ res, err := tg.Execute(ctx, nil, target, "query", nil, 0, 0)
require.NoError(t, err)
require.Equal(t, res, sqlResult1)
@@ -114,7 +114,7 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) {
// execute the query in a go routine since it should be buffered, and check that it eventually succeed
queryChan := make(chan struct{})
go func() {
- res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ res, err = tg.Execute(ctx, nil, target, "query", nil, 0, 0)
queryChan <- struct{}{}
}()
@@ -186,7 +186,7 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) {
// run a query that we indeed get the result added to the sandbox connection back
// this also checks that the query reaches the primary tablet and not the replica
- res, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ res, err := tg.Execute(ctx, nil, target, "query", nil, 0, 0)
require.NoError(t, err)
require.Equal(t, res, sqlResult1)
@@ -224,7 +224,7 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) {
// execute the query in a go routine since it should be buffered, and check that it eventually succeed
queryChan := make(chan struct{})
go func() {
- res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ res, err = tg.Execute(ctx, nil, target, "query", nil, 0, 0)
queryChan <- struct{}{}
}()
@@ -332,7 +332,7 @@ func TestInconsistentStateDetectedBuffering(t *testing.T) {
var err error
queryChan := make(chan struct{})
go func() {
- res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ res, err = tg.Execute(ctx, nil, target, "query", nil, 0, 0)
queryChan <- struct{}{}
}()
diff --git a/go/vt/vtgate/tabletgateway_test.go b/go/vt/vtgate/tabletgateway_test.go
index cb3b48cd9f4..c7a7eebd98e 100644
--- a/go/vt/vtgate/tabletgateway_test.go
+++ b/go/vt/vtgate/tabletgateway_test.go
@@ -43,14 +43,14 @@ import (
func TestTabletGatewayExecute(t *testing.T) {
ctx := utils.LeakCheckContext(t)
testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
- _, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil)
+ _, err := tg.Execute(ctx, nil, target, "query", nil, 0, 0)
return err
},
func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
assert.Equal(t, want, sc.ExecCount.Load())
})
testTabletGatewayTransact(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
- _, err := tg.Execute(ctx, target, "query", nil, 1, 0, nil)
+ _, err := tg.Execute(ctx, nil, target, "query", nil, 1, 0)
return err
})
}
@@ -58,7 +58,7 @@ func TestTabletGatewayExecute(t *testing.T) {
func TestTabletGatewayExecuteStream(t *testing.T) {
ctx := utils.LeakCheckContext(t)
testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
- err := tg.StreamExecute(ctx, target, "query", nil, 0, 0, nil, func(qr *sqltypes.Result) error {
+ err := tg.StreamExecute(ctx, nil, target, "query", nil, 0, 0, func(qr *sqltypes.Result) error {
return nil
})
return err
@@ -71,7 +71,7 @@ func TestTabletGatewayExecuteStream(t *testing.T) {
func TestTabletGatewayBegin(t *testing.T) {
ctx := utils.LeakCheckContext(t)
testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
- _, err := tg.Begin(ctx, target, nil)
+ _, err := tg.Begin(ctx, nil, target)
return err
},
func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
@@ -98,7 +98,7 @@ func TestTabletGatewayRollback(t *testing.T) {
func TestTabletGatewayBeginExecute(t *testing.T) {
ctx := utils.LeakCheckContext(t)
testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
- _, _, err := tg.BeginExecute(ctx, target, nil, "query", nil, 0, nil)
+ _, _, err := tg.BeginExecute(ctx, nil, target, nil, "query", nil, 0)
return err
},
func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
@@ -190,7 +190,7 @@ func TestTabletGatewayReplicaTransactionError(t *testing.T) {
defer tg.Close(ctx)
_ = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
- _, err := tg.Execute(ctx, target, "query", nil, 1, 0, nil)
+ _, err := tg.Execute(ctx, nil, target, "query", nil, 1, 0)
verifyContainsError(t, err, "query service can only be used for non-transactional queries on replicas", vtrpcpb.Code_INTERNAL)
}
@@ -392,7 +392,7 @@ func TestWithRetry(t *testing.T) {
}
for _, tt := range testcases {
t.Run(tt.name, func(t *testing.T) {
- err := tg.withRetry(ctx, tt.target, nil, "", tt.inTransaction, tt.inner)
+ err := tg.withRetry(ctx, tt.target, nil, "", queryservice.WrapOpts{InTransaction: tt.inTransaction}, tt.inner)
if tt.expectedErr == "" {
require.NoError(t, err)
} else {
diff --git a/go/vt/vttablet/endtoend/framework/client.go b/go/vt/vttablet/endtoend/framework/client.go
index 59def25ab61..94bb3ccd154 100644
--- a/go/vt/vttablet/endtoend/framework/client.go
+++ b/go/vt/vttablet/endtoend/framework/client.go
@@ -28,6 +28,7 @@ import (
querypb "vitess.io/vitess/go/vt/proto/query"
topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)
@@ -99,11 +100,9 @@ func (client *QueryClient) Begin(clientFoundRows bool) error {
if client.transactionID != 0 {
return errors.New("already in transaction")
}
- var options *querypb.ExecuteOptions
- if clientFoundRows {
- options = &querypb.ExecuteOptions{ClientFoundRows: clientFoundRows}
- }
- state, err := client.server.Begin(client.ctx, client.target, options)
+
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{ClientFoundRows: clientFoundRows}}
+ state, err := client.server.Begin(client.ctx, session, client.target)
if err != nil {
return err
}
@@ -202,14 +201,16 @@ func (client *QueryClient) BeginExecute(query string, bindvars map[string]*query
if client.transactionID != 0 {
return nil, errors.New("already in transaction")
}
+
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
state, qr, err := client.server.BeginExecute(
client.ctx,
+ session,
client.target,
preQueries,
query,
bindvars,
client.reservedID,
- &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL},
)
client.transactionID = state.TransactionID
client.sessionStateChanges = state.SessionStateChanges
@@ -221,14 +222,15 @@ func (client *QueryClient) BeginExecute(query string, bindvars map[string]*query
// ExecuteWithOptions executes a query using 'options'.
func (client *QueryClient) ExecuteWithOptions(query string, bindvars map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+ session := &vtgatepb.Session{Options: options}
return client.server.Execute(
client.ctx,
+ session,
client.target,
query,
bindvars,
client.transactionID,
client.reservedID,
- options,
)
}
@@ -240,13 +242,14 @@ func (client *QueryClient) StreamExecute(query string, bindvars map[string]*quer
// StreamExecuteWithOptions executes a query & returns the results using 'options'.
func (client *QueryClient) StreamExecuteWithOptions(query string, bindvars map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
result := &sqltypes.Result{}
+ session := &vtgatepb.Session{Options: options}
err := client.server.StreamExecute(client.ctx,
+ session,
client.target,
query,
bindvars,
client.transactionID,
client.reservedID,
- options,
func(res *sqltypes.Result) error {
if result.Fields == nil {
result.Fields = res.Fields
@@ -263,14 +266,15 @@ func (client *QueryClient) StreamExecuteWithOptions(query string, bindvars map[s
// StreamBeginExecuteWithOptions starts a tx and executes a query using 'options', returning the results .
func (client *QueryClient) StreamBeginExecuteWithOptions(query string, preQueries []string, bindvars map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
result := &sqltypes.Result{}
+ session := &vtgatepb.Session{Options: options}
state, err := client.server.BeginStreamExecute(
client.ctx,
+ session,
client.target,
preQueries,
query,
bindvars,
client.reservedID,
- options,
func(res *sqltypes.Result) error {
if result.Fields == nil {
result.Fields = res.Fields
@@ -289,7 +293,8 @@ func (client *QueryClient) StreamBeginExecuteWithOptions(query string, preQuerie
// Stream streams the results of a query.
func (client *QueryClient) Stream(query string, bindvars map[string]*querypb.BindVariable, sendFunc func(*sqltypes.Result) error) error {
- return client.server.StreamExecute(client.ctx, client.target, query, bindvars, 0, 0, &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}, sendFunc)
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
+ return client.server.StreamExecute(client.ctx, session, client.target, query, bindvars, 0, 0, sendFunc)
}
// MessageStream streams messages from the message table.
@@ -316,7 +321,8 @@ func (client *QueryClient) ReserveExecute(query string, preQueries []string, bin
if client.reservedID != 0 {
return nil, errors.New("already reserved a connection")
}
- state, qr, err := client.server.ReserveExecute(client.ctx, client.target, preQueries, query, bindvars, client.transactionID, &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL})
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
+ state, qr, err := client.server.ReserveExecute(client.ctx, session, client.target, preQueries, query, bindvars, client.transactionID)
client.reservedID = state.ReservedID
if err != nil {
return nil, err
@@ -330,7 +336,8 @@ func (client *QueryClient) ReserveStreamExecute(query string, preQueries []strin
return nil, errors.New("already reserved a connection")
}
result := &sqltypes.Result{}
- state, err := client.server.ReserveStreamExecute(client.ctx, client.target, preQueries, query, bindvars, client.transactionID, &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL},
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
+ state, err := client.server.ReserveStreamExecute(client.ctx, session, client.target, preQueries, query, bindvars, client.transactionID,
func(res *sqltypes.Result) error {
if result.Fields == nil {
result.Fields = res.Fields
@@ -353,7 +360,8 @@ func (client *QueryClient) ReserveBeginExecute(query string, preQueries []string
if client.transactionID != 0 {
return nil, errors.New("already in transaction")
}
- state, qr, err := client.server.ReserveBeginExecute(client.ctx, client.target, preQueries, postBeginQueries, query, bindvars, &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL})
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
+ state, qr, err := client.server.ReserveBeginExecute(client.ctx, session, client.target, preQueries, postBeginQueries, query, bindvars)
client.transactionID = state.TransactionID
client.reservedID = state.ReservedID
client.sessionStateChanges = state.SessionStateChanges
@@ -372,7 +380,8 @@ func (client *QueryClient) ReserveBeginStreamExecute(query string, preQueries []
return nil, errors.New("already in transaction")
}
result := &sqltypes.Result{}
- state, err := client.server.ReserveBeginStreamExecute(client.ctx, client.target, preQueries, postBeginQueries, query, bindvars, &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL},
+ session := &vtgatepb.Session{Options: &querypb.ExecuteOptions{IncludedFields: querypb.ExecuteOptions_ALL}}
+ state, err := client.server.ReserveBeginStreamExecute(client.ctx, session, client.target, preQueries, postBeginQueries, query, bindvars,
func(res *sqltypes.Result) error {
if result.Fields == nil {
result.Fields = res.Fields
diff --git a/go/vt/vttablet/grpcqueryservice/server.go b/go/vt/vttablet/grpcqueryservice/server.go
index e3c179ce856..ced6dcea0bd 100644
--- a/go/vt/vttablet/grpcqueryservice/server.go
+++ b/go/vt/vttablet/grpcqueryservice/server.go
@@ -25,11 +25,13 @@ import (
"vitess.io/vitess/go/vt/callerid"
"vitess.io/vitess/go/vt/callinfo"
"vitess.io/vitess/go/vt/vterrors"
+ "vitess.io/vitess/go/vt/vtgate/executorcontext"
"vitess.io/vitess/go/vt/vttablet/queryservice"
binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
querypb "vitess.io/vitess/go/vt/proto/query"
queryservicepb "vitess.io/vitess/go/vt/proto/queryservice"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
)
// query is the gRPC query service implementation.
@@ -48,7 +50,10 @@ func (q *query) Execute(ctx context.Context, request *querypb.ExecuteRequest) (r
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- result, err := q.server.Execute(ctx, request.Target, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.ReservedId, request.Options)
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ result, err := q.server.Execute(ctx, session, request.Target, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.ReservedId)
if err != nil {
return nil, vterrors.ToGRPC(err)
}
@@ -64,7 +69,10 @@ func (q *query) StreamExecute(request *querypb.StreamExecuteRequest, stream quer
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- err = q.server.StreamExecute(ctx, request.Target, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.ReservedId, request.Options, func(reply *sqltypes.Result) error {
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ err = q.server.StreamExecute(ctx, session, request.Target, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.ReservedId, func(reply *sqltypes.Result) error {
return stream.Send(&querypb.StreamExecuteResponse{
Result: sqltypes.ResultToProto3(reply),
})
@@ -79,7 +87,10 @@ func (q *query) Begin(ctx context.Context, request *querypb.BeginRequest) (respo
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, err := q.server.Begin(ctx, request.Target, request.Options)
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, err := q.server.Begin(ctx, session, request.Target)
if err != nil {
return nil, vterrors.ToGRPC(err)
}
@@ -252,7 +263,10 @@ func (q *query) BeginExecute(ctx context.Context, request *querypb.BeginExecuteR
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, result, err := q.server.BeginExecute(ctx, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.ReservedId, request.Options)
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, result, err := q.server.BeginExecute(ctx, session, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.ReservedId)
if err != nil {
// if we have a valid transactionID, return the error in-band
if state.TransactionID != 0 {
@@ -279,7 +293,10 @@ func (q *query) BeginStreamExecute(request *querypb.BeginStreamExecuteRequest, s
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, err := q.server.BeginStreamExecute(ctx, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.ReservedId, request.Options, func(reply *sqltypes.Result) error {
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, err := q.server.BeginStreamExecute(ctx, session, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.ReservedId, func(reply *sqltypes.Result) error {
return stream.Send(&querypb.BeginStreamExecuteResponse{
Result: sqltypes.ResultToProto3(reply),
})
@@ -392,7 +409,10 @@ func (q *query) ReserveExecute(ctx context.Context, request *querypb.ReserveExec
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, result, err := q.server.ReserveExecute(ctx, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.Options)
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, result, err := q.server.ReserveExecute(ctx, session, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.TransactionId)
if err != nil {
// if we have a valid reservedID, return the error in-band
if state.ReservedID != 0 {
@@ -418,7 +438,10 @@ func (q *query) ReserveStreamExecute(request *querypb.ReserveStreamExecuteReques
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, err := q.server.ReserveStreamExecute(ctx, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.TransactionId, request.Options, func(reply *sqltypes.Result) error {
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, err := q.server.ReserveStreamExecute(ctx, session, request.Target, request.PreQueries, request.Query.Sql, request.Query.BindVariables, request.TransactionId, func(reply *sqltypes.Result) error {
return stream.Send(&querypb.ReserveStreamExecuteResponse{
Result: sqltypes.ResultToProto3(reply),
})
@@ -442,7 +465,10 @@ func (q *query) ReserveBeginExecute(ctx context.Context, request *querypb.Reserv
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, result, err := q.server.ReserveBeginExecute(ctx, request.Target, request.PreQueries, request.PostBeginQueries, request.Query.Sql, request.Query.BindVariables, request.Options)
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, result, err := q.server.ReserveBeginExecute(ctx, session, request.Target, request.PreQueries, request.PostBeginQueries, request.Query.Sql, request.Query.BindVariables)
if err != nil {
// if we have a valid reservedID or transactionID, return the error in-band
if state.TransactionID != 0 || state.ReservedID != 0 {
@@ -472,7 +498,10 @@ func (q *query) ReserveBeginStreamExecute(request *querypb.ReserveBeginStreamExe
request.EffectiveCallerId,
request.ImmediateCallerId,
)
- state, err := q.server.ReserveBeginStreamExecute(ctx, request.Target, request.PreQueries, request.PostBeginQueries, request.Query.Sql, request.Query.BindVariables, request.Options, func(reply *sqltypes.Result) error {
+
+ session := executorcontext.NewSafeSession(&vtgatepb.Session{Options: request.Options})
+
+ state, err := q.server.ReserveBeginStreamExecute(ctx, session, request.Target, request.PreQueries, request.PostBeginQueries, request.Query.Sql, request.Query.BindVariables, func(reply *sqltypes.Result) error {
return stream.Send(&querypb.ReserveBeginStreamExecuteResponse{
Result: sqltypes.ResultToProto3(reply),
})
diff --git a/go/vt/vttablet/grpctabletconn/conn.go b/go/vt/vttablet/grpctabletconn/conn.go
index 0ef94031cf3..ce48d0d430e 100644
--- a/go/vt/vttablet/grpctabletconn/conn.go
+++ b/go/vt/vttablet/grpctabletconn/conn.go
@@ -112,7 +112,7 @@ func DialTablet(ctx context.Context, tablet *topodatapb.Tablet, failFast grpccli
}
// Execute sends the query to VTTablet.
-func (conn *gRPCQueryClient) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (conn *gRPCQueryClient) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
@@ -128,7 +128,7 @@ func (conn *gRPCQueryClient) Execute(ctx context.Context, target *querypb.Target
BindVariables: bindVars,
},
TransactionId: transactionID,
- Options: options,
+ Options: getOptions(session),
ReservedId: reservedID,
}
er, err := conn.c.Execute(ctx, req)
@@ -139,7 +139,7 @@ func (conn *gRPCQueryClient) Execute(ctx context.Context, target *querypb.Target
}
// StreamExecute executes the query and streams results back through callback.
-func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
+func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error {
// All streaming clients should follow the code pattern below.
// The first part of the function starts the stream while holding
// a lock on conn.mu. The second part receives the data and calls
@@ -166,7 +166,7 @@ func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, target *querypb.
Sql: query,
BindVariables: bindVars,
},
- Options: options,
+ Options: getOptions(session),
TransactionId: transactionID,
ReservedId: reservedID,
}
@@ -198,7 +198,7 @@ func (conn *gRPCQueryClient) StreamExecute(ctx context.Context, target *querypb.
}
// Begin starts a transaction.
-func (conn *gRPCQueryClient) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (state queryservice.TransactionState, err error) {
+func (conn *gRPCQueryClient) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (state queryservice.TransactionState, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
@@ -209,7 +209,7 @@ func (conn *gRPCQueryClient) Begin(ctx context.Context, target *querypb.Target,
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
- Options: options,
+ Options: getOptions(session),
}
br, err := conn.c.Begin(ctx, req)
if err != nil {
@@ -463,7 +463,7 @@ func (conn *gRPCQueryClient) UnresolvedTransactions(ctx context.Context, target
}
// BeginExecute starts a transaction and runs an Execute.
-func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (state queryservice.TransactionState, result *sqltypes.Result, err error) {
+func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64) (state queryservice.TransactionState, result *sqltypes.Result, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
@@ -480,7 +480,7 @@ func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, target *querypb.T
BindVariables: bindVars,
},
ReservedId: reservedID,
- Options: options,
+ Options: getOptions(session),
}
reply, err := conn.c.BeginExecute(ctx, req)
if err != nil {
@@ -496,7 +496,7 @@ func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, target *querypb.T
}
// BeginStreamExecute starts a transaction and runs an Execute.
-func (conn *gRPCQueryClient) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state queryservice.TransactionState, err error) {
+func (conn *gRPCQueryClient) BeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (state queryservice.TransactionState, err error) {
// Please see comments in StreamExecute to see how this works.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -524,7 +524,7 @@ func (conn *gRPCQueryClient) BeginStreamExecute(ctx context.Context, target *que
BindVariables: bindVars,
},
ReservedId: reservedID,
- Options: options,
+ Options: getOptions(session),
}
stream, err := conn.c.BeginStreamExecute(ctx, req)
if err != nil {
@@ -861,7 +861,7 @@ func (conn *gRPCQueryClient) HandlePanic(err *error) {
}
// ReserveBeginExecute implements the queryservice interface
-func (conn *gRPCQueryClient) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (state queryservice.ReservedTransactionState, result *sqltypes.Result, err error) {
+func (conn *gRPCQueryClient) ReserveBeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (state queryservice.ReservedTransactionState, result *sqltypes.Result, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
@@ -872,7 +872,7 @@ func (conn *gRPCQueryClient) ReserveBeginExecute(ctx context.Context, target *qu
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
- Options: options,
+ Options: getOptions(session),
PreQueries: preQueries,
PostBeginQueries: postBeginQueries,
Query: &querypb.BoundQuery{
@@ -896,7 +896,7 @@ func (conn *gRPCQueryClient) ReserveBeginExecute(ctx context.Context, target *qu
}
// ReserveBeginStreamExecute implements the queryservice interface
-func (conn *gRPCQueryClient) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state queryservice.ReservedTransactionState, err error) {
+func (conn *gRPCQueryClient) ReserveBeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (state queryservice.ReservedTransactionState, err error) {
// Please see comments in StreamExecute to see how this works.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -917,7 +917,7 @@ func (conn *gRPCQueryClient) ReserveBeginStreamExecute(ctx context.Context, targ
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
- Options: options,
+ Options: getOptions(session),
PreQueries: preQueries,
PostBeginQueries: postBeginQueries,
Query: &querypb.BoundQuery{
@@ -977,7 +977,7 @@ func (conn *gRPCQueryClient) ReserveBeginStreamExecute(ctx context.Context, targ
}
// ReserveExecute implements the queryservice interface
-func (conn *gRPCQueryClient) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (state queryservice.ReservedState, result *sqltypes.Result, err error) {
+func (conn *gRPCQueryClient) ReserveExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (state queryservice.ReservedState, result *sqltypes.Result, err error) {
conn.mu.RLock()
defer conn.mu.RUnlock()
if conn.cc == nil {
@@ -993,7 +993,7 @@ func (conn *gRPCQueryClient) ReserveExecute(ctx context.Context, target *querypb
BindVariables: bindVariables,
},
TransactionId: transactionID,
- Options: options,
+ Options: getOptions(session),
PreQueries: preQueries,
}
reply, err := conn.c.ReserveExecute(ctx, req)
@@ -1010,7 +1010,7 @@ func (conn *gRPCQueryClient) ReserveExecute(ctx context.Context, target *querypb
}
// ReserveStreamExecute implements the queryservice interface
-func (conn *gRPCQueryClient) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state queryservice.ReservedState, err error) {
+func (conn *gRPCQueryClient) ReserveStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (state queryservice.ReservedState, err error) {
// Please see comments in StreamExecute to see how this works.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -1031,7 +1031,7 @@ func (conn *gRPCQueryClient) ReserveStreamExecute(ctx context.Context, target *q
Target: target,
EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
- Options: options,
+ Options: getOptions(session),
PreQueries: preQueries,
Query: &querypb.BoundQuery{
Sql: sql,
@@ -1167,3 +1167,12 @@ func (conn *gRPCQueryClient) Close(ctx context.Context) error {
func (conn *gRPCQueryClient) Tablet() *topodatapb.Tablet {
return conn.tablet
}
+
+// getOptions safely extracts ExecuteOptions from a session, returning nil if session is nil.
+func getOptions(session queryservice.Session) *querypb.ExecuteOptions {
+ if session == nil {
+ return nil
+ }
+
+ return session.GetOptions()
+}
diff --git a/go/vt/vttablet/grpctabletconn/conn_test.go b/go/vt/vttablet/grpctabletconn/conn_test.go
index e20cf36a797..a7da2194d11 100644
--- a/go/vt/vttablet/grpctabletconn/conn_test.go
+++ b/go/vt/vttablet/grpctabletconn/conn_test.go
@@ -191,22 +191,22 @@ func TestGoRoutineLeakPrevention(t *testing.T) {
cc: &grpc.ClientConn{},
c: mqc,
}
- _ = qc.StreamExecute(context.Background(), nil, "", nil, 0, 0, nil, func(result *sqltypes.Result) error {
+ _ = qc.StreamExecute(context.Background(), nil, nil, "", nil, 0, 0, func(result *sqltypes.Result) error {
return nil
})
require.Error(t, mqc.lastCallCtx.Err())
- _, _ = qc.BeginStreamExecute(context.Background(), nil, nil, "", nil, 0, nil, func(result *sqltypes.Result) error {
+ _, _ = qc.BeginStreamExecute(context.Background(), nil, nil, nil, "", nil, 0, func(result *sqltypes.Result) error {
return nil
})
require.Error(t, mqc.lastCallCtx.Err())
- _, _ = qc.ReserveBeginStreamExecute(context.Background(), nil, nil, nil, "", nil, nil, func(result *sqltypes.Result) error {
+ _, _ = qc.ReserveBeginStreamExecute(context.Background(), nil, nil, nil, nil, "", nil, func(result *sqltypes.Result) error {
return nil
})
require.Error(t, mqc.lastCallCtx.Err())
- _, _ = qc.ReserveStreamExecute(context.Background(), nil, nil, "", nil, 0, nil, func(result *sqltypes.Result) error {
+ _, _ = qc.ReserveStreamExecute(context.Background(), nil, nil, nil, "", nil, 0, func(result *sqltypes.Result) error {
return nil
})
require.Error(t, mqc.lastCallCtx.Err())
diff --git a/go/vt/vttablet/queryservice/fakes/error_query_service.go b/go/vt/vttablet/queryservice/fakes/error_query_service.go
index 93db80ee2a8..bfb3c2ff5b8 100644
--- a/go/vt/vttablet/queryservice/fakes/error_query_service.go
+++ b/go/vt/vttablet/queryservice/fakes/error_query_service.go
@@ -28,7 +28,7 @@ import (
// ErrorQueryService is an object that returns an error for all methods.
var ErrorQueryService = queryservice.Wrap(
nil,
- func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, inTransaction bool, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {
+ func(ctx context.Context, target *querypb.Target, conn queryservice.QueryService, name string, opts queryservice.WrapOpts, inner func(context.Context, *querypb.Target, queryservice.QueryService) (bool, error)) error {
return errors.New("ErrorQueryService does not implement any method")
},
)
diff --git a/go/vt/vttablet/queryservice/fakes/stream_health_query_service.go b/go/vt/vttablet/queryservice/fakes/stream_health_query_service.go
index e992c12baca..7d259503eee 100644
--- a/go/vt/vttablet/queryservice/fakes/stream_health_query_service.go
+++ b/go/vt/vttablet/queryservice/fakes/stream_health_query_service.go
@@ -56,12 +56,12 @@ func NewStreamHealthQueryService(target *querypb.Target) *StreamHealthQueryServi
}
// Begin implemented as a no op
-func (q *StreamHealthQueryService) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (q *StreamHealthQueryService) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (queryservice.TransactionState, error) {
return queryservice.TransactionState{}, nil
}
// Execute implemented as a no op
-func (q *StreamHealthQueryService) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (q *StreamHealthQueryService) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
return &sqltypes.Result{}, nil
}
diff --git a/go/vt/vttablet/queryservice/queryservice.go b/go/vt/vttablet/queryservice/queryservice.go
index d6972bfb6a3..75d09b1c05f 100644
--- a/go/vt/vttablet/queryservice/queryservice.go
+++ b/go/vt/vttablet/queryservice/queryservice.go
@@ -19,16 +19,25 @@ limitations under the License.
package queryservice
import (
- topodatapb "vitess.io/vitess/go/vt/proto/topodata"
-
"context"
+ topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+
"vitess.io/vitess/go/sqltypes"
binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
querypb "vitess.io/vitess/go/vt/proto/query"
)
+// Session represents the current session.
+type Session interface {
+ // GetSessionUUID returns the session's UUID.
+ GetSessionUUID() string
+
+ // GetOptions returns the current execute options.
+ GetOptions() *querypb.ExecuteOptions
+}
+
// QueryService is the interface implemented by the tablet's query service.
// All streaming methods accept a callback function that will be called for
// each response. If the callback returns an error, that error is returned
@@ -41,7 +50,7 @@ type QueryService interface {
// Transaction management
// Begin returns the transaction id to use for further operations
- Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (TransactionState, error)
+ Begin(ctx context.Context, session Session, target *querypb.Target) (TransactionState, error)
// Commit commits the current transaction
Commit(ctx context.Context, target *querypb.Target, transactionID int64) (int64, error)
@@ -80,16 +89,16 @@ type QueryService interface {
UnresolvedTransactions(ctx context.Context, target *querypb.Target, abandonAgeSeconds int64) ([]*querypb.TransactionMetadata, error)
// Execute for query execution
- Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error)
+ Execute(ctx context.Context, session Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error)
// StreamExecute for query execution with streaming
- StreamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error
+ StreamExecute(ctx context.Context, session Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error
// Combo methods, they also return the transactionID from the
// Begin part. If err != nil, the transactionID may still be
// non-zero, and needs to be propagated back (like for a DB
// Integrity Error)
- BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (TransactionState, *sqltypes.Result, error)
- BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (TransactionState, error)
+ BeginExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64) (TransactionState, *sqltypes.Result, error)
+ BeginStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (TransactionState, error)
// Messaging methods.
MessageStream(ctx context.Context, target *querypb.Target, name string, callback func(*sqltypes.Result) error) error
@@ -114,13 +123,13 @@ type QueryService interface {
// HandlePanic will be called if any of the functions panic.
HandlePanic(err *error)
- ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (ReservedTransactionState, *sqltypes.Result, error)
+ ReserveBeginExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (ReservedTransactionState, *sqltypes.Result, error)
- ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (ReservedTransactionState, error)
+ ReserveBeginStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (ReservedTransactionState, error)
- ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (ReservedState, *sqltypes.Result, error)
+ ReserveExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (ReservedState, *sqltypes.Result, error)
- ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (ReservedState, error)
+ ReserveStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (ReservedState, error)
Release(ctx context.Context, target *querypb.Target, transactionID, reservedID int64) error
diff --git a/go/vt/vttablet/queryservice/wrapped.go b/go/vt/vttablet/queryservice/wrapped.go
index 2e31c66dba2..4534170d399 100644
--- a/go/vt/vttablet/queryservice/wrapped.go
+++ b/go/vt/vttablet/queryservice/wrapped.go
@@ -34,7 +34,14 @@ var _ QueryService = &wrappedService{}
// The inner function returns err and canRetry.
// If canRetry is true, the error is specific to the current vttablet and can be retried elsewhere.
// The flag will be false if there was no error.
-type WrapperFunc func(ctx context.Context, target *querypb.Target, conn QueryService, name string, inTransaction bool, inner func(context.Context, *querypb.Target, QueryService) (canRetry bool, err error)) error
+type WrapperFunc func(ctx context.Context, target *querypb.Target, conn QueryService, name string, opts WrapOpts, inner func(context.Context, *querypb.Target, QueryService) (canRetry bool, err error)) error
+
+// WrapOpts is the options passed to [WrapperFunc].
+type WrapOpts struct {
+ InTransaction bool
+
+ Session Session
+}
// Wrap returns a wrapped version of the original QueryService implementation.
// This lets you avoid repeating boiler-plate code by consolidating it in the
@@ -109,10 +116,11 @@ type wrappedService struct {
wrapper WrapperFunc
}
-func (ws *wrappedService) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (state TransactionState, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "Begin", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+func (ws *wrappedService) Begin(ctx context.Context, session Session, target *querypb.Target) (state TransactionState, err error) {
+ opts := WrapOpts{InTransaction: false, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "Begin", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- state, innerErr = conn.Begin(ctx, target, options)
+ state, innerErr = conn.Begin(ctx, session, target)
return canRetry(ctx, innerErr), innerErr
})
return state, wrapFatalTxErrorInVTError(err, true, vterrors.VT15001)
@@ -120,7 +128,8 @@ func (ws *wrappedService) Begin(ctx context.Context, target *querypb.Target, opt
func (ws *wrappedService) Commit(ctx context.Context, target *querypb.Target, transactionID int64) (int64, error) {
var rID int64
- err := ws.wrapper(ctx, target, ws.impl, "Commit", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err := ws.wrapper(ctx, target, ws.impl, "Commit", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
rID, innerErr = conn.Commit(ctx, target, transactionID)
return canRetry(ctx, innerErr), innerErr
@@ -133,7 +142,8 @@ func (ws *wrappedService) Commit(ctx context.Context, target *querypb.Target, tr
func (ws *wrappedService) Rollback(ctx context.Context, target *querypb.Target, transactionID int64) (int64, error) {
var rID int64
- err := ws.wrapper(ctx, target, ws.impl, "Rollback", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err := ws.wrapper(ctx, target, ws.impl, "Rollback", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
rID, innerErr = conn.Rollback(ctx, target, transactionID)
return canRetry(ctx, innerErr), innerErr
@@ -145,7 +155,8 @@ func (ws *wrappedService) Rollback(ctx context.Context, target *querypb.Target,
}
func (ws *wrappedService) Prepare(ctx context.Context, target *querypb.Target, transactionID int64, dtid string) error {
- err := ws.wrapper(ctx, target, ws.impl, "Prepare", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err := ws.wrapper(ctx, target, ws.impl, "Prepare", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.Prepare(ctx, target, transactionID, dtid)
return canRetry(ctx, innerErr), innerErr
})
@@ -153,7 +164,8 @@ func (ws *wrappedService) Prepare(ctx context.Context, target *querypb.Target, t
}
func (ws *wrappedService) CommitPrepared(ctx context.Context, target *querypb.Target, dtid string) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "CommitPrepared", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "CommitPrepared", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.CommitPrepared(ctx, target, dtid)
return canRetry(ctx, innerErr), innerErr
})
@@ -161,7 +173,8 @@ func (ws *wrappedService) CommitPrepared(ctx context.Context, target *querypb.Ta
}
func (ws *wrappedService) RollbackPrepared(ctx context.Context, target *querypb.Target, dtid string, originalID int64) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "RollbackPrepared", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "RollbackPrepared", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.RollbackPrepared(ctx, target, dtid, originalID)
return canRetry(ctx, innerErr), innerErr
})
@@ -169,7 +182,8 @@ func (ws *wrappedService) RollbackPrepared(ctx context.Context, target *querypb.
}
func (ws *wrappedService) CreateTransaction(ctx context.Context, target *querypb.Target, dtid string, participants []*querypb.Target) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "CreateTransaction", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "CreateTransaction", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.CreateTransaction(ctx, target, dtid, participants)
return canRetry(ctx, innerErr), innerErr
})
@@ -177,7 +191,8 @@ func (ws *wrappedService) CreateTransaction(ctx context.Context, target *querypb
}
func (ws *wrappedService) StartCommit(ctx context.Context, target *querypb.Target, transactionID int64, dtid string) (state querypb.StartCommitState, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "StartCommit", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "StartCommit", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
state, innerErr = conn.StartCommit(ctx, target, transactionID, dtid)
return canRetry(ctx, innerErr), innerErr
@@ -186,7 +201,8 @@ func (ws *wrappedService) StartCommit(ctx context.Context, target *querypb.Targe
}
func (ws *wrappedService) SetRollback(ctx context.Context, target *querypb.Target, dtid string, transactionID int64) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "SetRollback", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "SetRollback", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.SetRollback(ctx, target, dtid, transactionID)
return canRetry(ctx, innerErr), innerErr
})
@@ -194,7 +210,8 @@ func (ws *wrappedService) SetRollback(ctx context.Context, target *querypb.Targe
}
func (ws *wrappedService) ConcludeTransaction(ctx context.Context, target *querypb.Target, dtid string) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "ConcludeTransaction", true, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: true}
+ err = ws.wrapper(ctx, target, ws.impl, "ConcludeTransaction", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.ConcludeTransaction(ctx, target, dtid)
return canRetry(ctx, innerErr), innerErr
})
@@ -202,7 +219,8 @@ func (ws *wrappedService) ConcludeTransaction(ctx context.Context, target *query
}
func (ws *wrappedService) ReadTransaction(ctx context.Context, target *querypb.Target, dtid string) (metadata *querypb.TransactionMetadata, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "ReadTransaction", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ err = ws.wrapper(ctx, target, ws.impl, "ReadTransaction", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
metadata, innerErr = conn.ReadTransaction(ctx, target, dtid)
return canRetry(ctx, innerErr), innerErr
@@ -211,7 +229,8 @@ func (ws *wrappedService) ReadTransaction(ctx context.Context, target *querypb.T
}
func (ws *wrappedService) UnresolvedTransactions(ctx context.Context, target *querypb.Target, abandonAgeSeconds int64) (transactions []*querypb.TransactionMetadata, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "UnresolvedTransactions", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ err = ws.wrapper(ctx, target, ws.impl, "UnresolvedTransactions", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
transactions, innerErr = conn.UnresolvedTransactions(ctx, target, abandonAgeSeconds)
return canRetry(ctx, innerErr), innerErr
@@ -219,11 +238,12 @@ func (ws *wrappedService) UnresolvedTransactions(ctx context.Context, target *qu
return transactions, err
}
-func (ws *wrappedService) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (qr *sqltypes.Result, err error) {
+func (ws *wrappedService) Execute(ctx context.Context, session Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64) (qr *sqltypes.Result, err error) {
inDedicatedConn := transactionID != 0 || reservedID != 0
- err = ws.wrapper(ctx, target, ws.impl, "Execute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "Execute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- qr, innerErr = conn.Execute(ctx, target, query, bindVars, transactionID, reservedID, options)
+ qr, innerErr = conn.Execute(ctx, session, target, query, bindVars, transactionID, reservedID)
// You cannot retry if you're in a transaction.
retryable := canRetry(ctx, innerErr) && (!inDedicatedConn)
return retryable, innerErr
@@ -232,11 +252,12 @@ func (ws *wrappedService) Execute(ctx context.Context, target *querypb.Target, q
}
// StreamExecute implements the QueryService interface
-func (ws *wrappedService) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
+func (ws *wrappedService) StreamExecute(ctx context.Context, session Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error {
inDedicatedConn := transactionID != 0 || reservedID != 0
- err := ws.wrapper(ctx, target, ws.impl, "StreamExecute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err := ws.wrapper(ctx, target, ws.impl, "StreamExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
streamingStarted := false
- innerErr := conn.StreamExecute(ctx, target, query, bindVars, transactionID, reservedID, options, func(qr *sqltypes.Result) error {
+ innerErr := conn.StreamExecute(ctx, session, target, query, bindVars, transactionID, reservedID, func(qr *sqltypes.Result) error {
streamingStarted = true
return callback(qr)
})
@@ -247,36 +268,40 @@ func (ws *wrappedService) StreamExecute(ctx context.Context, target *querypb.Tar
return wrapFatalTxErrorInVTError(err, transactionID != 0, vterrors.VT15001)
}
-func (ws *wrappedService) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (state TransactionState, qr *sqltypes.Result, err error) {
+func (ws *wrappedService) BeginExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64) (state TransactionState, qr *sqltypes.Result, err error) {
inDedicatedConn := reservedID != 0
- err = ws.wrapper(ctx, target, ws.impl, "BeginExecute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "BeginExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- state, qr, innerErr = conn.BeginExecute(ctx, target, preQueries, query, bindVars, reservedID, options)
+ state, qr, innerErr = conn.BeginExecute(ctx, session, target, preQueries, query, bindVars, reservedID)
return canRetry(ctx, innerErr) && !inDedicatedConn, innerErr
})
return state, qr, wrapFatalTxErrorInVTError(err, true, vterrors.VT15001)
}
// BeginStreamExecute implements the QueryService interface
-func (ws *wrappedService) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state TransactionState, err error) {
+func (ws *wrappedService) BeginStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (state TransactionState, err error) {
inDedicatedConn := reservedID != 0
- err = ws.wrapper(ctx, target, ws.impl, "BeginStreamExecute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "BeginStreamExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- state, innerErr = conn.BeginStreamExecute(ctx, target, preQueries, query, bindVars, reservedID, options, callback)
+ state, innerErr = conn.BeginStreamExecute(ctx, session, target, preQueries, query, bindVars, reservedID, callback)
return canRetry(ctx, innerErr) && !inDedicatedConn, innerErr
})
return state, wrapFatalTxErrorInVTError(err, true, vterrors.VT15001)
}
func (ws *wrappedService) MessageStream(ctx context.Context, target *querypb.Target, name string, callback func(*sqltypes.Result) error) error {
- return ws.wrapper(ctx, target, ws.impl, "MessageStream", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, target, ws.impl, "MessageStream", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.MessageStream(ctx, target, name, callback)
return canRetry(ctx, innerErr), innerErr
})
}
func (ws *wrappedService) MessageAck(ctx context.Context, target *querypb.Target, name string, ids []*querypb.Value) (count int64, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "MessageAck", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ err = ws.wrapper(ctx, target, ws.impl, "MessageAck", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
count, innerErr = conn.MessageAck(ctx, target, name, ids)
return canRetry(ctx, innerErr), innerErr
@@ -285,35 +310,40 @@ func (ws *wrappedService) MessageAck(ctx context.Context, target *querypb.Target
}
func (ws *wrappedService) VStream(ctx context.Context, request *binlogdatapb.VStreamRequest, send func([]*binlogdatapb.VEvent) error) error {
- return ws.wrapper(ctx, request.Target, ws.impl, "VStream", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, request.Target, ws.impl, "VStream", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.VStream(ctx, request, send)
return false, innerErr
})
}
func (ws *wrappedService) VStreamRows(ctx context.Context, request *binlogdatapb.VStreamRowsRequest, send func(*binlogdatapb.VStreamRowsResponse) error) error {
- return ws.wrapper(ctx, request.Target, ws.impl, "VStreamRows", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, request.Target, ws.impl, "VStreamRows", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.VStreamRows(ctx, request, send)
return false, innerErr
})
}
func (ws *wrappedService) VStreamTables(ctx context.Context, request *binlogdatapb.VStreamTablesRequest, send func(response *binlogdatapb.VStreamTablesResponse) error) error {
- return ws.wrapper(ctx, request.Target, ws.impl, "VStreamTables", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, request.Target, ws.impl, "VStreamTables", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.VStreamTables(ctx, request, send)
return false, innerErr
})
}
func (ws *wrappedService) VStreamResults(ctx context.Context, target *querypb.Target, query string, send func(*binlogdatapb.VStreamResultsResponse) error) error {
- return ws.wrapper(ctx, target, ws.impl, "VStreamResults", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, target, ws.impl, "VStreamResults", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.VStreamResults(ctx, target, query, send)
return false, innerErr
})
}
func (ws *wrappedService) StreamHealth(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
- return ws.wrapper(ctx, nil, ws.impl, "StreamHealth", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, nil, ws.impl, "StreamHealth", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.StreamHealth(ctx, callback)
return canRetry(ctx, innerErr), innerErr
})
@@ -324,10 +354,11 @@ func (ws *wrappedService) HandlePanic(err *error) {
}
// ReserveBeginExecute implements the QueryService interface
-func (ws *wrappedService) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (state ReservedTransactionState, res *sqltypes.Result, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "ReserveBeginExecute", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+func (ws *wrappedService) ReserveBeginExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (state ReservedTransactionState, res *sqltypes.Result, err error) {
+ opts := WrapOpts{InTransaction: false, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "ReserveBeginExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var err error
- state, res, err = conn.ReserveBeginExecute(ctx, target, preQueries, postBeginQueries, sql, bindVariables, options)
+ state, res, err = conn.ReserveBeginExecute(ctx, session, target, preQueries, postBeginQueries, sql, bindVariables)
return canRetry(ctx, err), err
})
@@ -335,21 +366,23 @@ func (ws *wrappedService) ReserveBeginExecute(ctx context.Context, target *query
}
// ReserveBeginStreamExecute implements the QueryService interface
-func (ws *wrappedService) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state ReservedTransactionState, err error) {
- err = ws.wrapper(ctx, target, ws.impl, "ReserveBeginStreamExecute", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+func (ws *wrappedService) ReserveBeginStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (state ReservedTransactionState, err error) {
+ opts := WrapOpts{InTransaction: false, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "ReserveBeginStreamExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- state, innerErr = conn.ReserveBeginStreamExecute(ctx, target, preQueries, postBeginQueries, sql, bindVariables, options, callback)
+ state, innerErr = conn.ReserveBeginStreamExecute(ctx, session, target, preQueries, postBeginQueries, sql, bindVariables, callback)
return canRetry(ctx, innerErr), innerErr
})
return state, err
}
// ReserveExecute implements the QueryService interface
-func (ws *wrappedService) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (state ReservedState, res *sqltypes.Result, err error) {
+func (ws *wrappedService) ReserveExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (state ReservedState, res *sqltypes.Result, err error) {
inDedicatedConn := transactionID != 0
- err = ws.wrapper(ctx, target, ws.impl, "ReserveExecute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "ReserveExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var err error
- state, res, err = conn.ReserveExecute(ctx, target, preQueries, sql, bindVariables, transactionID, options)
+ state, res, err = conn.ReserveExecute(ctx, session, target, preQueries, sql, bindVariables, transactionID)
return canRetry(ctx, err) && !inDedicatedConn, err
})
@@ -357,11 +390,12 @@ func (ws *wrappedService) ReserveExecute(ctx context.Context, target *querypb.Ta
}
// ReserveStreamExecute implements the QueryService interface
-func (ws *wrappedService) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (state ReservedState, err error) {
+func (ws *wrappedService) ReserveStreamExecute(ctx context.Context, session Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (state ReservedState, err error) {
inDedicatedConn := transactionID != 0
- err = ws.wrapper(ctx, target, ws.impl, "ReserveStreamExecute", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn, Session: session}
+ err = ws.wrapper(ctx, target, ws.impl, "ReserveStreamExecute", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
var innerErr error
- state, innerErr = conn.ReserveStreamExecute(ctx, target, preQueries, sql, bindVariables, transactionID, options, callback)
+ state, innerErr = conn.ReserveStreamExecute(ctx, session, target, preQueries, sql, bindVariables, transactionID, callback)
return canRetry(ctx, innerErr) && !inDedicatedConn, innerErr
})
return state, err
@@ -369,14 +403,16 @@ func (ws *wrappedService) ReserveStreamExecute(ctx context.Context, target *quer
func (ws *wrappedService) Release(ctx context.Context, target *querypb.Target, transactionID, reservedID int64) error {
inDedicatedConn := transactionID != 0 || reservedID != 0
- return ws.wrapper(ctx, target, ws.impl, "Release", inDedicatedConn, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: inDedicatedConn}
+ return ws.wrapper(ctx, target, ws.impl, "Release", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
// No point retrying Release.
return false, conn.Release(ctx, target, transactionID, reservedID)
})
}
func (ws *wrappedService) GetSchema(ctx context.Context, target *querypb.Target, tableType querypb.SchemaTableType, tableNames []string, callback func(schemaRes *querypb.GetSchemaResponse) error) (err error) {
- err = ws.wrapper(ctx, target, ws.impl, "GetSchema", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ err = ws.wrapper(ctx, target, ws.impl, "GetSchema", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
innerErr := conn.GetSchema(ctx, target, tableType, tableNames, callback)
return canRetry(ctx, innerErr), innerErr
})
@@ -384,7 +420,8 @@ func (ws *wrappedService) GetSchema(ctx context.Context, target *querypb.Target,
}
func (ws *wrappedService) Close(ctx context.Context) error {
- return ws.wrapper(ctx, nil, ws.impl, "Close", false, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
+ opts := WrapOpts{InTransaction: false}
+ return ws.wrapper(ctx, nil, ws.impl, "Close", opts, func(ctx context.Context, target *querypb.Target, conn QueryService) (bool, error) {
// No point retrying Close.
return false, conn.Close(ctx)
})
diff --git a/go/vt/vttablet/sandboxconn/sandboxconn.go b/go/vt/vttablet/sandboxconn/sandboxconn.go
index e6a13199204..92f96f114de 100644
--- a/go/vt/vttablet/sandboxconn/sandboxconn.go
+++ b/go/vt/vttablet/sandboxconn/sandboxconn.go
@@ -264,7 +264,7 @@ func (sbc *SandboxConn) SetSchemaResult(r []SchemaResult) {
}
// Execute is part of the QueryService interface.
-func (sbc *SandboxConn) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (sbc *SandboxConn) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
sbc.panicIfNeeded()
sbc.execMu.Lock()
defer sbc.execMu.Unlock()
@@ -283,7 +283,7 @@ func (sbc *SandboxConn) Execute(ctx context.Context, target *querypb.Target, que
Sql: query,
BindVariables: bv,
})
- sbc.Options = append(sbc.Options, options)
+ sbc.Options = append(sbc.Options, getOptions(session))
if err := sbc.getError(); err != nil {
return nil, err
}
@@ -297,7 +297,7 @@ func (sbc *SandboxConn) Execute(ctx context.Context, target *querypb.Target, que
}
// StreamExecute is part of the QueryService interface.
-func (sbc *SandboxConn) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
+func (sbc *SandboxConn) StreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error {
sbc.panicIfNeeded()
sbc.sExecMu.Lock()
sbc.ExecCount.Add(1)
@@ -309,7 +309,7 @@ func (sbc *SandboxConn) StreamExecute(ctx context.Context, target *querypb.Targe
Sql: query,
BindVariables: bv,
})
- sbc.Options = append(sbc.Options, options)
+ sbc.Options = append(sbc.Options, getOptions(session))
err := sbc.getError()
if err != nil {
sbc.sExecMu.Unlock()
@@ -338,12 +338,12 @@ func (sbc *SandboxConn) StreamExecute(ctx context.Context, target *querypb.Targe
}
// Begin is part of the QueryService interface.
-func (sbc *SandboxConn) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (sbc *SandboxConn) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (queryservice.TransactionState, error) {
sbc.panicIfNeeded()
- return sbc.begin(ctx, target, nil, 0, options)
+ return sbc.begin(ctx, session, target, nil, 0)
}
-func (sbc *SandboxConn) begin(ctx context.Context, target *querypb.Target, preQueries []string, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (sbc *SandboxConn) begin(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, reservedID int64) (queryservice.TransactionState, error) {
sbc.BeginCount.Add(1)
err := sbc.getError()
if err != nil {
@@ -355,7 +355,7 @@ func (sbc *SandboxConn) begin(ctx context.Context, target *querypb.Target, preQu
transactionID = sbc.TransactionID.Add(1)
}
for _, preQuery := range preQueries {
- _, err := sbc.Execute(ctx, target, preQuery, nil, transactionID, reservedID, options)
+ _, err := sbc.Execute(ctx, session, target, preQuery, nil, transactionID, reservedID)
if err != nil {
return queryservice.TransactionState{}, err
}
@@ -494,30 +494,30 @@ func (sbc *SandboxConn) UnresolvedTransactions(context.Context, *querypb.Target,
}
// BeginExecute is part of the QueryService interface.
-func (sbc *SandboxConn) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
+func (sbc *SandboxConn) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64) (queryservice.TransactionState, *sqltypes.Result, error) {
sbc.panicIfNeeded()
- state, err := sbc.begin(ctx, target, preQueries, reservedID, options)
+ state, err := sbc.begin(ctx, session, target, preQueries, reservedID)
if state.TransactionID != 0 {
sbc.setTxReservedID(state.TransactionID, reservedID)
}
if err != nil {
return queryservice.TransactionState{}, nil, err
}
- result, err := sbc.Execute(ctx, target, query, bindVars, state.TransactionID, reservedID, options)
+ result, err := sbc.Execute(ctx, session, target, query, bindVars, state.TransactionID, reservedID)
return state, result, err
}
// BeginStreamExecute is part of the QueryService interface.
-func (sbc *SandboxConn) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
+func (sbc *SandboxConn) BeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
sbc.panicIfNeeded()
- state, err := sbc.begin(ctx, target, preQueries, reservedID, options)
+ state, err := sbc.begin(ctx, session, target, preQueries, reservedID)
if state.TransactionID != 0 {
sbc.setTxReservedID(state.TransactionID, reservedID)
}
if err != nil {
return queryservice.TransactionState{}, err
}
- err = sbc.StreamExecute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options, callback)
+ err = sbc.StreamExecute(ctx, session, target, sql, bindVariables, state.TransactionID, reservedID, callback)
return state, err
}
@@ -669,10 +669,10 @@ func (sbc *SandboxConn) HandlePanic(err *error) {
}
// ReserveBeginExecute implements the QueryService interface
-func (sbc *SandboxConn) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
+func (sbc *SandboxConn) ReserveBeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
sbc.panicIfNeeded()
- reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, 0, options)
- state, result, err := sbc.BeginExecute(ctx, target, postBeginQueries, sql, bindVariables, reservedID, options)
+ reservedID := sbc.reserve(ctx, session, target, preQueries, bindVariables, 0)
+ state, result, err := sbc.BeginExecute(ctx, session, target, postBeginQueries, sql, bindVariables, reservedID)
if state.TransactionID != 0 {
sbc.setTxReservedID(state.TransactionID, reservedID)
}
@@ -684,10 +684,10 @@ func (sbc *SandboxConn) ReserveBeginExecute(ctx context.Context, target *querypb
}
// ReserveBeginStreamExecute is part of the QueryService interface.
-func (sbc *SandboxConn) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
+func (sbc *SandboxConn) ReserveBeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
sbc.panicIfNeeded()
- reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, 0, options)
- state, err := sbc.BeginStreamExecute(ctx, target, postBeginQueries, sql, bindVariables, reservedID, options, callback)
+ reservedID := sbc.reserve(ctx, session, target, preQueries, bindVariables, 0)
+ state, err := sbc.BeginStreamExecute(ctx, session, target, postBeginQueries, sql, bindVariables, reservedID, callback)
if state.TransactionID != 0 {
sbc.setTxReservedID(state.TransactionID, reservedID)
}
@@ -699,10 +699,10 @@ func (sbc *SandboxConn) ReserveBeginStreamExecute(ctx context.Context, target *q
}
// ReserveExecute implements the QueryService interface
-func (sbc *SandboxConn) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (queryservice.ReservedState, *sqltypes.Result, error) {
+func (sbc *SandboxConn) ReserveExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (queryservice.ReservedState, *sqltypes.Result, error) {
sbc.panicIfNeeded()
- reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, transactionID, options)
- result, err := sbc.Execute(ctx, target, sql, bindVariables, transactionID, reservedID, options)
+ reservedID := sbc.reserve(ctx, session, target, preQueries, bindVariables, transactionID)
+ result, err := sbc.Execute(ctx, session, target, sql, bindVariables, transactionID, reservedID)
if transactionID != 0 {
sbc.setTxReservedID(transactionID, reservedID)
}
@@ -713,10 +713,10 @@ func (sbc *SandboxConn) ReserveExecute(ctx context.Context, target *querypb.Targ
}
// ReserveStreamExecute is part of the QueryService interface.
-func (sbc *SandboxConn) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
+func (sbc *SandboxConn) ReserveStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
sbc.panicIfNeeded()
- reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, transactionID, options)
- err := sbc.StreamExecute(ctx, target, sql, bindVariables, transactionID, reservedID, options, callback)
+ reservedID := sbc.reserve(ctx, session, target, preQueries, bindVariables, transactionID)
+ err := sbc.StreamExecute(ctx, session, target, sql, bindVariables, transactionID, reservedID, callback)
if transactionID != 0 {
sbc.setTxReservedID(transactionID, reservedID)
}
@@ -726,10 +726,10 @@ func (sbc *SandboxConn) ReserveStreamExecute(ctx context.Context, target *queryp
}, err
}
-func (sbc *SandboxConn) reserve(ctx context.Context, target *querypb.Target, preQueries []string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) int64 {
+func (sbc *SandboxConn) reserve(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, bindVariables map[string]*querypb.BindVariable, transactionID int64) int64 {
sbc.ReserveCount.Add(1)
for _, query := range preQueries {
- sbc.Execute(ctx, target, query, bindVariables, transactionID, 0, options)
+ sbc.Execute(ctx, session, target, query, bindVariables, transactionID, 0)
}
if transactionID != 0 {
return transactionID
@@ -894,3 +894,12 @@ func (sbc *SandboxConn) panicIfNeeded() {
panic(sbc.panicThis)
}
}
+
+// getOptions safely extracts ExecuteOptions from a session, returning nil if session is nil.
+func getOptions(session queryservice.Session) *querypb.ExecuteOptions {
+ if session == nil {
+ return nil
+ }
+
+ return session.GetOptions()
+}
diff --git a/go/vt/vttablet/tabletconntest/fakequeryservice.go b/go/vt/vttablet/tabletconntest/fakequeryservice.go
index 13ec5838fe3..6ba49809b6a 100644
--- a/go/vt/vttablet/tabletconntest/fakequeryservice.go
+++ b/go/vt/vttablet/tabletconntest/fakequeryservice.go
@@ -33,6 +33,7 @@ import (
binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
querypb "vitess.io/vitess/go/vt/proto/query"
topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)
@@ -104,10 +105,12 @@ var TestVTGateCallerID = &querypb.VTGateCallerID{
Username: "test_username",
}
-// TestExecuteOptions is a test execute options.
-var TestExecuteOptions = &querypb.ExecuteOptions{
- IncludedFields: querypb.ExecuteOptions_TYPE_ONLY,
- ClientFoundRows: true,
+// TestSession is a test execute options.
+var TestSession = &vtgatepb.Session{
+ Options: &querypb.ExecuteOptions{
+ IncludedFields: querypb.ExecuteOptions_TYPE_ONLY,
+ ClientFoundRows: true,
+ },
}
// TestAsTransaction is a test 'asTransaction' flag.
@@ -141,7 +144,7 @@ func (f *FakeQueryService) checkTargetCallerID(ctx context.Context, name string,
const beginTransactionID int64 = 9990
// Begin is part of the queryservice.QueryService interface
-func (f *FakeQueryService) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (f *FakeQueryService) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (queryservice.TransactionState, error) {
if f.HasBeginError {
return queryservice.TransactionState{}, f.TabletError
}
@@ -149,8 +152,8 @@ func (f *FakeQueryService) Begin(ctx context.Context, target *querypb.Target, op
panic(errors.New("test-triggered panic"))
}
f.checkTargetCallerID(ctx, "Begin", target)
- if !proto.Equal(options, TestExecuteOptions) {
- f.t.Errorf("invalid Execute.ExecuteOptions: got %v expected %v", options, TestExecuteOptions)
+ if !proto.Equal(getOptions(session), TestSession.GetOptions()) {
+ f.t.Errorf("invalid Execute.ExecuteOptions: got %v expected %v", getOptions(session), TestSession.GetOptions())
}
return queryservice.TransactionState{TransactionID: beginTransactionID, TabletAlias: TestAlias}, nil
}
@@ -412,7 +415,7 @@ var ExecuteQueryResult = sqltypes.Result{
}
// Execute is part of the queryservice.QueryService interface
-func (f *FakeQueryService) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (f *FakeQueryService) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
if f.HasError {
return nil, f.TabletError
}
@@ -425,8 +428,8 @@ func (f *FakeQueryService) Execute(ctx context.Context, target *querypb.Target,
if !sqltypes.BindVariablesEqual(bindVariables, ExecuteBindVars) {
f.t.Errorf("invalid Execute.BindVariables: got %v expected %v", bindVariables, ExecuteBindVars)
}
- if !proto.Equal(options, TestExecuteOptions) {
- f.t.Errorf("invalid Execute.ExecuteOptions: got %v expected %v", options, TestExecuteOptions)
+ if !proto.Equal(getOptions(session), TestSession.GetOptions()) {
+ f.t.Errorf("invalid Execute.ExecuteOptions: got %v expected %v", getOptions(session), TestSession.GetOptions())
}
f.checkTargetCallerID(ctx, "Execute", target)
if transactionID != f.ExpectedTransactionID {
@@ -472,7 +475,7 @@ var StreamExecuteQueryResult2 = sqltypes.Result{
}
// StreamExecute is part of the queryservice.QueryService interface
-func (f *FakeQueryService) StreamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
+func (f *FakeQueryService) StreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error {
if f.Panics && f.StreamExecutePanicsEarly {
panic(errors.New("test-triggered panic early"))
}
@@ -482,8 +485,8 @@ func (f *FakeQueryService) StreamExecute(ctx context.Context, target *querypb.Ta
if !sqltypes.BindVariablesEqual(bindVariables, StreamExecuteBindVars) {
f.t.Errorf("invalid StreamExecute.BindVariables: got %v expected %v", bindVariables, StreamExecuteBindVars)
}
- if !proto.Equal(options, TestExecuteOptions) {
- f.t.Errorf("invalid StreamExecute.ExecuteOptions: got %v expected %v", options, TestExecuteOptions)
+ if !proto.Equal(getOptions(session), TestSession.GetOptions()) {
+ f.t.Errorf("invalid StreamExecute.ExecuteOptions: got %v expected %v", getOptions(session), TestSession.GetOptions())
}
f.checkTargetCallerID(ctx, "StreamExecute", target)
if err := callback(&StreamExecuteQueryResult1); err != nil {
@@ -567,32 +570,32 @@ var ExecuteBatchQueryResultList = []sqltypes.Result{
}
// BeginExecute combines Begin and Execute.
-func (f *FakeQueryService) BeginExecute(ctx context.Context, target *querypb.Target, _ []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
- state, err := f.Begin(ctx, target, options)
+func (f *FakeQueryService) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, _ []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64) (queryservice.TransactionState, *sqltypes.Result, error) {
+ state, err := f.Begin(ctx, session, target)
if err != nil {
return state, nil, err
}
// TODO(deepthi): what alias should we actually return here?
- result, err := f.Execute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options)
+ result, err := f.Execute(ctx, session, target, sql, bindVariables, state.TransactionID, reservedID)
return state, result, err
}
// BeginStreamExecute combines Begin and StreamExecute.
-func (f *FakeQueryService) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
- state, err := f.Begin(ctx, target, options)
+func (f *FakeQueryService) BeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
+ state, err := f.Begin(ctx, session, target)
if err != nil {
return state, err
}
for _, preQuery := range preQueries {
- _, err := f.Execute(ctx, target, preQuery, nil, state.TransactionID, reservedID, options)
+ _, err := f.Execute(ctx, session, target, preQuery, nil, state.TransactionID, reservedID)
if err != nil {
return state, err
}
}
- err = f.StreamExecute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options, callback)
+ err = f.StreamExecute(ctx, session, target, sql, bindVariables, state.TransactionID, reservedID, callback)
return state, err
}
@@ -732,27 +735,27 @@ func (f *FakeQueryService) GetServingKeyspaces() []string {
}
// ReserveBeginExecute satisfies the Gateway interface
-func (f *FakeQueryService) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
+func (f *FakeQueryService) ReserveBeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
panic("implement me")
}
// ReserveBeginStreamExecute satisfies the Gateway interface
-func (f *FakeQueryService) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
+func (f *FakeQueryService) ReserveBeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
panic("implement me")
}
// ReserveExecute implements the QueryService interface
-func (f *FakeQueryService) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (queryservice.ReservedState, *sqltypes.Result, error) {
+func (f *FakeQueryService) ReserveExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (queryservice.ReservedState, *sqltypes.Result, error) {
panic("implement me")
}
// ReserveStreamExecute satisfies the Gateway interface
-func (f *FakeQueryService) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
+func (f *FakeQueryService) ReserveStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
state, err := f.reserve(transactionID)
if err != nil {
return state, err
}
- err = f.StreamExecute(ctx, target, sql, bindVariables, transactionID, state.ReservedID, options, callback)
+ err = f.StreamExecute(ctx, session, target, sql, bindVariables, transactionID, state.ReservedID, callback)
return state, err
}
@@ -784,3 +787,12 @@ func CreateFakeServer(t testing.TB) *FakeQueryService {
t: t,
}
}
+
+// getOptions safely extracts ExecuteOptions from a session, returning nil if session is nil.
+func getOptions(session queryservice.Session) *querypb.ExecuteOptions {
+ if session == nil {
+ return nil
+ }
+
+ return session.GetOptions()
+}
diff --git a/go/vt/vttablet/tabletconntest/tabletconntest.go b/go/vt/vttablet/tabletconntest/tabletconntest.go
index 91aeceb0dcd..a10b2b42bff 100644
--- a/go/vt/vttablet/tabletconntest/tabletconntest.go
+++ b/go/vt/vttablet/tabletconntest/tabletconntest.go
@@ -102,7 +102,7 @@ func testBegin(t *testing.T, conn queryservice.QueryService, f *FakeQueryService
t.Log("testBegin")
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- state, err := conn.Begin(ctx, TestTarget, TestExecuteOptions)
+ state, err := conn.Begin(ctx, TestSession, TestTarget)
if err != nil {
t.Fatalf("Begin failed: %v", err)
}
@@ -116,7 +116,7 @@ func testBeginError(t *testing.T, conn queryservice.QueryService, f *FakeQuerySe
t.Log("testBeginError")
f.HasBeginError = true
testErrorHelper(t, f, "Begin", func(ctx context.Context) error {
- _, err := conn.Begin(ctx, TestTarget, nil)
+ _, err := conn.Begin(ctx, nil, TestTarget)
return err
})
f.HasBeginError = false
@@ -125,7 +125,7 @@ func testBeginError(t *testing.T, conn queryservice.QueryService, f *FakeQuerySe
func testBeginPanics(t *testing.T, conn queryservice.QueryService, f *FakeQueryService) {
t.Log("testBeginPanics")
testPanicHelper(t, f, "Begin", func(ctx context.Context) error {
- _, err := conn.Begin(ctx, TestTarget, nil)
+ _, err := conn.Begin(ctx, nil, TestTarget)
return err
})
}
@@ -436,7 +436,7 @@ func testExecute(t *testing.T, conn queryservice.QueryService, f *FakeQueryServi
f.ExpectedTransactionID = ExecuteTransactionID
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- qr, err := conn.Execute(ctx, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID, TestExecuteOptions)
+ qr, err := conn.Execute(ctx, TestSession, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
@@ -449,7 +449,7 @@ func testExecuteError(t *testing.T, conn queryservice.QueryService, f *FakeQuery
t.Log("testExecuteError")
f.HasError = true
testErrorHelper(t, f, "Execute", func(ctx context.Context) error {
- _, err := conn.Execute(ctx, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID, TestExecuteOptions)
+ _, err := conn.Execute(ctx, TestSession, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID)
return err
})
f.HasError = false
@@ -458,7 +458,7 @@ func testExecuteError(t *testing.T, conn queryservice.QueryService, f *FakeQuery
func testExecutePanics(t *testing.T, conn queryservice.QueryService, f *FakeQueryService) {
t.Log("testExecutePanics")
testPanicHelper(t, f, "Execute", func(ctx context.Context) error {
- _, err := conn.Execute(ctx, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID, TestExecuteOptions)
+ _, err := conn.Execute(ctx, TestSession, TestTarget, ExecuteQuery, ExecuteBindVars, ExecuteTransactionID, ReserveConnectionID)
return err
})
}
@@ -468,7 +468,7 @@ func testBeginExecute(t *testing.T, conn queryservice.QueryService, f *FakeQuery
f.ExpectedTransactionID = beginTransactionID
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- state, qr, err := conn.BeginExecute(ctx, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID, TestExecuteOptions)
+ state, qr, err := conn.BeginExecute(ctx, TestSession, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID)
if err != nil {
t.Fatalf("BeginExecute failed: %v", err)
}
@@ -485,7 +485,7 @@ func testBeginExecuteErrorInBegin(t *testing.T, conn queryservice.QueryService,
t.Log("testBeginExecuteErrorInBegin")
f.HasBeginError = true
testErrorHelper(t, f, "BeginExecute.Begin", func(ctx context.Context) error {
- state, _, err := conn.BeginExecute(ctx, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID, TestExecuteOptions)
+ state, _, err := conn.BeginExecute(ctx, TestSession, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID)
if state.TransactionID != 0 {
t.Errorf("Unexpected transactionID from BeginExecute: got %v wanted 0", state.TransactionID)
}
@@ -499,7 +499,7 @@ func testBeginExecuteErrorInExecute(t *testing.T, conn queryservice.QueryService
f.HasError = true
testErrorHelper(t, f, "BeginExecute.Execute", func(ctx context.Context) error {
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- state, _, err := conn.BeginExecute(ctx, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID, TestExecuteOptions)
+ state, _, err := conn.BeginExecute(ctx, TestSession, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID)
if state.TransactionID != beginTransactionID {
t.Errorf("Unexpected transactionID from BeginExecute: got %v wanted %v", state.TransactionID, beginTransactionID)
}
@@ -511,7 +511,7 @@ func testBeginExecuteErrorInExecute(t *testing.T, conn queryservice.QueryService
func testBeginExecutePanics(t *testing.T, conn queryservice.QueryService, f *FakeQueryService) {
t.Log("testBeginExecutePanics")
testPanicHelper(t, f, "BeginExecute", func(ctx context.Context) error {
- _, _, err := conn.BeginExecute(ctx, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID, TestExecuteOptions)
+ _, _, err := conn.BeginExecute(ctx, TestSession, TestTarget, nil, ExecuteQuery, ExecuteBindVars, ReserveConnectionID)
return err
})
}
@@ -521,7 +521,7 @@ func testStreamExecute(t *testing.T, conn queryservice.QueryService, f *FakeQuer
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
i := 0
- err := conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ err := conn.StreamExecute(ctx, TestSession, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, func(qr *sqltypes.Result) error {
switch i {
case 0:
if len(qr.Rows) == 0 {
@@ -557,7 +557,7 @@ func testStreamExecuteError(t *testing.T, conn queryservice.QueryService, f *Fak
testErrorHelper(t, f, "StreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- return conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ return conn.StreamExecute(ctx, TestSession, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.ErrorWait:
@@ -586,7 +586,7 @@ func testStreamExecutePanics(t *testing.T, conn queryservice.QueryService, f *Fa
f.StreamExecutePanicsEarly = true
testPanicHelper(t, f, "StreamExecute.Early", func(ctx context.Context) error {
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- return conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ return conn.StreamExecute(ctx, TestSession, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, func(qr *sqltypes.Result) error {
return nil
})
})
@@ -596,7 +596,7 @@ func testStreamExecutePanics(t *testing.T, conn queryservice.QueryService, f *Fa
testPanicHelper(t, f, "StreamExecute.Late", func(ctx context.Context) error {
f.PanicWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- return conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ return conn.StreamExecute(ctx, TestSession, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.PanicWait:
@@ -621,7 +621,7 @@ func testBeginStreamExecute(t *testing.T, conn queryservice.QueryService, f *Fak
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
i := 0
- _, err := conn.BeginStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ _, err := conn.BeginStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
switch i {
case 0:
if len(qr.Rows) == 0 {
@@ -656,7 +656,7 @@ func testReserveStreamExecute(t *testing.T, conn queryservice.QueryService, f *F
ctx := context.Background()
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
i := 0
- _, err := conn.ReserveStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ _, err := conn.ReserveStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
switch i {
case 0:
if len(qr.Rows) == 0 {
@@ -692,7 +692,7 @@ func testBeginStreamExecuteErrorInBegin(t *testing.T, conn queryservice.QuerySer
testErrorHelper(t, f, "StreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- _, err := conn.BeginStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ _, err := conn.BeginStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.ErrorWait:
@@ -720,7 +720,7 @@ func testBeginStreamExecuteErrorInExecute(t *testing.T, conn queryservice.QueryS
testErrorHelper(t, f, "StreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- state, err := conn.BeginStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ state, err := conn.BeginStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.ErrorWait:
@@ -749,7 +749,7 @@ func testReserveStreamExecuteErrorInReserve(t *testing.T, conn queryservice.Quer
testErrorHelper(t, f, "ReserveStreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- _, err := conn.ReserveStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ _, err := conn.ReserveStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.ErrorWait:
@@ -777,7 +777,7 @@ func testReserveStreamExecuteErrorInExecute(t *testing.T, conn queryservice.Quer
testErrorHelper(t, f, "ReserveStreamExecute", func(ctx context.Context) error {
f.ErrorWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- state, err := conn.ReserveStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ state, err := conn.ReserveStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.ErrorWait:
@@ -808,7 +808,7 @@ func testBeginStreamExecutePanics(t *testing.T, conn queryservice.QueryService,
f.StreamExecutePanicsEarly = true
testPanicHelper(t, f, "StreamExecute.Early", func(ctx context.Context) error {
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- return conn.StreamExecute(ctx, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ return conn.StreamExecute(ctx, TestSession, TestTarget, StreamExecuteQuery, StreamExecuteBindVars, 0, 0, func(qr *sqltypes.Result) error {
return nil
})
})
@@ -818,7 +818,7 @@ func testBeginStreamExecutePanics(t *testing.T, conn queryservice.QueryService,
testPanicHelper(t, f, "StreamExecute.Late", func(ctx context.Context) error {
f.PanicWait = make(chan struct{})
ctx = callerid.NewContext(ctx, TestCallerID, TestVTGateCallerID)
- _, err := conn.BeginStreamExecute(ctx, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, TestExecuteOptions, func(qr *sqltypes.Result) error {
+ _, err := conn.BeginStreamExecute(ctx, TestSession, TestTarget, nil, StreamExecuteQuery, StreamExecuteBindVars, 0, func(qr *sqltypes.Result) error {
// For some errors, the call can be retried.
select {
case <-f.PanicWait:
diff --git a/go/vt/vttablet/tabletmanager/framework_test.go b/go/vt/vttablet/tabletmanager/framework_test.go
index c3ca31e4862..8f4cc38b59b 100644
--- a/go/vt/vttablet/tabletmanager/framework_test.go
+++ b/go/vt/vttablet/tabletmanager/framework_test.go
@@ -217,7 +217,7 @@ type fakeTabletConn struct {
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
+func (ftc *fakeTabletConn) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (queryservice.TransactionState, error) {
return queryservice.TransactionState{
TransactionID: 1,
}, nil
@@ -274,24 +274,24 @@ func (ftc *fakeTabletConn) ReadTransaction(ctx context.Context, target *querypb.
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
+func (ftc *fakeTabletConn) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (*sqltypes.Result, error) {
return nil, nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) StreamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
+func (ftc *fakeTabletConn) StreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) error {
return nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
+func (ftc *fakeTabletConn) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64) (queryservice.TransactionState, *sqltypes.Result, error) {
return queryservice.TransactionState{
TransactionID: 1,
}, nil, nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
+func (ftc *fakeTabletConn) BeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
return queryservice.TransactionState{
TransactionID: 1,
}, nil
@@ -327,28 +327,28 @@ func (ftc *fakeTabletConn) HandlePanic(err *error) {
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
+func (ftc *fakeTabletConn) ReserveBeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
return queryservice.ReservedTransactionState{
ReservedID: 1,
}, nil, nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
+func (ftc *fakeTabletConn) ReserveBeginStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
return queryservice.ReservedTransactionState{
ReservedID: 1,
}, nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (queryservice.ReservedState, *sqltypes.Result, error) {
+func (ftc *fakeTabletConn) ReserveExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (queryservice.ReservedState, *sqltypes.Result, error) {
return queryservice.ReservedState{
ReservedID: 1,
}, nil, nil
}
// fakeTabletConn implements the QueryService interface.
-func (ftc *fakeTabletConn) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
+func (ftc *fakeTabletConn) ReserveStreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
return queryservice.ReservedState{
ReservedID: 1,
}, nil
diff --git a/go/vt/vttablet/tabletmanager/rpc_query.go b/go/vt/vttablet/tabletmanager/rpc_query.go
index b390ed86ba6..3a3fdc39b2b 100644
--- a/go/vt/vttablet/tabletmanager/rpc_query.go
+++ b/go/vt/vttablet/tabletmanager/rpc_query.go
@@ -293,6 +293,6 @@ func (tm *TabletManager) ExecuteQuery(ctx context.Context, req *tabletmanagerdat
if err != nil {
return nil, err
}
- result, err := tm.QueryServiceControl.QueryService().Execute(ctx, target, uq, nil, 0, 0, nil)
+ result, err := tm.QueryServiceControl.QueryService().Execute(ctx, nil, target, uq, nil, 0, 0)
return sqltypes.ResultToProto3(result), err
}
diff --git a/go/vt/vttablet/tabletserver/bench_test.go b/go/vt/vttablet/tabletserver/bench_test.go
index fd2d86c2812..d2cd12c853f 100644
--- a/go/vt/vttablet/tabletserver/bench_test.go
+++ b/go/vt/vttablet/tabletserver/bench_test.go
@@ -36,8 +36,10 @@ import (
// BenchmarkExecuteVarBinary-4 100 14610045 ns/op
// BenchmarkExecuteExpression-4 1000 1047798 ns/op
-var benchQuery = "select a from test_table where v = :vtg1 and v0 = :vtg2 and v1 = :vtg3 and v2 = :vtg4 and v3 = :vtg5 and v4 = :vtg6 and v5 = :vtg7 and v6 = :vtg8 and v7 = :vtg9 and v8 = :vtg10 and v9 = :vtg11"
-var benchVarValue []byte
+var (
+ benchQuery = "select a from test_table where v = :vtg1 and v0 = :vtg2 and v1 = :vtg3 and v2 = :vtg4 and v3 = :vtg5 and v4 = :vtg6 and v5 = :vtg7 and v6 = :vtg8 and v7 = :vtg9 and v8 = :vtg10 and v9 = :vtg11"
+ benchVarValue []byte
+)
func init() {
// benchQuerySize is the approximate size of the query.
@@ -70,8 +72,8 @@ func BenchmarkExecuteVarBinary(b *testing.B) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
db.SetAllowAll(true)
- for i := 0; i < b.N; i++ {
- if _, err := tsv.Execute(ctx, &target, benchQuery, bv, 0, 0, nil); err != nil {
+ for b.Loop() {
+ if _, err := tsv.Execute(ctx, nil, &target, benchQuery, bv, 0, 0); err != nil {
panic(err)
}
}
@@ -97,8 +99,8 @@ func BenchmarkExecuteExpression(b *testing.B) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
db.SetAllowAll(true)
- for i := 0; i < b.N; i++ {
- if _, err := tsv.Execute(ctx, &target, benchQuery, bv, 0, 0, nil); err != nil {
+ for b.Loop() {
+ if _, err := tsv.Execute(ctx, nil, &target, benchQuery, bv, 0, 0); err != nil {
panic(err)
}
}
diff --git a/go/vt/vttablet/tabletserver/dt_executor_test.go b/go/vt/vttablet/tabletserver/dt_executor_test.go
index e1a7c80656a..5f24ef2a0f4 100644
--- a/go/vt/vttablet/tabletserver/dt_executor_test.go
+++ b/go/vt/vttablet/tabletserver/dt_executor_test.go
@@ -771,7 +771,7 @@ func newNoTwopcExecutor(t *testing.T, ctx context.Context) (txe *DTExecutor, tsv
func newTxForPrep(ctx context.Context, tsv *TabletServer) int64 {
txid := newTransaction(tsv, nil)
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- _, err := tsv.Execute(ctx, &target, "update test_table set name = 2 where pk = 1", nil, txid, 0, nil)
+ _, err := tsv.Execute(ctx, nil, &target, "update test_table set name = 2 where pk = 1", nil, txid, 0)
if err != nil {
panic(err)
}
diff --git a/go/vt/vttablet/tabletserver/query_executor_test.go b/go/vt/vttablet/tabletserver/query_executor_test.go
index ae590d45b0e..2e94f0eaa33 100644
--- a/go/vt/vttablet/tabletserver/query_executor_test.go
+++ b/go/vt/vttablet/tabletserver/query_executor_test.go
@@ -55,6 +55,7 @@ import (
querypb "vitess.io/vitess/go/vt/proto/query"
tableaclpb "vitess.io/vitess/go/vt/proto/tableacl"
topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)
@@ -95,229 +96,230 @@ func TestQueryExecutorPlans(t *testing.T) {
outsideTxErr bool
// TxThrottler allows the test case to override the transaction throttler
txThrottler txthrottler.TxThrottler
- }{{
- input: "select * from t",
- dbResponses: []dbResponse{{
- query: "select * from t limit 10001",
- result: selectResult,
- }},
- resultWant: selectResult,
- planWant: "Select",
- logWant: "select * from t limit 10001",
- inTxWant: "select * from t limit 10001",
- }, {
- input: "select * from t limit 1",
- dbResponses: []dbResponse{{
- query: "select * from t limit 1",
- result: selectResult,
- }},
- resultWant: selectResult,
- planWant: "Select",
- logWant: "select * from t limit 1",
- inTxWant: "select * from t limit 1",
- }, {
- input: "show engines",
- dbResponses: []dbResponse{{
- query: "show engines",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "Show",
- logWant: "show engines",
- }, {
- input: "repair t",
- dbResponses: []dbResponse{{
- query: "repair t",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "OtherAdmin",
- logWant: "repair t",
- }, {
- input: "insert into test_table(a) values(1)",
- dbResponses: []dbResponse{{
- query: "insert into test_table(a) values (1)",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "Insert",
- logWant: "insert into test_table(a) values (1)",
- }, {
- input: "replace into test_table(a) values(1)",
- dbResponses: []dbResponse{{
- query: "replace into test_table(a) values (1)",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "Insert",
- logWant: "replace into test_table(a) values (1)",
- }, {
- input: "update test_table set a=1",
- dbResponses: []dbResponse{{
- query: "update test_table set a = 1 limit 10001",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "UpdateLimit",
- // The UpdateLimit query will not use autocommit because
- // it needs to roll back on failure.
- logWant: "begin; update test_table set a = 1 limit 10001; commit",
- inTxWant: "update test_table set a = 1 limit 10001",
- }, {
- input: "select a, b from test_table",
- passThrough: true,
- inDMLExec: true,
- dbResponses: []dbResponse{{
- query: "select a, b from test_table",
- result: selectResult,
- }},
- resultWant: selectResult,
- planWant: "SelectNoLimit",
- logWant: "select a, b from test_table",
- outsideTxErr: true,
- errorWant: "[BUG] SelectNoLimit unexpected plan type",
- }, {
- input: "update test_table set a=1",
- passThrough: true,
- dbResponses: []dbResponse{{
- query: "update test_table set a = 1",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "Update",
- logWant: "update test_table set a = 1",
- }, {
- input: "delete from test_table",
- dbResponses: []dbResponse{{
- query: "delete from test_table limit 10001",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "DeleteLimit",
- // The DeleteLimit query will not use autocommit because
- // it needs to roll back on failure.
- logWant: "begin; delete from test_table limit 10001; commit",
- inTxWant: "delete from test_table limit 10001",
- }, {
- input: "delete from test_table",
- passThrough: true,
- dbResponses: []dbResponse{{
- query: "delete from test_table",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "Delete",
- logWant: "delete from test_table",
- }, {
- input: "alter table test_table add zipcode int",
- dbResponses: []dbResponse{{
- query: "alter table test_table add column zipcode int",
- result: dmlResult,
- }},
- resultWant: dmlResult,
- planWant: "DDL",
- logWant: "alter table test_table add column zipcode int",
- onlyInTxErr: true,
- errorWant: "DDL statement executed inside a transaction",
- }, {
- input: "savepoint a",
- dbResponses: []dbResponse{{
- query: "savepoint a",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "Savepoint",
- logWant: "savepoint a",
- inTxWant: "savepoint a",
- }, {
- input: "create index a on user(id)",
- dbResponses: []dbResponse{{
- query: "alter table `user` add key a (id)",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "DDL",
- logWant: "alter table `user` add key a (id)",
- inTxWant: "alter table `user` add key a (id)",
- onlyInTxErr: true,
- errorWant: "DDL statement executed inside a transaction",
- }, {
- input: "create index a on user(id1 + id2)",
- dbResponses: []dbResponse{{
- query: "create index a on user(id1 + id2)",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "DDL",
- logWant: "create index a on user(id1 + id2)",
- inTxWant: "create index a on user(id1 + id2)",
- onlyInTxErr: true,
- errorWant: "DDL statement executed inside a transaction",
- }, {
- input: "ROLLBACK work to SAVEPOINT a",
- dbResponses: []dbResponse{{
- query: "ROLLBACK work to SAVEPOINT a",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "RollbackSavepoint",
- logWant: "ROLLBACK work to SAVEPOINT a",
- inTxWant: "ROLLBACK work to SAVEPOINT a",
- }, {
- input: "RELEASE savepoint a",
- dbResponses: []dbResponse{{
- query: "RELEASE savepoint a",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "Release",
- logWant: "RELEASE savepoint a",
- inTxWant: "RELEASE savepoint a",
- }, {
- input: "show create database db_name",
- dbResponses: []dbResponse{{
- query: "show create database ks",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "Show",
- logWant: "show create database ks",
- }, {
- input: "show create database mysql",
- dbResponses: []dbResponse{{
- query: "show create database mysql",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "Show",
- logWant: "show create database mysql",
- }, {
- input: "show create table mysql.user",
- dbResponses: []dbResponse{{
- query: "show create table mysql.`user`",
- result: emptyResult,
- }},
- resultWant: emptyResult,
- planWant: "Show",
- logWant: "show create table mysql.`user`",
- }, {
- input: "update test_table set a=1",
- dbResponses: []dbResponse{{
- query: "update test_table set a = 1 limit 10001",
- result: dmlResult,
- }},
- errorWant: "Transaction throttled",
- txThrottler: &mockTxThrottler{true},
- }, {
- input: "update test_table set a=1",
- passThrough: true,
- dbResponses: []dbResponse{{
- query: "update test_table set a = 1 limit 10001",
- result: dmlResult,
- }},
- errorWant: "Transaction throttled",
- txThrottler: &mockTxThrottler{true},
- },
+ }{
+ {
+ input: "select * from t",
+ dbResponses: []dbResponse{{
+ query: "select * from t limit 10001",
+ result: selectResult,
+ }},
+ resultWant: selectResult,
+ planWant: "Select",
+ logWant: "select * from t limit 10001",
+ inTxWant: "select * from t limit 10001",
+ }, {
+ input: "select * from t limit 1",
+ dbResponses: []dbResponse{{
+ query: "select * from t limit 1",
+ result: selectResult,
+ }},
+ resultWant: selectResult,
+ planWant: "Select",
+ logWant: "select * from t limit 1",
+ inTxWant: "select * from t limit 1",
+ }, {
+ input: "show engines",
+ dbResponses: []dbResponse{{
+ query: "show engines",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "Show",
+ logWant: "show engines",
+ }, {
+ input: "repair t",
+ dbResponses: []dbResponse{{
+ query: "repair t",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "OtherAdmin",
+ logWant: "repair t",
+ }, {
+ input: "insert into test_table(a) values(1)",
+ dbResponses: []dbResponse{{
+ query: "insert into test_table(a) values (1)",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "Insert",
+ logWant: "insert into test_table(a) values (1)",
+ }, {
+ input: "replace into test_table(a) values(1)",
+ dbResponses: []dbResponse{{
+ query: "replace into test_table(a) values (1)",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "Insert",
+ logWant: "replace into test_table(a) values (1)",
+ }, {
+ input: "update test_table set a=1",
+ dbResponses: []dbResponse{{
+ query: "update test_table set a = 1 limit 10001",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "UpdateLimit",
+ // The UpdateLimit query will not use autocommit because
+ // it needs to roll back on failure.
+ logWant: "begin; update test_table set a = 1 limit 10001; commit",
+ inTxWant: "update test_table set a = 1 limit 10001",
+ }, {
+ input: "select a, b from test_table",
+ passThrough: true,
+ inDMLExec: true,
+ dbResponses: []dbResponse{{
+ query: "select a, b from test_table",
+ result: selectResult,
+ }},
+ resultWant: selectResult,
+ planWant: "SelectNoLimit",
+ logWant: "select a, b from test_table",
+ outsideTxErr: true,
+ errorWant: "[BUG] SelectNoLimit unexpected plan type",
+ }, {
+ input: "update test_table set a=1",
+ passThrough: true,
+ dbResponses: []dbResponse{{
+ query: "update test_table set a = 1",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "Update",
+ logWant: "update test_table set a = 1",
+ }, {
+ input: "delete from test_table",
+ dbResponses: []dbResponse{{
+ query: "delete from test_table limit 10001",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "DeleteLimit",
+ // The DeleteLimit query will not use autocommit because
+ // it needs to roll back on failure.
+ logWant: "begin; delete from test_table limit 10001; commit",
+ inTxWant: "delete from test_table limit 10001",
+ }, {
+ input: "delete from test_table",
+ passThrough: true,
+ dbResponses: []dbResponse{{
+ query: "delete from test_table",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "Delete",
+ logWant: "delete from test_table",
+ }, {
+ input: "alter table test_table add zipcode int",
+ dbResponses: []dbResponse{{
+ query: "alter table test_table add column zipcode int",
+ result: dmlResult,
+ }},
+ resultWant: dmlResult,
+ planWant: "DDL",
+ logWant: "alter table test_table add column zipcode int",
+ onlyInTxErr: true,
+ errorWant: "DDL statement executed inside a transaction",
+ }, {
+ input: "savepoint a",
+ dbResponses: []dbResponse{{
+ query: "savepoint a",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "Savepoint",
+ logWant: "savepoint a",
+ inTxWant: "savepoint a",
+ }, {
+ input: "create index a on user(id)",
+ dbResponses: []dbResponse{{
+ query: "alter table `user` add key a (id)",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "DDL",
+ logWant: "alter table `user` add key a (id)",
+ inTxWant: "alter table `user` add key a (id)",
+ onlyInTxErr: true,
+ errorWant: "DDL statement executed inside a transaction",
+ }, {
+ input: "create index a on user(id1 + id2)",
+ dbResponses: []dbResponse{{
+ query: "create index a on user(id1 + id2)",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "DDL",
+ logWant: "create index a on user(id1 + id2)",
+ inTxWant: "create index a on user(id1 + id2)",
+ onlyInTxErr: true,
+ errorWant: "DDL statement executed inside a transaction",
+ }, {
+ input: "ROLLBACK work to SAVEPOINT a",
+ dbResponses: []dbResponse{{
+ query: "ROLLBACK work to SAVEPOINT a",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "RollbackSavepoint",
+ logWant: "ROLLBACK work to SAVEPOINT a",
+ inTxWant: "ROLLBACK work to SAVEPOINT a",
+ }, {
+ input: "RELEASE savepoint a",
+ dbResponses: []dbResponse{{
+ query: "RELEASE savepoint a",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "Release",
+ logWant: "RELEASE savepoint a",
+ inTxWant: "RELEASE savepoint a",
+ }, {
+ input: "show create database db_name",
+ dbResponses: []dbResponse{{
+ query: "show create database ks",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "Show",
+ logWant: "show create database ks",
+ }, {
+ input: "show create database mysql",
+ dbResponses: []dbResponse{{
+ query: "show create database mysql",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "Show",
+ logWant: "show create database mysql",
+ }, {
+ input: "show create table mysql.user",
+ dbResponses: []dbResponse{{
+ query: "show create table mysql.`user`",
+ result: emptyResult,
+ }},
+ resultWant: emptyResult,
+ planWant: "Show",
+ logWant: "show create table mysql.`user`",
+ }, {
+ input: "update test_table set a=1",
+ dbResponses: []dbResponse{{
+ query: "update test_table set a = 1 limit 10001",
+ result: dmlResult,
+ }},
+ errorWant: "Transaction throttled",
+ txThrottler: &mockTxThrottler{true},
+ }, {
+ input: "update test_table set a=1",
+ passThrough: true,
+ dbResponses: []dbResponse{{
+ query: "update test_table set a = 1 limit 10001",
+ result: dmlResult,
+ }},
+ errorWant: "Transaction throttled",
+ txThrottler: &mockTxThrottler{true},
+ },
}
for _, tcase := range testcases {
t.Run(tcase.input, func(t *testing.T) {
@@ -352,7 +354,7 @@ func TestQueryExecutorPlans(t *testing.T) {
// Test inside a transaction.
target := tsv.sm.Target()
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
if !tcase.outsideTxErr && tcase.errorWant != "" && !tcase.onlyInTxErr {
require.EqualError(t, err, tcase.errorWant)
return
@@ -449,7 +451,7 @@ func TestQueryExecutorQueryAnnotation(t *testing.T) {
// Test inside a transaction.
target := tsv.sm.Target()
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
require.NoError(t, err)
require.NotNil(t, state.TabletAlias, "alias should not be nil")
assert.Equal(t, tsv.alias, state.TabletAlias, "Wrong alias returned by Begin")
@@ -516,7 +518,7 @@ func TestQueryExecutorSelectImpossible(t *testing.T) {
assert.Equal(t, tcase.planWant, qre.logStats.PlanType, tcase.input)
assert.Equal(t, tcase.logWant, qre.logStats.RewrittenSQL(), tcase.input)
target := tsv.sm.Target()
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
require.NoError(t, err)
require.NotNil(t, state.TabletAlias, "alias should not be nil")
assert.Equal(t, tsv.alias, state.TabletAlias, "Wrong tablet alias from Begin")
@@ -647,7 +649,7 @@ func TestQueryExecutorLimitFailure(t *testing.T) {
// Test inside a transaction.
target := tsv.sm.Target()
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
require.NoError(t, err)
require.NotNil(t, state.TabletAlias, "alias should not be nil")
assert.Equal(t, tsv.alias, state.TabletAlias, "Wrong tablet alias from Begin")
@@ -1638,7 +1640,7 @@ func newTestTabletServer(ctx context.Context, flags executorFlags, db *fakesqldb
func newTransaction(tsv *TabletServer, options *querypb.ExecuteOptions) int64 {
target := tsv.sm.Target()
- state, err := tsv.Begin(context.Background(), target, options)
+ state, err := tsv.Begin(context.Background(), &vtgatepb.Session{Options: options}, target)
if err != nil {
panic(vterrors.Wrap(err, "failed to start a transaction"))
}
diff --git a/go/vt/vttablet/tabletserver/tabletserver.go b/go/vt/vttablet/tabletserver/tabletserver.go
index 12a33bad575..60f5ab2d436 100644
--- a/go/vt/vttablet/tabletserver/tabletserver.go
+++ b/go/vt/vttablet/tabletserver/tabletserver.go
@@ -464,11 +464,11 @@ func (tsv *TabletServer) IsHealthy() error {
_, err := tsv.Execute(
tabletenv.LocalContext(),
nil,
+ nil,
"/* health */ select 1 from dual",
nil,
0,
0,
- nil,
)
return err
}
@@ -535,8 +535,8 @@ func (tsv *TabletServer) SchemaEngine() *schema.Engine {
}
// Begin starts a new transaction. This is allowed only if the state is StateServing.
-func (tsv *TabletServer) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (state queryservice.TransactionState, err error) {
- return tsv.begin(ctx, target, nil, 0, nil, options)
+func (tsv *TabletServer) Begin(ctx context.Context, session queryservice.Session, target *querypb.Target) (state queryservice.TransactionState, err error) {
+ return tsv.begin(ctx, target, nil, 0, nil, getOptions(session))
}
func (tsv *TabletServer) begin(
@@ -903,7 +903,7 @@ func (tsv *TabletServer) UnresolvedTransactions(ctx context.Context, target *que
}
// Execute executes the query and returns the result as response.
-func (tsv *TabletServer) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (result *sqltypes.Result, err error) {
+func (tsv *TabletServer) Execute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64) (result *sqltypes.Result, err error) {
span, ctx := trace.NewSpan(ctx, "TabletServer.Execute")
trace.AnnotateSQL(span, sqlparser.Preview(sql))
defer span.Finish()
@@ -912,7 +912,7 @@ func (tsv *TabletServer) Execute(ctx context.Context, target *querypb.Target, sq
return nil, vterrors.New(vtrpcpb.Code_INTERNAL, "[BUG] transactionID and reserveID must match if both are non-zero")
}
- return tsv.execute(ctx, target, sql, bindVariables, transactionID, reservedID, nil, options)
+ return tsv.execute(ctx, target, sql, bindVariables, transactionID, reservedID, nil, getOptions(session))
}
func (tsv *TabletServer) execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, settings []string, options *querypb.ExecuteOptions) (result *sqltypes.Result, err error) {
@@ -1009,12 +1009,12 @@ func smallerTimeout(t1, t2 time.Duration) time.Duration {
// StreamExecute executes the query and streams the result.
// The first QueryResult will have Fields set (and Rows nil).
// The subsequent QueryResult will have Rows set (and Fields nil).
-func (tsv *TabletServer) StreamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (err error) {
+func (tsv *TabletServer) StreamExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, callback func(*sqltypes.Result) error) (err error) {
if transactionID != 0 && reservedID != 0 && transactionID != reservedID {
return vterrors.New(vtrpcpb.Code_INTERNAL, "[BUG] transactionID and reserveID must match if both are non-zero")
}
- return tsv.streamExecute(ctx, target, sql, bindVariables, transactionID, reservedID, nil, options, callback)
+ return tsv.streamExecute(ctx, target, sql, bindVariables, transactionID, reservedID, nil, getOptions(session), callback)
}
func (tsv *TabletServer) streamExecute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, reservedID int64, settings []string, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
@@ -1077,7 +1077,9 @@ func (tsv *TabletServer) streamExecute(ctx context.Context, target *querypb.Targ
}
// BeginExecute combines Begin and Execute.
-func (tsv *TabletServer) BeginExecute(ctx context.Context, target *querypb.Target, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
+func (tsv *TabletServer) BeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64) (queryservice.TransactionState, *sqltypes.Result, error) {
+ options := getOptions(session)
+
// Disable hot row protection in case of reserve connection.
if tsv.enableHotRowProtection && reservedID == 0 {
txDone, err := tsv.beginWaitForSameRangeTransactions(ctx, target, options, sql, bindVariables)
@@ -1094,27 +1096,27 @@ func (tsv *TabletServer) BeginExecute(ctx context.Context, target *querypb.Targe
return state, nil, err
}
- result, err := tsv.Execute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options)
+ result, err := tsv.Execute(ctx, session, target, sql, bindVariables, state.TransactionID, reservedID)
return state, result, err
}
// BeginStreamExecute combines Begin and StreamExecute.
func (tsv *TabletServer) BeginStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
postBeginQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
reservedID int64,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (queryservice.TransactionState, error) {
- state, err := tsv.begin(ctx, target, postBeginQueries, reservedID, nil, options)
+ state, err := tsv.begin(ctx, target, postBeginQueries, reservedID, nil, getOptions(session))
if err != nil {
return state, err
}
- err = tsv.StreamExecute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options, callback)
+ err = tsv.StreamExecute(ctx, session, target, sql, bindVariables, state.TransactionID, reservedID, callback)
return state, err
}
@@ -1271,7 +1273,7 @@ func (tsv *TabletServer) execDML(ctx context.Context, target *querypb.Target, qu
return 0, err
}
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
if err != nil {
return 0, err
}
@@ -1282,7 +1284,7 @@ func (tsv *TabletServer) execDML(ctx context.Context, target *querypb.Target, qu
tsv.Rollback(ctx, target, state.TransactionID)
}
}()
- qr, err := tsv.Execute(ctx, target, query, bv, state.TransactionID, 0, nil)
+ qr, err := tsv.Execute(ctx, nil, target, query, bv, state.TransactionID, 0)
if err != nil {
return 0, err
}
@@ -1335,7 +1337,9 @@ func (tsv *TabletServer) VStreamResults(ctx context.Context, target *querypb.Tar
}
// ReserveBeginExecute implements the QueryService interface
-func (tsv *TabletServer) ReserveBeginExecute(ctx context.Context, target *querypb.Target, settings []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (state queryservice.ReservedTransactionState, result *sqltypes.Result, err error) {
+func (tsv *TabletServer) ReserveBeginExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, settings []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable) (state queryservice.ReservedTransactionState, result *sqltypes.Result, err error) {
+ options := getOptions(session)
+
state, result, err = tsv.beginExecuteWithSettings(ctx, target, settings, postBeginQueries, sql, bindVariables, options)
// If there is an error and the error message is about allowing query in reserved connection only,
// then we do not return an error from here and continue to use the reserved connection path.
@@ -1396,7 +1400,6 @@ func (tsv *TabletServer) ReserveBeginExecute(ctx context.Context, target *queryp
return nil
},
)
-
if err != nil {
return state, nil, err
}
@@ -1411,14 +1414,16 @@ func (tsv *TabletServer) ReserveBeginExecute(ctx context.Context, target *queryp
// ReserveBeginStreamExecute combines Begin and StreamExecute.
func (tsv *TabletServer) ReserveBeginStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
settings []string,
postBeginQueries []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (state queryservice.ReservedTransactionState, err error) {
+ options := getOptions(session)
+
txState, err := tsv.begin(ctx, target, postBeginQueries, 0, settings, options)
if err != nil {
return txToReserveState(txState), err
@@ -1429,7 +1434,9 @@ func (tsv *TabletServer) ReserveBeginStreamExecute(
}
// ReserveExecute implements the QueryService interface
-func (tsv *TabletServer) ReserveExecute(ctx context.Context, target *querypb.Target, settings []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (state queryservice.ReservedState, result *sqltypes.Result, err error) {
+func (tsv *TabletServer) ReserveExecute(ctx context.Context, session queryservice.Session, target *querypb.Target, settings []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64) (state queryservice.ReservedState, result *sqltypes.Result, err error) {
+ options := getOptions(session)
+
result, err = tsv.executeWithSettings(ctx, target, settings, sql, bindVariables, transactionID, options)
// If there is an error and the error message is about allowing query in reserved connection only,
// then we do not return an error from here and continue to use the reserved connection path.
@@ -1464,7 +1471,6 @@ func (tsv *TabletServer) ReserveExecute(ctx context.Context, target *querypb.Tar
return nil
},
)
-
if err != nil {
return state, nil, err
}
@@ -1476,15 +1482,15 @@ func (tsv *TabletServer) ReserveExecute(ctx context.Context, target *querypb.Tar
// ReserveStreamExecute combines Begin and StreamExecute.
func (tsv *TabletServer) ReserveStreamExecute(
ctx context.Context,
+ session queryservice.Session,
target *querypb.Target,
settings []string,
sql string,
bindVariables map[string]*querypb.BindVariable,
transactionID int64,
- options *querypb.ExecuteOptions,
callback func(*sqltypes.Result) error,
) (state queryservice.ReservedState, err error) {
- return state, tsv.streamExecute(ctx, target, sql, bindVariables, transactionID, 0, settings, options, callback)
+ return state, tsv.streamExecute(ctx, target, sql, bindVariables, transactionID, 0, settings, getOptions(session), callback)
}
// Release implements the QueryService interface
@@ -2114,3 +2120,12 @@ func skipQueryPlanCache(options *querypb.ExecuteOptions) bool {
func (tsv *TabletServer) getShard() string {
return tsv.sm.Target().Shard
}
+
+// getOptions safely extracts ExecuteOptions from a session, returning nil if session is nil.
+func getOptions(session queryservice.Session) *querypb.ExecuteOptions {
+ if session == nil {
+ return nil
+ }
+
+ return session.GetOptions()
+}
diff --git a/go/vt/vttablet/tabletserver/tabletserver_test.go b/go/vt/vttablet/tabletserver/tabletserver_test.go
index 62ae91a922a..a98c691af00 100644
--- a/go/vt/vttablet/tabletserver/tabletserver_test.go
+++ b/go/vt/vttablet/tabletserver/tabletserver_test.go
@@ -59,6 +59,7 @@ import (
querypb "vitess.io/vitess/go/vt/proto/query"
topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+ vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)
@@ -137,7 +138,7 @@ func TestBeginOnReplica(t *testing.T) {
options := querypb.ExecuteOptions{
TransactionIsolation: querypb.ExecuteOptions_CONSISTENT_SNAPSHOT_READ_ONLY,
}
- state, err := tsv.Begin(ctx, &target, &options)
+ state, err := tsv.Begin(ctx, &vtgatepb.Session{Options: &options}, &target)
require.NoError(t, err, "failed to create read only tx on replica")
assert.Equal(t, tsv.alias, state.TabletAlias, "Wrong tablet alias from Begin")
_, err = tsv.Rollback(ctx, &target, state.TransactionID)
@@ -145,7 +146,7 @@ func TestBeginOnReplica(t *testing.T) {
// test that we can still create transactions even in read-only mode
options = querypb.ExecuteOptions{}
- state, err = tsv.Begin(ctx, &target, &options)
+ state, err = tsv.Begin(ctx, &vtgatepb.Session{Options: &options}, &target)
require.NoError(t, err, "expected write tx to be allowed")
_, err = tsv.Rollback(ctx, &target, state.TransactionID)
require.NoError(t, err)
@@ -163,14 +164,14 @@ func TestTabletServerPrimaryToReplica(t *testing.T) {
tsv.te.shutdownGracePeriod = 1
tsv.sm.shutdownGracePeriod = 1
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state1, err := tsv.Begin(ctx, &target, nil)
+ state1, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
- _, err = tsv.Execute(ctx, &target, "update test_table set `name` = 2 where pk = 1", nil, state1.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, "update test_table set `name` = 2 where pk = 1", nil, state1.TransactionID, 0)
require.NoError(t, err)
err = tsv.Prepare(ctx, &target, state1.TransactionID, "aa")
require.NoError(t, err)
- state2, err := tsv.Begin(ctx, &target, nil)
+ state2, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
// This makes txid2 busy
@@ -246,7 +247,8 @@ func TestTabletServerRedoLogIsKeptBetweenRestarts(t *testing.T) {
got := tsv.te.preparedPool.conns["dtid0"].TxProperties().Queries
want := []tx.Query{{
Sql: "update test_table set `name` = 2 where pk = 1 limit 10001",
- Tables: []string{"test_table"}}}
+ Tables: []string{"test_table"},
+ }}
utils.MustMatch(t, want, got, "Prepared queries")
turnOffTxEngine()
assert.Empty(t, tsv.te.preparedPool.conns, "tsv.te.preparedPool.conns")
@@ -286,7 +288,8 @@ func TestTabletServerRedoLogIsKeptBetweenRestarts(t *testing.T) {
got = tsv.te.preparedPool.conns["a:b:10"].TxProperties().Queries
want = []tx.Query{{
Sql: "update test_table set `name` = 2 where pk = 1 limit 10001",
- Tables: []string{"test_table"}}}
+ Tables: []string{"test_table"},
+ }}
utils.MustMatch(t, want, got, "Prepared queries")
wantFailed := map[string]error{
"bogus": errPrepFailed, // The query is rejected by database so added to failed list.
@@ -474,8 +477,8 @@ func TestTabletServerBeginFail(t *testing.T) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
ctx, cancel = context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
- tsv.Begin(ctx, &target, nil)
- _, err := tsv.Begin(ctx, &target, nil)
+ tsv.Begin(ctx, nil, &target)
+ _, err := tsv.Begin(ctx, nil, &target)
require.EqualError(t, err, "transaction pool aborting request due to already expired context", "Begin err")
}
@@ -498,9 +501,9 @@ func TestTabletServerCommitTransaction(t *testing.T) {
db.AddQuery(executeSQL, executeSQLResult)
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, err := tsv.Begin(ctx, &target, nil)
+ state, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
- _, err = tsv.Execute(ctx, &target, executeSQL, nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, executeSQL, nil, state.TransactionID, 0)
require.NoError(t, err)
_, err = tsv.Commit(ctx, &target, state.TransactionID)
require.NoError(t, err)
@@ -540,12 +543,12 @@ func TestTabletServerRollback(t *testing.T) {
db.AddQuery(executeSQL, executeSQLResult)
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, err := tsv.Begin(ctx, &target, nil)
+ state, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
if err != nil {
t.Fatalf("call TabletServer.Begin failed: %v", err)
}
- _, err = tsv.Execute(ctx, &target, executeSQL, nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, executeSQL, nil, state.TransactionID, 0)
require.NoError(t, err)
_, err = tsv.Rollback(ctx, &target, state.TransactionID)
require.NoError(t, err)
@@ -558,9 +561,9 @@ func TestTabletServerPrepare(t *testing.T) {
_, tsv, _, closer := newTestTxExecutor(t, ctx)
defer closer()
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, err := tsv.Begin(ctx, &target, nil)
+ state, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
- _, err = tsv.Execute(ctx, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0)
require.NoError(t, err)
defer tsv.RollbackPrepared(ctx, &target, "aa", 0)
err = tsv.Prepare(ctx, &target, state.TransactionID, "aa")
@@ -574,9 +577,9 @@ func TestTabletServerCommitPrepared(t *testing.T) {
_, tsv, _, closer := newTestTxExecutor(t, ctx)
defer closer()
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, err := tsv.Begin(ctx, &target, nil)
+ state, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
- _, err = tsv.Execute(ctx, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0)
require.NoError(t, err)
err = tsv.Prepare(ctx, &target, state.TransactionID, "aa")
require.NoError(t, err)
@@ -623,12 +626,12 @@ func TestTabletServerWithNilTarget(t *testing.T) {
expectedCount := tsv.stats.QueryTimingsByTabletType.Counts()[fullKey]
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
require.NoError(t, err)
expectedCount++
require.Equal(t, expectedCount, tsv.stats.QueryTimingsByTabletType.Counts()[fullKey])
- _, err = tsv.Execute(ctx, target, executeSQL, nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, target, executeSQL, nil, state.TransactionID, 0)
require.NoError(t, err)
expectedCount++
require.Equal(t, expectedCount, tsv.stats.QueryTimingsByTabletType.Counts()[fullKey])
@@ -638,7 +641,7 @@ func TestTabletServerWithNilTarget(t *testing.T) {
expectedCount++
require.Equal(t, expectedCount, tsv.stats.QueryTimingsByTabletType.Counts()[fullKey])
- state, err = tsv.Begin(ctx, target, nil)
+ state, err = tsv.Begin(ctx, nil, target)
require.NoError(t, err)
expectedCount++
require.Equal(t, expectedCount, tsv.stats.QueryTimingsByTabletType.Counts()[fullKey])
@@ -651,7 +654,7 @@ func TestTabletServerWithNilTarget(t *testing.T) {
// Finally be sure that we return an error now as expected when NOT
// using a local context but passing a nil target.
nonLocalCtx := context.Background()
- _, err = tsv.Begin(nonLocalCtx, target, nil)
+ _, err = tsv.Begin(nonLocalCtx, nil, target)
require.True(t, errors.Is(err, ErrNoTarget))
_, err = tsv.resolveTargetType(nonLocalCtx, target)
require.True(t, errors.Is(err, ErrNoTarget))
@@ -773,11 +776,11 @@ func TestTabletServerReserveConnection(t *testing.T) {
options := &querypb.ExecuteOptions{}
// reserve a connection
- state, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, 0, options)
+ state, _, err := tsv.ReserveExecute(ctx, &vtgatepb.Session{Options: options}, &target, nil, "set sql_mode = ''", nil, 0)
require.NoError(t, err)
// run a query in it
- _, err = tsv.Execute(ctx, &target, "select 42", nil, 0, state.ReservedID, options)
+ _, err = tsv.Execute(ctx, &vtgatepb.Session{Options: options}, &target, "select 42", nil, 0, state.ReservedID)
require.NoError(t, err)
// release the connection
@@ -797,7 +800,7 @@ func TestTabletServerExecNonExistentConnection(t *testing.T) {
options := &querypb.ExecuteOptions{}
// run a query with a non-existent reserved id
- _, err := tsv.Execute(ctx, &target, "select 42", nil, 0, 123456, options)
+ _, err := tsv.Execute(ctx, &vtgatepb.Session{Options: options}, &target, "select 42", nil, 0, 123456)
require.Error(t, err)
}
@@ -828,7 +831,7 @@ func TestMakeSureToCloseDbConnWhenBeginQueryFails(t *testing.T) {
options := &querypb.ExecuteOptions{}
// run a query with a non-existent reserved id
- _, _, err := tsv.ReserveBeginExecute(ctx, &target, []string{}, nil, "select 42", nil, options)
+ _, _, err := tsv.ReserveBeginExecute(ctx, &vtgatepb.Session{Options: options}, &target, []string{}, nil, "select 42", nil)
require.Error(t, err)
}
@@ -844,7 +847,7 @@ func TestTabletServerReserveAndBeginCommit(t *testing.T) {
options := &querypb.ExecuteOptions{}
// reserve a connection and a transaction
- state, _, err := tsv.ReserveBeginExecute(ctx, &target, nil, nil, "set sql_mode = ''", nil, options)
+ state, _, err := tsv.ReserveBeginExecute(ctx, &vtgatepb.Session{Options: options}, &target, nil, nil, "set sql_mode = ''", nil)
require.NoError(t, err)
defer func() {
// fallback so the test finishes quickly
@@ -852,13 +855,13 @@ func TestTabletServerReserveAndBeginCommit(t *testing.T) {
}()
// run a query in it
- _, err = tsv.Execute(ctx, &target, "select 42", nil, state.TransactionID, state.ReservedID, options)
+ _, err = tsv.Execute(ctx, &vtgatepb.Session{Options: options}, &target, "select 42", nil, state.TransactionID, state.ReservedID)
require.NoError(t, err)
// run a query in a non-existent connection
- _, err = tsv.Execute(ctx, &target, "select 42", nil, state.TransactionID, state.ReservedID+100, options)
+ _, err = tsv.Execute(ctx, &vtgatepb.Session{Options: options}, &target, "select 42", nil, state.TransactionID, state.ReservedID+100)
require.Error(t, err)
- _, err = tsv.Execute(ctx, &target, "select 42", nil, state.TransactionID+100, state.ReservedID, options)
+ _, err = tsv.Execute(ctx, &vtgatepb.Session{Options: options}, &target, "select 42", nil, state.TransactionID+100, state.ReservedID)
require.Error(t, err)
// commit
@@ -868,7 +871,7 @@ func TestTabletServerReserveAndBeginCommit(t *testing.T) {
rID := newRID
// begin and rollback
- beginState, _, err := tsv.BeginExecute(ctx, &target, nil, "select 42", nil, rID, options)
+ beginState, _, err := tsv.BeginExecute(ctx, &vtgatepb.Session{Options: options}, &target, nil, "select 42", nil, rID)
require.NoError(t, err)
assert.Equal(t, newRID, beginState.TransactionID)
rID = newRID
@@ -894,9 +897,9 @@ func TestTabletServerRollbackPrepared(t *testing.T) {
_, tsv, _, closer := newTestTxExecutor(t, ctx)
defer closer()
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, err := tsv.Begin(ctx, &target, nil)
+ state, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
- _, err = tsv.Execute(ctx, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0, nil)
+ _, err = tsv.Execute(ctx, nil, &target, "update test_table set `name` = 2 where pk = 1", nil, state.TransactionID, 0)
require.NoError(t, err)
err = tsv.Prepare(ctx, &target, state.TransactionID, "aa")
require.NoError(t, err)
@@ -924,7 +927,7 @@ func TestTabletServerStreamExecute(t *testing.T) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
callback := func(*sqltypes.Result) error { return nil }
- if err := tsv.StreamExecute(ctx, &target, executeSQL, nil, 0, 0, nil, callback); err != nil {
+ if err := tsv.StreamExecute(ctx, nil, &target, executeSQL, nil, 0, 0, callback); err != nil {
t.Fatalf("TabletServer.StreamExecute should success: %s, but get error: %v",
executeSQL, err)
}
@@ -954,7 +957,7 @@ func TestTabletServerStreamExecuteComments(t *testing.T) {
ch := tabletenv.StatsLogger.Subscribe("test stats logging")
defer tabletenv.StatsLogger.Unsubscribe(ch)
- if err := tsv.StreamExecute(ctx, &target, executeSQL, nil, 0, 0, nil, callback); err != nil {
+ if err := tsv.StreamExecute(ctx, nil, &target, executeSQL, nil, 0, 0, callback); err != nil {
t.Fatalf("TabletServer.StreamExecute should success: %s, but get error: %v",
executeSQL, err)
}
@@ -990,7 +993,7 @@ func TestTabletServerBeginStreamExecute(t *testing.T) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
callback := func(*sqltypes.Result) error { return nil }
- state, err := tsv.BeginStreamExecute(ctx, &target, nil, executeSQL, nil, 0, nil, callback)
+ state, err := tsv.BeginStreamExecute(ctx, nil, &target, nil, executeSQL, nil, 0, callback)
if err != nil {
t.Fatalf("TabletServer.BeginStreamExecute should success: %s, but get error: %v",
executeSQL, err)
@@ -1023,7 +1026,7 @@ func TestTabletServerBeginStreamExecuteWithError(t *testing.T) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
callback := func(*sqltypes.Result) error { return nil }
- state, err := tsv.BeginStreamExecute(ctx, &target, nil, executeSQL, nil, 0, nil, callback)
+ state, err := tsv.BeginStreamExecute(ctx, nil, &target, nil, executeSQL, nil, 0, callback)
require.Error(t, err)
err = tsv.Release(ctx, &target, state.TransactionID, 0)
require.NoError(t, err)
@@ -1090,7 +1093,7 @@ func TestSerializeTransactionsSameRow(t *testing.T) {
go func() {
defer wg.Done()
- state1, _, err := tsv.BeginExecute(ctx, &target, nil, q1, bvTx1, 0, nil)
+ state1, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q1, bvTx1, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q1, err)
}
@@ -1105,7 +1108,7 @@ func TestSerializeTransactionsSameRow(t *testing.T) {
defer wg.Done()
<-tx1Started
- state2, _, err := tsv.BeginExecute(ctx, &target, nil, q2, bvTx2, 0, nil)
+ state2, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q2, bvTx2, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q2, err)
}
@@ -1125,7 +1128,7 @@ func TestSerializeTransactionsSameRow(t *testing.T) {
defer wg.Done()
<-tx1Started
- state3, _, err := tsv.BeginExecute(ctx, &target, nil, q3, bvTx3, 0, nil)
+ state3, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q3, bvTx3, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q3, err)
}
@@ -1160,7 +1163,7 @@ func TestDMLQueryWithoutWhereClause(t *testing.T) {
db.AddQuery(q+" limit 10001", &sqltypes.Result{})
- state, _, err := tsv.BeginExecute(ctx, &target, nil, q, nil, 0, nil)
+ state, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q, nil, 0)
require.NoError(t, err)
_, err = tsv.Commit(ctx, &target, state.TransactionID)
require.NoError(t, err)
@@ -1234,7 +1237,7 @@ func TestSerializeTransactionsSameRow_ConcurrentTransactions(t *testing.T) {
go func() {
defer wg.Done()
- state1, _, err := tsv.BeginExecute(ctx, &target, nil, q1, bvTx1, 0, nil)
+ state1, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q1, bvTx1, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q1, err)
}
@@ -1253,7 +1256,7 @@ func TestSerializeTransactionsSameRow_ConcurrentTransactions(t *testing.T) {
// In that case, we would see less than 3 pending transactions.
<-tx1Started
- state2, _, err := tsv.BeginExecute(ctx, &target, nil, q2, bvTx2, 0, nil)
+ state2, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q2, bvTx2, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q2, err)
}
@@ -1272,7 +1275,7 @@ func TestSerializeTransactionsSameRow_ConcurrentTransactions(t *testing.T) {
// In that case, we would see less than 3 pending transactions.
<-tx1Started
- state3, _, err := tsv.BeginExecute(ctx, &target, nil, q3, bvTx3, 0, nil)
+ state3, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q3, bvTx3, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q3, err)
}
@@ -1370,7 +1373,7 @@ func TestSerializeTransactionsSameRow_TooManyPendingRequests(t *testing.T) {
go func() {
defer wg.Done()
- state1, _, err := tsv.BeginExecute(ctx, &target, nil, q1, bvTx1, 0, nil)
+ state1, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q1, bvTx1, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q1, err)
}
@@ -1386,7 +1389,7 @@ func TestSerializeTransactionsSameRow_TooManyPendingRequests(t *testing.T) {
defer close(tx2Failed)
<-tx1Started
- _, _, err := tsv.BeginExecute(ctx, &target, nil, q2, bvTx2, 0, nil)
+ _, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q2, bvTx2, 0)
if err == nil || vterrors.Code(err) != vtrpcpb.Code_RESOURCE_EXHAUSTED || err.Error() != "hot row protection: too many queued transactions (1 >= 1) for the same row (table + WHERE clause: 'test_table where pk = 1 and `name` = 1')" {
t.Errorf("tx2 should have failed because there are too many pending requests: %v", err)
}
@@ -1459,7 +1462,7 @@ func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) {
go func() {
defer wg.Done()
- state1, _, err := tsv.BeginExecute(ctx, &target, nil, q1, bvTx1, 0, nil)
+ state1, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q1, bvTx1, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q1, err)
}
@@ -1479,7 +1482,7 @@ func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) {
// Wait until tx1 has started to make the test deterministic.
<-tx1Started
- _, _, err := tsv.BeginExecute(ctxTx2, &target, nil, q2, bvTx2, 0, nil)
+ _, _, err := tsv.BeginExecute(ctxTx2, nil, &target, nil, q2, bvTx2, 0)
if err == nil || vterrors.Code(err) != vtrpcpb.Code_CANCELED || err.Error() != "context canceled" {
t.Errorf("tx2 should have failed because the context was canceled: %v", err)
}
@@ -1496,7 +1499,7 @@ func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) {
t.Error(err)
}
- state3, _, err := tsv.BeginExecute(ctx, &target, nil, q3, bvTx3, 0, nil)
+ state3, _, err := tsv.BeginExecute(ctx, nil, &target, nil, q3, bvTx3, 0)
if err != nil {
t.Errorf("failed to execute query: %s: %s", q3, err)
}
@@ -2146,6 +2149,7 @@ var aclJSON1 = `{
}
]
}`
+
var aclJSON2 = `{
"table_groups": [
{
@@ -2156,6 +2160,7 @@ var aclJSON2 = `{
}
]
}`
+
var aclJSONOverlapError = `{
"table_groups": [
{
@@ -2316,7 +2321,7 @@ func TestReserveBeginExecute(t *testing.T) {
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
db.AddQueryPattern("set @@sql_mode = ''", &sqltypes.Result{})
- state, _, err := tsv.ReserveBeginExecute(ctx, &target, nil, nil, "set @@sql_mode = ''", nil, &querypb.ExecuteOptions{})
+ state, _, err := tsv.ReserveBeginExecute(ctx, nil, &target, nil, nil, "set @@sql_mode = ''", nil)
require.NoError(t, err)
assert.Greater(t, state.TransactionID, int64(0), "transactionID")
@@ -2343,7 +2348,7 @@ func TestReserveExecute_WithoutTx(t *testing.T) {
db.AddQueryPattern("set sql_mode = ''", &sqltypes.Result{})
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- state, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, 0, &querypb.ExecuteOptions{})
+ state, _, err := tsv.ReserveExecute(ctx, nil, &target, nil, "set sql_mode = ''", nil, 0)
require.NoError(t, err)
assert.NotEqual(t, int64(0), state.ReservedID, "reservedID should not be zero")
expected := []string{
@@ -2367,12 +2372,12 @@ func TestReserveExecute_WithTx(t *testing.T) {
db.AddQueryPattern("set sql_mode = ''", &sqltypes.Result{})
target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
- beginState, err := tsv.Begin(ctx, &target, &querypb.ExecuteOptions{})
+ beginState, err := tsv.Begin(ctx, nil, &target)
require.NoError(t, err)
require.NotEqual(t, int64(0), beginState.TransactionID)
db.ResetQueryLog()
- reserveState, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, beginState.TransactionID, &querypb.ExecuteOptions{})
+ reserveState, _, err := tsv.ReserveExecute(ctx, nil, &target, nil, "set sql_mode = ''", nil, beginState.TransactionID)
require.NoError(t, err)
defer tsv.Release(ctx, &target, beginState.TransactionID, reserveState.ReservedID)
assert.Equal(t, beginState.TransactionID, reserveState.ReservedID, "reservedID should be equal to transactionID")
@@ -2431,19 +2436,19 @@ func TestRelease(t *testing.T) {
switch {
case test.begin && test.reserve:
- state, _, err := tsv.ReserveBeginExecute(ctx, &target, nil, nil, "set sql_mode = ''", nil, &querypb.ExecuteOptions{})
+ state, _, err := tsv.ReserveBeginExecute(ctx, nil, &target, nil, nil, "set sql_mode = ''", nil)
require.NoError(t, err)
transactionID = state.TransactionID
reservedID = state.ReservedID
require.NotEqual(t, int64(0), transactionID)
require.NotEqual(t, int64(0), reservedID)
case test.begin:
- state, _, err := tsv.BeginExecute(ctx, &target, nil, "select 42", nil, 0, &querypb.ExecuteOptions{})
+ state, _, err := tsv.BeginExecute(ctx, nil, &target, nil, "select 42", nil, 0)
require.NoError(t, err)
transactionID = state.TransactionID
require.NotEqual(t, int64(0), transactionID)
case test.reserve:
- state, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, 0, &querypb.ExecuteOptions{})
+ state, _, err := tsv.ReserveExecute(ctx, nil, &target, nil, "set sql_mode = ''", nil, 0)
require.NoError(t, err)
reservedID = state.ReservedID
require.NotEqual(t, int64(0), reservedID)
@@ -2478,27 +2483,27 @@ func TestReserveStats(t *testing.T) {
ctx = callerid.NewContext(ctx, nil, callerID)
// Starts reserved connection and transaction
- rbeState, _, err := tsv.ReserveBeginExecute(ctx, &target, nil, nil, "set sql_mode = ''", nil, &querypb.ExecuteOptions{})
+ rbeState, _, err := tsv.ReserveBeginExecute(ctx, nil, &target, nil, nil, "set sql_mode = ''", nil)
require.NoError(t, err)
assert.EqualValues(t, 1, tsv.te.txPool.env.Stats().UserActiveReservedCount.Counts()["test"])
// Starts reserved connection
- reState, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, 0, &querypb.ExecuteOptions{})
+ reState, _, err := tsv.ReserveExecute(ctx, nil, &target, nil, "set sql_mode = ''", nil, 0)
require.NoError(t, err)
assert.EqualValues(t, 2, tsv.te.txPool.env.Stats().UserActiveReservedCount.Counts()["test"])
// Use previous reserved connection to start transaction
- reBeState, _, err := tsv.BeginExecute(ctx, &target, nil, "select 42", nil, reState.ReservedID, &querypb.ExecuteOptions{})
+ reBeState, _, err := tsv.BeginExecute(ctx, nil, &target, nil, "select 42", nil, reState.ReservedID)
require.NoError(t, err)
assert.EqualValues(t, 2, tsv.te.txPool.env.Stats().UserActiveReservedCount.Counts()["test"])
// Starts transaction.
- beState, _, err := tsv.BeginExecute(ctx, &target, nil, "select 42", nil, 0, &querypb.ExecuteOptions{})
+ beState, _, err := tsv.BeginExecute(ctx, nil, &target, nil, "select 42", nil, 0)
require.NoError(t, err)
assert.EqualValues(t, 2, tsv.te.txPool.env.Stats().UserActiveReservedCount.Counts()["test"])
// Reserved the connection on previous transaction
- beReState, _, err := tsv.ReserveExecute(ctx, &target, nil, "set sql_mode = ''", nil, beState.TransactionID, &querypb.ExecuteOptions{})
+ beReState, _, err := tsv.ReserveExecute(ctx, nil, &target, nil, "set sql_mode = ''", nil, beState.TransactionID)
require.NoError(t, err)
assert.EqualValues(t, 3, tsv.te.txPool.env.Stats().UserActiveReservedCount.Counts()["test"])
@@ -2544,11 +2549,11 @@ func TestDatabaseNameReplaceByKeyspaceNameExecuteMethod(t *testing.T) {
target := tsv.sm.target
// Testing Execute Method
- state, err := tsv.Begin(ctx, target, nil)
+ state, err := tsv.Begin(ctx, nil, target)
require.NoError(t, err)
- res, err := tsv.Execute(ctx, target, executeSQL, nil, state.TransactionID, 0, &querypb.ExecuteOptions{
+ res, err := tsv.Execute(ctx, &vtgatepb.Session{Options: &querypb.ExecuteOptions{
IncludedFields: querypb.ExecuteOptions_ALL,
- })
+ }}, target, executeSQL, nil, state.TransactionID, 0)
require.NoError(t, err)
for _, field := range res.Fields {
require.Equal(t, "keyspaceName", field.Database)
@@ -2590,9 +2595,9 @@ func TestDatabaseNameReplaceByKeyspaceNameStreamExecuteMethod(t *testing.T) {
}
return nil
}
- err := tsv.StreamExecute(ctx, target, executeSQL, nil, 0, 0, &querypb.ExecuteOptions{
+ err := tsv.StreamExecute(ctx, &vtgatepb.Session{Options: &querypb.ExecuteOptions{
IncludedFields: querypb.ExecuteOptions_ALL,
- }, callback)
+ }}, target, executeSQL, nil, 0, 0, callback)
require.NoError(t, err)
}
@@ -2621,9 +2626,9 @@ func TestDatabaseNameReplaceByKeyspaceNameBeginExecuteMethod(t *testing.T) {
target := tsv.sm.target
// Test BeginExecute Method
- state, res, err := tsv.BeginExecute(ctx, target, nil, executeSQL, nil, 0, &querypb.ExecuteOptions{
+ state, res, err := tsv.BeginExecute(ctx, &vtgatepb.Session{Options: &querypb.ExecuteOptions{
IncludedFields: querypb.ExecuteOptions_ALL,
- })
+ }}, target, nil, executeSQL, nil, 0)
require.NoError(t, err)
for _, field := range res.Fields {
require.Equal(t, "keyspaceName", field.Database)
@@ -2661,9 +2666,9 @@ func TestDatabaseNameReplaceByKeyspaceNameReserveExecuteMethod(t *testing.T) {
target := tsv.sm.target
// Test ReserveExecute
- _, res, err := tsv.ReserveExecute(ctx, target, nil, executeSQL, nil, 0, &querypb.ExecuteOptions{
+ _, res, err := tsv.ReserveExecute(ctx, &vtgatepb.Session{Options: &querypb.ExecuteOptions{
IncludedFields: querypb.ExecuteOptions_ALL,
- })
+ }}, target, nil, executeSQL, nil, 0)
require.NoError(t, err)
for _, field := range res.Fields {
require.Equal(t, "keyspaceName", field.Database)
@@ -2695,9 +2700,9 @@ func TestDatabaseNameReplaceByKeyspaceNameReserveBeginExecuteMethod(t *testing.T
target := tsv.sm.target
// Test for ReserveBeginExecute
- state, res, err := tsv.ReserveBeginExecute(ctx, target, nil, nil, executeSQL, nil, &querypb.ExecuteOptions{
+ state, res, err := tsv.ReserveBeginExecute(ctx, &vtgatepb.Session{Options: &querypb.ExecuteOptions{
IncludedFields: querypb.ExecuteOptions_ALL,
- })
+ }}, target, nil, nil, executeSQL, nil)
require.NoError(t, err)
for _, field := range res.Fields {
require.Equal(t, "keyspaceName", field.Database)
diff --git a/go/vtbench/client.go b/go/vtbench/client.go
index 70c9615b321..e141dba605d 100644
--- a/go/vtbench/client.go
+++ b/go/vtbench/client.go
@@ -55,7 +55,6 @@ func (c *mysqlClientConn) connect(ctx context.Context, cp ConnParams) error {
Pass: cp.Password,
UnixSocket: cp.UnixSocket,
})
-
if err != nil {
return err
}
@@ -162,5 +161,5 @@ func (c *grpcVttabletConn) connect(ctx context.Context, cp ConnParams) error {
}
func (c *grpcVttabletConn) execute(ctx context.Context, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
- return c.qs.Execute(ctx, &c.target, query, bindVars, 0, 0, nil)
+ return c.qs.Execute(ctx, nil, &c.target, query, bindVars, 0, 0)
}