sql: deal with retriable errors when using a new txn

In cockroachdb#46588 a bug was introduced when a retriable error was encountered while
using a new transaction for preparing. Prior to that commit, all error were
treated as not retriable. This was sort of a bummer. Retriable errors can
occur due to read within uncertainty. Before this PR, those retriable errors
would make their way to the client. Now we'll handle those retry errors
internally underneath `connExecutor.prepare`

Fixes cockroachdb#43251

Release note: None

pkg/sql/conn_executor_prepare.go

@@ -15,7 +15,6 @@ import (
 	"fmt"
 
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
-	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase"
@@ -150,33 +149,33 @@ func (ex *connExecutor) prepare(
 	// Preparing needs a transaction because it needs to retrieve db/table
 	// descriptors for type checking. If we already have an open transaction for
 	// this planner, use it. Using the user's transaction here is critical for
-	// proper deadlock detection. At the time of writing it is the case that any
+	// proper deadlock detection. At the time of writing, it is the case that any
 	// data read on behalf of this transaction is not cached for use in other
 	// transactions. It's critical that this fact remain true but nothing really
 	// enforces it. If we create a new transaction (newTxn is true), we'll need to
 	// finish it before we return.
-	newTxn, txn := false, ex.state.mu.txn
-	if txn == nil || txn.Serialize().Status != roachpb.PENDING {
-		newTxn, txn = true, client.NewTxn(ctx, ex.server.cfg.DB, ex.server.cfg.NodeID.Get(), client.RootTxn)
+
+	var flags planFlags
+	prepare := func(ctx context.Context, txn *client.Txn) (err error) {
+		ex.statsCollector.reset(&ex.server.sqlStats, ex.appStats, &ex.phaseTimes)
+		p := &ex.planner
+		ex.resetPlanner(ctx, p, txn,
+			ex.server.cfg.Clock.PhysicalTime() /* stmtTS */, stmt.NumAnnotations)
+		p.stmt = &stmt
+		p.semaCtx.Annotations = tree.MakeAnnotations(stmt.NumAnnotations)
+		flags, err = ex.populatePrepared(ctx, txn, placeholderHints, p)
+		return err
 	}
 
-	ex.statsCollector.reset(&ex.server.sqlStats, ex.appStats, &ex.phaseTimes)
-	p := &ex.planner
-	ex.resetPlanner(ctx, p, txn, ex.server.cfg.Clock.PhysicalTime() /* stmtTS */, stmt.NumAnnotations)
-	p.stmt = &stmt
-	flags, err := ex.populatePrepared(ctx, txn, placeholderHints, p)
-	if err != nil {
-		// NB: if this is not a new transaction then let the connExecutor state
-		// machine decide whether we should clean up intents; we may be restarting
-		// and want to leave them in place.
-		if newTxn {
-			txn.CleanupOnError(ctx, err)
+	if txn := ex.state.mu.txn; txn != nil && !txn.Serialize().Status.IsFinalized() {
+		// Use the existing transaction.
+		if err := prepare(ctx, txn); err != nil {
+			return nil, err
 		}
-		return nil, err
-	}
-	if newTxn {
-		// Clean up the newly created transaction if we made one.
-		if err := txn.CommitOrCleanup(ctx); err != nil {
+	} else {
+		// Use a new transaction. This will handle retriable errors here rather
+		// than bubbling them up to the connExecutor state machine.
+		if err := ex.server.cfg.DB.Txn(ctx, prepare); err != nil {
 			return nil, err
 		}
 	}
@@ -194,6 +193,11 @@ func (ex *connExecutor) prepare(
 func (ex *connExecutor) populatePrepared(
 	ctx context.Context, txn *client.Txn, placeholderHints tree.PlaceholderTypes, p *planner,
 ) (planFlags, error) {
+	if before := ex.server.cfg.TestingKnobs.BeforePrepare; before != nil {
+		if err := before(ctx, ex.planner.stmt.String(), txn); err != nil {
+			return 0, err
+		}
+	}
 	stmt := p.stmt
 	if err := p.semaCtx.Placeholders.Init(stmt.NumPlaceholders, placeholderHints); err != nil {
 		return 0, err

pkg/sql/conn_executor_test.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
@@ -465,6 +466,37 @@ func TestPrepareInExplicitTransactionDoesNotDeadlock(t *testing.T) {
 	}
 }
 
+// TestRetriableErrorDuringPrepare ensures that when preparing and using a new
+// transaction, retriable errors are handled properly and do not propagate to
+// the user's transaction.
+func TestRetriableErrorDuringPrepare(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	const uniqueString = "'a very unique string'"
+	var failed int64
+	const numToFail = 2 // only fail on the first two attempts
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			SQLExecutor: &sql.ExecutorTestingKnobs{
+				BeforePrepare: func(ctx context.Context, stmt string, txn *client.Txn) error {
+					if strings.Contains(stmt, uniqueString) && atomic.AddInt64(&failed, 1) <= numToFail {
+						return roachpb.NewTransactionRetryWithProtoRefreshError("boom",
+							txn.ID(), *txn.Serialize())
+					}
+					return nil
+				},
+			},
+		},
+	})
+	defer s.Stopper().Stop(context.Background())
+
+	testDB := sqlutils.MakeSQLRunner(sqlDB)
+	testDB.Exec(t, "CREATE TABLE foo (i INT PRIMARY KEY)")
+
+	stmt, err := sqlDB.Prepare("SELECT " + uniqueString)
+	require.NoError(t, err)
+	defer func() { _ = stmt.Close() }()
+}
+
 // This test ensures that when in an explicit transaction and statement
 // preparation uses the user's transaction, errors during those planning queries
 // are handled correctly.

pkg/sql/exec_util.go

@@ -547,6 +547,10 @@ type ExecutorTestingKnobs struct {
 	// statement has been executed.
 	StatementFilter StatementFilter
 
+	// BeforePrepare can be used to trap execution of SQL statement preparation.
+	// If a nil error is returned, planning continues as usual.
+	BeforePrepare func(ctx context.Context, stmt string, txn *client.Txn) error
+
 	// BeforeExecute is called by the Executor before plan execution. It is useful
 	// for synchronizing statement execution.
 	BeforeExecute func(ctx context.Context, stmt string)

pkg/sql/logictest/testdata/logic_test/schema_change_in_txn

@@ -1771,4 +1771,3 @@ ALTER TABLE t42508 ADD COLUMN y INT DEFAULT nextval('s42508')
 
 statement error pgcode XXA00 unimplemented: cannot backfill such sequence operation.*\nHINT.*\n.*42508
 COMMIT
-