add optional flag to enable multi shard autocommit by default

demmer · demmer · commit 320f960d21c8 · 2025-01-27T13:43:06.000-08:00
Adds a new vtgate flag --default-multi-shard-autocommit which, as the name
implies, opts the query engine into using multi-shard autocommit semantics by
default, even if the plan does not contain the query directive
MULTI_SHARD_AUTOCOMMIT.

Signed-off-by: Michael Demmer &lt;mdemmer@slack-corp.com&gt;
diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go
@@ -75,7 +75,7 @@ func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShar
 	var schemaTracker vtgate.SchemaInfo // no schema tracker for these tests
 	queryLogBufferSize := 10
 	plans := theine.NewStore[vtgate.PlanCacheKey, *engine.Plan](4*1024*1024, false)
-	vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, plans, schemaTracker, false, opts.PlannerVersion, 0)
+	vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, plans, schemaTracker /* noScatter */, false /* defaultMultiShardAutocommit */, false, opts.PlannerVersion, 0)
 	vte.vtgateExecutor.SetQueryLogger(streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize))
 
 	return nil
diff --git a/go/vt/vtgate/engine/dml.go b/go/vt/vtgate/engine/dml.go
@@ -130,7 +130,7 @@ func allowOnlyPrimary(rss ...*srvtopo.ResolvedShard) error {
 }
 
 func execMultiShard(ctx context.Context, primitive Primitive, vcursor VCursor, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, multiShardAutoCommit bool) (*sqltypes.Result, error) {
-	autocommit := (len(rss) == 1 || multiShardAutoCommit) && vcursor.AutocommitApproval()
+	autocommit := (len(rss) == 1 || vcursor.DefaultMultiShardAutocommit() || multiShardAutoCommit) && vcursor.AutocommitApproval()
 	result, errs := vcursor.ExecuteMultiShard(ctx, primitive, rss, queries, true /* rollbackOnError */, autocommit)
 	return result, vterrors.Aggregate(errs)
 }
diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go
@@ -373,6 +373,10 @@ func (t *noopVCursor) GetDBDDLPluginName() string {
 	panic("unimplemented")
 }
 
+func (t *noopVCursor) DefaultMultiShardAutocommit() bool {
+	panic("unimplemented")
+}
+
 var _ VCursor = (*loggingVCursor)(nil)
 var _ SessionActions = (*loggingVCursor)(nil)
 
@@ -775,6 +779,10 @@ func (f *loggingVCursor) SetPriority(string) {
 	panic("implement me")
 }
 
+func (f *loggingVCursor) DefaultMultiShardAutocommit() bool {
+	return false
+}
+
 func (f *loggingVCursor) FindRoutedTable(tbl sqlparser.TableName) (*vindexes.Table, error) {
 	f.log = append(f.log, fmt.Sprintf("FindTable(%s)", sqlparser.String(tbl)))
 	return f.tableRoutes.tbl, nil
diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go
@@ -165,7 +165,7 @@ func (ins *Insert) executeInsertQueries(
 	queries []*querypb.BoundQuery,
 	insertID uint64,
 ) (*sqltypes.Result, error) {
-	autocommit := (len(rss) == 1 || ins.MultiShardAutocommit) && vcursor.AutocommitApproval()
+	autocommit := (len(rss) == 1 || vcursor.DefaultMultiShardAutocommit() || ins.MultiShardAutocommit) && vcursor.AutocommitApproval()
 	err := allowOnlyPrimary(rss...)
 	if err != nil {
 		return nil, err
diff --git a/go/vt/vtgate/engine/insert_select.go b/go/vt/vtgate/engine/insert_select.go
@@ -201,7 +201,7 @@ func (ins *InsertSelect) executeInsertQueries(
 	queries []*querypb.BoundQuery,
 	insertID uint64,
 ) (*sqltypes.Result, error) {
-	autocommit := (len(rss) == 1 || ins.MultiShardAutocommit) && vcursor.AutocommitApproval()
+	autocommit := (len(rss) == 1 || vcursor.DefaultMultiShardAutocommit() || ins.MultiShardAutocommit) && vcursor.AutocommitApproval()
 	err := allowOnlyPrimary(rss...)
 	if err != nil {
 		return nil, err
diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
@@ -129,6 +129,9 @@ type (
 
 		// CloneForReplicaWarming clones the VCursor for re-use in warming queries to replicas
 		CloneForReplicaWarming(ctx context.Context) VCursor
+
+		// DefaultMultiShardAutocommit returns true if multi shard autocommit semantics are enabled by default
+		DefaultMultiShardAutocommit() bool
 	}
 
 	// SessionActions gives primitives ability to interact with the session state
diff --git a/go/vt/vtgate/engine/send.go b/go/vt/vtgate/engine/send.go
@@ -140,7 +140,7 @@ func (s *Send) checkAndReturnShards(ctx context.Context, vcursor VCursor) ([]*sr
 
 func (s *Send) canAutoCommit(vcursor VCursor, rss []*srvtopo.ResolvedShard) bool {
 	if s.IsDML {
-		return (len(rss) == 1 || s.MultishardAutocommit) && vcursor.AutocommitApproval()
+		return (len(rss) == 1 || vcursor.DefaultMultiShardAutocommit() || s.MultishardAutocommit) && vcursor.AutocommitApproval()
 	}
 	return false
 }
diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
@@ -117,6 +117,9 @@ type Executor struct {
 	// allowScatter will fail planning if set to false and a plan contains any scatter queries
 	allowScatter bool
 
+	// defaultMultiShardAutocommit will by default execute multi shard DML statements with autocommit
+	defaultMultiShardAutocommit bool
+
 	// queryLogger is passed in for logging from this vtgate executor.
 	queryLogger *streamlog.StreamLogger[*logstats.LogStats]
 
@@ -149,7 +152,7 @@ func NewExecutor(
 	streamSize int,
 	plans *PlanCache,
 	schemaTracker SchemaInfo,
-	noScatter bool,
+	noScatter, defaultMultiShardAutocommit bool,
 	pv plancontext.PlannerVersion,
 	warmingReadsPercent int,
 ) *Executor {
@@ -168,6 +171,8 @@ func NewExecutor(
 		plans:               plans,
 		warmingReadsPercent: warmingReadsPercent,
 		warmingReadsChannel: make(chan bool, warmingReadsConcurrency),
+
+		defaultMultiShardAutocommit: defaultMultiShardAutocommit,
 	}
 
 	vschemaacl.Init()
diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go
@@ -184,7 +184,7 @@ func createExecutorEnvCallback(t testing.TB, eachShard func(shard, ks string, ta
 	// one-off queries from thrashing the cache. Disable the doorkeeper in the tests to prevent flakiness.
 	plans := theine.NewStore[PlanCacheKey, *engine.Plan](queryPlanCacheMemory, false)
 
-	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	executor.SetQueryLogger(queryLogger)
 
 	key.AnyShardPicker = DestinationAnyShardPickerFirstShard{}
@@ -231,7 +231,7 @@ func createCustomExecutor(t testing.TB, vschema string) (executor *Executor, sbc
 
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()
-	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	executor.SetQueryLogger(queryLogger)
 
 	t.Cleanup(func() {
@@ -268,7 +268,7 @@ func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqlty
 	sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil)
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()
-	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	executor.SetQueryLogger(queryLogger)
 
 	t.Cleanup(func() {
@@ -293,7 +293,7 @@ func createExecutorEnvWithPrimaryReplicaConn(t testing.TB, ctx context.Context,
 	replica = hc.AddTestTablet(cell, "0-replica", 1, KsTestUnsharded, "0", topodatapb.TabletType_REPLICA, true, 1, nil)
 
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
-	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, DefaultPlanCache(), nil, false, querypb.ExecuteOptions_Gen4, warmingReadsPercent)
+	executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, DefaultPlanCache(), nil, false, false, querypb.ExecuteOptions_Gen4, warmingReadsPercent)
 	executor.SetQueryLogger(queryLogger)
 
 	t.Cleanup(func() {
diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go
@@ -1563,7 +1563,7 @@ func TestStreamSelectIN(t *testing.T) {
 func createExecutor(ctx context.Context, serv *sandboxTopo, cell string, resolver *Resolver) *Executor {
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()
-	ex := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	ex := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	ex.SetQueryLogger(queryLogger)
 	return ex
 }
@@ -3189,7 +3189,7 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) {
 	}
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()
-	executor := NewExecutor(ctx, serv, cell, resolver, true, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	executor := NewExecutor(ctx, serv, cell, resolver, true, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	executor.SetQueryLogger(queryLogger)
 	defer executor.Close()
 	// some sleep for all goroutines to start
diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go
@@ -68,7 +68,7 @@ func TestStreamSQLSharded(t *testing.T) {
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()
 
-	executor := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, 0)
+	executor := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, false, querypb.ExecuteOptions_Gen4, 0)
 	executor.SetQueryLogger(queryLogger)
 
 	defer executor.Close()
diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go
@@ -121,6 +121,8 @@ type vcursorImpl struct {
 
 	warmingReadsPercent int
 	warmingReadsChannel chan bool
+
+	defaultMultiShardAutocommit bool
 }
 
 // newVcursorImpl creates a vcursorImpl. Before creating this object, you have to separate out any marginComments that came with
@@ -167,9 +169,11 @@ func newVCursorImpl(
 
 	warmingReadsPct := 0
 	var warmingReadsChan chan bool
+	defaultMultiShardAutocommit := false
 	if executor != nil {
 		warmingReadsPct = executor.warmingReadsPercent
 		warmingReadsChan = executor.warmingReadsChannel
+		defaultMultiShardAutocommit = executor.defaultMultiShardAutocommit
 	}
 	return &vcursorImpl{
 		safeSession:         safeSession,
@@ -188,6 +192,8 @@ func newVCursorImpl(
 		pv:                  pv,
 		warmingReadsPercent: warmingReadsPct,
 		warmingReadsChannel: warmingReadsChan,
+
+		defaultMultiShardAutocommit: defaultMultiShardAutocommit,
 	}, nil
 }
 
@@ -1361,3 +1367,8 @@ func (vc *vcursorImpl) UpdateForeignKeyChecksState(fkStateFromQuery *bool) {
 func (vc *vcursorImpl) GetForeignKeyChecksState() *bool {
 	return vc.fkChecksState
 }
+
+// GetForeignKeyChecksState gets the stored foreign key checks state in the vcursor.
+func (vc *vcursorImpl) DefaultMultiShardAutocommit() bool {
+	return vc.defaultMultiShardAutocommit
+}
diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go
@@ -74,6 +74,8 @@ var (
 	noScatter          bool
 	enableShardRouting bool
 
+	defaultMultiShardAutocommit bool
+
 	// healthCheckRetryDelay is the time to wait before retrying healthcheck
 	healthCheckRetryDelay = 2 * time.Millisecond
 	// healthCheckTimeout is the timeout on the RPC call to tablets
@@ -151,6 +153,8 @@ func registerFlags(fs *pflag.FlagSet) {
 	fs.IntVar(&warmingReadsPercent, "warming-reads-percent", 0, "Percentage of reads on the primary to forward to replicas. Useful for keeping buffer pools warm")
 	fs.IntVar(&warmingReadsConcurrency, "warming-reads-concurrency", 500, "Number of concurrent warming reads allowed")
 	fs.DurationVar(&warmingReadsQueryTimeout, "warming-reads-query-timeout", 5*time.Second, "Timeout of warming read queries")
+	fs.BoolVar(&defaultMultiShardAutocommit, "default-multi-shard-autocommit", defaultMultiShardAutocommit, "By default execute multi-shard DML statements with autocommit, even without the MULTI_SHARD_AUTOCOMMIT directive")
+
 }
 
 func init() {
@@ -317,6 +321,7 @@ func Init(
 		plans,
 		si,
 		noScatter,
+		defaultMultiShardAutocommit,
 		pv,
 		warmingReadsPercent,
 	)

Original file line number	Diff line number	Diff line change
`@@ -130,7 +130,7 @@ func allowOnlyPrimary(rss ...*srvtopo.ResolvedShard) error {`
`130`	`130`	`}`
`131`	`131`
`132`	`132`	`func execMultiShard(ctx context.Context, primitive Primitive, vcursor VCursor, rss []srvtopo.ResolvedShard, queries []querypb.BoundQuery, multiShardAutoCommit bool) (*sqltypes.Result, error) {`
`133`		`- autocommit := (len(rss) == 1 \|\| multiShardAutoCommit) && vcursor.AutocommitApproval()`
	`133`	`+ autocommit := (len(rss) == 1 \|\| vcursor.DefaultMultiShardAutocommit() \|\| multiShardAutoCommit) && vcursor.AutocommitApproval()`
`134`	`134`	`result, errs := vcursor.ExecuteMultiShard(ctx, primitive, rss, queries, true /* rollbackOnError */, autocommit)`
`135`	`135`	`return result, vterrors.Aggregate(errs)`
`136`	`136`	`}`
Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,9 @@ type (`
`129`	`129`
`130`	`130`	`// CloneForReplicaWarming clones the VCursor for re-use in warming queries to replicas`
`131`	`131`	`CloneForReplicaWarming(ctx context.Context) VCursor`
	`132`	`+`
	`133`	`+ // DefaultMultiShardAutocommit returns true if multi shard autocommit semantics are enabled by default`
	`134`	`+ DefaultMultiShardAutocommit() bool`
`132`	`135`	`}`
`133`	`136`
`134`	`137`	`// SessionActions gives primitives ability to interact with the session state`
Original file line number	Diff line number	Diff line change
`@@ -140,7 +140,7 @@ func (s Send) checkAndReturnShards(ctx context.Context, vcursor VCursor) ([]sr`
`140`	`140`
`141`	`141`	`func (s Send) canAutoCommit(vcursor VCursor, rss []srvtopo.ResolvedShard) bool {`
`142`	`142`	`if s.IsDML {`
`143`		`- return (len(rss) == 1 \|\| s.MultishardAutocommit) && vcursor.AutocommitApproval()`
	`143`	`+ return (len(rss) == 1 \|\| vcursor.DefaultMultiShardAutocommit() \|\| s.MultishardAutocommit) && vcursor.AutocommitApproval()`
`144`	`144`	`}`
`145`	`145`	`return false`
`146`	`146`	`}`