diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go index d511e2d2ea0..f9ae8be3820 100644 --- a/go/vt/vtexplain/vtexplain_vtgate.go +++ b/go/vt/vtexplain/vtexplain_vtgate.go @@ -38,6 +38,7 @@ import ( "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate" "vitess.io/vitess/go/vt/vtgate/engine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/vt/vtgate/vindexes" "vitess.io/vitess/go/vt/vttablet/queryservice" @@ -235,7 +236,7 @@ func (vte *VTExplain) vtgateExecute(sql string) ([]*engine.Plan, map[string]*Tab // This will ensure that the commit/rollback order is predictable. vte.sortShardSession() - _, err := vte.vtgateExecutor.Execute(context.Background(), nil, "VtexplainExecute", vtgate.NewSafeSession(vte.vtgateSession), sql, nil) + _, err := vte.vtgateExecutor.Execute(context.Background(), nil, "VtexplainExecute", econtext.NewSafeSession(vte.vtgateSession), sql, nil) if err != nil { for _, tc := range vte.explainTopo.TabletConns { tc.tabletQueries = nil diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 1ba99c01ef2..2e65cefbabe 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -23,10 +23,10 @@ import ( "github.com/stretchr/testify/require" "vitess.io/vitess/go/sqltypes" - querypb "vitess.io/vitess/go/vt/proto/query" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" ) // This file contains tests for all the autocommit code paths @@ -382,7 +382,7 @@ func TestAutocommitTransactionStarted(t *testing.T) { // single shard query - no savepoint needed sql := "update `user` set a = 2 where id = 1" - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) require.Len(t, sbc1.Queries, 1) require.Equal(t, sql, sbc1.Queries[0].Sql) @@ -394,7 +394,7 @@ func TestAutocommitTransactionStarted(t *testing.T) { // multi shard query - savepoint needed sql = "update `user` set a = 2 where id in (1, 4)" expectedSql := "update `user` set a = 2 where id in ::__vals" - _, err = executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + _, err = executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) require.Len(t, sbc1.Queries, 2) require.Contains(t, sbc1.Queries[0].Sql, "savepoint") @@ -413,7 +413,7 @@ func TestAutocommitDirectTarget(t *testing.T) { } sql := "insert into `simple`(val) values ('val')" - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) assertQueries(t, sbclookup, []*querypb.BoundQuery{{ @@ -434,7 +434,7 @@ func TestAutocommitDirectRangeTarget(t *testing.T) { } sql := "delete from sharded_user_msgs limit 1000" - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) assertQueries(t, sbc1, []*querypb.BoundQuery{{ @@ -451,5 +451,5 @@ func autocommitExec(executor *Executor, sql string) (*sqltypes.Result, error) { TransactionMode: vtgatepb.TransactionMode_MULTI, } - return executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + return executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) } diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go index 928f42fca30..106751194a8 100644 --- a/go/vt/vtgate/executor.go +++ b/go/vt/vtgate/executor.go @@ -57,6 +57,7 @@ import ( "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/engine" "vitess.io/vitess/go/vt/vtgate/evalengine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/vt/vtgate/planbuilder" "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" @@ -223,7 +224,7 @@ func NewExecutor( } // Execute executes a non-streaming query. -func (e *Executor) Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, method string, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable) (result *sqltypes.Result, err error) { +func (e *Executor) Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, method string, safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable) (result *sqltypes.Result, err error) { span, ctx := trace.NewSpan(ctx, "executor.Execute") span.Annotate("method", method) trace.AnnotateSQL(span, sqlparser.Preview(sql)) @@ -286,7 +287,7 @@ func (e *Executor) StreamExecute( ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, method string, - safeSession *SafeSession, + safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error, @@ -411,12 +412,12 @@ func canReturnRows(stmtType sqlparser.StatementType) bool { } } -func saveSessionStats(safeSession *SafeSession, stmtType sqlparser.StatementType, rowsAffected, insertID uint64, rowsReturned int, err error) { +func saveSessionStats(safeSession *econtext.SafeSession, stmtType sqlparser.StatementType, rowsAffected, insertID uint64, rowsReturned int, err error) { safeSession.RowCount = -1 if err != nil { return } - if !safeSession.foundRowsHandled { + if !safeSession.IsFoundRowsHandled() { safeSession.FoundRows = uint64(rowsReturned) } if insertID > 0 { @@ -430,7 +431,7 @@ func saveSessionStats(safeSession *SafeSession, stmtType sqlparser.StatementType } } -func (e *Executor) execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) (sqlparser.StatementType, *sqltypes.Result, error) { +func (e *Executor) execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) (sqlparser.StatementType, *sqltypes.Result, error) { var err error var qr *sqltypes.Result var stmtType sqlparser.StatementType @@ -448,7 +449,7 @@ func (e *Executor) execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConn } // addNeededBindVars adds bind vars that are needed by the plan -func (e *Executor) addNeededBindVars(vcursor *vcursorImpl, bindVarNeeds *sqlparser.BindVarNeeds, bindVars map[string]*querypb.BindVariable, session *SafeSession) error { +func (e *Executor) addNeededBindVars(vcursor *vcursorImpl, bindVarNeeds *sqlparser.BindVarNeeds, bindVars map[string]*querypb.BindVariable, session *econtext.SafeSession) error { for _, funcName := range bindVarNeeds.NeedFunctionResult { switch funcName { case sqlparser.DBVarName: @@ -572,21 +573,21 @@ func (e *Executor) addNeededBindVars(vcursor *vcursorImpl, bindVarNeeds *sqlpars return nil } -func ifOptionsExist(session *SafeSession, f func(*querypb.ExecuteOptions)) { +func ifOptionsExist(session *econtext.SafeSession, f func(*querypb.ExecuteOptions)) { options := session.GetOptions() if options != nil { f(options) } } -func ifReadAfterWriteExist(session *SafeSession, f func(*vtgatepb.ReadAfterWrite)) { +func ifReadAfterWriteExist(session *econtext.SafeSession, f func(*vtgatepb.ReadAfterWrite)) { raw := session.ReadAfterWrite if raw != nil { f(raw) } } -func (e *Executor) handleBegin(ctx context.Context, safeSession *SafeSession, logStats *logstats.LogStats, stmt sqlparser.Statement) (*sqltypes.Result, error) { +func (e *Executor) handleBegin(ctx context.Context, safeSession *econtext.SafeSession, logStats *logstats.LogStats, stmt sqlparser.Statement) (*sqltypes.Result, error) { execStart := time.Now() logStats.PlanTime = execStart.Sub(logStats.StartTime) @@ -599,7 +600,7 @@ func (e *Executor) handleBegin(ctx context.Context, safeSession *SafeSession, lo return &sqltypes.Result{}, err } -func (e *Executor) handleCommit(ctx context.Context, safeSession *SafeSession, logStats *logstats.LogStats) (*sqltypes.Result, error) { +func (e *Executor) handleCommit(ctx context.Context, safeSession *econtext.SafeSession, logStats *logstats.LogStats) (*sqltypes.Result, error) { execStart := time.Now() logStats.PlanTime = execStart.Sub(logStats.StartTime) logStats.ShardQueries = uint64(len(safeSession.ShardSessions)) @@ -611,11 +612,11 @@ func (e *Executor) handleCommit(ctx context.Context, safeSession *SafeSession, l } // Commit commits the existing transactions -func (e *Executor) Commit(ctx context.Context, safeSession *SafeSession) error { +func (e *Executor) Commit(ctx context.Context, safeSession *econtext.SafeSession) error { return e.txConn.Commit(ctx, safeSession) } -func (e *Executor) handleRollback(ctx context.Context, safeSession *SafeSession, logStats *logstats.LogStats) (*sqltypes.Result, error) { +func (e *Executor) handleRollback(ctx context.Context, safeSession *econtext.SafeSession, logStats *logstats.LogStats) (*sqltypes.Result, error) { execStart := time.Now() logStats.PlanTime = execStart.Sub(logStats.StartTime) logStats.ShardQueries = uint64(len(safeSession.ShardSessions)) @@ -625,7 +626,7 @@ func (e *Executor) handleRollback(ctx context.Context, safeSession *SafeSession, return &sqltypes.Result{}, err } -func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession, sql string, planType string, logStats *logstats.LogStats, nonTxResponse func(query string) (*sqltypes.Result, error), ignoreMaxMemoryRows bool) (*sqltypes.Result, error) { +func (e *Executor) handleSavepoint(ctx context.Context, safeSession *econtext.SafeSession, sql string, planType string, logStats *logstats.LogStats, nonTxResponse func(query string) (*sqltypes.Result, error), ignoreMaxMemoryRows bool) (*sqltypes.Result, error) { execStart := time.Now() logStats.PlanTime = execStart.Sub(logStats.StartTime) logStats.ShardQueries = uint64(len(safeSession.ShardSessions)) @@ -637,7 +638,7 @@ func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession // If no transaction exists on any of the shard sessions, // then savepoint does not need to be executed, it will be only stored in the session // and later will be executed when a transaction is started. - if !safeSession.isTxOpen() { + if !safeSession.IsTxOpen() { if safeSession.InTransaction() { // Storing, as this needs to be executed just after starting transaction on the shard. safeSession.StoreSavepoint(sql) @@ -645,7 +646,7 @@ func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession } return nonTxResponse(sql) } - orig := safeSession.commitOrder + orig := safeSession.GetCommitOrder() qr, err := e.executeSPInAllSessions(ctx, safeSession, sql, ignoreMaxMemoryRows) safeSession.SetCommitOrder(orig) if err != nil { @@ -657,7 +658,7 @@ func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession // executeSPInAllSessions function executes the savepoint query in all open shard sessions (pre, normal and post) // which has non-zero transaction id (i.e. an open transaction on the shard connection). -func (e *Executor) executeSPInAllSessions(ctx context.Context, safeSession *SafeSession, sql string, ignoreMaxMemoryRows bool) (*sqltypes.Result, error) { +func (e *Executor) executeSPInAllSessions(ctx context.Context, safeSession *econtext.SafeSession, sql string, ignoreMaxMemoryRows bool) (*sqltypes.Result, error) { var qr *sqltypes.Result var errs []error for _, co := range []vtgatepb.CommitOrder{vtgatepb.CommitOrder_PRE, vtgatepb.CommitOrder_NORMAL, vtgatepb.CommitOrder_POST} { @@ -665,7 +666,7 @@ func (e *Executor) executeSPInAllSessions(ctx context.Context, safeSession *Safe var rss []*srvtopo.ResolvedShard var queries []*querypb.BoundQuery - for _, shardSession := range safeSession.getSessions() { + for _, shardSession := range safeSession.GetSessions() { // This will avoid executing savepoint on reserved connections // which has no open transaction. if shardSession.TransactionId == 0 { @@ -718,7 +719,7 @@ func (e *Executor) handleKill(ctx context.Context, mysqlCtx vtgateservice.MySQLC // CloseSession releases the current connection, which rollbacks open transactions and closes reserved connections. // It is called then the MySQL servers closes the connection to its client. -func (e *Executor) CloseSession(ctx context.Context, safeSession *SafeSession) error { +func (e *Executor) CloseSession(ctx context.Context, safeSession *econtext.SafeSession) error { return e.txConn.ReleaseAll(ctx, safeSession) } @@ -1088,11 +1089,6 @@ func (e *Executor) ParseDestinationTarget(targetString string) (string, topodata return destKeyspace, destTabletType, dest, err } -type iQueryOption interface { - cachePlan() bool - getSelectLimit() int -} - // getPlan computes the plan for the given query. If one is in // the cache, it reuses it. func (e *Executor) getPlan( @@ -1136,7 +1132,7 @@ func (e *Executor) getPlan( bindVars, parameterize, vcursor.keyspace, - vcursor.safeSession.getSelectLimit(), + vcursor.safeSession.GetSelectLimit(), setVarComment, vcursor.safeSession.SystemVariables, vcursor.GetForeignKeyChecksState(), @@ -1195,7 +1191,7 @@ func (e *Executor) cacheAndBuildStatement( bindVarNeeds *sqlparser.BindVarNeeds, logStats *logstats.LogStats, ) (*engine.Plan, error) { - planCachable := sqlparser.CachePlan(stmt) && vcursor.safeSession.cachePlan() + planCachable := sqlparser.CachePlan(stmt) && vcursor.safeSession.CachePlan() if planCachable { planKey := e.hashPlan(ctx, vcursor, query) @@ -1354,7 +1350,7 @@ func isValidPayloadSize(query string) bool { } // Prepare executes a prepare statements. -func (e *Executor) Prepare(ctx context.Context, method string, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable) (fld []*querypb.Field, err error) { +func (e *Executor) Prepare(ctx context.Context, method string, safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable) (fld []*querypb.Field, err error) { logStats := logstats.NewLogStats(ctx, method, sql, safeSession.GetSessionUUID(), bindVars) fld, err = e.prepare(ctx, safeSession, sql, bindVars, logStats) logStats.Error = err @@ -1373,7 +1369,7 @@ func (e *Executor) Prepare(ctx context.Context, method string, safeSession *Safe return fld, err } -func (e *Executor) prepare(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) ([]*querypb.Field, error) { +func (e *Executor) prepare(ctx context.Context, safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) ([]*querypb.Field, error) { // Start an implicit transaction if necessary. if !safeSession.Autocommit && !safeSession.InTransaction() { if err := e.txConn.Begin(ctx, safeSession, nil); err != nil { @@ -1409,7 +1405,7 @@ func (e *Executor) prepare(ctx context.Context, safeSession *SafeSession, sql st return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unrecognized prepare statement: %s", sql) } -func (e *Executor) handlePrepare(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) ([]*querypb.Field, error) { +func (e *Executor) handlePrepare(ctx context.Context, safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) ([]*querypb.Field, error) { query, comments := sqlparser.SplitMarginComments(sql) vcursor, _ := newVCursorImpl(safeSession, comments, e, logStats, e.vm, e.VSchema(), e.resolver.resolver, e.serv, e.warnShardedOnly, e.pv) @@ -1460,17 +1456,17 @@ func parseAndValidateQuery(query string, parser *sqlparser.Parser) (sqlparser.St } // ExecuteMultiShard implements the IExecutor interface -func (e *Executor) ExecuteMultiShard(ctx context.Context, primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool, ignoreMaxMemoryRows bool, resultsObserver resultsObserver) (qr *sqltypes.Result, errs []error) { +func (e *Executor) ExecuteMultiShard(ctx context.Context, primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *econtext.SafeSession, autocommit bool, ignoreMaxMemoryRows bool, resultsObserver resultsObserver) (qr *sqltypes.Result, errs []error) { return e.scatterConn.ExecuteMultiShard(ctx, primitive, rss, queries, session, autocommit, ignoreMaxMemoryRows, resultsObserver) } // StreamExecuteMulti implements the IExecutor interface -func (e *Executor) StreamExecuteMulti(ctx context.Context, primitive engine.Primitive, query string, rss []*srvtopo.ResolvedShard, vars []map[string]*querypb.BindVariable, session *SafeSession, autocommit bool, callback func(reply *sqltypes.Result) error, resultsObserver resultsObserver) []error { +func (e *Executor) StreamExecuteMulti(ctx context.Context, primitive engine.Primitive, query string, rss []*srvtopo.ResolvedShard, vars []map[string]*querypb.BindVariable, session *econtext.SafeSession, autocommit bool, callback func(reply *sqltypes.Result) error, resultsObserver resultsObserver) []error { return e.scatterConn.StreamExecuteMulti(ctx, primitive, query, rss, vars, session, autocommit, callback, resultsObserver) } // ExecuteLock implements the IExecutor interface -func (e *Executor) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) { +func (e *Executor) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *econtext.SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) { return e.scatterConn.ExecuteLock(ctx, rs, query, session, lockFuncType) } @@ -1581,7 +1577,7 @@ func getTabletThrottlerStatus(tabletHostPort string) (string, error) { } // ReleaseLock implements the IExecutor interface -func (e *Executor) ReleaseLock(ctx context.Context, session *SafeSession) error { +func (e *Executor) ReleaseLock(ctx context.Context, session *econtext.SafeSession) error { return e.txConn.ReleaseLock(ctx, session) } diff --git a/go/vt/vtgate/executor_ddl_test.go b/go/vt/vtgate/executor_ddl_test.go index 3274fd94475..770b61c16db 100644 --- a/go/vt/vtgate/executor_ddl_test.go +++ b/go/vt/vtgate/executor_ddl_test.go @@ -21,6 +21,7 @@ import ( "testing" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "github.com/stretchr/testify/require" ) @@ -56,7 +57,7 @@ func TestDDLFlags(t *testing.T) { for _, testcase := range testcases { t.Run(fmt.Sprintf("%s-%v-%v", testcase.sql, testcase.enableDirectDDL, testcase.enableOnlineDDL), func(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) enableDirectDDL = testcase.enableDirectDDL enableOnlineDDL = testcase.enableOnlineDDL _, err := executor.Execute(ctx, nil, "TestDDLFlags", session, testcase.sql, nil) diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index 3dce4e212ef..09ef6032533 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -25,6 +25,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/mysql/config" "vitess.io/vitess/go/mysql/sqlerror" "vitess.io/vitess/go/sqltypes" @@ -234,7 +236,7 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { )} executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) - safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + safeSession := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where nv_lu_col = 2", @@ -296,7 +298,7 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { )} executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) - safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + safeSession := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where erl_lu_col = 2", @@ -358,7 +360,7 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { )} executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) - safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + safeSession := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where srl_lu_col = 2", @@ -420,7 +422,7 @@ func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { )} executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) - safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + safeSession := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where nrl_lu_col = 2", @@ -2066,7 +2068,7 @@ func TestInsertPartialFail1(t *testing.T) { context.Background(), nil, "TestExecute", - NewSafeSession(&vtgatepb.Session{InTransaction: true}), + econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}), "insert into user(id, v, name) values (1, 2, 'myname')", nil, ) @@ -2082,7 +2084,7 @@ func TestInsertPartialFail2(t *testing.T) { // Make the second DML fail, it should result in a rollback. sbc1.MustFailExecute[sqlparser.StmtInsert] = 1 - safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + safeSession := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executor.Execute( context.Background(), nil, @@ -2656,7 +2658,7 @@ func TestReservedConnDML(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestReservedConnDML") defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true}) _, err := executor.Execute(ctx, nil, "TestReservedConnDML", session, "use "+KsTestUnsharded, nil) require.NoError(t, err) @@ -2708,7 +2710,7 @@ func TestStreamingDML(t *testing.T) { logChan := executor.queryLogger.Subscribe(method) defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) tcases := []struct { query string @@ -2792,7 +2794,7 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) require.True(t, session.GetAutocommit()) require.False(t, session.InTransaction()) @@ -2845,7 +2847,7 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) require.True(t, session.GetAutocommit()) require.False(t, session.InTransaction()) @@ -2886,7 +2888,7 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { func TestMultiInternalSavepoint(t *testing.T) { executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) _, err := executorExecSession(ctx, executor, "begin", nil, session.Session) require.NoError(t, err) @@ -2935,7 +2937,7 @@ func TestInsertSelectFromDual(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) query := "insert into user(id, v, name) select 1, 2, 'myname' from dual" wantQueries := []*querypb.BoundQuery{{ @@ -2990,7 +2992,7 @@ func TestInsertSelectFromTable(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) query := "insert into user(id, name) select c1, c2 from music" wantQueries := []*querypb.BoundQuery{{ diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index 332139c4a78..2ee3425209f 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -28,6 +28,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/cache/theine" "vitess.io/vitess/go/constants/sidecar" "vitess.io/vitess/go/sqltypes" @@ -307,7 +309,7 @@ func executorExecSession(ctx context.Context, executor *Executor, sql string, bv ctx, nil, "TestExecute", - NewSafeSession(session), + econtext.NewSafeSession(session), sql, bv) } @@ -320,7 +322,7 @@ func executorPrepare(ctx context.Context, executor *Executor, session *vtgatepb. return executor.Prepare( ctx, "TestExecute", - NewSafeSession(session), + econtext.NewSafeSession(session), sql, bv) } @@ -331,7 +333,7 @@ func executorStream(ctx context.Context, executor *Executor, sql string) (qr *sq ctx, nil, "TestExecuteStream", - NewSafeSession(nil), + econtext.NewSafeSession(nil), sql, nil, func(qr *sqltypes.Result) error { diff --git a/go/vt/vtgate/executor_scatter_stats_test.go b/go/vt/vtgate/executor_scatter_stats_test.go index 84dd2744e8b..b665f850a23 100644 --- a/go/vt/vtgate/executor_scatter_stats_test.go +++ b/go/vt/vtgate/executor_scatter_stats_test.go @@ -24,12 +24,13 @@ import ( "github.com/stretchr/testify/require" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" ) func TestScatterStatsWithNoScatterQuery(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from main1", nil) require.NoError(t, err) @@ -41,7 +42,7 @@ func TestScatterStatsWithNoScatterQuery(t *testing.T) { func TestScatterStatsWithSingleScatterQuery(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) require.NoError(t, err) @@ -53,7 +54,7 @@ func TestScatterStatsWithSingleScatterQuery(t *testing.T) { func TestScatterStatsHttpWriting(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) require.NoError(t, err) diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 8ba89d25daf..9fd112dbdb4 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -30,6 +30,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + _flag "vitess.io/vitess/go/internal/flag" "vitess.io/vitess/go/mysql/collations" "vitess.io/vitess/go/sqltypes" @@ -59,7 +61,7 @@ func TestSelectNext(t *testing.T) { }} // Autocommit - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) _, err := executor.Execute(context.Background(), nil, "TestSelectNext", session, query, bv) require.NoError(t, err) @@ -69,7 +71,7 @@ func TestSelectNext(t *testing.T) { sbclookup.Queries = nil // Txn - session = NewAutocommitSession(&vtgatepb.Session{}) + session = econtext.NewAutocommitSession(&vtgatepb.Session{}) session.Session.InTransaction = true _, err = executor.Execute(context.Background(), nil, "TestSelectNext", session, query, bv) require.NoError(t, err) @@ -80,7 +82,7 @@ func TestSelectNext(t *testing.T) { sbclookup.Queries = nil // Reserve - session = NewAutocommitSession(&vtgatepb.Session{}) + session = econtext.NewAutocommitSession(&vtgatepb.Session{}) session.Session.InReservedConn = true _, err = executor.Execute(context.Background(), nil, "TestSelectNext", session, query, bv) require.NoError(t, err) @@ -91,7 +93,7 @@ func TestSelectNext(t *testing.T) { sbclookup.Queries = nil // Reserve and Txn - session = NewAutocommitSession(&vtgatepb.Session{}) + session = econtext.NewAutocommitSession(&vtgatepb.Session{}) session.Session.InReservedConn = true session.Session.InTransaction = true _, err = executor.Execute(context.Background(), nil, "TestSelectNext", session, query, bv) @@ -107,7 +109,7 @@ func TestSelectDBA(t *testing.T) { query := "select * from INFORMATION_SCHEMA.foo" _, err := executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -117,7 +119,7 @@ func TestSelectDBA(t *testing.T) { sbc1.Queries = nil query = "SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES ist WHERE ist.table_schema = 'performance_schema' AND ist.table_name = 'foo'" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -133,7 +135,7 @@ func TestSelectDBA(t *testing.T) { sbc1.Queries = nil query = "select 1 from information_schema.table_constraints where constraint_schema = 'vt_ks' and table_name = 'user'" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -149,7 +151,7 @@ func TestSelectDBA(t *testing.T) { sbc1.Queries = nil query = "select 1 from information_schema.table_constraints where constraint_schema = 'vt_ks'" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -167,7 +169,7 @@ func TestSystemVariablesMySQLBelow80(t *testing.T) { executor.normalize = true setVarEnabled = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -204,7 +206,7 @@ func TestSystemVariablesWithSetVarDisabled(t *testing.T) { defer func() { setVarEnabled = true }() - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -237,7 +239,7 @@ func TestSetSystemVariablesTx(t *testing.T) { executor, sbc1, _, _, _ := createCustomExecutor(t, "{}", "8.0.1") executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: "TestExecutor"}) _, err := executor.Execute(context.Background(), nil, "TestBegin", session, "begin", map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -283,7 +285,7 @@ func TestSetSystemVariables(t *testing.T) { executor, _, _, lookup, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded, SystemVariables: map[string]string{}}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded, SystemVariables: map[string]string{}}) // Set @@sql_mode and execute a select statement. We should have SET_VAR in the select statement @@ -394,7 +396,7 @@ func TestSetSystemVariablesWithReservedConnection(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -437,7 +439,7 @@ func TestSelectVindexFunc(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) query := "select * from hash_index where id = 1" - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) _, err := executor.Execute(context.Background(), nil, "TestSelectVindexFunc", session, query, nil) require.ErrorContains(t, err, "VT09005: no database selected") @@ -450,7 +452,7 @@ func TestCreateTableValidTimestamp(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor", SystemVariables: map[string]string{"sql_mode": "ALLOW_INVALID_DATES"}}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor", SystemVariables: map[string]string{"sql_mode": "ALLOW_INVALID_DATES"}}) query := "create table aa(t timestamp default 0)" _, err := executor.Execute(context.Background(), nil, "TestSelect", session, query, map[string]*querypb.BindVariable{}) @@ -472,7 +474,7 @@ func TestGen4SelectDBA(t *testing.T) { query := "select * from INFORMATION_SCHEMA.TABLE_CONSTRAINTS" _, err := executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -483,7 +485,7 @@ func TestGen4SelectDBA(t *testing.T) { sbc1.Queries = nil query = "SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES ist WHERE ist.table_schema = 'performance_schema' AND ist.table_name = 'foo'" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -501,7 +503,7 @@ func TestGen4SelectDBA(t *testing.T) { sbc1.Queries = nil query = "select 1 from information_schema.table_constraints where constraint_schema = 'vt_ks' and table_name = 'user'" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -519,7 +521,7 @@ func TestGen4SelectDBA(t *testing.T) { sbc1.Queries = nil query = "select 1 from information_schema.table_constraints where constraint_schema = 'vt_ks'" - _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}) + _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}), query, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select :vtg1 /* INT64 */ from information_schema.table_constraints where constraint_schema = :__vtschemaname /* VARCHAR */", @@ -534,7 +536,7 @@ func TestGen4SelectDBA(t *testing.T) { sbc1.Queries = nil query = "select t.table_schema,t.table_name,c.column_name,c.column_type from tables t join columns c on c.table_schema = t.table_schema and c.table_name = t.table_name where t.table_schema = 'TestExecutor' and c.table_schema = 'TestExecutor' order by t.table_schema,t.table_name,c.column_name" _, err = executor.Execute(context.Background(), nil, "TestSelectDBA", - NewSafeSession(&vtgatepb.Session{TargetString: "information_schema"}), + econtext.NewSafeSession(&vtgatepb.Session{TargetString: "information_schema"}), query, map[string]*querypb.BindVariable{}, ) require.NoError(t, err) @@ -651,7 +653,7 @@ func TestStreamBuffering(t *testing.T) { context.Background(), nil, "TestStreamBuffering", - NewSafeSession(session), + econtext.NewSafeSession(session), "select id from music_user_map where id = 1", nil, func(qr *sqltypes.Result) error { @@ -723,7 +725,7 @@ func TestStreamLimitOffset(t *testing.T) { context.Background(), nil, "TestStreamLimitOffset", - NewSafeSession(session), + econtext.NewSafeSession(session), "select id, textcol from user order by id limit 2 offset 2", nil, func(qr *sqltypes.Result) error { @@ -1083,7 +1085,7 @@ func TestSelectDatabase(t *testing.T) { newSession := &vtgatepb.Session{ TargetString: "@primary", } - session := NewSafeSession(newSession) + session := econtext.NewSafeSession(newSession) session.TargetString = "TestExecutor@primary" result, err := executor.Execute( context.Background(), @@ -2951,7 +2953,7 @@ func TestSubQueryAndQueryWithLimit(t *testing.T) { sbc1.SetResults(result1) sbc2.SetResults(result2) - exec(executor, NewSafeSession(&vtgatepb.Session{ + exec(executor, econtext.NewSafeSession(&vtgatepb.Session{ TargetString: "@primary", }), "select id1, id2 from t1 where id1 >= ( select id1 from t1 order by id1 asc limit 1) limit 100") require.Equal(t, 2, len(sbc1.Queries)) @@ -3000,7 +3002,7 @@ func TestSelectUsingMultiEqualOnLookupColumn(t *testing.T) { }}, }}) - result, err := exec(executor, NewSafeSession(&vtgatepb.Session{ + result, err := exec(executor, econtext.NewSafeSession(&vtgatepb.Session{ TargetString: KsTestSharded, }), "select nv_lu_col, other from t2_lookup WHERE (nv_lu_col = 1 AND other = 'bar') OR (nv_lu_col = 2 AND other = 'baz') OR (nv_lu_col = 3 AND other = 'qux') OR (nv_lu_col = 4 AND other = 'brz') OR (nv_lu_col = 5 AND other = 'brz')") @@ -3197,7 +3199,7 @@ func TestSelectWithUnionAll(t *testing.T) { func TestSelectLock(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) - session := NewSafeSession(nil) + session := econtext.NewSafeSession(nil) session.Session.InTransaction = true session.ShardSessions = []*vtgatepb.Session_ShardSession{{ Target: &querypb.Target{ @@ -3265,7 +3267,7 @@ func TestLockReserve(t *testing.T) { "select release_lock('lock name') from dual", } - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) for _, sql := range tcases { t.Run(sql, func(t *testing.T) { @@ -3283,7 +3285,7 @@ func TestLockReserve(t *testing.T) { func TestSelectFromInformationSchema(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) - session := NewSafeSession(nil) + session := econtext.NewSafeSession(nil) // check failure when trying to query two keyspaces _, err := exec(executor, session, "SELECT B.TABLE_NAME FROM INFORMATION_SCHEMA.TABLES AS A, INFORMATION_SCHEMA.COLUMNS AS B WHERE A.TABLE_SCHEMA = 'TestExecutor' AND A.TABLE_SCHEMA = 'TestXBadSharding'") @@ -3411,7 +3413,7 @@ func TestGen4SelectStraightJoin(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select u.id from user u straight_join user2 u2 on u.id = u2.id" _, err := executor.Execute(context.Background(), nil, "TestGen4SelectStraightJoin", @@ -3434,7 +3436,7 @@ func TestGen4MultiColumnVindexEqual(t *testing.T) { executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select * from user_region where cola = 1 and colb = 2" _, err := executor.Execute(context.Background(), nil, "TestGen4MultiColumnVindex", session, query, map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -3473,7 +3475,7 @@ func TestGen4MultiColumnVindexIn(t *testing.T) { executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select * from user_region where cola IN (1,17984) and colb IN (2,3,4)" _, err := executor.Execute(context.Background(), nil, "TestGen4MultiColumnVindex", session, query, map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -3512,7 +3514,7 @@ func TestGen4MultiColMixedColComparision(t *testing.T) { executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select * from user_region where colb = 2 and cola IN (1,17984)" _, err := executor.Execute(context.Background(), nil, "TestGen4MultiColMixedColComparision", session, query, map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -3549,7 +3551,7 @@ func TestGen4MultiColBestVindexSel(t *testing.T) { executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select * from user_region where colb = 2 and cola IN (1,17984) and cola = 1" _, err := executor.Execute(context.Background(), nil, "TestGen4MultiColBestVindexSel", session, query, map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -3595,7 +3597,7 @@ func TestGen4MultiColMultiEqual(t *testing.T) { executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) query := "select * from user_region where (cola,colb) in ((17984,2),(17984,3))" _, err := executor.Execute(context.Background(), nil, "TestGen4MultiColMultiEqual", session, query, map[string]*querypb.BindVariable{}) require.NoError(t, err) @@ -3769,7 +3771,7 @@ func TestRegionRange(t *testing.T) { for _, tcase := range tcases { t.Run(strconv.Itoa(tcase.regionID), func(t *testing.T) { sql := fmt.Sprintf("select * from user_region where cola = %d", tcase.regionID) - _, err := executor.Execute(context.Background(), nil, "TestRegionRange", NewAutocommitSession(&vtgatepb.Session{}), sql, nil) + _, err := executor.Execute(context.Background(), nil, "TestRegionRange", econtext.NewAutocommitSession(&vtgatepb.Session{}), sql, nil) require.NoError(t, err) count := 0 for _, sbc := range conns { @@ -3817,7 +3819,7 @@ func TestMultiCol(t *testing.T) { shards: []string{"20a0-"}, }} - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) for _, tcase := range tcases { t.Run(fmt.Sprintf("%d_%d_%d", tcase.cola, tcase.colb, tcase.colc), func(t *testing.T) { @@ -3907,7 +3909,7 @@ func TestMultiColPartial(t *testing.T) { shards: []string{"20a0c0-"}, }} - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) for _, tcase := range tcases { t.Run(tcase.where, func(t *testing.T) { @@ -4197,7 +4199,7 @@ func TestSelectAggregationRandom(t *testing.T) { executor := createExecutor(ctx, serv, cell, resolver) defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) rs, err := executor.Execute(context.Background(), nil, "TestSelectCFC", session, "select /*vt+ PLANNER=gen4 */ A.a, A.b, (A.a / A.b) as c from (select sum(a) as a, sum(b) as b from user) A", nil) require.NoError(t, err) @@ -4207,7 +4209,7 @@ func TestSelectAggregationRandom(t *testing.T) { func TestSelectDateTypes(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) qr, err := executor.Execute(context.Background(), nil, "TestSelectDateTypes", session, "select '2020-01-01' + interval month(date_sub(FROM_UNIXTIME(1234), interval 1 month))-1 month", nil) require.NoError(t, err) @@ -4218,7 +4220,7 @@ func TestSelectDateTypes(t *testing.T) { func TestSelectHexAndBit(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) qr, err := executor.Execute(context.Background(), nil, "TestSelectHexAndBit", session, "select 0b1001, b'1001', 0x9, x'09'", nil) require.NoError(t, err) @@ -4234,7 +4236,7 @@ func TestSelectHexAndBit(t *testing.T) { func TestSelectCFC(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) _, err := executor.Execute(context.Background(), nil, "TestSelectCFC", session, "select /*vt+ PLANNER=gen4 */ c2 from tbl_cfc where c1 like 'A%'", nil) require.NoError(t, err) @@ -4263,7 +4265,7 @@ func TestSelectView(t *testing.T) { require.NoError(t, err) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) _, err = executor.Execute(context.Background(), nil, "TestSelectView", session, "select * from user_details_view", nil) require.NoError(t, err) @@ -4304,7 +4306,7 @@ func TestWarmingReads(t *testing.T) { executor, primary, replica := createExecutorEnvWithPrimaryReplicaConn(t, ctx, 100) executor.normalize = true - session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) // Since queries on the replica will run in a separate go-routine, we need synchronization for the Queries field in the sandboxconn. replica.RequireQueriesLocking() @@ -4428,7 +4430,7 @@ func TestStreamJoinQuery(t *testing.T) { func TestSysVarGlobalAndSession(t *testing.T) { executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("innodb_lock_wait_timeout", "uint64"), "20"), diff --git a/go/vt/vtgate/executor_set_test.go b/go/vt/vtgate/executor_set_test.go index 12e8e272bd7..62101639a11 100644 --- a/go/vt/vtgate/executor_set_test.go +++ b/go/vt/vtgate/executor_set_test.go @@ -22,6 +22,7 @@ import ( "vitess.io/vitess/go/mysql/sqlerror" querypb "vitess.io/vitess/go/vt/proto/query" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/test/utils" @@ -266,7 +267,7 @@ func TestExecutorSet(t *testing.T) { }} for i, tcase := range testcases { t.Run(fmt.Sprintf("%d-%s", i, tcase.in), func(t *testing.T) { - session := NewSafeSession(&vtgatepb.Session{Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: true}) _, err := executorEnv.Execute(ctx, nil, "TestExecute", session, tcase.in, nil) if tcase.err == "" { require.NoError(t, err) @@ -374,7 +375,7 @@ func TestExecutorSetOp(t *testing.T) { }} for _, tcase := range testcases { t.Run(tcase.in, func(t *testing.T) { - session := NewAutocommitSession(&vtgatepb.Session{ + session := econtext.NewAutocommitSession(&vtgatepb.Session{ TargetString: "@primary", }) session.TargetString = KsTestUnsharded @@ -392,7 +393,7 @@ func TestExecutorSetMetadata(t *testing.T) { t.Run("Session 1", func(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_keyspace_v1= '1'" _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) @@ -406,7 +407,7 @@ func TestExecutorSetMetadata(t *testing.T) { }() executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_keyspace_v1= '1'" _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) @@ -469,7 +470,7 @@ func TestPlanExecutorSetUDV(t *testing.T) { }} for _, tcase := range testcases { t.Run(tcase.in, func(t *testing.T) { - session := NewSafeSession(&vtgatepb.Session{Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: true}) _, err := executor.Execute(ctx, nil, "TestExecute", session, tcase.in, nil) if err != nil { require.EqualError(t, err, tcase.err) @@ -515,7 +516,7 @@ func TestSetVar(t *testing.T) { executor, _, _, sbc, ctx := createCustomExecutor(t, "{}", "8.0.0") executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded}) sbc.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("orig|new", "varchar|varchar"), @@ -554,7 +555,7 @@ func TestSetVarShowVariables(t *testing.T) { executor, _, _, sbc, ctx := createCustomExecutor(t, "{}", "8.0.0") executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, TargetString: KsTestUnsharded}) sbc.SetResults([]*sqltypes.Result{ // select query result for checking any change in system settings @@ -597,7 +598,7 @@ func TestExecutorSetAndSelect(t *testing.T) { sysVar: "tx_isolation", exp: `[[VARCHAR("READ-UNCOMMITTED")]]`, // this returns the value set in previous query. }} - session := NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded, EnableSystemSettings: true}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded, EnableSystemSettings: true}) for _, tcase := range testcases { t.Run(fmt.Sprintf("%s-%s", tcase.sysVar, tcase.val), func(t *testing.T) { sbc.ExecCount.Store(0) // reset the value @@ -631,7 +632,7 @@ func TestExecutorSetAndSelect(t *testing.T) { func TestExecutorTimeZone(t *testing.T) { e, _, _, _, ctx := createExecutorEnv(t) - session := NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded, EnableSystemSettings: true}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded, EnableSystemSettings: true}) session.SetSystemVariable("time_zone", "'+08:00'") qr, err := e.Execute(ctx, nil, "TestExecutorSetAndSelect", session, "select now()", nil) diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index b8cfeaf3cd5..a8500dd59c4 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -31,6 +31,7 @@ import ( topodatapb "vitess.io/vitess/go/vt/proto/topodata" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "vitess.io/vitess/go/vt/vtenv" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/logstats" _ "vitess.io/vitess/go/vt/vtgate/vindexes" "vitess.io/vitess/go/vt/vttablet/sandboxconn" @@ -102,7 +103,7 @@ func executorStreamMessages(executor *Executor, sql string) (qr *sqltypes.Result ctx, nil, "TestExecuteStream", - NewSafeSession(session), + econtext.NewSafeSession(session), sql, nil, func(qr *sqltypes.Result) error { diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index 3732a37d1d1..717367b5125 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -36,6 +36,8 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/mysql/collations" "vitess.io/vitess/go/mysql/sqlerror" "vitess.io/vitess/go/sqltypes" @@ -64,7 +66,7 @@ func TestExecutorResultsExceeded(t *testing.T) { warnMemoryRows = 3 defer func() { warnMemoryRows = save }() - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) initial := warnings.Counts()["ResultsExceeded"] @@ -88,7 +90,7 @@ func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { maxMemoryRows = 3 defer func() { maxMemoryRows = save }() - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) result := sqltypes.MakeTestResult(sqltypes.MakeTestFields("col", "int64"), "1", "2", "3", "4") fn := func(r *sqltypes.Result) error { return nil @@ -122,7 +124,7 @@ func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { func TestExecutorTransactionsNoAutoCommit(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"}) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -188,7 +190,7 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { } // Prevent use of non-primary if in_transaction is on. - session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary", InTransaction: true}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", InTransaction: true}) _, err = executor.Execute(ctx, nil, "TestExecute", session, "use @replica", nil) require.EqualError(t, err, `can't execute the given command because you have an active transaction`) } @@ -205,7 +207,7 @@ func TestDirectTargetRewrites(t *testing.T) { } sql := "select database()" - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(session), sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) assertQueries(t, sbclookup, []*querypb.BoundQuery{{ Sql: "select :__vtdbname as `database()` from dual", @@ -216,7 +218,7 @@ func TestDirectTargetRewrites(t *testing.T) { func TestExecutorTransactionsAutoCommit(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true, SessionUUID: "suuid"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true, SessionUUID: "suuid"}) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -270,7 +272,7 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) oltpOptions := &querypb.ExecuteOptions{Workload: querypb.ExecuteOptions_OLTP} - session := NewSafeSession(&vtgatepb.Session{ + session := econtext.NewSafeSession(&vtgatepb.Session{ TargetString: "@primary", Autocommit: true, Options: oltpOptions, @@ -339,7 +341,7 @@ func TestExecutorDeleteMetadata(t *testing.T) { }() executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_v1= '1'" _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) @@ -367,7 +369,7 @@ func TestExecutorDeleteMetadata(t *testing.T) { func TestExecutorAutocommit(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -446,7 +448,7 @@ func TestExecutorAutocommit(t *testing.T) { // transition autocommit from 0 to 1 in the middle of a transaction. startCount = sbclookup.CommitCount.Load() - session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err = executor.Execute(ctx, nil, "TestExecute", session, "begin", nil) require.NoError(t, err) _, err = executor.Execute(ctx, nil, "TestExecute", session, "update main1 set id=1", nil) @@ -468,7 +470,7 @@ func TestExecutorAutocommit(t *testing.T) { func TestExecutorShowColumns(t *testing.T) { executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: ""}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ""}) queries := []string{ "SHOW COLUMNS FROM `user` in `TestExecutor`", @@ -520,7 +522,7 @@ func assertMatchesNoOrder(t *testing.T, expected, got string) { func TestExecutorShow(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) for _, query := range []string{"show vitess_keyspaces", "show keyspaces"} { qr, err := executor.Execute(ctx, nil, "TestExecute", session, query, nil) @@ -1024,7 +1026,7 @@ func TestExecutorShow(t *testing.T) { utils.MustMatch(t, wantqr, qr, fmt.Sprintf("%s, with a bad keyspace", query)) query = "show vschema tables" - session = NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) qr, err = executor.Execute(ctx, nil, "TestExecute", session, query, nil) require.NoError(t, err) wantqr = &sqltypes.Result{ @@ -1050,7 +1052,7 @@ func TestExecutorShow(t *testing.T) { utils.MustMatch(t, wantqr, qr, query) query = "show vschema tables" - session = NewSafeSession(&vtgatepb.Session{}) + session = econtext.NewSafeSession(&vtgatepb.Session{}) _, err = executor.Execute(ctx, nil, "TestExecute", session, query, nil) want = errNoKeyspace.Error() assert.EqualError(t, err, want, query) @@ -1061,7 +1063,7 @@ func TestExecutorShow(t *testing.T) { assert.EqualError(t, err, want, query) query = "show vschema tables" - session = NewSafeSession(&vtgatepb.Session{TargetString: "no_such_keyspace"}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: "no_such_keyspace"}) _, err = executor.Execute(ctx, nil, "TestExecute", session, query, nil) want = "VT05003: unknown database 'no_such_keyspace' in vschema" assert.EqualError(t, err, want, query) @@ -1080,7 +1082,7 @@ func TestExecutorShow(t *testing.T) { func TestExecutorShowTargeted(t *testing.T) { executor, _, sbc2, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor/40-60"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor/40-60"}) queries := []string{ "show databases", @@ -1107,7 +1109,7 @@ func TestExecutorShowTargeted(t *testing.T) { func TestExecutorShowFromSystemSchema(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "mysql"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "mysql"}) _, err := executor.Execute(ctx, nil, "TestExecutorShowFromSystemSchema", session, "show tables", nil) require.NoError(t, err) @@ -1116,7 +1118,7 @@ func TestExecutorShowFromSystemSchema(t *testing.T) { func TestExecutorUse(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary"}) stmts := []string{ "use TestExecutor", @@ -1135,13 +1137,13 @@ func TestExecutorUse(t *testing.T) { utils.MustMatch(t, wantSession, session.Session, "session does not match") } - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "use 1", nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{}), "use 1", nil) wantErr := "syntax error at position 6 near '1'" if err == nil || err.Error() != wantErr { t.Errorf("got: %v, want %v", err, wantErr) } - _, err = executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "use UnexistentKeyspace", nil) + _, err = executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{}), "use UnexistentKeyspace", nil) require.EqualError(t, err, "VT05003: unknown database 'UnexistentKeyspace' in vschema") } @@ -1155,7 +1157,7 @@ func TestExecutorComment(t *testing.T) { wantResult := &sqltypes.Result{} for _, stmt := range stmts { - gotResult, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), stmt, nil) + gotResult, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), stmt, nil) if err != nil { t.Error(err) } @@ -1240,7 +1242,7 @@ func TestExecutorDDL(t *testing.T) { sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) stmtType := "DDL" - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) if tc.hasNoKeyspaceErr { require.EqualError(t, err, errNoKeyspace.Error(), "expect query to fail: %q", stmt) stmtType = "" // For error case, plan is not generated to query log will not contain any stmtType. @@ -1278,7 +1280,7 @@ func TestExecutorDDL(t *testing.T) { sbc1.ExecCount.Store(0) sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: ""}), stmt.input, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: ""}), stmt.input, nil) if stmt.hasErr { require.EqualError(t, err, errNoKeyspace.Error(), "expect query to fail") testQueryLog(t, executor, logChan, "TestExecute", "", stmt.input, 0) @@ -1302,7 +1304,7 @@ func TestExecutorDDLFk(t *testing.T) { executor, _, _, sbc, ctx := createExecutorEnv(t) sbc.ExecCount.Store(0) foreignKeyMode = fkMode - _, err := executor.Execute(ctx, nil, mName, NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), stmt, nil) + _, err := executor.Execute(ctx, nil, mName, econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), stmt, nil) if fkMode == "allow" { require.NoError(t, err) require.EqualValues(t, 1, sbc.ExecCount.Load()) @@ -1322,7 +1324,7 @@ func TestExecutorAlterVSchemaKeyspace(t *testing.T) { }() executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) executor.serv.WatchSrvVSchema(ctx, executor.cell, func(vschema *vschemapb.SrvVSchema, err error) bool { @@ -1364,7 +1366,7 @@ func TestExecutorCreateVindexDDL(t *testing.T) { t.Fatalf("test_vindex should not exist in original vschema") } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema create vindex test_vindex using hash" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) @@ -1388,7 +1390,7 @@ func TestExecutorCreateVindexDDL(t *testing.T) { // Create a new vschema keyspace implicitly by creating a vindex with a different // target in the session // ksNew := "test_new_keyspace" - session = NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt = "alter vschema create vindex test_vindex2 using hash" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) if err != nil { @@ -1439,7 +1441,7 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { vschemaTables = append(vschemaTables, t) } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add table test_table" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) @@ -1451,7 +1453,7 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { _ = waitForVschemaTables(t, ks, append([]string{"test_table", "test_table2"}, vschemaTables...), executor) // Should fail adding a table on a sharded keyspace - session = NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) stmt = "alter vschema add table test_table" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "add vschema table: unsupported on sharded keyspace TestExecutor") @@ -1470,7 +1472,7 @@ func TestExecutorVindexDDLACL(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) ctxRedUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "redUser"}) ctxBlueUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "blueUser"}) @@ -1515,7 +1517,7 @@ func TestExecutorVindexDDLACL(t *testing.T) { func TestExecutorUnrecognized(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "invalid statement", nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{}), "invalid statement", nil) require.Error(t, err, "unrecognized statement: invalid statement'") } @@ -1525,7 +1527,7 @@ func TestExecutorDeniedErrorNoBuffer(t *testing.T) { vschemaWaitTimeout = 500 * time.Millisecond - session := NewAutocommitSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{TargetString: "@primary"}) startExec := time.Now() _, err := executor.Execute(ctx, nil, "TestExecutorDeniedErrorNoBuffer", session, "select * from user", nil) require.NoError(t, err, "enforce denied tables not buffered") @@ -1560,8 +1562,8 @@ var pv = querypb.ExecuteOptions_Gen4 func TestGetPlanUnnormalized(t *testing.T) { r, _, _, _, ctx := createExecutorEnv(t) - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + emptyvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + unshardedvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" plan1, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) @@ -1625,7 +1627,7 @@ func assertCacheContains(t *testing.T, e *Executor, vc *vcursorImpl, sql string) func getPlanCached(t *testing.T, ctx context.Context, e *Executor, vcursor *vcursorImpl, sql string, comments sqlparser.MarginComments, bindVars map[string]*querypb.BindVariable, skipQueryPlanCache bool) (*engine.Plan, *logstats.LogStats) { logStats := logstats.NewLogStats(ctx, "Test", "", "", nil) - vcursor.safeSession = &SafeSession{ + vcursor.safeSession = &econtext.SafeSession{ Session: &vtgatepb.Session{ Options: &querypb.ExecuteOptions{SkipQueryPlanCache: skipQueryPlanCache}}, } @@ -1644,7 +1646,7 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { t.Run("Cache", func(t *testing.T) { r, _, _, _, ctx := createExecutorEnv(t) - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + emptyvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true) @@ -1668,7 +1670,7 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { // Skip cache using directive r, _, _, _, ctx := createExecutorEnv(t) - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + unshardedvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) @@ -1679,12 +1681,12 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { assertCacheSize(t, r.plans, 1) // the target string will be resolved and become part of the plan cache key, which adds a new entry - ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + ksIDVc1, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above - ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + ksIDVc2, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) }) @@ -1694,7 +1696,7 @@ func TestGetPlanCacheNormalized(t *testing.T) { t.Run("Cache", func(t *testing.T) { r, _, _, _, ctx := createExecutorEnv(t) r.normalize = true - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + emptyvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true /* skipQueryPlanCache */) @@ -1711,7 +1713,7 @@ func TestGetPlanCacheNormalized(t *testing.T) { // Skip cache using directive r, _, _, _, ctx := createExecutorEnv(t) r.normalize = true - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + unshardedvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) @@ -1722,12 +1724,12 @@ func TestGetPlanCacheNormalized(t *testing.T) { assertCacheSize(t, r.plans, 1) // the target string will be resolved and become part of the plan cache key, which adds a new entry - ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + ksIDVc1, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above - ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + ksIDVc2, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) }) @@ -1737,8 +1739,8 @@ func TestGetPlanNormalized(t *testing.T) { r, _, _, _, ctx := createExecutorEnv(t) r.normalize = true - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + emptyvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + unshardedvc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" query2 := "select * from music_user_map where id = 2" @@ -1785,7 +1787,7 @@ func TestGetPlanPriority(t *testing.T) { {name: "empty priority", sql: "select * from music_user_map", expectedPriority: "", expectedError: nil}, } - session := NewSafeSession(&vtgatepb.Session{TargetString: "@unknown", Options: &querypb.ExecuteOptions{}}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@unknown", Options: &querypb.ExecuteOptions{}}) for _, aTestCase := range testCases { testCase := aTestCase @@ -1966,7 +1968,7 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) warningCount := warnings.Counts()["WarnPayloadSizeExceeded"] testMaxPayloadSizeExceeded := []string{ "select * from main1", @@ -2014,7 +2016,7 @@ func TestOlapSelectDatabase(t *testing.T) { cbInvoked = true return nil } - err := executor.StreamExecute(context.Background(), nil, "TestExecute", NewSafeSession(session), sql, nil, cb) + err := executor.StreamExecute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(session), sql, nil, cb) assert.NoError(t, err) assert.True(t, cbInvoked) } @@ -2022,7 +2024,7 @@ func TestOlapSelectDatabase(t *testing.T) { func TestExecutorClearsWarnings(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{ + session := econtext.NewSafeSession(&vtgatepb.Session{ Warnings: []*querypb.QueryWarning{{Code: 234, Message: "oh noes"}}, }) _, err := executor.Execute(context.Background(), nil, "TestExecute", session, "select 42", nil) @@ -2058,7 +2060,7 @@ func TestServingKeyspaces(t *testing.T) { }) require.ElementsMatch(t, []string{"TestExecutor", "TestUnsharded"}, gw.GetServingKeyspaces()) - result, err := executor.Execute(ctx, nil, "TestServingKeyspaces", NewSafeSession(&vtgatepb.Session{}), "select keyspace_name from dual", nil) + result, err := executor.Execute(ctx, nil, "TestServingKeyspaces", econtext.NewSafeSession(&vtgatepb.Session{}), "select keyspace_name from dual", nil) require.NoError(t, err) require.Equal(t, `[[VARCHAR("TestExecutor")]]`, fmt.Sprintf("%v", result.Rows)) @@ -2074,7 +2076,7 @@ func TestServingKeyspaces(t *testing.T) { // Clear plan cache, to force re-planning of the query. executor.ClearPlans() require.ElementsMatch(t, []string{"TestUnsharded"}, gw.GetServingKeyspaces()) - result, err = executor.Execute(ctx, nil, "TestServingKeyspaces", NewSafeSession(&vtgatepb.Session{}), "select keyspace_name from dual", nil) + result, err = executor.Execute(ctx, nil, "TestServingKeyspaces", econtext.NewSafeSession(&vtgatepb.Session{}), "select keyspace_name from dual", nil) require.NoError(t, err) require.Equal(t, `[[VARCHAR("TestUnsharded")]]`, fmt.Sprintf("%v", result.Rows)) } @@ -2150,7 +2152,7 @@ func TestExecutorOther(t *testing.T) { sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) if tc.hasNoKeyspaceErr { assert.Error(t, err, errNoKeyspace) } else if tc.hasDestinationShardErr { @@ -2206,7 +2208,7 @@ func TestExecutorAnalyze(t *testing.T) { sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) require.NoError(t, err) utils.MustMatch(t, tc.wantCnts, cnts{ @@ -2270,7 +2272,7 @@ func TestExecutorExplainStmt(t *testing.T) { sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) - _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) assert.NoError(t, err) utils.MustMatch(t, tc.wantCnts, cnts{ @@ -2360,7 +2362,7 @@ func TestExecutorOtherAdmin(t *testing.T) { sbc2.ExecCount.Store(0) sbclookup.ExecCount.Store(0) - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil) if tc.hasNoKeyspaceErr { assert.Error(t, err, errNoKeyspace) } else if tc.hasDestinationShardErr { @@ -2387,7 +2389,7 @@ func TestExecutorSavepointInTx(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) - session := NewSafeSession(&vtgatepb.Session{Autocommit: false, TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: false, TargetString: "@primary"}) _, err := exec(executor, session, "savepoint a") require.NoError(t, err) _, err = exec(executor, session, "rollback to a") @@ -2470,7 +2472,7 @@ func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) - session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "TestExecutor", EnableSystemSettings: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "TestExecutor", EnableSystemSettings: true}) sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("orig|new", "varchar|varchar"), "a|"), }) @@ -2537,7 +2539,7 @@ func TestExecutorSavepointWithoutTx(t *testing.T) { logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) - session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary", InTransaction: false}) + session := econtext.NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary", InTransaction: false}) _, err := exec(executor, session, "savepoint a") require.NoError(t, err) _, err = exec(executor, session, "rollback to a") @@ -2622,7 +2624,7 @@ func TestExecutorCallProc(t *testing.T) { sbc2.ExecCount.Store(0) sbcUnsharded.ExecCount.Store(0) - _, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), "CALL proc()", nil) + _, err := executor.Execute(context.Background(), nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), "CALL proc()", nil) if tc.hasNoKeyspaceErr { assert.EqualError(t, err, errNoKeyspace.Error()) } else if tc.unshardedOnlyErr { @@ -2646,7 +2648,7 @@ func TestExecutorTempTable(t *testing.T) { initialWarningsCount := warnings.Counts()["WarnUnshardedOnly"] executor.warnShardedOnly = true creatQuery := "create temporary table temp_t(id bigint primary key)" - session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) _, err := executor.Execute(ctx, nil, "TestExecutorTempTable", session, creatQuery, nil) require.NoError(t, err) assert.EqualValues(t, 1, sbcUnsharded.ExecCount.Load()) @@ -2665,7 +2667,7 @@ func TestExecutorShowVitessMigrations(t *testing.T) { executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) showQuery := "show vitess_migrations" - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) require.NoError(t, err) assert.Contains(t, sbc1.StringQueries(), "show vitess_migrations") @@ -2676,7 +2678,7 @@ func TestExecutorDescHash(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) showQuery := "desc hash_index" - session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) require.NoError(t, err) } @@ -2684,7 +2686,7 @@ func TestExecutorDescHash(t *testing.T) { func TestExecutorVExplainQueries(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) sbclookup.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("name|user_id", "varchar|int64"), "apa|1", "apa|2"), @@ -2697,7 +2699,7 @@ func TestExecutorVExplainQueries(t *testing.T) { // Test the streaming side as well var results []sqltypes.Row - session = NewAutocommitSession(&vtgatepb.Session{}) + session = econtext.NewAutocommitSession(&vtgatepb.Session{}) err = executor.StreamExecute(ctx, nil, "TestExecutorVExplainQueries", session, "vexplain queries select * from user where name = 'apa'", nil, func(result *sqltypes.Result) error { results = append(results, result.Rows...) return nil @@ -2710,7 +2712,7 @@ func TestExecutorVExplainQueries(t *testing.T) { func TestExecutorStartTxnStmt(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) tcases := []struct { beginSQL string @@ -2757,7 +2759,7 @@ func TestExecutorPrepareExecute(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true - session := NewAutocommitSession(&vtgatepb.Session{}) + session := econtext.NewAutocommitSession(&vtgatepb.Session{}) // prepare statement. _, err := executor.Execute(context.Background(), nil, "TestExecutorPrepareExecute", session, "prepare prep_user from 'select * from user where id = ?'", nil) @@ -2834,7 +2836,7 @@ func TestExecutorSettingsInTwoPC(t *testing.T) { sbc2.SetResults(tcase.testRes) // create a new session - session := NewSafeSession(&vtgatepb.Session{ + session := econtext.NewSafeSession(&vtgatepb.Session{ TargetString: KsTestSharded, TransactionMode: vtgatepb.TransactionMode_TWOPC, EnableSystemSettings: true, @@ -2892,7 +2894,7 @@ func TestExecutorRejectTwoPC(t *testing.T) { sbc2.SetResults(tcase.testRes) // create a new session - session := NewSafeSession(&vtgatepb.Session{ + session := econtext.NewSafeSession(&vtgatepb.Session{ TargetString: KsTestSharded, TransactionMode: vtgatepb.TransactionMode_TWOPC, EnableSystemSettings: true, @@ -2922,7 +2924,7 @@ func TestExecutorTruncateErrors(t *testing.T) { truncateErrorLen = 32 defer func() { truncateErrorLen = save }() - session := NewSafeSession(&vtgatepb.Session{}) + session := econtext.NewSafeSession(&vtgatepb.Session{}) fn := func(r *sqltypes.Result) error { return nil } @@ -2982,7 +2984,7 @@ func TestExecutorFlushStmt(t *testing.T) { for _, tc := range tcs { t.Run(tc.query+tc.targetStr, func(t *testing.T) { - _, err := executor.Execute(context.Background(), nil, "TestExecutorFlushStmt", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), tc.query, nil) + _, err := executor.Execute(context.Background(), nil, "TestExecutorFlushStmt", econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), tc.query, nil) if tc.expectedErr == "" { require.NoError(t, err) } else { @@ -3029,7 +3031,7 @@ func TestExecutorKillStmt(t *testing.T) { allowKillStmt = !tc.disallow t.Run("execute:"+tc.query+tc.errStr, func(t *testing.T) { mysqlCtx := &fakeMysqlConnection{ErrMsg: tc.errStr} - _, err := executor.Execute(context.Background(), mysqlCtx, "TestExecutorKillStmt", NewAutocommitSession(&vtgatepb.Session{}), tc.query, nil) + _, err := executor.Execute(context.Background(), mysqlCtx, "TestExecutorKillStmt", econtext.NewAutocommitSession(&vtgatepb.Session{}), tc.query, nil) if tc.errStr != "" { require.ErrorContains(t, err, tc.errStr) } else { @@ -3039,7 +3041,7 @@ func TestExecutorKillStmt(t *testing.T) { }) t.Run("stream:"+tc.query+tc.errStr, func(t *testing.T) { mysqlCtx := &fakeMysqlConnection{ErrMsg: tc.errStr} - err := executor.StreamExecute(context.Background(), mysqlCtx, "TestExecutorKillStmt", NewAutocommitSession(&vtgatepb.Session{}), tc.query, nil, func(result *sqltypes.Result) error { + err := executor.StreamExecute(context.Background(), mysqlCtx, "TestExecutorKillStmt", econtext.NewAutocommitSession(&vtgatepb.Session{}), tc.query, nil, func(result *sqltypes.Result) error { return nil }) if tc.errStr != "" { @@ -3075,7 +3077,7 @@ func (f *fakeMysqlConnection) KillConnection(ctx context.Context, connID uint32) var _ vtgateservice.MySQLConnection = (*fakeMysqlConnection)(nil) -func exec(executor *Executor, session *SafeSession, sql string) (*sqltypes.Result, error) { +func exec(executor *Executor, session *econtext.SafeSession, sql string) (*sqltypes.Result, error) { return executor.Execute(context.Background(), nil, "TestExecute", session, sql, nil) } diff --git a/go/vt/vtgate/executor_vexplain_test.go b/go/vt/vtgate/executor_vexplain_test.go index 99eb77c7ed4..a9516492f1b 100644 --- a/go/vt/vtgate/executor_vexplain_test.go +++ b/go/vt/vtgate/executor_vexplain_test.go @@ -26,6 +26,8 @@ import ( "github.com/stretchr/testify/assert" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "github.com/stretchr/testify/require" "vitess.io/vitess/go/sqltypes" @@ -135,7 +137,7 @@ func TestVExplainKeys(t *testing.T) { for _, tt := range tests { t.Run(tt.Query, func(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) gotResult, err := executor.Execute(context.Background(), nil, "Execute", session, "vexplain keys "+tt.Query, nil) require.NoError(t, err) diff --git a/go/vt/vtgate/executor_vschema_ddl_test.go b/go/vt/vtgate/executor_vschema_ddl_test.go index 1c912ed0d62..825b65ab8f3 100644 --- a/go/vt/vtgate/executor_vschema_ddl_test.go +++ b/go/vt/vtgate/executor_vschema_ddl_test.go @@ -25,6 +25,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/callerid" @@ -138,7 +140,7 @@ func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { vschemaacl.AuthorizedDDLUsers = "" }() executor, _, _, _, ctx := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { @@ -180,7 +182,7 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { t.Fatalf("test_vindex should not exist in original vschema") } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema create vindex test_vindex using hash" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) @@ -222,7 +224,7 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { t.Fatalf("test_vindex should not exist in original vschema") } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema drop vindex test_vindex" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "vindex test_vindex does not exists in keyspace TestExecutor" @@ -296,7 +298,7 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { vschemaTables = append(vschemaTables, t) } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add table test_table" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) @@ -308,7 +310,7 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { _ = waitForVschemaTables(t, ks, append([]string{"test_table", "test_table2"}, vschemaTables...), executor) // Should fail adding a table on a sharded keyspace - session = NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) stmt = "alter vschema add table test_table" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "add vschema table: unsupported on sharded keyspace TestExecutor" @@ -343,7 +345,7 @@ func TestExecutorAddSequenceDDL(t *testing.T) { vschemaTables = append(vschemaTables, t) } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add sequence test_seq" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) @@ -357,7 +359,7 @@ func TestExecutorAddSequenceDDL(t *testing.T) { // Should fail adding a table on a sharded keyspace ksSharded := "TestExecutor" - session = NewSafeSession(&vtgatepb.Session{TargetString: ksSharded}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: ksSharded}) stmt = "alter vschema add sequence sequence_table" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) @@ -403,7 +405,7 @@ func TestExecutorDropSequenceDDL(t *testing.T) { t.Fatalf("test_seq should not exist in original vschema") } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) // add test sequence stmt := "alter vschema add sequence test_seq" @@ -428,7 +430,7 @@ func TestExecutorDropSequenceDDL(t *testing.T) { } // Should fail dropping a non-existing test sequence - session = NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session = econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt = "alter vschema drop sequence test_seq" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) @@ -447,7 +449,7 @@ func TestExecutorDropAutoIncDDL(t *testing.T) { executor, _, _, _, ctx := createExecutorEnv(t) ks := KsTestUnsharded - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add table test_table" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) @@ -488,7 +490,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { }() executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) ks := "TestExecutor" - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema @@ -706,7 +708,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { require.EqualError(t, err, "table TestExecutor.nonexistent not defined in vschema") stmt = "alter vschema on nonexistent drop vindex test_lookup" - _, err = executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: "InvalidKeyspace"}), stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", econtext.NewSafeSession(&vtgatepb.Session{TargetString: "InvalidKeyspace"}), stmt, nil) require.EqualError(t, err, "VT05003: unknown database 'InvalidKeyspace' in vschema") stmt = "alter vschema on nowhere.nohow drop vindex test_lookup" @@ -731,7 +733,7 @@ func TestPlanExecutorVindexDDLACL(t *testing.T) { // t.Skip("not yet planned") executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) + session := econtext.NewSafeSession(&vtgatepb.Session{TargetString: ks}) ctxRedUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "redUser"}) ctxBlueUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "blueUser"}) diff --git a/go/vt/vtgate/executor_vstream_test.go b/go/vt/vtgate/executor_vstream_test.go index 5466e9e8f3f..22fb7ee1034 100644 --- a/go/vt/vtgate/executor_vstream_test.go +++ b/go/vt/vtgate/executor_vstream_test.go @@ -21,6 +21,7 @@ import ( "time" "vitess.io/vitess/go/vt/vtgate/engine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" querypb "vitess.io/vitess/go/vt/proto/query" @@ -76,7 +77,7 @@ func TestVStreamSQLUnsharded(t *testing.T) { results := make(chan *sqltypes.Result, 20) go func() { - err := executor.StreamExecute(ctx, nil, "TestExecuteStream", NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), sql, nil, func(qr *sqltypes.Result) error { + err := executor.StreamExecute(ctx, nil, "TestExecuteStream", econtext.NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), sql, nil, func(qr *sqltypes.Result) error { results <- qr return nil }) diff --git a/go/vt/vtgate/safe_session.go b/go/vt/vtgate/executorcontext/safe_session.go similarity index 91% rename from go/vt/vtgate/safe_session.go rename to go/vt/vtgate/executorcontext/safe_session.go index 1d57c63ef35..6abe0a69329 100644 --- a/go/vt/vtgate/safe_session.go +++ b/go/vt/vtgate/executorcontext/safe_session.go @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. */ -package vtgate +package executorcontext import ( "fmt" @@ -23,22 +23,19 @@ import ( "sync" "time" - "vitess.io/vitess/go/sqltypes" - "google.golang.org/protobuf/proto" "vitess.io/vitess/go/mysql/datetime" - + "vitess.io/vitess/go/sqltypes" + 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" "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/srvtopo" "vitess.io/vitess/go/vt/sysvars" "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/engine" - - 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" ) type ( @@ -58,7 +55,7 @@ type ( rollbackOnPartialExec string savepointName string - // this is a signal that found_rows has already been handles by the primitives, + // this is a signal that found_rows has already been handled by the primitives, // and doesn't have to be updated by the executor foundRowsHandled bool @@ -66,12 +63,12 @@ type ( // as the query that started a new transaction on the shard belong to a vindex. queryFromVindex bool - logging *executeLogger + logging *ExecuteLogger *vtgatepb.Session } - executeLogger struct { + ExecuteLogger struct { mu sync.Mutex entries []engine.ExecuteEntry lastID int @@ -127,6 +124,8 @@ const ( savepointRollback ) +const TxRollback = "Rollback Transaction" + // NewSafeSession returns a new SafeSession based on the Session func NewSafeSession(sessn *vtgatepb.Session) *SafeSession { if sessn == nil { @@ -205,6 +204,50 @@ func (session *SafeSession) resetCommonLocked() { } } +// NewAutocommitSession returns a SafeSession based on the original +// session, but with autocommit enabled. +func (session *SafeSession) NewAutocommitSession() *SafeSession { + ss := NewAutocommitSession(session.Session) + ss.logging = session.logging + return ss +} + +// IsFoundRowsHandled returns the foundRowsHandled. +func (session *SafeSession) IsFoundRowsHandled() bool { + session.mu.Lock() + defer session.mu.Unlock() + return session.foundRowsHandled +} + +// SetFoundRows set the found rows value. +func (session *SafeSession) SetFoundRows(value uint64) { + session.mu.Lock() + defer session.mu.Unlock() + session.FoundRows = value + session.foundRowsHandled = true +} + +// GetRollbackOnPartialExec returns the rollbackOnPartialExec value. +func (session *SafeSession) GetRollbackOnPartialExec() string { + session.mu.Lock() + defer session.mu.Unlock() + return session.rollbackOnPartialExec +} + +// SetQueryFromVindex set the queryFromVindex value. +func (session *SafeSession) SetQueryFromVindex(value bool) { + session.mu.Lock() + defer session.mu.Unlock() + session.queryFromVindex = value +} + +// GetQueryFromVindex returns the queryFromVindex value. +func (session *SafeSession) GetQueryFromVindex() bool { + session.mu.Lock() + defer session.mu.Unlock() + return session.queryFromVindex +} + // SetQueryTimeout sets the query timeout func (session *SafeSession) SetQueryTimeout(queryTimeout int64) { session.mu.Lock() @@ -312,7 +355,7 @@ func (session *SafeSession) SetRollbackCommand() { if session.savepointState == savepointSet { session.rollbackOnPartialExec = fmt.Sprintf("rollback to %s", session.savepointName) } else { - session.rollbackOnPartialExec = txRollback + session.rollbackOnPartialExec = TxRollback } session.savepointState = savepointRollbackSet } @@ -340,6 +383,18 @@ func (session *SafeSession) SetCommitOrder(co vtgatepb.CommitOrder) { session.commitOrder = co } +// GetCommitOrder returns the commit order. +func (session *SafeSession) GetCommitOrder() vtgatepb.CommitOrder { + session.mu.Lock() + defer session.mu.Unlock() + return session.commitOrder +} + +// GetLogger returns executor logger. +func (session *SafeSession) GetLogger() *ExecuteLogger { + return session.logging +} + // InTransaction returns true if we are in a transaction func (session *SafeSession) InTransaction() bool { session.mu.Lock() @@ -678,12 +733,11 @@ func (session *SafeSession) UpdateLockHeartbeat() { session.LastLockHeartbeat = time.Now().Unix() } -// TriggerLockHeartBeat returns if it time to trigger next lock heartbeat -func (session *SafeSession) TriggerLockHeartBeat() bool { +// GetLockHeartbeat returns last time the lock heartbeat was sent. +func (session *SafeSession) GetLockHeartbeat() int64 { session.mu.Lock() defer session.mu.Unlock() - now := time.Now().Unix() - return now-session.LastLockHeartbeat >= int64(lockHeartbeatTime.Seconds()) + return session.LastLockHeartbeat } // InLockSession returns whether locking is used on this session. @@ -836,9 +890,7 @@ func (session *SafeSession) GetOrCreateOptions() *querypb.ExecuteOptions { return session.Session.Options } -var _ iQueryOption = (*SafeSession)(nil) - -func (session *SafeSession) cachePlan() bool { +func (session *SafeSession) CachePlan() bool { if session == nil || session.Options == nil { return true } @@ -849,7 +901,7 @@ func (session *SafeSession) cachePlan() bool { return !(session.Options.SkipQueryPlanCache || session.Options.HasCreatedTempTables) } -func (session *SafeSession) getSelectLimit() int { +func (session *SafeSession) GetSelectLimit() int { if session == nil || session.Options == nil { return -1 } @@ -860,16 +912,16 @@ func (session *SafeSession) getSelectLimit() int { return int(session.Options.SqlSelectLimit) } -// isTxOpen returns true if there is open connection to any of the shard. -func (session *SafeSession) isTxOpen() bool { +// IsTxOpen returns true if there is open connection to any of the shard. +func (session *SafeSession) IsTxOpen() bool { session.mu.Lock() defer session.mu.Unlock() return len(session.ShardSessions) > 0 || len(session.PreSessions) > 0 || len(session.PostSessions) > 0 } -// getSessions returns the shard session for the current commit order. -func (session *SafeSession) getSessions() []*vtgatepb.Session_ShardSession { +// GetSessions returns the shard session for the current commit order. +func (session *SafeSession) GetSessions() []*vtgatepb.Session_ShardSession { session.mu.Lock() defer session.mu.Unlock() @@ -956,7 +1008,7 @@ func (session *SafeSession) EnableLogging(parser *sqlparser.Parser) { session.mu.Lock() defer session.mu.Unlock() - session.logging = &executeLogger{ + session.logging = &ExecuteLogger{ parser: parser, } } @@ -994,7 +1046,15 @@ func (session *SafeSession) GetPrepareData(name string) *vtgatepb.PrepareData { return session.PrepareStatement[name] } -func (l *executeLogger) log(primitive engine.Primitive, target *querypb.Target, gateway srvtopo.Gateway, query string, begin bool, bv map[string]*querypb.BindVariable) { +func (session *SafeSession) Log(primitive engine.Primitive, target *querypb.Target, gateway srvtopo.Gateway, query string, begin bool, bv map[string]*querypb.BindVariable) { + session.logging.Log(primitive, target, gateway, query, begin, bv) +} + +func (session *SafeSession) GetLogs() []engine.ExecuteEntry { + return session.logging.GetLogs() +} + +func (l *ExecuteLogger) Log(primitive engine.Primitive, target *querypb.Target, gateway srvtopo.Gateway, query string, begin bool, bv map[string]*querypb.BindVariable) { if l == nil { return } @@ -1033,7 +1093,10 @@ func (l *executeLogger) log(primitive engine.Primitive, target *querypb.Target, }) } -func (l *executeLogger) GetLogs() []engine.ExecuteEntry { +func (l *ExecuteLogger) GetLogs() []engine.ExecuteEntry { + if l == nil { + return nil + } l.mu.Lock() defer l.mu.Unlock() result := make([]engine.ExecuteEntry, len(l.entries)) diff --git a/go/vt/vtgate/safe_session_test.go b/go/vt/vtgate/executorcontext/safe_session_test.go similarity index 99% rename from go/vt/vtgate/safe_session_test.go rename to go/vt/vtgate/executorcontext/safe_session_test.go index ce681fe7fd3..9f41365dd87 100644 --- a/go/vt/vtgate/safe_session_test.go +++ b/go/vt/vtgate/executorcontext/safe_session_test.go @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. */ -package vtgate +package executorcontext import ( "reflect" diff --git a/go/vt/vtgate/legacy_scatter_conn_test.go b/go/vt/vtgate/legacy_scatter_conn_test.go index 4512fc0724e..3a10852bbc7 100644 --- a/go/vt/vtgate/legacy_scatter_conn_test.go +++ b/go/vt/vtgate/legacy_scatter_conn_test.go @@ -26,6 +26,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/mysql/collations" "vitess.io/vitess/go/sqltypes" @@ -99,7 +101,7 @@ func TestLegacyExecuteFailOnAutocommit(t *testing.T) { }, Autocommit: false, } - _, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) + _, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, econtext.NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) err := vterrors.Aggregate(errs) require.Error(t, err) require.Contains(t, err.Error(), "in autocommit mode, transactionID should be zero but was: 123") @@ -123,7 +125,7 @@ func TestScatterConnExecuteMulti(t *testing.T) { } } - qr, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, NewSafeSession(nil), false /*autocommit*/, false, nullResultsObserver{}) + qr, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, econtext.NewSafeSession(nil), false /*autocommit*/, false, nullResultsObserver{}) return qr, vterrors.Aggregate(errs) }) } @@ -138,7 +140,7 @@ func TestScatterConnStreamExecuteMulti(t *testing.T) { bvs := make([]map[string]*querypb.BindVariable, len(rss)) qr := new(sqltypes.Result) var mu sync.Mutex - errors := sc.StreamExecuteMulti(ctx, nil, "query", rss, bvs, NewSafeSession(&vtgatepb.Session{InTransaction: true}), true /* autocommit */, func(r *sqltypes.Result) error { + errors := sc.StreamExecuteMulti(ctx, nil, "query", rss, bvs, econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}), true /* autocommit */, func(r *sqltypes.Result) error { mu.Lock() defer mu.Unlock() qr.AppendResult(r) @@ -280,7 +282,7 @@ func TestMaxMemoryRows(t *testing.T) { []key.Destination{key.DestinationShard("0"), key.DestinationShard("1")}) require.NoError(t, err) - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) queries := []*querypb.BoundQuery{{ Sql: "query1", BindVariables: map[string]*querypb.BindVariable{}, @@ -328,7 +330,7 @@ func TestLegaceHealthCheckFailsOnReservedConnections(t *testing.T) { res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") - session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} rss, _, err := res.ResolveDestinations(ctx, keyspace, topodatapb.TabletType_REPLICA, nil, destinations) require.NoError(t, err) @@ -346,12 +348,12 @@ func TestLegaceHealthCheckFailsOnReservedConnections(t *testing.T) { require.Error(t, vterrors.Aggregate(errs)) } -func executeOnShards(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) { +func executeOnShards(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *econtext.SafeSession, destinations []key.Destination) { t.Helper() require.Empty(t, executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations)) } -func executeOnShardsReturnsErr(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) error { +func executeOnShardsReturnsErr(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *econtext.SafeSession, destinations []key.Destination) error { t.Helper() rss, _, err := res.ResolveDestinations(ctx, keyspace, topodatapb.TabletType_REPLICA, nil, destinations) require.NoError(t, err) @@ -429,7 +431,7 @@ func TestMultiExecs(t *testing.T) { observer := recordingResultsObserver{} - session := NewSafeSession(&vtgatepb.Session{}) + session := econtext.NewSafeSession(&vtgatepb.Session{}) _, err := sc.ExecuteMultiShard(ctx, nil, rss, queries, session, false, false, &observer) require.NoError(t, vterrors.Aggregate(err)) if len(sbc0.Queries) == 0 || len(sbc1.Queries) == 0 { @@ -511,7 +513,7 @@ func TestScatterConnSingleDB(t *testing.T) { want := "multi-db transaction attempted" // TransactionMode_SINGLE in session - session := NewSafeSession(&vtgatepb.Session{InTransaction: true, TransactionMode: vtgatepb.TransactionMode_SINGLE}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true, TransactionMode: vtgatepb.TransactionMode_SINGLE}) queries := []*querypb.BoundQuery{{Sql: "query1"}} _, errors := sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) require.Empty(t, errors) @@ -521,7 +523,7 @@ func TestScatterConnSingleDB(t *testing.T) { // TransactionMode_SINGLE in txconn sc.txConn.mode = vtgatepb.TransactionMode_SINGLE - session = NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session = econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, errors = sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) require.Empty(t, errors) _, errors = sc.ExecuteMultiShard(ctx, nil, rss1, queries, session, false, false, nullResultsObserver{}) @@ -530,7 +532,7 @@ func TestScatterConnSingleDB(t *testing.T) { // TransactionMode_MULTI in txconn. Should not fail. sc.txConn.mode = vtgatepb.TransactionMode_MULTI - session = NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session = econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, errors = sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) require.Empty(t, errors) _, errors = sc.ExecuteMultiShard(ctx, nil, rss1, queries, session, false, false, nullResultsObserver{}) @@ -601,7 +603,7 @@ func TestReservePrequeries(t *testing.T) { res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") - session := NewSafeSession(&vtgatepb.Session{ + session := econtext.NewSafeSession(&vtgatepb.Session{ InTransaction: false, InReservedConn: true, SystemVariables: map[string]string{ diff --git a/go/vt/vtgate/plan_execute.go b/go/vt/vtgate/plan_execute.go index 1c0915470ef..7439df4d699 100644 --- a/go/vt/vtgate/plan_execute.go +++ b/go/vt/vtgate/plan_execute.go @@ -29,6 +29,7 @@ import ( "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/engine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/vt/vtgate/vtgateservice" ) @@ -59,7 +60,7 @@ func waitForNewerVSchema(ctx context.Context, e *Executor, lastVSchemaCreated ti func (e *Executor) newExecute( ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, - safeSession *SafeSession, + safeSession *econtext.SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats, @@ -225,7 +226,7 @@ func (e *Executor) newExecute( func (e *Executor) handleTransactions( ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, - safeSession *SafeSession, + safeSession *econtext.SafeSession, plan *engine.Plan, logStats *logstats.LogStats, vcursor *vcursorImpl, @@ -267,7 +268,7 @@ func (e *Executor) handleTransactions( return nil, nil } -func (e *Executor) startTxIfNecessary(ctx context.Context, safeSession *SafeSession) error { +func (e *Executor) startTxIfNecessary(ctx context.Context, safeSession *econtext.SafeSession) error { if !safeSession.Autocommit && !safeSession.InTransaction() { if err := e.txConn.Begin(ctx, safeSession, nil); err != nil { return err @@ -276,7 +277,7 @@ func (e *Executor) startTxIfNecessary(ctx context.Context, safeSession *SafeSess return nil } -func (e *Executor) insideTransaction(ctx context.Context, safeSession *SafeSession, logStats *logstats.LogStats, execPlan func() error) error { +func (e *Executor) insideTransaction(ctx context.Context, safeSession *econtext.SafeSession, logStats *logstats.LogStats, execPlan func() error) error { mustCommit := false if safeSession.Autocommit && !safeSession.InTransaction() { mustCommit = true @@ -320,7 +321,7 @@ func (e *Executor) insideTransaction(ctx context.Context, safeSession *SafeSessi func (e *Executor) executePlan( ctx context.Context, - safeSession *SafeSession, + safeSession *econtext.SafeSession, plan *engine.Plan, vcursor *vcursorImpl, bindVars map[string]*querypb.BindVariable, @@ -342,7 +343,7 @@ func (e *Executor) executePlan( } // rollbackExecIfNeeded rollbacks the partial execution if earlier it was detected that it needs partial query execution to be rolled back. -func (e *Executor) rollbackExecIfNeeded(ctx context.Context, safeSession *SafeSession, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats, err error) error { +func (e *Executor) rollbackExecIfNeeded(ctx context.Context, safeSession *econtext.SafeSession, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats, err error) error { if safeSession.InTransaction() && safeSession.IsRollbackSet() { rErr := e.rollbackPartialExec(ctx, safeSession, bindVars, logStats) return vterrors.Wrap(err, rErr.Error()) @@ -353,7 +354,7 @@ func (e *Executor) rollbackExecIfNeeded(ctx context.Context, safeSession *SafeSe // rollbackPartialExec rollbacks to the savepoint or rollbacks transaction based on the value set on SafeSession.rollbackOnPartialExec. // Once, it is used the variable is reset. // If it fails to rollback to the previous savepoint then, the transaction is forced to be rolled back. -func (e *Executor) rollbackPartialExec(ctx context.Context, safeSession *SafeSession, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) error { +func (e *Executor) rollbackPartialExec(ctx context.Context, safeSession *econtext.SafeSession, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) error { var err error var errMsg strings.Builder @@ -367,8 +368,8 @@ func (e *Executor) rollbackPartialExec(ctx context.Context, safeSession *SafeSes } // needs to rollback only once. - rQuery := safeSession.rollbackOnPartialExec - if rQuery != txRollback { + rQuery := safeSession.GetRollbackOnPartialExec() + if rQuery != econtext.TxRollback { safeSession.SavepointRollback() _, _, err = e.execute(ctx, nil, safeSession, rQuery, bindVars, logStats) // If no error, the revert is successful with the savepoint. Notify the reason as error to the client. diff --git a/go/vt/vtgate/scatter_conn.go b/go/vt/vtgate/scatter_conn.go index f7db598127e..1029fb3d63a 100644 --- a/go/vt/vtgate/scatter_conn.go +++ b/go/vt/vtgate/scatter_conn.go @@ -24,26 +24,25 @@ import ( "sync/atomic" "time" - "vitess.io/vitess/go/mysql/sqlerror" - "vitess.io/vitess/go/vt/sqlparser" - "google.golang.org/protobuf/proto" + "vitess.io/vitess/go/mysql/sqlerror" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/stats" "vitess.io/vitess/go/vt/concurrency" "vitess.io/vitess/go/vt/discovery" "vitess.io/vitess/go/vt/log" + 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" + "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/srvtopo" "vitess.io/vitess/go/vt/topo/topoproto" "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/engine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vttablet/queryservice" - - 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" ) // ScatterConn is used for executing queries across @@ -108,7 +107,7 @@ func (stc *ScatterConn) startAction(name string, target *querypb.Target) (time.T return startTime, statsKey } -func (stc *ScatterConn) endAction(startTime time.Time, allErrors *concurrency.AllErrorRecorder, statsKey []string, err *error, session *SafeSession) { +func (stc *ScatterConn) endAction(startTime time.Time, allErrors *concurrency.AllErrorRecorder, statsKey []string, err *error, session *econtext.SafeSession) { if *err != nil { allErrors.RecordError(*err) // Don't increment the error counter for duplicate @@ -152,7 +151,7 @@ func (stc *ScatterConn) ExecuteMultiShard( primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, - session *SafeSession, + session *econtext.SafeSession, autocommit bool, ignoreMaxMemoryRows bool, resultsObserver resultsObserver, @@ -166,7 +165,7 @@ func (stc *ScatterConn) ExecuteMultiShard( var mu sync.Mutex qr = new(sqltypes.Result) - if session.InLockSession() && session.TriggerLockHeartBeat() { + if session.InLockSession() && triggerLockHeartBeat(session) { go stc.runLockQuery(ctx, session) } @@ -260,7 +259,7 @@ func (stc *ScatterConn) ExecuteMultiShard( default: return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unexpected actionNeeded on query execution: %v", info.actionNeeded) } - session.logging.log(primitive, rs.Target, rs.Gateway, queries[i].Sql, info.actionNeeded == begin || info.actionNeeded == reserveBegin, queries[i].BindVariables) + session.Log(primitive, rs.Target, rs.Gateway, queries[i].Sql, info.actionNeeded == begin || info.actionNeeded == reserveBegin, queries[i].BindVariables) // We need to new shard info irrespective of the error. newInfo := info.updateTransactionAndReservedID(transactionID, reservedID, alias) @@ -289,7 +288,13 @@ func (stc *ScatterConn) ExecuteMultiShard( return qr, allErrors.GetErrors() } -func (stc *ScatterConn) runLockQuery(ctx context.Context, session *SafeSession) { +func triggerLockHeartBeat(session *econtext.SafeSession) bool { + now := time.Now().Unix() + lastHeartbeat := session.GetLockHeartbeat() + return now-lastHeartbeat >= int64(lockHeartbeatTime.Seconds()) +} + +func (stc *ScatterConn) runLockQuery(ctx context.Context, session *econtext.SafeSession) { rs := &srvtopo.ResolvedShard{Target: session.LockSession.Target, Gateway: stc.gateway} query := &querypb.BoundQuery{Sql: "select 1", BindVariables: nil} _, lockErr := stc.ExecuteLock(ctx, rs, query, session, sqlparser.IsUsedLock) @@ -298,7 +303,7 @@ func (stc *ScatterConn) runLockQuery(ctx context.Context, session *SafeSession) } } -func checkAndResetShardSession(info *shardActionInfo, err error, session *SafeSession, target *querypb.Target) reset { +func checkAndResetShardSession(info *shardActionInfo, err error, session *econtext.SafeSession, target *querypb.Target) reset { retry := none if info.reservedID != 0 && info.transactionID == 0 { if wasConnectionClosed(err) { @@ -314,7 +319,7 @@ func checkAndResetShardSession(info *shardActionInfo, err error, session *SafeSe return retry } -func getQueryService(ctx context.Context, rs *srvtopo.ResolvedShard, info *shardActionInfo, session *SafeSession, skipReset bool) (queryservice.QueryService, error) { +func getQueryService(ctx context.Context, rs *srvtopo.ResolvedShard, info *shardActionInfo, session *econtext.SafeSession, skipReset bool) (queryservice.QueryService, error) { if info.alias == nil { return rs.Gateway, nil } @@ -365,12 +370,12 @@ func (stc *ScatterConn) StreamExecuteMulti( query string, rss []*srvtopo.ResolvedShard, bindVars []map[string]*querypb.BindVariable, - session *SafeSession, + session *econtext.SafeSession, autocommit bool, callback func(reply *sqltypes.Result) error, resultsObserver resultsObserver, ) []error { - if session.InLockSession() && session.TriggerLockHeartBeat() { + if session.InLockSession() && triggerLockHeartBeat(session) { go stc.runLockQuery(ctx, session) } @@ -469,7 +474,7 @@ func (stc *ScatterConn) StreamExecuteMulti( default: return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unexpected actionNeeded on query execution: %v", info.actionNeeded) } - session.logging.log(primitive, rs.Target, rs.Gateway, query, info.actionNeeded == begin || info.actionNeeded == reserveBegin, bindVars[i]) + session.Log(primitive, rs.Target, rs.Gateway, query, info.actionNeeded == begin || info.actionNeeded == reserveBegin, bindVars[i]) // We need to new shard info irrespective of the error. newInfo := info.updateTransactionAndReservedID(transactionID, reservedID, alias) @@ -604,7 +609,7 @@ func (stc *ScatterConn) multiGo( startTime, statsKey := stc.startAction(name, rs.Target) // Send a dummy session. // TODO(sougou): plumb a real session through this call. - defer stc.endAction(startTime, allErrors, statsKey, &err, NewSafeSession(nil)) + defer stc.endAction(startTime, allErrors, statsKey, &err, econtext.NewSafeSession(nil)) err = action(rs, i) } @@ -646,7 +651,7 @@ func (stc *ScatterConn) multiGoTransaction( ctx context.Context, name string, rss []*srvtopo.ResolvedShard, - session *SafeSession, + session *econtext.SafeSession, autocommit bool, action shardActionTransactionFunc, ) (allErrors *concurrency.AllErrorRecorder) { @@ -727,7 +732,7 @@ func (stc *ScatterConn) multiGoTransaction( // It returns an error recorder in which each shard error is recorded positionally, // i.e. if rss[2] had an error, then the error recorder will store that error // in the second position. -func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) { +func (stc *ScatterConn) ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *econtext.SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) { var ( qr *sqltypes.Result @@ -830,7 +835,7 @@ func requireNewQS(err error, target *querypb.Target) bool { } // actionInfo looks at the current session, and returns information about what needs to be done for this tablet -func actionInfo(ctx context.Context, target *querypb.Target, session *SafeSession, autocommit bool, txMode vtgatepb.TransactionMode) (*shardActionInfo, error) { +func actionInfo(ctx context.Context, target *querypb.Target, session *econtext.SafeSession, autocommit bool, txMode vtgatepb.TransactionMode) (*shardActionInfo, error) { if !(session.InTransaction() || session.InReservedConn()) { return &shardActionInfo{}, nil } @@ -869,7 +874,7 @@ func actionInfo(ctx context.Context, target *querypb.Target, session *SafeSessio } // lockInfo looks at the current session, and returns information about what needs to be done for this tablet -func lockInfo(target *querypb.Target, session *SafeSession, lockFuncType sqlparser.LockingFuncType) (*shardActionInfo, error) { +func lockInfo(target *querypb.Target, session *econtext.SafeSession, lockFuncType sqlparser.LockingFuncType) (*shardActionInfo, error) { info := &shardActionInfo{actionNeeded: nothing} if session.LockSession != nil { if !proto.Equal(target, session.LockSession.Target) { diff --git a/go/vt/vtgate/scatter_conn_test.go b/go/vt/vtgate/scatter_conn_test.go index c5d4f350433..ab8680ca5e6 100644 --- a/go/vt/vtgate/scatter_conn_test.go +++ b/go/vt/vtgate/scatter_conn_test.go @@ -21,6 +21,7 @@ import ( "testing" "vitess.io/vitess/go/vt/log" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/mysql/sqlerror" vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" @@ -100,7 +101,7 @@ func TestExecuteFailOnAutocommit(t *testing.T) { }, Autocommit: false, } - _, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) + _, errs := sc.ExecuteMultiShard(ctx, nil, rss, queries, econtext.NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) err := vterrors.Aggregate(errs) require.Error(t, err) require.Contains(t, err.Error(), "in autocommit mode, transactionID should be zero but was: 123") @@ -183,7 +184,7 @@ func TestExecutePanic(t *testing.T) { require.Contains(t, logMessage, "(*ScatterConn).multiGoTransaction") }() - _, _ = sc.ExecuteMultiShard(ctx, nil, rss, queries, NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) + _, _ = sc.ExecuteMultiShard(ctx, nil, rss, queries, econtext.NewSafeSession(session), true /*autocommit*/, false, nullResultsObserver{}) } @@ -204,7 +205,7 @@ func TestReservedOnMultiReplica(t *testing.T) { res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") - session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} for i := 0; i < 10; i++ { executeOnShards(t, ctx, res, keyspace, sc, session, destinations) @@ -351,7 +352,7 @@ func TestReservedBeginTableDriven(t *testing.T) { res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") t.Run(test.name, func(t *testing.T) { - session := NewSafeSession(&vtgatepb.Session{}) + session := econtext.NewSafeSession(&vtgatepb.Session{}) for _, action := range test.actions { session.Session.InTransaction = action.transaction session.Session.InReservedConn = action.reserved @@ -384,7 +385,7 @@ func TestReservedConnFail(t *testing.T) { _ = hc.AddTestTablet("aa", "1", 1, keyspace, "1", topodatapb.TabletType_REPLICA, true, 1, nil) res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") - session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} executeOnShards(t, ctx, res, keyspace, sc, session, destinations) diff --git a/go/vt/vtgate/tx_conn.go b/go/vt/vtgate/tx_conn.go index 968b41d38d3..ac5cfd78d6e 100644 --- a/go/vt/vtgate/tx_conn.go +++ b/go/vt/vtgate/tx_conn.go @@ -33,6 +33,7 @@ import ( vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/vterrors" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vttablet/queryservice" ) @@ -72,7 +73,7 @@ const ( // Begin begins a new transaction. If one is already in progress, it commits it // and starts a new one. -func (txc *TxConn) Begin(ctx context.Context, session *SafeSession, txAccessModes []sqlparser.TxAccessMode) error { +func (txc *TxConn) Begin(ctx context.Context, session *econtext.SafeSession, txAccessModes []sqlparser.TxAccessMode) error { if session.InTransaction() { if err := txc.Commit(ctx, session); err != nil { return err @@ -94,7 +95,7 @@ func (txc *TxConn) Begin(ctx context.Context, session *SafeSession, txAccessMode // Commit commits the current transaction. The type of commit can be // best effort or 2pc depending on the session setting. -func (txc *TxConn) Commit(ctx context.Context, session *SafeSession) error { +func (txc *TxConn) Commit(ctx context.Context, session *econtext.SafeSession) error { defer session.ResetTx() if !session.InTransaction() { return nil @@ -115,7 +116,7 @@ func (txc *TxConn) Commit(ctx context.Context, session *SafeSession) error { return txc.commitNormal(ctx, session) } -func recordCommitTime(session *SafeSession, twopc bool, startTime time.Time) { +func recordCommitTime(session *econtext.SafeSession, twopc bool, startTime time.Time) { switch { case len(session.ShardSessions) == 0: // No-op @@ -135,7 +136,7 @@ func (txc *TxConn) queryService(ctx context.Context, alias *topodatapb.TabletAli return txc.tabletGateway.QueryServiceByAlias(ctx, alias, nil) } -func (txc *TxConn) commitShard(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error { +func (txc *TxConn) commitShard(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { if s.TransactionId == 0 { return nil } @@ -151,19 +152,19 @@ func (txc *TxConn) commitShard(ctx context.Context, s *vtgatepb.Session_ShardSes } s.TransactionId = 0 s.ReservedId = reservedID - logging.log(nil, s.Target, nil, "commit", false, nil) + logging.Log(nil, s.Target, nil, "commit", false, nil) return nil } -func (txc *TxConn) commitNormal(ctx context.Context, session *SafeSession) error { - if err := txc.runSessions(ctx, session.PreSessions, session.logging, txc.commitShard); err != nil { +func (txc *TxConn) commitNormal(ctx context.Context, session *econtext.SafeSession) error { + if err := txc.runSessions(ctx, session.PreSessions, session.GetLogger(), txc.commitShard); err != nil { _ = txc.Release(ctx, session) return err } // Retain backward compatibility on commit order for the normal session. for i, shardSession := range session.ShardSessions { - if err := txc.commitShard(ctx, shardSession, session.logging); err != nil { + if err := txc.commitShard(ctx, shardSession, session.GetLogger()); err != nil { if i > 0 { nShards := i elipsis := false @@ -189,7 +190,7 @@ func (txc *TxConn) commitNormal(ctx context.Context, session *SafeSession) error } } - if err := txc.runSessions(ctx, session.PostSessions, session.logging, txc.commitShard); err != nil { + if err := txc.runSessions(ctx, session.PostSessions, session.GetLogger(), txc.commitShard); err != nil { // If last commit fails, there will be nothing to rollback. session.RecordWarning(&querypb.QueryWarning{Message: fmt.Sprintf("post-operation transaction had an error: %v", err)}) // With reserved connection we should release them. @@ -201,7 +202,7 @@ func (txc *TxConn) commitNormal(ctx context.Context, session *SafeSession) error } // commit2PC will not used the pinned tablets - to make sure we use the current source, we need to use the gateway's queryservice -func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err error) { +func (txc *TxConn) commit2PC(ctx context.Context, session *econtext.SafeSession) (err error) { // If the number of participants is one or less, then it's a normal commit. if len(session.ShardSessions) <= 1 { return txc.commitNormal(ctx, session) @@ -240,7 +241,7 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err err } txPhase = Commit2pcPrepare - prepareAction := func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error { + prepareAction := func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { if DebugTwoPc { // Test code to simulate a failure during RM prepare if terr := checkTestFailure(ctx, "RMPrepare_-40_FailNow", s.Target); terr != nil { return terr @@ -248,7 +249,7 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err err } return txc.tabletGateway.Prepare(ctx, s.Target, s.TransactionId, dtid) } - if err = txc.runSessions(ctx, rmShards, session.logging, prepareAction); err != nil { + if err = txc.runSessions(ctx, rmShards, session.GetLogger(), prepareAction); err != nil { return err } @@ -271,7 +272,7 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err err } txPhase = Commit2pcPrepareCommit - prepareCommitAction := func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error { + prepareCommitAction := func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { if DebugTwoPc { // Test code to simulate a failure during RM prepare if terr := checkTestFailure(ctx, "RMCommit_-40_FailNow", s.Target); terr != nil { return terr @@ -279,7 +280,7 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err err } return txc.tabletGateway.CommitPrepared(ctx, s.Target, dtid) } - if err = txc.runSessions(ctx, rmShards, session.logging, prepareCommitAction); err != nil { + if err = txc.runSessions(ctx, rmShards, session.GetLogger(), prepareCommitAction); err != nil { return err } @@ -291,14 +292,14 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) (err err return nil } -func (txc *TxConn) checkValidCondition(session *SafeSession) error { +func (txc *TxConn) checkValidCondition(session *econtext.SafeSession) error { if len(session.PreSessions) != 0 || len(session.PostSessions) != 0 { return vterrors.VT12001("atomic distributed transaction commit with consistent lookup vindex") } return nil } -func (txc *TxConn) errActionAndLogWarn(ctx context.Context, session *SafeSession, txPhase commitPhase, dtid string, mmShard *vtgatepb.Session_ShardSession, rmShards []*vtgatepb.Session_ShardSession) { +func (txc *TxConn) errActionAndLogWarn(ctx context.Context, session *econtext.SafeSession, txPhase commitPhase, dtid string, mmShard *vtgatepb.Session_ShardSession, rmShards []*vtgatepb.Session_ShardSession) { switch txPhase { case Commit2pcCreateTransaction: // Normal rollback is safe because nothing was prepared yet. @@ -307,7 +308,7 @@ func (txc *TxConn) errActionAndLogWarn(ctx context.Context, session *SafeSession } case Commit2pcPrepare: // Rollback the prepared and unprepared transactions. - if resumeErr := txc.rollbackTx(ctx, dtid, mmShard, rmShards, session.logging); resumeErr != nil { + if resumeErr := txc.rollbackTx(ctx, dtid, mmShard, rmShards, session.GetLogger()); resumeErr != nil { log.Warningf("Rollback failed after Prepare failure: %v", resumeErr) } case Commit2pcStartCommit, Commit2pcPrepareCommit: @@ -336,7 +337,7 @@ func createWarningMessage(dtid string, txPhase commitPhase) string { } // Rollback rolls back the current transaction. There are no retries on this operation. -func (txc *TxConn) Rollback(ctx context.Context, session *SafeSession) error { +func (txc *TxConn) Rollback(ctx context.Context, session *econtext.SafeSession) error { if !session.InTransaction() { return nil } @@ -345,7 +346,7 @@ func (txc *TxConn) Rollback(ctx context.Context, session *SafeSession) error { allsessions := append(session.PreSessions, session.ShardSessions...) allsessions = append(allsessions, session.PostSessions...) - err := txc.runSessions(ctx, allsessions, session.logging, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error { + err := txc.runSessions(ctx, allsessions, session.GetLogger(), func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { if s.TransactionId == 0 { return nil } @@ -359,7 +360,7 @@ func (txc *TxConn) Rollback(ctx context.Context, session *SafeSession) error { } s.TransactionId = 0 s.ReservedId = reservedID - logging.log(nil, s.Target, nil, "rollback", false, nil) + logging.Log(nil, s.Target, nil, "rollback", false, nil) return nil }) if err != nil { @@ -372,7 +373,7 @@ func (txc *TxConn) Rollback(ctx context.Context, session *SafeSession) error { } // Release releases the reserved connection and/or rollbacks the transaction -func (txc *TxConn) Release(ctx context.Context, session *SafeSession) error { +func (txc *TxConn) Release(ctx context.Context, session *econtext.SafeSession) error { if !session.InTransaction() && !session.InReservedConn() { return nil } @@ -381,7 +382,7 @@ func (txc *TxConn) Release(ctx context.Context, session *SafeSession) error { allsessions := append(session.PreSessions, session.ShardSessions...) allsessions = append(allsessions, session.PostSessions...) - return txc.runSessions(ctx, allsessions, session.logging, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error { + return txc.runSessions(ctx, allsessions, session.GetLogger(), func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { if s.ReservedId == 0 && s.TransactionId == 0 { return nil } @@ -400,7 +401,7 @@ func (txc *TxConn) Release(ctx context.Context, session *SafeSession) error { } // ReleaseLock releases the reserved connection used for locking. -func (txc *TxConn) ReleaseLock(ctx context.Context, session *SafeSession) error { +func (txc *TxConn) ReleaseLock(ctx context.Context, session *econtext.SafeSession) error { if !session.InLockSession() { return nil } @@ -419,7 +420,7 @@ func (txc *TxConn) ReleaseLock(ctx context.Context, session *SafeSession) error } // ReleaseAll releases all the shard sessions and lock session. -func (txc *TxConn) ReleaseAll(ctx context.Context, session *SafeSession) error { +func (txc *TxConn) ReleaseAll(ctx context.Context, session *econtext.SafeSession) error { if !session.InTransaction() && !session.InReservedConn() && !session.InLockSession() { return nil } @@ -431,7 +432,7 @@ func (txc *TxConn) ReleaseAll(ctx context.Context, session *SafeSession) error { allsessions = append(allsessions, session.LockSession) } - return txc.runSessions(ctx, allsessions, session.logging, func(ctx context.Context, s *vtgatepb.Session_ShardSession, loggging *executeLogger) error { + return txc.runSessions(ctx, allsessions, session.GetLogger(), func(ctx context.Context, s *vtgatepb.Session_ShardSession, loggging *econtext.ExecuteLogger) error { if s.ReservedId == 0 && s.TransactionId == 0 { return nil } @@ -503,12 +504,12 @@ func (txc *TxConn) resolveTx(ctx context.Context, target *querypb.Target, transa // rollbackTx rollbacks the specified distributed transaction. // Rollbacks happens on the metadata manager and all participants irrespective of the failure. -func (txc *TxConn) rollbackTx(ctx context.Context, dtid string, mmShard *vtgatepb.Session_ShardSession, participants []*vtgatepb.Session_ShardSession, logging *executeLogger) error { +func (txc *TxConn) rollbackTx(ctx context.Context, dtid string, mmShard *vtgatepb.Session_ShardSession, participants []*vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger) error { var errs []error if mmErr := txc.rollbackMM(ctx, dtid, mmShard); mmErr != nil { errs = append(errs, mmErr) } - if rmErr := txc.runSessions(ctx, participants, logging, func(ctx context.Context, session *vtgatepb.Session_ShardSession, logger *executeLogger) error { + if rmErr := txc.runSessions(ctx, participants, logging, func(ctx context.Context, session *vtgatepb.Session_ShardSession, logger *econtext.ExecuteLogger) error { return txc.tabletGateway.RollbackPrepared(ctx, session.Target, dtid, session.TransactionId) }); rmErr != nil { errs = append(errs, rmErr) @@ -549,7 +550,7 @@ func (txc *TxConn) resumeCommit(ctx context.Context, target *querypb.Target, tra } // runSessions executes the action for all shardSessions in parallel and returns a consolidated error. -func (txc *TxConn) runSessions(ctx context.Context, shardSessions []*vtgatepb.Session_ShardSession, logging *executeLogger, action func(context.Context, *vtgatepb.Session_ShardSession, *executeLogger) error) error { +func (txc *TxConn) runSessions(ctx context.Context, shardSessions []*vtgatepb.Session_ShardSession, logging *econtext.ExecuteLogger, action func(context.Context, *vtgatepb.Session_ShardSession, *econtext.ExecuteLogger) error) error { // Fastpath. if len(shardSessions) == 1 { return action(ctx, shardSessions[0], logging) diff --git a/go/vt/vtgate/tx_conn_test.go b/go/vt/vtgate/tx_conn_test.go index dd1415bce87..c508967e502 100644 --- a/go/vt/vtgate/tx_conn_test.go +++ b/go/vt/vtgate/tx_conn_test.go @@ -26,6 +26,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/event/syslogger" "vitess.io/vitess/go/mysql/sqlerror" "vitess.io/vitess/go/test/utils" @@ -51,7 +53,7 @@ func TestTxConnBegin(t *testing.T) { session := &vtgatepb.Session{} // begin - safeSession := NewSafeSession(session) + safeSession := econtext.NewSafeSession(session) err := sc.txConn.Begin(ctx, safeSession, nil) require.NoError(t, err) wantSession := vtgatepb.Session{InTransaction: true} @@ -75,7 +77,7 @@ func TestTxConnCommitFailure(t *testing.T) { // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rssm[0], queries, session, false, false, nullResultsObserver{}) wantSession := vtgatepb.Session{ InTransaction: true, @@ -176,7 +178,7 @@ func TestTxConnCommitFailureAfterNonAtomicCommitMaxShards(t *testing.T) { // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) wantSession := vtgatepb.Session{ InTransaction: true, ShardSessions: []*vtgatepb.Session_ShardSession{}, @@ -230,7 +232,7 @@ func TestTxConnCommitFailureBeforeNonAtomicCommitMaxShards(t *testing.T) { // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) wantSession := vtgatepb.Session{ InTransaction: true, ShardSessions: []*vtgatepb.Session_ShardSession{}, @@ -282,7 +284,7 @@ func TestTxConnCommitSuccess(t *testing.T) { sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) wantSession := vtgatepb.Session{ InTransaction: true, @@ -335,7 +337,7 @@ func TestTxConnReservedCommitSuccess(t *testing.T) { sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) wantSession := vtgatepb.Session{ InTransaction: true, @@ -420,7 +422,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) { sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure shard session order - session := NewSafeSession(&vtgatepb.Session{InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InReservedConn: true}) // this will create reserved connections against all tablets _, errs := sc.ExecuteMultiShard(ctx, nil, rss1, queries, session, false, false, nullResultsObserver{}) @@ -515,7 +517,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) { sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure shard session order - session := NewSafeSession(&vtgatepb.Session{InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InReservedConn: true}) // this will create reserved connections against all tablets _, errs := sc.ExecuteMultiShard(ctx, nil, rss1, queries, session, false, false, nullResultsObserver{}) @@ -611,7 +613,7 @@ func TestTxConnCommitOrderFailure1(t *testing.T) { queries := []*querypb.BoundQuery{{Sql: "query1"}} // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) session.SetCommitOrder(vtgatepb.CommitOrder_PRE) @@ -646,7 +648,7 @@ func TestTxConnCommitOrderFailure2(t *testing.T) { }} // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(context.Background(), nil, rss1, queries, session, false, false, nullResultsObserver{}) session.SetCommitOrder(vtgatepb.CommitOrder_PRE) @@ -680,7 +682,7 @@ func TestTxConnCommitOrderFailure3(t *testing.T) { }} // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) session.SetCommitOrder(vtgatepb.CommitOrder_PRE) @@ -722,7 +724,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) { }} // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) wantSession := vtgatepb.Session{ InTransaction: true, @@ -820,7 +822,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) { }} // Sequence the executes to ensure commit order - session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) wantSession := vtgatepb.Session{ InTransaction: true, @@ -957,7 +959,7 @@ func TestTxConnCommit2PC(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PC") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) session.TransactionMode = vtgatepb.TransactionMode_TWOPC @@ -974,7 +976,7 @@ func TestTxConnCommit2PCOneParticipant(t *testing.T) { ctx := utils.LeakCheckContext(t) sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCOneParticipant") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) session.TransactionMode = vtgatepb.TransactionMode_TWOPC require.NoError(t, @@ -987,7 +989,7 @@ func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) { sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCreateTransactionFail") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss1, queries, session, false, false, nullResultsObserver{}) @@ -1011,7 +1013,7 @@ func TestTxConnCommit2PCPrepareFail(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCPrepareFail") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) @@ -1039,7 +1041,7 @@ func TestTxConnCommit2PCStartCommitFail(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCStartCommitFail") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) @@ -1061,7 +1063,7 @@ func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCommitPreparedFail") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) @@ -1083,7 +1085,7 @@ func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCConcludeTransactionFail") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) @@ -1103,7 +1105,7 @@ func TestTxConnRollback(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnRollback") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) require.NoError(t, @@ -1119,7 +1121,7 @@ func TestTxConnReservedRollback(t *testing.T) { sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) require.NoError(t, @@ -1156,7 +1158,7 @@ func TestTxConnReservedRollbackFailure(t *testing.T) { sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") - session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) + session := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false, nullResultsObserver{}) sc.ExecuteMultiShard(ctx, nil, rss01, twoQueries, session, false, false, nullResultsObserver{}) @@ -1435,7 +1437,7 @@ func TestTxConnMultiGoSessions(t *testing.T) { Keyspace: "0", }, }} - err := txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *executeLogger) error { + err := txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *econtext.ExecuteLogger) error { return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "err %s", s.Target.Keyspace) }) want := "err 0" @@ -1450,7 +1452,7 @@ func TestTxConnMultiGoSessions(t *testing.T) { Keyspace: "1", }, }} - err = txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *executeLogger) error { + err = txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *econtext.ExecuteLogger) error { return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "err %s", s.Target.Keyspace) }) want = "err 0\nerr 1" @@ -1458,7 +1460,7 @@ func TestTxConnMultiGoSessions(t *testing.T) { wantCode := vtrpcpb.Code_INTERNAL assert.Equal(t, wantCode, vterrors.Code(err), "error code") - err = txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *executeLogger) error { + err = txc.runSessions(ctx, input, nil, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logger *econtext.ExecuteLogger) error { return nil }) require.NoError(t, err) @@ -1501,7 +1503,7 @@ func TestTxConnAccessModeReset(t *testing.T) { tcases := []struct { name string - f func(ctx context.Context, session *SafeSession) error + f func(ctx context.Context, session *econtext.SafeSession) error }{{ name: "begin-commit", f: sc.txConn.Commit, @@ -1518,7 +1520,7 @@ func TestTxConnAccessModeReset(t *testing.T) { for _, tcase := range tcases { t.Run(tcase.name, func(t *testing.T) { - safeSession := NewSafeSession(&vtgatepb.Session{ + safeSession := econtext.NewSafeSession(&vtgatepb.Session{ Options: &querypb.ExecuteOptions{ TransactionAccessMode: []querypb.ExecuteOptions_TransactionAccessMode{querypb.ExecuteOptions_READ_ONLY}, }, diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go index e9b1d3d7712..079e7e0fab1 100644 --- a/go/vt/vtgate/vcursor_impl.go +++ b/go/vt/vtgate/vcursor_impl.go @@ -50,6 +50,7 @@ import ( "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vtgate/buffer" "vitess.io/vitess/go/vt/vtgate/engine" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" "vitess.io/vitess/go/vt/vtgate/semantics" @@ -68,14 +69,14 @@ var ( type ( // vcursor_impl needs these facilities to be able to be able to execute queries for vindexes iExecute interface { - Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, method string, session *SafeSession, s string, vars map[string]*querypb.BindVariable) (*sqltypes.Result, error) - ExecuteMultiShard(ctx context.Context, primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool, ignoreMaxMemoryRows bool, resultsObserver resultsObserver) (qr *sqltypes.Result, errs []error) - StreamExecuteMulti(ctx context.Context, primitive engine.Primitive, query string, rss []*srvtopo.ResolvedShard, vars []map[string]*querypb.BindVariable, session *SafeSession, autocommit bool, callback func(reply *sqltypes.Result) error, observer resultsObserver) []error - ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) - Commit(ctx context.Context, safeSession *SafeSession) error + Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, method string, session *econtext.SafeSession, s string, vars map[string]*querypb.BindVariable) (*sqltypes.Result, error) + ExecuteMultiShard(ctx context.Context, primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *econtext.SafeSession, autocommit bool, ignoreMaxMemoryRows bool, resultsObserver resultsObserver) (qr *sqltypes.Result, errs []error) + StreamExecuteMulti(ctx context.Context, primitive engine.Primitive, query string, rss []*srvtopo.ResolvedShard, vars []map[string]*querypb.BindVariable, session *econtext.SafeSession, autocommit bool, callback func(reply *sqltypes.Result) error, observer resultsObserver) []error + ExecuteLock(ctx context.Context, rs *srvtopo.ResolvedShard, query *querypb.BoundQuery, session *econtext.SafeSession, lockFuncType sqlparser.LockingFuncType) (*sqltypes.Result, error) + Commit(ctx context.Context, safeSession *econtext.SafeSession) error ExecuteMessageStream(ctx context.Context, rss []*srvtopo.ResolvedShard, name string, callback func(*sqltypes.Result) error) error ExecuteVStream(ctx context.Context, rss []*srvtopo.ResolvedShard, filter *binlogdatapb.Filter, gtid string, callback func(evs []*binlogdatapb.VEvent) error) error - ReleaseLock(ctx context.Context, session *SafeSession) error + ReleaseLock(ctx context.Context, session *econtext.SafeSession) error showVitessReplicationStatus(ctx context.Context, filter *sqlparser.ShowFilter) (*sqltypes.Result, error) showShards(ctx context.Context, filter *sqlparser.ShowFilter, destTabletType topodatapb.TabletType) (*sqltypes.Result, error) @@ -102,7 +103,7 @@ type ( // vcursorImpl implements the VCursor functionality used by dependent // packages to call back into VTGate. vcursorImpl struct { - safeSession *SafeSession + safeSession *econtext.SafeSession keyspace string tabletType topodatapb.TabletType destination key.Destination @@ -145,7 +146,7 @@ type ( // including as identifying markers. So, they have to be added back to all queries that are executed // on behalf of the original query. func newVCursorImpl( - safeSession *SafeSession, + safeSession *econtext.SafeSession, marginComments sqlparser.MarginComments, executor *Executor, logStats *logstats.LogStats, @@ -612,8 +613,7 @@ func (vc *vcursorImpl) Execute(ctx context.Context, method string, query string, session := vc.safeSession if co == vtgatepb.CommitOrder_AUTOCOMMIT { // For autocommit, we have to create an independent session. - session = NewAutocommitSession(vc.safeSession.Session) - session.logging = vc.safeSession.logging + session = vc.safeSession.NewAutocommitSession() rollbackOnError = false } else { session.SetCommitOrder(co) @@ -648,8 +648,6 @@ func (vc *vcursorImpl) markSavepoint(ctx context.Context, needsRollbackOnParialE return nil } -const txRollback = "Rollback Transaction" - // ExecuteMultiShard is part of the engine.VCursor interface. func (vc *vcursorImpl) ExecuteMultiShard(ctx context.Context, primitive engine.Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, rollbackOnError, canAutocommit bool) (*sqltypes.Result, []error) { noOfShards := len(rss) @@ -699,7 +697,7 @@ func (vc *vcursorImpl) ExecuteStandalone(ctx context.Context, primitive engine.P } // The autocommit flag is always set to false because we currently don't // execute DMLs through ExecuteStandalone. - qr, errs := vc.executor.ExecuteMultiShard(ctx, primitive, rss, bqs, NewAutocommitSession(vc.safeSession.Session), false /* autocommit */, vc.ignoreMaxMemoryRows, vc.resultsObserver) + qr, errs := vc.executor.ExecuteMultiShard(ctx, primitive, rss, bqs, econtext.NewAutocommitSession(vc.safeSession.Session), false /* autocommit */, vc.ignoreMaxMemoryRows, vc.resultsObserver) vc.logShardsQueried(primitive, len(rss)) return qr, vterrors.Aggregate(errs) } @@ -722,9 +720,9 @@ func (vc *vcursorImpl) ExecuteKeyspaceID(ctx context.Context, keyspace string, k // This creates a transaction but that transaction is for locking purpose only and should not cause multi-db transaction error. // This fields helps in to ignore multi-db transaction error when it states `queryFromVindex`. if !rollbackOnError { - vc.safeSession.queryFromVindex = true + vc.safeSession.SetQueryFromVindex(true) defer func() { - vc.safeSession.queryFromVindex = false + vc.safeSession.SetQueryFromVindex(false) }() } qr, errs := vc.ExecuteMultiShard(ctx, nil, rss, queries, rollbackOnError, autocommit) @@ -1023,8 +1021,7 @@ func (vc *vcursorImpl) SetWorkloadName(workloadName string) { // SetFoundRows implements the SessionActions interface func (vc *vcursorImpl) SetFoundRows(foundRows uint64) { - vc.safeSession.FoundRows = foundRows - vc.safeSession.foundRowsHandled = true + vc.safeSession.SetFoundRows(foundRows) } // SetDDLStrategy implements the SessionActions interface @@ -1387,8 +1384,7 @@ func (vc *vcursorImpl) ReleaseLock(ctx context.Context) error { } func (vc *vcursorImpl) cloneWithAutocommitSession() *vcursorImpl { - safeSession := NewAutocommitSession(vc.safeSession.Session) - safeSession.logging = vc.safeSession.logging + safeSession := vc.safeSession.NewAutocommitSession() return &vcursorImpl{ safeSession: safeSession, keyspace: vc.keyspace, @@ -1413,7 +1409,7 @@ func (vc *vcursorImpl) VExplainLogging() { } func (vc *vcursorImpl) GetVExplainLogs() []engine.ExecuteEntry { - return vc.safeSession.logging.GetLogs() + return vc.safeSession.GetLogs() } func (vc *vcursorImpl) FindRoutedShard(keyspace, shard string) (keyspaceName string, err error) { @@ -1460,7 +1456,7 @@ func (vc *vcursorImpl) CloneForReplicaWarming(ctx context.Context) engine.VCurso clonedCtx := callerid.NewContext(timedCtx, callerId, immediateCallerId) v := &vcursorImpl{ - safeSession: NewAutocommitSession(vc.safeSession.Session), + safeSession: econtext.NewAutocommitSession(vc.safeSession.Session), keyspace: vc.keyspace, tabletType: topodatapb.TabletType_REPLICA, destination: vc.destination, @@ -1492,7 +1488,7 @@ func (vc *vcursorImpl) CloneForMirroring(ctx context.Context) engine.VCursor { clonedCtx := callerid.NewContext(ctx, callerId, immediateCallerId) v := &vcursorImpl{ - safeSession: NewAutocommitSession(vc.safeSession.Session), + safeSession: econtext.NewAutocommitSession(vc.safeSession.Session), keyspace: vc.keyspace, tabletType: vc.tabletType, destination: vc.destination, diff --git a/go/vt/vtgate/vcursor_impl_test.go b/go/vt/vtgate/vcursor_impl_test.go index 95d9a18078d..1f91af36d05 100644 --- a/go/vt/vtgate/vcursor_impl_test.go +++ b/go/vt/vtgate/vcursor_impl_test.go @@ -12,6 +12,8 @@ import ( "github.com/stretchr/testify/require" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" + "vitess.io/vitess/go/vt/key" "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/srvtopo" @@ -190,7 +192,7 @@ func TestDestinationKeyspace(t *testing.T) { r, _, _, _, _ := createExecutorEnv(t) for i, tc := range tests { t.Run(strconv.Itoa(i)+tc.targetString, func(t *testing.T) { - impl, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: tc.targetString}), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, nil, nil, false, querypb.ExecuteOptions_Gen4) + impl, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{TargetString: tc.targetString}), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, nil, nil, false, querypb.ExecuteOptions_Gen4) impl.vschema = tc.vschema dest, keyspace, tabletType, err := impl.TargetDestination(tc.qualifier) if tc.expectedError == "" { @@ -253,7 +255,7 @@ func TestSetTarget(t *testing.T) { r, _, _, _, _ := createExecutorEnv(t) for i, tc := range tests { t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { - vc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{InTransaction: true}), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, nil, nil, false, querypb.ExecuteOptions_Gen4) + vc, _ := newVCursorImpl(econtext.NewSafeSession(&vtgatepb.Session{InTransaction: true}), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, nil, nil, false, querypb.ExecuteOptions_Gen4) vc.vschema = tc.vschema err := vc.SetTarget(tc.targetString) if tc.expectedError == "" { @@ -302,7 +304,7 @@ func TestKeyForPlan(t *testing.T) { r, _, _, _, _ := createExecutorEnv(t) for i, tc := range tests { t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { - ss := NewSafeSession(&vtgatepb.Session{InTransaction: false}) + ss := econtext.NewSafeSession(&vtgatepb.Session{InTransaction: false}) ss.SetTargetString(tc.targetString) vc, err := newVCursorImpl(ss, sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&fakeTopoServer{}, nil, ""), nil, false, querypb.ExecuteOptions_Gen4) require.NoError(t, err) @@ -328,7 +330,7 @@ func TestFirstSortedKeyspace(t *testing.T) { } r, _, _, _, _ := createExecutorEnv(t) - vc, err := newVCursorImpl(NewSafeSession(nil), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: vschemaWith2KS}, vschemaWith2KS, srvtopo.NewResolver(&fakeTopoServer{}, nil, ""), nil, false, querypb.ExecuteOptions_Gen4) + vc, err := newVCursorImpl(econtext.NewSafeSession(nil), sqlparser.MarginComments{}, r, nil, &fakeVSchemaOperator{vschema: vschemaWith2KS}, vschemaWith2KS, srvtopo.NewResolver(&fakeTopoServer{}, nil, ""), nil, false, querypb.ExecuteOptions_Gen4) require.NoError(t, err) ks, err := vc.FirstSortedKeyspace() require.NoError(t, err) @@ -339,7 +341,7 @@ func TestFirstSortedKeyspace(t *testing.T) { // Validates the timeout value is set based on override rule. func TestSetExecQueryTimeout(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) - safeSession := NewSafeSession(nil) + safeSession := econtext.NewSafeSession(nil) vc, err := newVCursorImpl(safeSession, sqlparser.MarginComments{}, executor, nil, nil, &vindexes.VSchema{}, nil, nil, false, querypb.ExecuteOptions_Gen4) require.NoError(t, err) @@ -382,7 +384,7 @@ func TestSetExecQueryTimeout(t *testing.T) { func TestRecordMirrorStats(t *testing.T) { executor, _, _, _, _ := createExecutorEnv(t) - safeSession := NewSafeSession(nil) + safeSession := econtext.NewSafeSession(nil) logStats := logstats.NewLogStats(context.Background(), t.Name(), "select 1", "", nil) vc, err := newVCursorImpl(safeSession, sqlparser.MarginComments{}, executor, logStats, nil, &vindexes.VSchema{}, nil, nil, false, querypb.ExecuteOptions_Gen4) require.NoError(t, err) diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go index e9e7cd65011..d9d9063aef9 100644 --- a/go/vt/vtgate/vtgate.go +++ b/go/vt/vtgate/vtgate.go @@ -51,6 +51,7 @@ import ( "vitess.io/vitess/go/vt/topo/topoproto" "vitess.io/vitess/go/vt/vtenv" "vitess.io/vitess/go/vt/vterrors" + econtext "vitess.io/vitess/go/vt/vtgate/executorcontext" "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" vtschema "vitess.io/vitess/go/vt/vtgate/schema" "vitess.io/vitess/go/vt/vtgate/txresolver" @@ -469,7 +470,7 @@ func (vtg *VTGate) Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConn if bvErr := sqltypes.ValidateBindVariables(bindVariables); bvErr != nil { err = vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "%v", bvErr) } else { - safeSession := NewSafeSession(session) + safeSession := econtext.NewSafeSession(session) qr, err = vtg.executor.Execute(ctx, mysqlCtx, "Execute", safeSession, sql, bindVariables) safeSession.RemoveInternalSavepoint() } @@ -526,7 +527,7 @@ func (vtg *VTGate) StreamExecute(ctx context.Context, mysqlCtx vtgateservice.MyS defer vtg.timings.Record(statsKey, time.Now()) - safeSession := NewSafeSession(session) + safeSession := econtext.NewSafeSession(session) var err error if bvErr := sqltypes.ValidateBindVariables(bindVariables); bvErr != nil { err = vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "%v", bvErr) @@ -560,7 +561,7 @@ func (vtg *VTGate) StreamExecute(ctx context.Context, mysqlCtx vtgateservice.MyS // same effect as if a "rollback" statement was executed, but does not affect the query // statistics. func (vtg *VTGate) CloseSession(ctx context.Context, session *vtgatepb.Session) error { - return vtg.executor.CloseSession(ctx, NewSafeSession(session)) + return vtg.executor.CloseSession(ctx, econtext.NewSafeSession(session)) } // Prepare supports non-streaming prepare statement query with multi shards @@ -575,7 +576,7 @@ func (vtg *VTGate) Prepare(ctx context.Context, session *vtgatepb.Session, sql s goto handleError } - fld, err = vtg.executor.Prepare(ctx, "Prepare", NewSafeSession(session), sql, bindVariables) + fld, err = vtg.executor.Prepare(ctx, "Prepare", econtext.NewSafeSession(session), sql, bindVariables) if err == nil { return session, fld, nil }