From 468740e06e72fcc8e71c5a6dd75bf0beff0aab80 Mon Sep 17 00:00:00 2001 From: Timo Furrer Date: Fri, 21 Aug 2026 10:54:45 +0200 Subject: [PATCH] Add Conn.TxFromCurrentTransaction and TxOptions.BeginSQL Allows a caller to begin a transaction as part of a Batch or a pgconn pipeline and still get a Tx for the rest of the unit of work, saving the round trip that Conn.BeginTx spends on its own begin query. TxOptions.BeginSQL exports the statement Conn.BeginTx would send, so callers queueing it themselves do not have to hand-roll it. Refs #1667 --- tx.go | 34 +++++++++++ tx_test.go | 166 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+) diff --git a/tx.go b/tx.go index dcb0feb43..c3ee80481 100644 --- a/tx.go +++ b/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 == "" { diff --git a/tx_test.go b/tx_test.go index 4896941e5..0dfccee78 100644 --- a/tx_test.go +++ b/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) +}