Repository: vitessio/vitess
Pull request #6430 (2020-07-09 20:54:07)

Add query comment to ignore max memory row limit
### Overview
We currently have warn and max thresholds for the number of rows that can be returned by a given query. However, there exists no mechanism for ignoring the `max_memory_rows` limit for a given query when this flag is enabled. This PR introduces a new query comment directive that will allow us to ignore the `max_memory_rows` limit, similar to that for `max_payload_size` in [this PR](https://github.com/vitessio/vitess/pull/6143).

**Relevant Issue:** https://github.com/vitessio/vitess/issues/6213
**Relevant PR:** https://github.com/vitessio/vitess/pull/6143

### Implementation
1. Define a new comment directive `DirectiveIgnoreMaxMemoryRows`.
2. Extract the `DirectiveIgnoreMaxMemoryRows` value at the executor step and pass that into vcursor for later validation.
2. Skip `MaxMemoryRows` validation if `DirectiveIgnoreMaxMemoryRows` is set.

### Testing
Testing was done via unit testing and manual testing. The following manual test cases were covered:
- Inserts/updates/deletes continue to work as expected
- Selects within the `max_memory_rows` limit continue to work as expected
- Selects that exceed the `max_memory_rows` limit fail as expected
- Selects that exceed the `max_memory_rows` limit that have the `DirectiveIgnoreMaxMemoryRows` query comment will succeed as expected

deepthi: @spark4 can you resolve the conflicts?

Patch:
diff --git a/go/vt/sqlparser/comments.go b/go/vt/sqlparser/comments.go
index 9d5df6d31f..f4bc122b95 100644
--- a/go/vt/sqlparser/comments.go
+++ b/go/vt/sqlparser/comments.go
@@ -34,6 +34,8 @@ const (
 	DirectiveScatterErrorsAsWarnings = "SCATTER_ERRORS_AS_WARNINGS"
 	// DirectiveIgnoreMaxPayloadSize skips payload size validation when set.
 	DirectiveIgnoreMaxPayloadSize = "IGNORE_MAX_PAYLOAD_SIZE"
+	// DirectiveIgnoreMaxMemoryRows skips memory row validation when set.
+	DirectiveIgnoreMaxMemoryRows = "IGNORE_MAX_MEMORY_ROWS"
 )
 
 func isNonSpace(r rune) bool {
@@ -321,3 +323,24 @@ func IgnoreMaxPayloadSizeDirective(stmt Statement) bool {
 		return false
 	}
 }
+
+// IgnoreMaxMaxMemoryRowsDirective returns true if the max memory rows override
+// directive is set to true.
+func IgnoreMaxMaxMemoryRowsDirective(stmt Statement) bool {
+	switch stmt := stmt.(type) {
+	case *Select:
+		directives := ExtractCommentDirectives(stmt.Comments)
+		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
+	case *Insert:
+		directives := ExtractCommentDirectives(stmt.Comments)
+		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
+	case *Update:
+		directives := ExtractCommentDirectives(stmt.Comments)
+		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
+	case *Delete:
+		directives := ExtractCommentDirectives(stmt.Comments)
+		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
+	default:
+		return false
+	}
+}
diff --git a/go/vt/vtgate/engine/join.go b/go/vt/vtgate/engine/join.go
index 7fd70ec26c..29d61e9706 100644
--- a/go/vt/vtgate/engine/join.go
+++ b/go/vt/vtgate/engine/join.go
@@ -88,7 +88,7 @@ func (jn *Join) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariab
 		} else {
 			result.RowsAffected += uint64(len(rresult.Rows))
 		}
-		if len(result.Rows) > vcursor.MaxMemoryRows() {
+		if vcursor.ExceedsMaxMemoryRows(len(result.Rows)) {
 			return nil, fmt.Errorf("in-memory row count exceeded allowed limit of %d", vcursor.MaxMemoryRows())
 		}
 	}
diff --git a/go/vt/vtgate/engine/memory_sort.go b/go/vt/vtgate/engine/memory_sort.go
index c243a31d60..8c8badbd3f 100644
--- a/go/vt/vtgate/engine/memory_sort.go
+++ b/go/vt/vtgate/engine/memory_sort.go
@@ -120,7 +120,7 @@ func (ms *MemorySort) StreamExecute(vcursor VCursor, bindVars map[string]*queryp
 		for len(sh.rows) > count {
 			_ = heap.Pop(sh)
 		}
-		if len(sh.rows) > vcursor.MaxMemoryRows() {
+		if vcursor.ExceedsMaxMemoryRows(len(sh.rows)) {
 			return fmt.Errorf("in-memory row count exceeded allowed limit of %d", vcursor.MaxMemoryRows())
 		}
 		return nil
diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
index 8003a7ab58..45470a270b 100644
--- a/go/vt/vtgate/engine/primitive.go
+++ b/go/vt/vtgate/engine/primitive.go
@@ -54,6 +54,11 @@ type (
 		// MaxMemoryRows returns the maxMemoryRows flag value.
 		MaxMemoryRows() int
 
+		// ExceedsMaxMemoryRows returns a boolean indicating whether
+		// the maxMemoryRows value has been exceeded. Returns false
+		// if the max memory rows override directive is set to true
+		ExceedsMaxMemoryRows(numRows int) bool
+
 		// SetContextTimeout updates the context and sets a timeout.
 		SetContextTimeout(timeout time.Duration) context.CancelFunc
 
diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
index 896922cc84..048af33dfb 100644
--- a/go/vt/vtgate/executor.go
+++ b/go/vt/vtgate/executor.go
@@ -268,8 +268,8 @@ func (e *Executor) addNeededBindVars(bindVarNeeds sqlparser.BindVarNeeds, bindVa
 	return nil
 }
 
-func (e *Executor) destinationExec(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, dest key.Destination, destKeyspace string, destTabletType topodatapb.TabletType, logStats *LogStats) (*sqltypes.Result, error) {
-	return e.resolver.Execute(ctx, sql, bindVars, destKeyspace, destTabletType, dest, safeSession, safeSession.Options, logStats, false /* canAutocommit */)
+func (e *Executor) destinationExec(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, dest key.Destination, destKeyspace string, destTabletType topodatapb.TabletType, logStats *LogStats, ignoreMaxMemoryRows bool) (*sqltypes.Result, error) {
+	return e.resolver.Execute(ctx, sql, bindVars, destKeyspace, destTabletType, dest, safeSession, safeSession.Options, logStats, false /* canAutocommit */, ignoreMaxMemoryRows)
 }
 
 func (e *Executor) handleBegin(ctx context.Context, safeSession *SafeSession, destTabletType topodatapb.TabletType, logStats *LogStats) (*sqltypes.Result, error) {
@@ -304,7 +304,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, nonTxResponse func(query string) (*sqltypes.Result, error)) (*sqltypes.Result, error) {
+func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession, sql string, planType string, 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 = uint32(len(safeSession.ShardSessions))
@@ -332,8 +332,7 @@ func (e *Executor) handleSavepoint(ctx context.Context, safeSession *SafeSession
 	for i := range rss {
 		queries[i] = &querypb.BoundQuery{Sql: sql}
 	}
-
-	qr, errs := e.ExecuteMultiShard(ctx, rss, queries, safeSession, false /*autocommit*/)
+	qr, errs := e.ExecuteMultiShard(ctx, rss, queries, safeSession, false /*autocommit*/, ignoreMaxMemoryRows)
 	err := vterrors.Aggregate(errs)
 	if err != nil {
 		return nil, err
@@ -709,6 +708,7 @@ func (e *Executor) handleShow(ctx context.Context, safeSession *SafeSession, sql
 		// This code is unreachable.
 		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unrecognized SHOW statement: %v", sql)
 	}
+	ignoreMaxMemoryRows := sqlparser.IgnoreMaxMaxMemoryRowsDirective(stmt)
 	execStart := time.Now()
 	defer func() { logStats.ExecuteTime = time.Since(execStart) }()
 	switch strings.ToLower(show.Type) {
@@ -725,7 +725,7 @@ func (e *Executor) handleShow(ctx context.Context, safeSession *SafeSession, sql
 			if len(keyspaces) == 0 {
 				return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "no keyspaces available")
 			}
-			return e.handleOther(ctx, safeSession, sql, bindVars, dest, keyspaces[0], destTabletType, logStats)
+			return e.handleOther(ctx, safeSession, sql, bindVars, dest, keyspaces[0], destTabletType, logStats, ignoreMaxMemoryRows)
 		}
 	// for STATUS, return empty result set
 	case sqlparser.KeywordString(sqlparser.STATUS):
@@ -1019,7 +1019,7 @@ func (e *Executor) handleShow(ctx context.Context, safeSession *SafeSession, sql
 	}
 
 	// Any other show statement is passed through
-	return e.handleOther(ctx, safeSession, sql, bindVars, dest, destKeyspace, destTabletType, logStats)
+	return e.handleOther(ctx, safeSession, sql, bindVars, dest, destKeyspace, destTabletType, logStats, ignoreMaxMemoryRows)
 }
 
 func (e *Executor) showTablets() (*sqltypes.Result, error) {
@@ -1084,7 +1084,7 @@ func (e *Executor) showTablets() (*sqltypes.Result, error) {
 	}, nil
 }
 
-func (e *Executor) handleOther(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, dest key.Destination, destKeyspace string, destTabletType topodatapb.TabletType, logStats *LogStats) (*sqltypes.Result, error) {
+func (e *Executor) handleOther(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, dest key.Destination, destKeyspace string, destTabletType topodatapb.TabletType, logStats *LogStats, ignoreMaxMemoryRows bool) (*sqltypes.Result, error) {
 	if destKeyspace == "" {
 		return nil, errNoKeyspace
 	}
@@ -1115,8 +1115,9 @@ func (e *Executor) handleOther(ctx context.Context, safeSession *SafeSession, sq
 	default:
 		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "Destination can only be a single shard for statement: %s, got: %v", sql, dest)
 	}
+
 	execStart := time.Now()
-	result, err := e.destinationExec(ctx, safeSession, sql, bindVars, dest, destKeyspace, destTabletType, logStats)
+	result, err := e.destinationExec(ctx, safeSession, sql, bindVars, dest, destKeyspace, destTabletType, logStats, ignoreMaxMemoryRows)
 
 	e.updateQueryCounts("Other", "", "", int64(logStats.ShardQueries))
 
@@ -1136,6 +1137,7 @@ func (e *Executor) StreamExecute(ctx context.Context, method string, safeSession
 	}
 	query, comments := sqlparser.SplitMarginComments(sql)
 	vcursor, _ := newVCursorImpl(ctx, safeSession, comments, e, logStats, e.vm, e.VSchema(), e.resolver.resolver)
+	vcursor.SetIgnoreMaxMemoryRows(true)
 
 	switch stmtType {
 	case sqlparser.StmtStream:
@@ -1312,21 +1314,24 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
 	if e.VSchema() == nil {
 		return nil, errors.New("vschema not initialized")
 	}
-	planKey := vcursor.planPrefixKey() + ":" + sql
-	if plan, ok := e.plans.Get(planKey); ok {
-		return plan.(*engine.Plan), nil
-	}
+
 	stmt, err := sqlparser.Parse(sql)
 	if err != nil {
 		return nil, err
 	}
-
 	query := sql
 	statement := stmt
 	bindVarNeeds := sqlparser.BindVarNeeds{}
 	if !sqlparser.IgnoreMaxPayloadSizeDirective(statement) && !isValidPayloadSize(query) {
 		return nil, vterrors.New(vtrpcpb.Code_RESOURCE_EXHAUSTED, "query payload size above threshold")
 	}
+	ignoreMaxMemoryRows := sqlparser.IgnoreMaxMaxMemoryRowsDirective(stmt)
+	vcursor.SetIgnoreMaxMemoryRows(ignoreMaxMemoryRows)
+
+	planKey := vcursor.planPrefixKey() + ":" + sql
+	if plan, ok := e.plans.Get(planKey); ok {
+		return plan.(*engine.Plan), nil
+	}
 
 	// Normalize if possible and retry.
 	if (e.normalize && sqlparser.CanNormalize(stmt)) || sqlparser.IsSetStatement(stmt) {
@@ -1653,8 +1658,8 @@ func (e *Executor) handlePrepare(ctx context.Context, safeSession *SafeSession,
 }
 
 // ExecuteMultiShard implements the IExecutor interface
-func (e *Executor) ExecuteMultiShard(ctx context.Context, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool) (qr *sqltypes.Result, errs []error) {
-	return e.scatterConn.ExecuteMultiShard(ctx, rss, queries, session, autocommit)
+func (e *Executor) ExecuteMultiShard(ctx context.Context, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool, ignoreMaxMemoryRows bool) (qr *sqltypes.Result, errs []error) {
+	return e.scatterConn.ExecuteMultiShard(ctx, rss, queries, session, autocommit, ignoreMaxMemoryRows)
 }
 
 // StreamExecuteMulti implements the IExecutor interface
diff --git a/go/vt/vtgate/plan_execute.go b/go/vt/vtgate/plan_execute.go
index e6f49d297d..f2180a6898 100644
--- a/go/vt/vtgate/plan_execute.go
+++ b/go/vt/vtgate/plan_execute.go
@@ -89,19 +89,19 @@ func (e *Executor) newExecute(ctx context.Context, safeSession *SafeSession, sql
 		qr, err := e.handleSavepoint(ctx, safeSession, plan.Original, "Savepoint", logStats, func(_ string) (*sqltypes.Result, error) {
 			// Safely to ignore as there is no transaction.
 			return &sqltypes.Result{}, nil
-		})
+		}, vcursor.ignoreMaxMemoryRows)
 		return sqlparser.StmtSavepoint, qr, err
 	case sqlparser.StmtSRollback:
 		qr, err := e.handleSavepoint(ctx, safeSession, plan.Original, "Rollback Savepoint", logStats, func(query string) (*sqltypes.Result, error) {
 			// Error as there is no transaction, so there is no savepoint that exists.
 			return nil, mysql.NewSQLError(mysql.ERSavepointNotExist, "42000", "SAVEPOINT does not exist: %s", query)
-		})
+		}, vcursor.ignoreMaxMemoryRows)
 		return sqlparser.StmtSRollback, qr, err
 	case sqlparser.StmtRelease:
 		qr, err := e.handleSavepoint(ctx, safeSession, plan.Original, "Release Savepoint", logStats, func(query string) (*sqltypes.Result, error) {
 			// Error as there is no transaction, so there is no savepoint that exists.
 			return nil, mysql.NewSQLError(mysql.ERSavepointNotExist, "42000", "SAVEPOINT does not exist: %s", query)
-		})
+		}, vcursor.ignoreMaxMemoryRows)
 		return sqlparser.StmtRelease, qr, err
 	}
 
diff --git a/go/vt/vtgate/resolver.go b/go/vt/vtgate/resolver.go
index 437ffd095a..c98b0eb664 100644
--- a/go/vt/vtgate/resolver.go
+++ b/go/vt/vtgate/resolver.go
@@ -66,6 +66,7 @@ func (res *Resolver) Execute(
 	options *querypb.ExecuteOptions,
 	logStats *LogStats,
 	canAutocommit bool,
+	ignoreMaxMemoryRows bool,
 ) (*sqltypes.Result, error) {
 	rss, err := res.resolver.ResolveDestination(ctx, keyspace, tabletType, destination)
 	if err != nil {
@@ -94,6 +95,7 @@ func (res *Resolver) Execute(
 			queries,
 			session,
 			autocommit,
+			ignoreMaxMemoryRows,
 		)
 		err = vterrors.Aggregate(errors)
 		if isRetryableError(err) {
diff --git a/go/vt/vtgate/scatter_conn.go b/go/vt/vtgate/scatter_conn.go
index c6a7fc7e79..14e24b166e 100644
--- a/go/vt/vtgate/scatter_conn.go
+++ b/go/vt/vtgate/scatter_conn.go
@@ -150,6 +150,7 @@ func (stc *ScatterConn) ExecuteMultiShard(
 	queries []*querypb.BoundQuery,
 	session *SafeSession,
 	autocommit bool,
+	ignoreMaxMemoryRows bool,
 ) (qr *sqltypes.Result, errs []error) {
 
 	if len(rss) != len(queries) {
@@ -226,15 +227,16 @@ func (stc *ScatterConn) ExecuteMultiShard(
 			}
 			mu.Lock()
 			defer mu.Unlock()
+
 			// Don't append more rows if row count is exceeded.
-			if len(qr.Rows) <= *maxMemoryRows {
+			if ignoreMaxMemoryRows || len(qr.Rows) <= *maxMemoryRows {
 				qr.AppendResult(innerqr)
 			}
 			return info.updateTransactionAndReservedID(transactionID, reservedID, alias), nil
 		},
 	)
 
-	if len(qr.Rows) > *maxMemoryRows {
+	if !ignoreMaxMemoryRows && len(qr.Rows) > *maxMemoryRows {
 		return nil, []error{vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "in-memory row count exceeded allowed limit of %d", *maxMemoryRows)}
 	}
 
diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go
index 45e20a3c00..d76f29d7c2 100644
--- a/go/vt/vtgate/vcursor_impl.go
+++ b/go/vt/vtgate/vcursor_impl.go
@@ -58,7 +58,7 @@ var _ iExecute = (*Executor)(nil)
 // vcursor_impl needs these facilities to be able to be able to execute queries for vindexes
 type iExecute interface {
 	Execute(ctx context.Context, method string, session *SafeSession, s string, vars map[string]*querypb.BindVariable) (*sqltypes.Result, error)
-	ExecuteMultiShard(ctx context.Context, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool) (qr *sqltypes.Result, errs []error)
+	ExecuteMultiShard(ctx context.Context, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, session *SafeSession, autocommit bool, ignoreMaxMemoryRows bool) (qr *sqltypes.Result, errs []error)
 	StreamExecuteMulti(ctx context.Context, s string, rss []*srvtopo.ResolvedShard, vars []map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(reply *sqltypes.Result) error) error
 
 	// TODO: remove when resolver is gone
@@ -88,6 +88,7 @@ type vcursorImpl struct {
 	// executed. If there was a subsequent failure, the transaction
 	// must be forced to rollback.
 	rollbackOnPartialExec bool
+	ignoreMaxMemoryRows   bool
 	vschema               *vindexes.VSchema
 	vm                    VSchemaOperator
 }
@@ -169,6 +170,17 @@ func (vc *vcursorImpl) MaxMemoryRows() int {
 	return *maxMemoryRows
 }
 
+// ExceedsMaxMemoryRows returns a boolean indicating whether the maxMemoryRows value has been exceeded.
+// Returns false if the max memory rows override directive is set to true.
+func (vc *vcursorImpl) ExceedsMaxMemoryRows(numRows int) bool {
+	return !vc.ignoreMaxMemoryRows && numRows > *maxMemoryRows
+}
+
+// SetIgnoreMaxMemoryRows sets the ignoreMaxMemoryRows value.
+func (vc *vcursorImpl) SetIgnoreMaxMemoryRows(ignoreMaxMemoryRows bool) {
+	vc.ignoreMaxMemoryRows = ignoreMaxMemoryRows
+}
+
 // SetContextTimeout updates context and sets a timeout.
 func (vc *vcursorImpl) SetContextTimeout(timeout time.Duration) context.CancelFunc {
 	ctx, cancel := context.WithTimeout(vc.ctx, timeout)
@@ -284,7 +296,7 @@ func (vc *vcursorImpl) Execute(method string, query string, bindVars map[string]
 // ExecuteMultiShard is part of the engine.VCursor interface.
 func (vc *vcursorImpl) ExecuteMultiShard(rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, rollbackOnError, autocommit bool) (*sqltypes.Result, []error) {
 	atomic.AddUint32(&vc.logStats.ShardQueries, uint32(len(queries)))
-	qr, errs := vc.executor.ExecuteMultiShard(vc.ctx, rss, commentedShardQueries(queries, vc.marginComments), vc.safeSession, autocommit)
+	qr, errs := vc.executor.ExecuteMultiShard(vc.ctx, rss, commentedShardQueries(queries, vc.marginComments), vc.safeSession, autocommit, vc.ignoreMaxMemoryRows)
 
 	if errs == nil && rollbackOnError {
 		vc.rollbackOnPartialExec = true
@@ -308,7 +320,7 @@ func (vc *vcursorImpl) ExecuteStandalone(query string, bindVars map[string]*quer
 	}
 	// The autocommit flag is always set to false because we currently don't
 	// execute DMLs through ExecuteStandalone.
-	qr, errs := vc.executor.ExecuteMultiShard(vc.ctx, rss, bqs, NewAutocommitSession(vc.safeSession.Session), false /* autocommit */)
+	qr, errs := vc.executor.ExecuteMultiShard(vc.ctx, rss, bqs, NewAutocommitSession(vc.safeSession.Session), false /* autocommit */, vc.ignoreMaxMemoryRows)
 	return qr, vterrors.Aggregate(errs)
 }
 


Test patch:
diff --git a/go/vt/sqlparser/comments_test.go b/go/vt/sqlparser/comments_test.go
index 8ec2a0e199..7dc0ba3511 100644
--- a/go/vt/sqlparser/comments_test.go
+++ b/go/vt/sqlparser/comments_test.go
@@ -404,6 +404,32 @@ func TestIgnoreMaxPayloadSizeDirective(t *testing.T) {
 	for _, test := range testCases {
 		stmt, _ := Parse(test.query)
 		got := IgnoreMaxPayloadSizeDirective(stmt)
-		assert.Equalf(t, test.expected, got, fmt.Sprintf("d.IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
+		assert.Equalf(t, test.expected, got, fmt.Sprintf("IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
+	}
+}
+
+func TestIgnoreMaxMaxMemoryRowsDirective(t *testing.T) {
+	testCases := []struct {
+		query    string
+		expected bool
+	}{
+		{"insert /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ into user(id) values (1), (2)", true},
+		{"insert into user(id) values (1), (2)", false},
+		{"update /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ users set name=1", true},
+		{"update users set name=1", false},
+		{"select /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ * from users", true},
+		{"select * from users", false},
+		{"delete /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ from users", true},
+		{"delete from users", false},
+		{"show /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ create table users", false},
+		{"show create table users", false},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.query, func(t *testing.T) {
+			stmt, _ := Parse(test.query)
+			got := IgnoreMaxMaxMemoryRowsDirective(stmt)
+			assert.Equalf(t, test.expected, got, fmt.Sprintf("IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
+		})
 	}
 }
diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go
index f0452eef37..2683863501 100644
--- a/go/vt/vtgate/engine/fake_vcursor_test.go
+++ b/go/vt/vtgate/engine/fake_vcursor_test.go
@@ -44,6 +44,7 @@ import (
 )
 
 var testMaxMemoryRows = 100
+var testIgnoreMaxMemoryRows = false
 
 var _ VCursor = (*noopVCursor)(nil)
 var _ SessionActions = (*noopVCursor)(nil)
@@ -82,6 +83,10 @@ func (t noopVCursor) MaxMemoryRows() int {
 	return testMaxMemoryRows
 }
 
+func (t noopVCursor) ExceedsMaxMemoryRows(numRows int) bool {
+	return !testIgnoreMaxMemoryRows && numRows > testMaxMemoryRows
+}
+
 func (t noopVCursor) SetContextTimeout(timeout time.Duration) context.CancelFunc {
 	return func() {}
 }
diff --git a/go/vt/vtgate/engine/join_test.go b/go/vt/vtgate/engine/join_test.go
index 7ae4582f10..1946dea6ec 100644
--- a/go/vt/vtgate/engine/join_test.go
+++ b/go/vt/vtgate/engine/join_test.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	"vitess.io/vitess/go/sqltypes"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -127,62 +128,77 @@ func TestJoinExecute(t *testing.T) {
 }
 
 func TestJoinExecuteMaxMemoryRows(t *testing.T) {
-	save := testMaxMemoryRows
+	saveMax := testMaxMemoryRows
+	saveIgnore := testIgnoreMaxMemoryRows
 	testMaxMemoryRows = 3
-	defer func() { testMaxMemoryRows = save }()
+	defer func() {
+		testMaxMemoryRows = saveMax
+		testIgnoreMaxMemoryRows = saveIgnore
+	}()
 
-	leftPrim := &fakePrimitive{
-		results: []*sqltypes.Result{
-			sqltypes.MakeTestResult(
-				sqltypes.MakeTestFields(
-					"col1|col2|col3",
-					"int64|varchar|varchar",
+	testCases := []struct {
+		ignoreMaxMemoryRows bool
+		err                 string
+	}{
+		{true, ""},
+		{false, "in-memory row count exceeded allowed limit of 3"},
+	}
+	for _, test := range testCases {
+		leftPrim := &fakePrimitive{
+			results: []*sqltypes.Result{
+				sqltypes.MakeTestResult(
+					sqltypes.MakeTestFields(
+						"col1|col2|col3",
+						"int64|varchar|varchar",
+					),
+					"1|a|aa",
+					"2|b|bb",
+					"3|c|cc",
 				),
-				"1|a|aa",
-				"2|b|bb",
-				"3|c|cc",
-			),
-		},
-	}
-	rightFields := sqltypes.MakeTestFields(
-		"col4|col5|col6",
-		"int64|varchar|varchar",
-	)
-	rightPrim := &fakePrimitive{
-		results: []*sqltypes.Result{
-			sqltypes.MakeTestResult(
-				rightFields,
-				"4|d|dd",
-			),
-			sqltypes.MakeTestResult(
-				rightFields,
-			),
-			sqltypes.MakeTestResult(
-				rightFields,
-				"5|e|ee",
-				"6|f|ff",
-				"7|g|gg",
-			),
-		},
-	}
-	bv := map[string]*querypb.BindVariable{
-		"a": sqltypes.Int64BindVariable(10),
-	}
+			},
+		}
+		rightFields := sqltypes.MakeTestFields(
+			"col4|col5|col6",
+			"int64|varchar|varchar",
+		)
+		rightPrim := &fakePrimitive{
+			results: []*sqltypes.Result{
+				sqltypes.MakeTestResult(
+					rightFields,
+					"4|d|dd",
+				),
+				sqltypes.MakeTestResult(
+					rightFields,
+				),
+				sqltypes.MakeTestResult(
+					rightFields,
+					"5|e|ee",
+					"6|f|ff",
+					"7|g|gg",
+				),
+			},
+		}
+		bv := map[string]*querypb.BindVariable{
+			"a": sqltypes.Int64BindVariable(10),
+		}
 
-	// Normal join
-	jn := &Join{
-		Opcode: NormalJoin,
-		Left:   leftPrim,
-		Right:  rightPrim,
-		Cols:   []int{-1, -2, 1, 2},
-		Vars: map[string]int{
-			"bv": 1,
-		},
-	}
-	_, err := jn.Execute(noopVCursor{}, bv, true)
-	want := "in-memory row count exceeded allowed limit of 3"
-	if err == nil || err.Error() != want {
-		t.Errorf("Execute(): %v, want %v", err, want)
+		// Normal join
+		jn := &Join{
+			Opcode: NormalJoin,
+			Left:   leftPrim,
+			Right:  rightPrim,
+			Cols:   []int{-1, -2, 1, 2},
+			Vars: map[string]int{
+				"bv": 1,
+			},
+		}
+		testIgnoreMaxMemoryRows = test.ignoreMaxMemoryRows
+		_, err := jn.Execute(noopVCursor{}, bv, true)
+		if testIgnoreMaxMemoryRows {
+			require.NoError(t, err)
+		} else {
+			require.EqualError(t, err, test.err)
+		}
 	}
 }
 
diff --git a/go/vt/vtgate/engine/memory_sort_test.go b/go/vt/vtgate/engine/memory_sort_test.go
index a9732f4e23..a2a9d68252 100644
--- a/go/vt/vtgate/engine/memory_sort_test.go
+++ b/go/vt/vtgate/engine/memory_sort_test.go
@@ -21,6 +21,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -340,38 +341,53 @@ func TestMemorySortMultiColumn(t *testing.T) {
 }
 
 func TestMemorySortMaxMemoryRows(t *testing.T) {
-	save := testMaxMemoryRows
+	saveMax := testMaxMemoryRows
+	saveIgnore := testIgnoreMaxMemoryRows
 	testMaxMemoryRows = 3
-	defer func() { testMaxMemoryRows = save }()
-
+	defer func() {
+		testMaxMemoryRows = saveMax
+		testIgnoreMaxMemoryRows = saveIgnore
+	}()
+
+	testCases := []struct {
+		ignoreMaxMemoryRows bool
+		err                 string
+	}{
+		{true, ""},
+		{false, "in-memory row count exceeded allowed limit of 3"},
+	}
 	fields := sqltypes.MakeTestFields(
 		"c1|c2",
 		"varbinary|decimal",
 	)
-	fp := &fakePrimitive{
-		results: []*sqltypes.Result{sqltypes.MakeTestResult(
-			fields,
-			"a|1",
-			"b|2",
-			"a|1",
-			"c|4",
-			"c|3",
-		)},
-	}
-
-	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
-			Col: 1,
-		}},
-		Input: fp,
-	}
-
-	err := ms.StreamExecute(noopVCursor{}, nil, false, func(qr *sqltypes.Result) error {
-		return nil
-	})
-	want := "in-memory row count exceeded allowed limit of 3"
-	if err == nil || err.Error() != want {
-		t.Errorf("StreamExecute err: %v, want %v", err, want)
+	for _, test := range testCases {
+		fp := &fakePrimitive{
+			results: []*sqltypes.Result{sqltypes.MakeTestResult(
+				fields,
+				"a|1",
+				"b|2",
+				"a|1",
+				"c|4",
+				"c|3",
+			)},
+		}
+
+		ms := &MemorySort{
+			OrderBy: []OrderbyParams{{
+				Col: 1,
+			}},
+			Input: fp,
+		}
+
+		testIgnoreMaxMemoryRows = test.ignoreMaxMemoryRows
+		err := ms.StreamExecute(noopVCursor{}, nil, false, func(qr *sqltypes.Result) error {
+			return nil
+		})
+		if testIgnoreMaxMemoryRows {
+			require.NoError(t, err)
+		} else {
+			require.EqualError(t, err, test.err)
+		}
 	}
 }
 
diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go
index 92e811e38d..f147927736 100644
--- a/go/vt/vtgate/executor_test.go
+++ b/go/vt/vtgate/executor_test.go
@@ -75,6 +75,44 @@ func TestExecutorResultsExceeded(t *testing.T) {
 	assert.Equal(t, initial+1, warnings.Counts()["ResultsExceeded"], "warnings count")
 }
 
+func TestExecutorMaxMemoryRowsExceeded(t *testing.T) {
+	save := *maxMemoryRows
+	*maxMemoryRows = 3
+	defer func() { *maxMemoryRows = save }()
+
+	executor, _, _, sbclookup := createExecutorEnv()
+	session := NewSafeSession(&vtgatepb.Session{TargetString: "@master"})
+	result := sqltypes.MakeTestResult(sqltypes.MakeTestFields("col", "int64"), "1", "2", "3", "4")
+	target := querypb.Target{}
+	fn := func(r *sqltypes.Result) error {
+		return nil
+	}
+	testCases := []struct {
+		query string
+		err   string
+	}{
+		{"select /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ * from main1", ""},
+		{"select * from main1", "in-memory row count exceeded allowed limit of 3"},
+	}
+
+	for _, test := range testCases {
+		sbclookup.SetResults([]*sqltypes.Result{result})
+		stmt, err := sqlparser.Parse(test.query)
+		require.NoError(t, err)
+
+		_, err = executor.Execute(ctx, "TestExecutorMaxMemoryRowsExceeded", session, test.query, nil)
+		if sqlparser.IgnoreMaxMaxMemoryRowsDirective(stmt) {
+			require.NoError(t, err, "no error when DirectiveIgnoreMaxMemoryRows is provided")
+		} else {
+			assert.EqualError(t, err, test.err, "maxMemoryRows limit exceeded")
+		}
+
+		sbclookup.SetResults([]*sqltypes.Result{result})
+		err = executor.StreamExecute(ctx, "TestExecutorMaxMemoryRowsExceeded", session, test.query, nil, target, fn)
+		require.NoError(t, err, "maxMemoryRows limit does not apply to StreamExecute")
+	}
+}
+
 func TestLegacyExecutorTransactionsNoAutoCommit(t *testing.T) {
 	executor, _, _, sbclookup := createExecutorEnv()
 	session := NewSafeSession(&vtgatepb.Session{TargetString: "@master"})
diff --git a/go/vt/vtgate/legacy_scatter_conn_test.go b/go/vt/vtgate/legacy_scatter_conn_test.go
index 93f0164e9b..381a8a60e5 100644
--- a/go/vt/vtgate/legacy_scatter_conn_test.go
+++ b/go/vt/vtgate/legacy_scatter_conn_test.go
@@ -97,7 +97,7 @@ func TestLegacyExecuteFailOnAutocommit(t *testing.T) {
 		},
 		Autocommit: false,
 	}
-	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(session), true /*autocommit*/)
+	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(session), true /*autocommit*/, false)
 	err := vterrors.Aggregate(errs)
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "in autocommit mode, transactionID should be zero but was: 123")
@@ -121,7 +121,7 @@ func TestScatterConnExecuteMulti(t *testing.T) {
 			}
 		}
 
-		qr, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(nil), false /*autocommit*/)
+		qr, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(nil), false /*autocommit*/, false)
 		return qr, vterrors.Aggregate(errs)
 	})
 }
@@ -287,26 +287,12 @@ func TestMaxMemoryRows(t *testing.T) {
 	sbc0 := hc.AddTestTablet("aa", "0", 1, "TestMaxMemoryRows", "0", topodatapb.TabletType_REPLICA, true, 1, nil)
 	sbc1 := hc.AddTestTablet("aa", "1", 1, "TestMaxMemoryRows", "1", topodatapb.TabletType_REPLICA, true, 1, nil)
 
-	tworows := &sqltypes.Result{
-		Rows: [][]sqltypes.Value{{
-			sqltypes.NewInt64(1),
-		}, {
-			sqltypes.NewInt64(1),
-		}},
-		RowsAffected: 1,
-		InsertID:     1,
-	}
-	sbc0.SetResults([]*sqltypes.Result{tworows, tworows})
-	sbc1.SetResults([]*sqltypes.Result{tworows, tworows})
-
 	res := srvtopo.NewResolver(&sandboxTopo{}, sc.gateway, "aa")
 	rss, _, err := res.ResolveDestinations(ctx, "TestMaxMemoryRows", topodatapb.TabletType_REPLICA, nil,
 		[]key.Destination{key.DestinationShard("0"), key.DestinationShard("1")})
 	require.NoError(t, err)
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-
-	want := "in-memory row count exceeded allowed limit of 3"
 	queries := []*querypb.BoundQuery{{
 		Sql:           "query1",
 		BindVariables: map[string]*querypb.BindVariable{},
@@ -314,8 +300,35 @@ func TestMaxMemoryRows(t *testing.T) {
 		Sql:           "query1",
 		BindVariables: map[string]*querypb.BindVariable{},
 	}}
-	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, session, false)
-	assert.EqualError(t, errs[0], want)
+	tworows := &sqltypes.Result{
+		Rows: [][]sqltypes.Value{{
+			sqltypes.NewInt64(1),
+		}, {
+			sqltypes.NewInt64(1),
+		}},
+		RowsAffected: 1,
+		InsertID:     1,
+	}
+
+	testCases := []struct {
+		ignoreMaxMemoryRows bool
+		err                 string
+	}{
+		{true, ""},
+		{false, "in-memory row count exceeded allowed limit of 3"},
+	}
+
+	for _, test := range testCases {
+		sbc0.SetResults([]*sqltypes.Result{tworows, tworows})
+		sbc1.SetResults([]*sqltypes.Result{tworows, tworows})
+
+		_, errs := sc.ExecuteMultiShard(ctx, rss, queries, session, false, test.ignoreMaxMemoryRows)
+		if test.ignoreMaxMemoryRows {
+			require.NoError(t, err)
+		} else {
+			assert.EqualError(t, errs[0], test.err)
+		}
+	}
 }
 
 func TestReservedBeginTableDriven(t *testing.T) {
@@ -515,7 +528,7 @@ func executeOnShards(t *testing.T, res *srvtopo.Resolver, keyspace string, sc *S
 		})
 	}
 
-	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, session, false)
+	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, session, false, false)
 	require.Empty(t, errs)
 }
 
@@ -557,7 +570,7 @@ func TestMultiExecs(t *testing.T) {
 		},
 	}
 
-	_, _ = sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(nil), false)
+	_, _ = sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(nil), false, false)
 	if len(sbc0.Queries) == 0 || len(sbc1.Queries) == 0 {
 		t.Fatalf("didn't get expected query")
 	}
@@ -651,27 +664,27 @@ func TestScatterConnSingleDB(t *testing.T) {
 	// TransactionMode_SINGLE in session
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true, TransactionMode: vtgatepb.TransactionMode_SINGLE})
 	queries := []*querypb.BoundQuery{{Sql: "query1"}}
-	_, errors := sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	_, errors := sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	require.Empty(t, errors)
-	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	require.Error(t, errors[0])
 	assert.Contains(t, errors[0].Error(), want)
 
 	// TransactionMode_SINGLE in txconn
 	sc.txConn.mode = vtgatepb.TransactionMode_SINGLE
 	session = NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	_, errors = sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	_, errors = sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	require.Empty(t, errors)
-	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	require.Error(t, errors[0])
 	assert.Contains(t, errors[0].Error(), want)
 
 	// TransactionMode_MULTI in txconn. Should not fail.
 	sc.txConn.mode = vtgatepb.TransactionMode_MULTI
 	session = NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	_, errors = sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	_, errors = sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	require.Empty(t, errors)
-	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	_, errors = sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	require.Empty(t, errors)
 }
 
diff --git a/go/vt/vtgate/scatter_conn_test.go b/go/vt/vtgate/scatter_conn_test.go
index f54c161a18..da415f90d0 100644
--- a/go/vt/vtgate/scatter_conn_test.go
+++ b/go/vt/vtgate/scatter_conn_test.go
@@ -88,7 +88,7 @@ func TestExecuteFailOnAutocommit(t *testing.T) {
 		},
 		Autocommit: false,
 	}
-	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(session), true /*autocommit*/)
+	_, errs := sc.ExecuteMultiShard(ctx, rss, queries, NewSafeSession(session), true /*autocommit*/, false)
 	err := vterrors.Aggregate(errs)
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "in autocommit mode, transactionID should be zero but was: 123")
diff --git a/go/vt/vtgate/tx_conn_test.go b/go/vt/vtgate/tx_conn_test.go
index 5f9b660dbc..89d820ea31 100644
--- a/go/vt/vtgate/tx_conn_test.go
+++ b/go/vt/vtgate/tx_conn_test.go
@@ -50,7 +50,7 @@ func TestTxConnBegin(t *testing.T) {
 	require.NoError(t, err)
 	wantSession := vtgatepb.Session{InTransaction: true}
 	utils.MustMatch(t, &wantSession, session, "Session")
-	_, errors := sc.ExecuteMultiShard(ctx, rss0, queries, safeSession, false)
+	_, errors := sc.ExecuteMultiShard(ctx, rss0, queries, safeSession, false, false)
 	require.Empty(t, errors)
 
 	// Begin again should cause a commit and a new begin.
@@ -66,7 +66,7 @@ func TestTxConnCommitSuccess(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession := vtgatepb.Session{
 		InTransaction: true,
 		ShardSessions: []*vtgatepb.Session_ShardSession{{
@@ -80,7 +80,7 @@ func TestTxConnCommitSuccess(t *testing.T) {
 		}},
 	}
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction: true,
 		ShardSessions: []*vtgatepb.Session_ShardSession{{
@@ -117,7 +117,7 @@ func TestTxConnReservedCommitSuccess(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession := vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -133,7 +133,7 @@ func TestTxConnReservedCommitSuccess(t *testing.T) {
 		}},
 	}
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -202,9 +202,9 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) {
 	session := NewSafeSession(&vtgatepb.Session{InReservedConn: true})
 
 	// this will create reserved connections against all tablets
-	_, errs := sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	_, errs := sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	require.Empty(t, errs)
-	_, errs = sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	_, errs = sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	require.Empty(t, errs)
 
 	wantSession := vtgatepb.Session{
@@ -232,7 +232,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) {
 	session.Session.InTransaction = true
 
 	// start a transaction against rss0
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -295,9 +295,9 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) {
 	session := NewSafeSession(&vtgatepb.Session{InReservedConn: true})
 
 	// this will create reserved connections against all tablets
-	_, errs := sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	_, errs := sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	require.Empty(t, errs)
-	_, errs = sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	_, errs = sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	require.Empty(t, errs)
 
 	wantSession := vtgatepb.Session{
@@ -325,7 +325,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) {
 	session.Session.InTransaction = true
 
 	// start a transaction against rss0
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -387,13 +387,13 @@ func TestTxConnCommitOrderFailure1(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_PRE)
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_POST)
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 
 	sbc0.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
 	err := sc.txConn.Commit(ctx, session)
@@ -420,13 +420,13 @@ func TestTxConnCommitOrderFailure2(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(context.Background(), rss1, queries, session, false)
+	sc.ExecuteMultiShard(context.Background(), rss1, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_PRE)
-	sc.ExecuteMultiShard(context.Background(), rss0, queries, session, false)
+	sc.ExecuteMultiShard(context.Background(), rss0, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_POST)
-	sc.ExecuteMultiShard(context.Background(), rss1, queries, session, false)
+	sc.ExecuteMultiShard(context.Background(), rss1, queries, session, false, false)
 
 	sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
 	err := sc.txConn.Commit(ctx, session)
@@ -452,13 +452,13 @@ func TestTxConnCommitOrderFailure3(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_PRE)
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_POST)
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 
 	sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
 	require.NoError(t,
@@ -487,7 +487,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession := vtgatepb.Session{
 		InTransaction: true,
 		ShardSessions: []*vtgatepb.Session_ShardSession{{
@@ -503,7 +503,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_PRE)
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction: true,
 		PreSessions: []*vtgatepb.Session_ShardSession{{
@@ -528,7 +528,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_POST)
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction: true,
 		PreSessions: []*vtgatepb.Session_ShardSession{{
@@ -562,7 +562,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	// Ensure nothing changes if we reuse a transaction.
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	require.NoError(t,
@@ -583,7 +583,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) {
 
 	// Sequence the executes to ensure commit order
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession := vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -601,7 +601,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_PRE)
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -629,7 +629,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	session.SetCommitOrder(vtgatepb.CommitOrder_POST)
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	wantSession = vtgatepb.Session{
 		InTransaction:  true,
 		InReservedConn: true,
@@ -667,7 +667,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) {
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	// Ensure nothing changes if we reuse a transaction.
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 	utils.MustMatch(t, &wantSession, session.Session, "Session")
 
 	require.NoError(t,
@@ -718,8 +718,8 @@ func TestTxConnCommit2PC(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PC")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
 	require.NoError(t,
 		sc.txConn.Commit(ctx, session))
@@ -733,7 +733,7 @@ func TestTxConnCommit2PC(t *testing.T) {
 func TestTxConnCommit2PCOneParticipant(t *testing.T) {
 	sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, "TestTxConnCommit2PCOneParticipant")
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
 	require.NoError(t,
 		sc.txConn.Commit(ctx, session))
@@ -744,8 +744,8 @@ func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) {
 	sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConnCommit2PCCreateTransactionFail")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss1, queries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss1, queries, session, false, false)
 
 	sbc0.MustFailCreateTransaction = 1
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
@@ -766,8 +766,8 @@ func TestTxConnCommit2PCPrepareFail(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCPrepareFail")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 
 	sbc1.MustFailPrepare = 1
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
@@ -786,8 +786,8 @@ func TestTxConnCommit2PCStartCommitFail(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCStartCommitFail")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 
 	sbc0.MustFailStartCommit = 1
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
@@ -806,8 +806,8 @@ func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCCommitPreparedFail")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 
 	sbc1.MustFailCommitPrepared = 1
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
@@ -826,8 +826,8 @@ func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCConcludeTransactionFail")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 
 	sbc0.MustFailConcludeTransaction = 1
 	session.TransactionMode = vtgatepb.TransactionMode_TWOPC
@@ -846,8 +846,8 @@ func TestTxConnRollback(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TxConnRollback")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 	require.NoError(t,
 		sc.txConn.Rollback(ctx, session))
 	wantSession := vtgatepb.Session{}
@@ -860,8 +860,8 @@ func TestTxConnReservedRollback(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TxConnReservedRollback")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 	require.NoError(t,
 		sc.txConn.Rollback(ctx, session))
 	wantSession := vtgatepb.Session{
@@ -895,8 +895,8 @@ func TestTxConnReservedRollbackFailure(t *testing.T) {
 	sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TxConnReservedRollback")
 
 	session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true})
-	sc.ExecuteMultiShard(ctx, rss0, queries, session, false)
-	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false)
+	sc.ExecuteMultiShard(ctx, rss0, queries, session, false, false)
+	sc.ExecuteMultiShard(ctx, rss01, twoQueries, session, false, false)
 
 	sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
 	require.Error(t,