Merge #153971

153971: sql: deflake TestReadCommittedImplicitPartitionUpsert r=fqazi a=fqazi

The test TestReadCommittedImplicitPartitionUpsert started flaking after we switched the implementation of TRUNCATE to the declarative schema changer. This happened because the testing knobs were incorrectly picking up internal queries. To address this, this patch modifies the testing knobs to include the SQL statement, so they can skip unrelated queries.

Fixes: #153644

Release note: None

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/sql/exec_util.go

@@ -2017,7 +2017,7 @@ type ExecutorTestingKnobs struct {
 
 	// AfterArbiterRead, if set, will be called after each row read from an arbiter index
 	// for an UPSERT or INSERT.
-	AfterArbiterRead func()
+	AfterArbiterRead func(query string)
 
 	// BeforeIndexSplitAndScatter is invoked with the split and scatter of an index
 	// occurs.

pkg/sql/insert.go

@@ -326,7 +326,7 @@ func (n *insertNode) BatchedNext(params runParams) (bool, error) {
 		if buildutil.CrdbTestBuild {
 			// This testing knob allows us to suspend execution to force a race condition.
 			if fn := params.ExecCfg().TestingKnobs.AfterArbiterRead; fn != nil {
-				fn()
+				fn(params.p.stmt.SQL)
 			}
 		}
 

pkg/sql/insert_fast_path.go

@@ -476,7 +476,7 @@ func (n *insertFastPathNode) BatchedNext(params runParams) (bool, error) {
 		if buildutil.CrdbTestBuild {
 			// This testing knob allows us to suspend execution to force a race condition.
 			if fn := params.ExecCfg().TestingKnobs.AfterArbiterRead; fn != nil {
-				fn()
+				fn(params.p.stmt.SQL)
 			}
 		}
 

pkg/sql/mutation_test.go

@@ -10,6 +10,7 @@ import (
 	gosql "database/sql"
 	"fmt"
 	"reflect"
+	"strings"
 	"sync"
 	"testing"
 
@@ -133,13 +134,20 @@ func TestReadCommittedImplicitPartitionUpsert(t *testing.T) {
 		mu.l.Unlock()
 	}
 
+	peekState := func() State {
+		mu.l.Lock()
+		defer mu.l.Unlock()
+		return mu.state
+	}
+
 	ctx := context.Background()
 	params, _ := createTestServerParamsAllowTenants()
 	// If test is in Ready state, transition to ReadDone and wait for conflict.
 	params.Knobs = base.TestingKnobs{
 		SQLExecutor: &sql.ExecutorTestingKnobs{
-			AfterArbiterRead: func() {
-				if mu.state != Ready {
+			AfterArbiterRead: func(query string) {
+				// Only wait for arbiter operations on the upsert table.
+				if peekState() != Ready || !strings.Contains(query, "d.upsert") {
 					return
 				}
 				setState(ReadDone)

pkg/sql/upsert.go

@@ -207,7 +207,7 @@ func (r *upsertRun) processSourceRow(params runParams, rowVals tree.Datums) erro
 	if buildutil.CrdbTestBuild {
 		// This testing knob allows us to suspend execution to force a race condition.
 		if fn := params.ExecCfg().TestingKnobs.AfterArbiterRead; fn != nil {
-			fn()
+			fn(params.p.stmt.SQL)
 		}
 	}