Add Conn.TxFromCurrentTransaction and TxOptions.BeginSQL

tx.go

@@ -3,6 +3,7 @@ package pgx
 import (
 	"context"
 	"errors"
+	"fmt"
 	"strconv"
 	"strings"
 
@@ -53,6 +54,14 @@ type TxOptions struct {
 
 var emptyTxOptions TxOptions
 
+// BeginSQL returns the SQL statement that [Conn.BeginTx] would send for
+// txOptions. It is exported so that callers that start a transaction themselves,
+// e.g. by queueing it as the first query of a [Batch] to save a round trip, can
+// use the same statement pgx would.
+func (txOptions TxOptions) BeginSQL() string {
+	return txOptions.beginSQL()
+}
+
 func (txOptions TxOptions) beginSQL() string {
 	if txOptions == emptyTxOptions {
 		return "begin"
@@ -115,6 +124,31 @@ func (c *Conn) BeginTx(ctx context.Context, txOptions TxOptions) (Tx, error) {
 	}, nil
 }
 
+// TxFromCurrentTransaction returns a [Tx] for the transaction c is already in,
+// without sending a begin query. It is for callers that begin the transaction
+// themselves in a way that avoids a dedicated round trip, such as queueing
+// [TxOptions.BeginSQL] as the first query of a [Batch] or of a pgconn pipeline.
+//
+// Only [TxOptions.CommitQuery] is read from txOptions. The isolation level,
+// access mode and deferrable mode must already have been established by the
+// begin query the caller sent.
+//
+// It returns an error if c is not currently in a transaction, which includes the
+// case where the caller's begin query has been sent but its results have not been
+// read yet. The returned Tx is otherwise identical to one from [Conn.BeginTx].
+func (c *Conn) TxFromCurrentTransaction(txOptions TxOptions) (Tx, error) {
+	switch status := c.PgConn().TxStatus(); status {
+	case 'T', 'E':
+	default:
+		return nil, fmt.Errorf("connection is not in a transaction (transaction status %q)", status)
+	}
+
+	return &dbTx{
+		conn:        c,
+		commitQuery: txOptions.CommitQuery,
+	}, nil
+}
+
 func isConnectionFatal(pgErr *pgconn.PgError) bool {
 	severity := pgErr.SeverityUnlocalized
 	if severity == "" {

tx_test.go

@@ -717,3 +717,169 @@ func TestBeginTxFatalErrorKillsConn(t *testing.T) {
 
 	require.True(t, conn.IsClosed())
 }
+
+func TestTxOptionsBeginSQL(t *testing.T) {
+	t.Parallel()
+
+	for _, tt := range []struct {
+		name      string
+		txOptions pgx.TxOptions
+		expected  string
+	}{
+		{
+			name:      "empty",
+			txOptions: pgx.TxOptions{},
+			expected:  "begin",
+		},
+		{
+			name:      "iso level",
+			txOptions: pgx.TxOptions{IsoLevel: pgx.Serializable},
+			expected:  "begin isolation level serializable",
+		},
+		{
+			name:      "access mode",
+			txOptions: pgx.TxOptions{AccessMode: pgx.ReadOnly},
+			expected:  "begin read only",
+		},
+		{
+			name: "all modes",
+			txOptions: pgx.TxOptions{
+				IsoLevel:       pgx.Serializable,
+				AccessMode:     pgx.ReadWrite,
+				DeferrableMode: pgx.NotDeferrable,
+			},
+			expected: "begin isolation level serializable read write not deferrable",
+		},
+		{
+			name:      "begin query overrides modes",
+			txOptions: pgx.TxOptions{IsoLevel: pgx.Serializable, BeginQuery: "begin priority high"},
+			expected:  "begin priority high",
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			require.Equal(t, tt.expected, tt.txOptions.BeginSQL())
+		})
+	}
+}
+
+func TestTxFromCurrentTransaction(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	_, err := conn.Exec(ctx, "create temporary table foo(id integer primary key)")
+	require.NoError(t, err)
+
+	txOptions := pgx.TxOptions{IsoLevel: pgx.RepeatableRead}
+
+	batch := &pgx.Batch{}
+	batch.Queue(txOptions.BeginSQL())
+	batch.Queue("insert into foo(id) values (1)")
+	require.NoError(t, conn.SendBatch(ctx, batch).Close())
+
+	tx, err := conn.TxFromCurrentTransaction(txOptions)
+	require.NoError(t, err)
+
+	var isoLevel string
+	err = tx.QueryRow(ctx, "select current_setting('transaction_isolation')").Scan(&isoLevel)
+	require.NoError(t, err)
+	require.Equal(t, string(pgx.RepeatableRead), isoLevel)
+
+	_, err = tx.Exec(ctx, "insert into foo(id) values (2)")
+	require.NoError(t, err)
+
+	require.NoError(t, tx.Commit(ctx))
+
+	var n int64
+	err = conn.QueryRow(ctx, "select count(*) from foo").Scan(&n)
+	require.NoError(t, err)
+	require.EqualValues(t, 2, n)
+}
+
+func TestTxFromCurrentTransactionRollback(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	_, err := conn.Exec(ctx, "create temporary table foo(id integer primary key)")
+	require.NoError(t, err)
+
+	batch := &pgx.Batch{}
+	batch.Queue(pgx.TxOptions{}.BeginSQL())
+	batch.Queue("insert into foo(id) values (1)")
+	require.NoError(t, conn.SendBatch(ctx, batch).Close())
+
+	tx, err := conn.TxFromCurrentTransaction(pgx.TxOptions{})
+	require.NoError(t, err)
+	require.NoError(t, tx.Rollback(ctx))
+
+	var n int64
+	err = conn.QueryRow(ctx, "select count(*) from foo").Scan(&n)
+	require.NoError(t, err)
+	require.EqualValues(t, 0, n)
+}
+
+func TestTxFromCurrentTransactionWhenNotInTransaction(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	tx, err := conn.TxFromCurrentTransaction(pgx.TxOptions{})
+	require.Error(t, err)
+	require.Nil(t, tx)
+
+	// The connection is untouched and remains usable.
+	var n int32
+	require.NoError(t, conn.QueryRow(ctx, "select 1").Scan(&n))
+	require.EqualValues(t, 1, n)
+}
+
+func TestTxFromCurrentTransactionWhenTransactionFailed(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	batch := &pgx.Batch{}
+	batch.Queue(pgx.TxOptions{}.BeginSQL())
+	batch.Queue("select 1/0")
+	require.Error(t, conn.SendBatch(ctx, batch).Close())
+
+	// The transaction is in the failed state, but a Tx is still needed to roll it back.
+	tx, err := conn.TxFromCurrentTransaction(pgx.TxOptions{})
+	require.NoError(t, err)
+	require.NoError(t, tx.Rollback(ctx))
+
+	var n int32
+	require.NoError(t, conn.QueryRow(ctx, "select 1").Scan(&n))
+	require.EqualValues(t, 1, n)
+}
+
+func TestTxFromCurrentTransactionCommitQuery(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	txOptions := pgx.TxOptions{CommitQuery: "commit /* custom */"}
+
+	batch := &pgx.Batch{}
+	batch.Queue(txOptions.BeginSQL())
+	batch.Queue("select 1")
+	require.NoError(t, conn.SendBatch(ctx, batch).Close())
+
+	tx, err := conn.TxFromCurrentTransaction(txOptions)
+	require.NoError(t, err)
+	require.NoError(t, tx.Commit(ctx))
+	require.ErrorIs(t, tx.Commit(ctx), pgx.ErrTxClosed)
+}