internal/database: Transact supports any isolation level

- Add an arg to Transact for the isolation level
- Remove TransactSerializable

This makes it possible to use other levels, and makes it easier to see
which level is being used for each transaction.

Change-Id: Iba5e2920b4139e5e2f0f8c6b331a658d7c84f60f
Reviewed-on: https://team-review.git.corp.google.com/c/golang/discovery/+/758942
CI-Result: Cloud Build <devtools-proctor-result-processor@system.gserviceaccount.com>
Reviewed-by: Julie Qiu <julieqiu@google.com>

2 files changed, 25 insertions, 24 deletions

diff --git a/internal/database/database.go b/internal/database/database.go
index 3e009e1c..2bd6b8cc 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -116,16 +116,26 @@ func (db *DB) RunQuery(ctx context.Context, query string, f func(*sql.Rows) erro
 	return rows.Err()
 }
 
-// Transact executes the given function in the context of a SQL transaction,
-// rolling back the transaction if the function panics or returns an error. It
-// uses the default transaction options.
+// Transact executes the given function in the context of a SQL transaction at
+// the given isolation level, rolling back the transaction if the function
+// panics or returns an error.
 //
 // The given function is called with a DB that is associated with a transaction.
 // The DB should be used only inside the function; if it is used to access the
 // database after the function returns, the calls will return errors.
-func (db *DB) Transact(ctx context.Context, txFunc func(*DB) error) (err error) {
-	defer derrors.Wrap(&err, "Transact")
-	return db.transact(ctx, nil, txFunc)
+//
+// If the isolation level requires it, Transact will retry the transaction upon
+// serialization failure, so txFunc may be called more than once.
+
+func (db *DB) Transact(ctx context.Context, iso sql.IsolationLevel, txFunc func(*DB) error) (err error) {
+	defer derrors.Wrap(&err, "Transact(%s)", iso)
+	// For the levels which require retry, see
+	// https://www.postgresql.org/docs/11/transaction-iso.html.
+	opts := &sql.TxOptions{Isolation: iso}
+	if iso == sql.LevelRepeatableRead || iso == sql.LevelSerializable {
+		return db.transactWithRetry(ctx, opts, txFunc)
+	}
+	return db.transact(ctx, opts, txFunc)
 }
 
 // serializationFailureCode is the Postgres error code returned when a serializable
@@ -133,23 +143,13 @@ func (db *DB) Transact(ctx context.Context, txFunc func(*DB) error) (err error)
 // See https://www.postgresql.org/docs/current/errcodes-appendix.html.
 const serializationFailureCode = "40001"
 
-// TransactSerializable executes the given function in the context of a SQL transaction,
-// rolling back the transaction if the function panics or returns an error. It uses
-// the given context and the Serializable transaction isolation level.
-//
-// The given function is called with a DB that is associated with a transaction.
-// The DB should be used only inside the function; if it is used to access the
-// database after the function returns, the calls will return errors.
-//
-// TransactSerializable will retry the transaction upon serialization failure, so txFunc
-// may be called more than once.
-func (db *DB) TransactSerializable(ctx context.Context, txFunc func(*DB) error) (err error) {
-	defer derrors.Wrap(&err, "TransactSerializable")
+func (db *DB) transactWithRetry(ctx context.Context, opts *sql.TxOptions, txFunc func(*DB) error) (err error) {
+	defer derrors.Wrap(&err, "transactWithRetry(%v)", opts)
 	// Retry on serialization failure, up to some max.
 	// See https://www.postgresql.org/docs/11/transaction-iso.html.
 	const maxRetries = 30
 	for i := 0; i <= maxRetries; i++ {
-		err = db.transact(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable}, txFunc)
+		err = db.transact(ctx, opts, txFunc)
 		var perr *pq.Error
 		if errors.As(err, &perr) && perr.Code == serializationFailureCode {
 			db.mu.Lock()
diff --git a/internal/database/database_test.go b/internal/database/database_test.go
index 7d95bd83..935b6707 100644
--- a/internal/database/database_test.go
+++ b/internal/database/database_test.go
@@ -193,7 +193,7 @@ func TestLargeBulkInsert(t *testing.T) {
 	for i := 0; i < size; i++ {
 		vals[i] = i + 1
 	}
-	if err := testDB.Transact(ctx, func(db *DB) error {
+	if err := testDB.Transact(ctx, sql.LevelDefault, func(db *DB) error {
 		return db.BulkInsert(ctx, "test_large_bulk", []string{"i"}, vals, "")
 	}); err != nil {
 		t.Fatal(err)
@@ -220,7 +220,7 @@ func TestLargeBulkInsert(t *testing.T) {
 func TestDBAfterTransactFails(t *testing.T) {
 	ctx := context.Background()
 	var tx *DB
-	err := testDB.Transact(ctx, func(d *DB) error {
+	err := testDB.Transact(ctx, sql.LevelDefault, func(d *DB) error {
 		tx = d
 		return nil
 	})
@@ -269,7 +269,7 @@ func TestBulkUpdate(t *testing.T) {
 	for i := 0; i < 50; i++ {
 		values = append(values, i, i)
 	}
-	err := testDB.Transact(ctx, func(tx *DB) error {
+	err := testDB.Transact(ctx, sql.LevelDefault, func(tx *DB) error {
 		return tx.BulkInsert(ctx, "bulk_update", cols, values, "")
 	})
 	if err != nil {
@@ -283,7 +283,7 @@ func TestBulkUpdate(t *testing.T) {
 		updateVals[1] = append(updateVals[1], -i)
 	}
 
-	err = testDB.Transact(ctx, func(tx *DB) error {
+	err = testDB.Transact(ctx, sql.LevelDefault, func(tx *DB) error {
 		return tx.BulkUpdate(ctx, "bulk_update", cols, []string{"INT", "INT"}, updateVals)
 	})
 	if err != nil {
@@ -349,7 +349,8 @@ func TestTransactSerializable(t *testing.T) {
 	for i := 0; i < numTransactions; i++ {
 		i := i
 		go func() {
-			errc <- testDB.TransactSerializable(ctx, func(tx *DB) error { return insertSum(tx, 1+i%2) })
+			errc <- testDB.Transact(ctx, sql.LevelSerializable,
+				func(tx *DB) error { return insertSum(tx, 1+i%2) })
 		}()
 	}
 	// None of the transactions should fail.