feat: add isolation level option for disabling internal retries

driver.go

@@ -1216,6 +1216,20 @@ func (c *conn) resetTransactionForRetry(ctx context.Context, errDuringCommit boo
 	return c.tx.resetForRetry(ctx)
 }
 
+type spannerIsolationLevel sql.IsolationLevel
+
+const (
+	levelNone spannerIsolationLevel = iota
+	levelDisableRetryAborts
+)
+
+// WithDisableRetryAborts returns a specific Spanner isolation level that contains
+// both the given standard isolation level and a custom Spanner isolation level that
+// disables internal retries for aborted transactions for a single transaction.
+func WithDisableRetryAborts(level sql.IsolationLevel) sql.IsolationLevel {
+	return sql.IsolationLevel(levelDisableRetryAborts)<<8 + level
+}
+
 func (c *conn) Begin() (driver.Tx, error) {
 	return c.BeginTx(context.Background(), driver.TxOptions{})
 }
@@ -1231,6 +1245,15 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 	if c.inBatch() {
 		return nil, status.Error(codes.FailedPrecondition, "This connection has an active batch. Run or abort the batch before starting a new transaction.")
 	}
+	disableRetryAborts := false
+	sil := opts.Isolation >> 8
+	opts.Isolation = opts.Isolation - sil
+	if sil > 0 {
+		switch spannerIsolationLevel(sil) {
+		case levelDisableRetryAborts:
+			disableRetryAborts = true
+		}
+	}
 
 	if opts.ReadOnly {
 		ro := c.client.ReadOnlyTransaction().WithTimestampBound(c.readOnlyStaleness)
@@ -1259,7 +1282,7 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 				c.commitTs = commitTs
 			}
 		},
-		retryAborts: c.retryAborts,
+		retryAborts: c.retryAborts && !disableRetryAborts,
 	}
 	c.commitTs = nil
 	return c.tx, nil

driver_with_mockserver_test.go

@@ -2922,6 +2922,78 @@ func TestRunTransactionCommitError(t *testing.T) {
 	}
 }
 
+func TestTransactionWithLevelDisableRetryAborts(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	tx, err := db.BeginTx(ctx, &sql.TxOptions{Isolation: WithDisableRetryAborts(sql.LevelSerializable)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	rows, err := tx.Query(testutil.SelectFooFromBar)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rows.Close()
+	for want := int64(1); rows.Next(); want++ {
+		cols, err := rows.Columns()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !cmp.Equal(cols, []string{"FOO"}) {
+			t.Fatalf("cols mismatch\nGot: %v\nWant: %v", cols, []string{"FOO"})
+		}
+		var got int64
+		err = rows.Scan(&got)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got != want {
+			t.Fatalf("value mismatch\nGot: %v\nWant: %v", got, want)
+		}
+	}
+	if err := rows.Err(); err != nil {
+		t.Fatal(err)
+	}
+	// Simulate that the transaction was aborted.
+	server.TestSpanner.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
+		Errors: []error{gstatus.Error(codes.Aborted, "Aborted")},
+	})
+	// Committing the transaction should fail, as we have disabled internal retries.
+	err = tx.Commit()
+	if err == nil {
+		t.Fatal("missing aborted error after commit")
+	}
+	code := spanner.ErrCode(err)
+	if w, g := code, codes.Aborted; w != g {
+		t.Fatalf("error code mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 1; g != w {
+		t.Fatalf("ExecuteSqlRequests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	req := sqlRequests[0].(*sppb.ExecuteSqlRequest)
+	if req.Transaction == nil {
+		t.Fatalf("missing transaction for ExecuteSqlRequest")
+	}
+	if req.Transaction.GetId() == nil {
+		t.Fatalf("missing id selector for ExecuteSqlRequest")
+	}
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 1; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	commitReq := commitRequests[0].(*sppb.CommitRequest)
+	if c, e := commitReq.GetTransactionId(), req.Transaction.GetId(); !cmp.Equal(c, e) {
+		t.Fatalf("transaction id mismatch\nCommit: %c\nExecute: %v", c, e)
+	}
+}
+
 func numeric(v string) big.Rat {
 	res, _ := big.NewRat(1, 1).SetString(v)
 	return *res