chore: allow tx opts to be supplied in a callback (#12779)

The database/sql driver uses inline-begin transaction by default for all
transaction shapes. It also allows setting transaction options after the
transaction has been started, but not yet activated. For example the
following script is allowed:

```
BEGIN TRANSACTION;
SET ISOLATION_LEVEL='repeatable_read';
UPDATE my_table SET my_col=1 WHERE id=1;
COMMIT;
```

The actual start of the transaction is the UPDATE statement. The isolation
level is set after the BEGIN TRANSACTION statement, and should be applied
to the current transaction. In order to enable this, statement-based
transactions should allow a callback to be registered that will be called
when the transaction is actually started. This callback should return any
updated transaction options since the BEGIN TRANSACTION call.

Author: olavloite

spanner/client_test.go

@@ -6782,6 +6782,42 @@ func TestClient_BatchWriteExcludeTxnFromChangeStreams(t *testing.T) {
 	}
 }
 
+func TestClient_DelayedExcludeTxnFromChangeStreams(t *testing.T) {
+	server, client, teardown := setupMockedTestServer(t)
+	defer teardown()
+	ctx := context.Background()
+
+	ms := []*Mutation{
+		Insert("Accounts", []string{"AccountId", "Nickname", "Balance"}, []interface{}{int64(1), "Foo", int64(50)}),
+		Insert("Accounts", []string{"AccountId", "Nickname", "Balance"}, []interface{}{int64(2), "Bar", int64(1)}),
+	}
+
+	// Randomly exclude the transaction from change streams.
+	exclude := time.Now().UnixNano()%2 == 0
+	// Start a transaction without an ExcludeTxnFromChangeStreams option, but register a callback that allows us to set
+	// it right before the transaction is actually started.
+	tx, err := NewReadWriteStmtBasedTransactionWithCallbackForOptions(ctx, client, TransactionOptions{}, func() TransactionOptions {
+		return TransactionOptions{ExcludeTxnFromChangeStreams: exclude}
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = tx.BufferWrite(ms)
+	if _, err := tx.Commit(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	beginRequests := requestsOfType(requests, reflect.TypeOf(&sppb.BeginTransactionRequest{}))
+	if g, w := len(beginRequests), 1; g != w {
+		t.Fatalf("num BeginTransaction requests mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	beginRequest := beginRequests[0].(*sppb.BeginTransactionRequest)
+	if g, w := beginRequest.Options.ExcludeTxnFromChangeStreams, exclude; g != w {
+		t.Fatalf("exclude option mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func TestParseServerTimingHeader(t *testing.T) {
 	tests := []struct {
 		name     string

spanner/transaction.go

@@ -1327,6 +1327,16 @@ type ReadWriteTransaction struct {
 	wb []*Mutation
 	// isLongRunningTransaction indicates whether the transaction is long-running or not.
 	isLongRunningTransaction bool
+	// getTransactionOptionsCallback is a callback function that is called right before the
+	// transaction is actually started (either inlined or with an explicit BeginTransaction RPC).
+	// This callback can be used for transactions that do not yet know all the options at the
+	// moment that the start of the transaction is registered. This allows the following type of
+	// scripts to be executed by the database/sql driver:
+	// BEGIN TRANSACTION
+	// SET ISOLATION_LEVEL='repeatable_read' -- This sets an additional option after the transaction was started.
+	// UPDATE my_table SET my_col=1 WHERE id=1 -- This triggers the actual creation of the transaction.
+	// COMMIT
+	getTransactionOptionsCallback func() TransactionOptions
 }
 
 func (t *ReadWriteTransaction) isDefaultInlinedBegin() bool {
@@ -1563,11 +1573,16 @@ func (t *ReadWriteTransaction) acquire(ctx context.Context) (*sessionHandle, *sp
 			// is accepted.
 			t.state = txInit
 			sh := t.sh
+			if t.getTransactionOptionsCallback != nil {
+				t.txOpts = t.txOpts.merge(t.getTransactionOptionsCallback())
+			}
 			ts := &sppb.TransactionSelector{
 				Selector: &sppb.TransactionSelector_Begin{
 					Begin: &sppb.TransactionOptions{
 						Mode: &sppb.TransactionOptions_ReadWrite_{
-							ReadWrite: &sppb.TransactionOptions_ReadWrite{},
+							ReadWrite: &sppb.TransactionOptions_ReadWrite{
+								ReadLockMode: t.txOpts.ReadLockMode,
+							},
 						},
 						ExcludeTxnFromChangeStreams: t.txOpts.ExcludeTxnFromChangeStreams,
 						IsolationLevel:              t.txOpts.IsolationLevel,
@@ -1622,6 +1637,9 @@ func (t *ReadWriteTransaction) getTransactionSelector() *sppb.TransactionSelecto
 			},
 		}
 	}
+	if t.getTransactionOptionsCallback != nil {
+		t.txOpts = t.txOpts.merge(t.getTransactionOptionsCallback())
+	}
 	mode := &sppb.TransactionOptions_ReadWrite_{
 		ReadWrite: &sppb.TransactionOptions_ReadWrite{
 			ReadLockMode: t.txOpts.ReadLockMode,
@@ -1786,6 +1804,9 @@ func (t *ReadWriteTransaction) begin(ctx context.Context, mutation *sppb.Mutatio
 		if sh != nil {
 			sh.updateLastUseTime()
 		}
+		if t.getTransactionOptionsCallback != nil {
+			t.txOpts = t.txOpts.merge(t.getTransactionOptionsCallback())
+		}
 		tx, precommitToken, err = beginTransaction(contextWithOutgoingMetadata(ctx, sh.getMetadata(), t.disableRouteToLeader), transactionBeginOptions{
 			multiplexEnabled: t.sp.isMultiplexedSessionForRWEnabled(),
 			sessionID:        sh.getID(),
@@ -2021,8 +2042,7 @@ type ReadWriteStmtBasedTransaction struct {
 	// ReadWriteTransaction contains methods for performing transactional reads.
 	ReadWriteTransaction
 
-	client  *Client
-	options TransactionOptions
+	client *Client
 }
 
 func (t *ReadWriteStmtBasedTransaction) isDefaultInlinedBegin() bool {
@@ -2060,10 +2080,29 @@ func NewReadWriteStmtBasedTransaction(ctx context.Context, c *Client) (*ReadWrit
 // NewReadWriteStmtBasedTransactionWithOptions is a configurable version of
 // NewReadWriteStmtBasedTransaction.
 func NewReadWriteStmtBasedTransactionWithOptions(ctx context.Context, c *Client, options TransactionOptions) (*ReadWriteStmtBasedTransaction, error) {
-	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, c, options, nil, nil)
+	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, c, options, nil, nil, nil)
+}
+
+// NewReadWriteStmtBasedTransactionWithCallbackForOptions starts a read-write
+// transaction with a callback that gives the actual transaction options.
+// Commit() or Rollback() must be called to end a transaction. If Commit() or
+// Rollback() is not called, the session that is used by the transaction will
+// not be returned to the pool and cause a session leak.
+//
+// ResetForRetry resets the transaction before a retry attempt. This function
+// returns a new transaction that should be used for the retry attempt. The
+// transaction that is returned by this function is assigned a higher priority
+// than the previous transaction, making it less probable to be aborted by
+// Spanner again during the retry.
+//
+// NewReadWriteStmtBasedTransactionWithCallbackForOptions is the same as
+// NewReadWriteStmtBasedTransactionWithOptions, but allows the caller to wait
+// with setting the actual transaction options until a later moment.
+func NewReadWriteStmtBasedTransactionWithCallbackForOptions(ctx context.Context, c *Client, opts TransactionOptions, callback func() TransactionOptions) (*ReadWriteStmtBasedTransaction, error) {
+	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, c, opts, nil, nil, callback)
 }
 
-func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *Client, options TransactionOptions, sh *sessionHandle, previousTransactionID transactionID) (*ReadWriteStmtBasedTransaction, error) {
+func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *Client, options TransactionOptions, sh *sessionHandle, previousTransactionID transactionID, callback func() TransactionOptions) (*ReadWriteStmtBasedTransaction, error) {
 	var (
 		err error
 		t   *ReadWriteStmtBasedTransaction
@@ -2081,7 +2120,8 @@ func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *C
 	}
 	t = &ReadWriteStmtBasedTransaction{
 		ReadWriteTransaction: ReadWriteTransaction{
-			txReadyOrClosed: make(chan struct{}),
+			txReadyOrClosed:               make(chan struct{}),
+			getTransactionOptionsCallback: callback,
 		},
 		client: c,
 	}
@@ -2120,7 +2160,6 @@ func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *C
 		return err
 	}
 
-	t.options = options
 	t.txOpts = c.txo.merge(options)
 	t.ct = c.ct
 	t.otConfig = c.otConfig
@@ -2202,8 +2241,8 @@ func (t *ReadWriteStmtBasedTransaction) ResetForRetry(ctx context.Context) (*Rea
 	// Create a new transaction that re-uses the current session if it is available.
 	// It should always use an explicit BeginTransaction RPC to ensure that the first
 	// statement is included in the transaction.
-	t.options.BeginTransactionOption = ExplicitBeginTransaction
-	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, t.client, t.options, t.sh, previousTransactionID)
+	t.txOpts.BeginTransactionOption = ExplicitBeginTransaction
+	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, t.client, t.txOpts, t.sh, previousTransactionID, t.getTransactionOptionsCallback)
 }
 
 // writeOnlyTransaction provides the most efficient way of doing write-only

spanner/transaction_test.go

@@ -1588,7 +1588,7 @@ func testReadWriteStmtBasedTransaction(t *testing.T, beginTransactionOption Begi
 		if err != nil {
 			return 0, attempts, fmt.Errorf("failed to begin a transaction: %v", err)
 		}
-		if g, w := tx.options.TransactionTag, "test"; g != w {
+		if g, w := tx.txOpts.TransactionTag, "test"; g != w {
 			t.Errorf("transaction tag mismatch\n Got: %v\nWant: %v", g, w)
 		}
 		rowCount, err = f(tx)