chore: pull retry_aborts_internally when needed (#518)

Pull the value of retry_aborts_internally once from the connecion state
when needed, instead of trying to set it (and keep it in sync) on the
current transaction of a connection. This also enables the use of
`SET LOCAL retry_aborts_internally = true|false`

Author: olavloite

conn.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"log/slog"
 	"slices"
+	"sync"
 	"time"
 
 	"cloud.google.com/go/spanner"
@@ -359,11 +360,6 @@ func (c *conn) SetRetryAbortsInternally(retry bool) error {
 }
 
 func (c *conn) setRetryAbortsInternally(retry bool) (driver.Result, error) {
-	if c.inTransaction() {
-		if _, err := c.tx.setRetryAbortsInternally(retry); err != nil {
-			return nil, err
-		}
-	}
 	if err := propertyRetryAbortsInternally.SetValue(c.state, retry, connectionstate.ContextUser); err != nil {
 		return nil, err
 	}
@@ -1175,9 +1171,8 @@ func (c *conn) BeginTx(ctx context.Context, driverOpts driver.TxOptions) (driver
 	if c.tempTransactionOptions != nil && c.tempTransactionOptions.close != nil {
 		tempCloseFunc = c.tempTransactionOptions.close
 	}
-	disableInternalRetries := !c.RetryAbortsInternally()
-	if c.tempTransactionOptions != nil {
-		disableInternalRetries = c.tempTransactionOptions.DisableInternalRetries
+	if !disableRetryAborts && c.tempTransactionOptions != nil {
+		disableRetryAborts = c.tempTransactionOptions.DisableInternalRetries
 	}
 
 	tx, err := spanner.NewReadWriteStmtBasedTransactionWithCallbackForOptions(ctx, c.client, opts, func() spanner.TransactionOptions {
@@ -1212,7 +1207,9 @@ func (c *conn) BeginTx(ctx context.Context, driverOpts driver.TxOptions) (driver
 			}
 		},
 		// Disable internal retries if any of these options have been set.
-		retryAborts: !disableInternalRetries && !disableRetryAborts,
+		retryAborts: sync.OnceValue(func() bool {
+			return c.RetryAbortsInternally() && !disableRetryAborts
+		}),
 	}
 	c.clearCommitResponse()
 	return c.tx, nil

conn_with_mockserver_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/googleapis/go-sql-spanner/connectionstate"
 	"github.com/googleapis/go-sql-spanner/testutil"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/emptypb"
@@ -809,6 +810,35 @@ func TestSetRetryAbortsInternallyInActiveTransaction(t *testing.T) {
 	_ = tx.Rollback()
 }
 
+func TestSetLocalRetryAbortsInternally(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+	ctx := context.Background()
+
+	tx, err := db.BeginTx(ctx, &sql.TxOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := tx.ExecContext(ctx, "set local retry_aborts_internally = false"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := tx.ExecContext(ctx, testutil.UpdateBarSetFoo); err != nil {
+		t.Fatal(err)
+	}
+	// Instruct the mock server to return Aborted when Commit is called.
+	// This should cause the transaction to fail, and as internal retries have been
+	// disabled, it should not be retried.
+	server.TestSpanner.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
+		Errors: []error{status.Error(codes.Aborted, "Aborted")},
+	})
+	err = tx.Commit()
+	if g, w := spanner.ErrCode(err), codes.Aborted; g != w {
+		t.Fatalf("error mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func TestSetAutocommitDMLMode(t *testing.T) {
 	t.Parallel()
 

driver_with_mockserver_test.go

@@ -2118,8 +2118,8 @@ func TestQueryWithDuplicateNamedParameterStartingWithUnderscore(t *testing.T) {
 		t.Fatal(err)
 	}
 	// Verify that 'bar' is used for both instances of the parameter @__name.
-	requests := drainRequestsFromServer(server.TestSpanner)
-	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	sqlRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
 	if len(sqlRequests) != 1 {
 		t.Fatalf("sql requests count mismatch\nGot: %v\nWant: %v", len(sqlRequests), 1)
 	}
@@ -2178,8 +2178,8 @@ func TestQueryWithReusedNamedParameterStartingWithUnderscore(t *testing.T) {
 		t.Fatal(err)
 	}
 	// Verify that 'foo' is used for both instances of the parameter @__name.
-	requests := drainRequestsFromServer(server.TestSpanner)
-	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	sqlRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
 	if len(sqlRequests) != 1 {
 		t.Fatalf("sql requests count mismatch\nGot: %v\nWant: %v", len(sqlRequests), 1)
 	}
@@ -4567,7 +4567,7 @@ func TestRunTransaction(t *testing.T) {
 		if g, w := rwTx.ctx, ctx; g != w {
 			return fmt.Errorf("getting the transaction through reflection failed")
 		}
-		if rwTx.retryAborts {
+		if rwTx.retryAborts() {
 			return fmt.Errorf("internal retries should be disabled during RunTransaction")
 		}
 
@@ -5028,7 +5028,7 @@ func TestBeginReadWriteTransaction(t *testing.T) {
 		if g, w := rwTx.ctx, ctx; g != w {
 			t.Fatal("getting the transaction through reflection failed")
 		}
-		if rwTx.retryAborts {
+		if rwTx.retryAborts() {
 			t.Fatal("internal retries should be disabled during this transaction")
 		}
 

transaction.go

@@ -54,8 +54,6 @@ type contextTransaction interface {
 	AbortBatch() (driver.Result, error)
 
 	BufferWrite(ms []*spanner.Mutation) error
-
-	setRetryAbortsInternally(retry bool) (driver.Result, error)
 }
 
 type rowIterator interface {
@@ -205,11 +203,6 @@ func (tx *readOnlyTransaction) BufferWrite([]*spanner.Mutation) error {
 	return spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "read-only transactions cannot write"))
 }
 
-func (tx *readOnlyTransaction) setRetryAbortsInternally(_ bool) (driver.Result, error) {
-	// no-op, ignore
-	return driver.ResultNoRows, nil
-}
-
 // ErrAbortedDueToConcurrentModification is returned by a read/write transaction
 // that was aborted by Cloud Spanner, and where the internal retry attempt
 // failed because it detected that the results during the retry were different
@@ -244,7 +237,7 @@ type readWriteTransaction struct {
 	close func(result txResult, commitResponse *spanner.CommitResponse, commitErr error)
 	// retryAborts indicates whether this transaction will automatically retry
 	// the transaction if it is aborted by Spanner. The default is true.
-	retryAborts bool
+	retryAborts func() bool
 
 	// statements contains the list of statements that has been executed on this
 	// transaction so far. These statements will be replayed on a new read write
@@ -415,7 +408,7 @@ func (tx *readWriteTransaction) Commit() (err error) {
 	}
 	var commitResponse spanner.CommitResponse
 	if tx.rwTx != nil {
-		if !tx.retryAborts {
+		if !tx.retryAborts() {
 			ts, err := tx.rwTx.CommitWithReturnResp(tx.ctx)
 			tx.close(txResultCommit, &ts, err)
 			return err
@@ -471,7 +464,7 @@ func (tx *readWriteTransaction) Query(ctx context.Context, stmt spanner.Statemen
 	}
 	// If internal retries have been disabled, we don't need to keep track of a
 	// running checksum for all results that we have seen.
-	if !tx.retryAborts {
+	if !tx.retryAborts() {
 		return &readOnlyRowIterator{tx.rwTx.QueryWithOptions(ctx, stmt, execOptions.QueryOptions)}, nil
 	}
 
@@ -509,7 +502,7 @@ func (tx *readWriteTransaction) ExecContext(ctx context.Context, stmt spanner.St
 		return &result{rowsAffected: updateCount}, nil
 	}
 
-	if !tx.retryAborts {
+	if !tx.retryAborts() {
 		return execTransactionalDML(ctx, tx.rwTx, stmt, statementInfo, options)
 	}
 
@@ -608,7 +601,7 @@ func (tx *readWriteTransaction) runDmlBatch(ctx context.Context) (*result, error
 	options := tx.batch.options
 	tx.batch = nil
 
-	if !tx.retryAborts {
+	if !tx.retryAborts() {
 		affected, err := tx.rwTx.BatchUpdateWithOptions(ctx, statements, options.QueryOptions)
 		return &result{rowsAffected: sum(affected)}, err
 	}
@@ -650,11 +643,3 @@ func errorsEqualForRetry(err1, err2 error) bool {
 	}
 	return false
 }
-
-func (tx *readWriteTransaction) setRetryAbortsInternally(retry bool) (driver.Result, error) {
-	if tx.active {
-		return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "cannot change retry mode while a transaction is active"))
-	}
-	tx.retryAborts = retry
-	return driver.ResultNoRows, nil
-}