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) +}