add WithinTransaction methods on Conn, Pool and PoolConnHEADmain

2 files changed, 49 insertions, 4 deletions

diff --git a/handle.go b/handle.go
index bd2462d..3cb0612 100644
--- a/handle.go
+++ b/handle.go
@@ -41,7 +41,13 @@ func (h connHandle) GSQLClose(ctx context.Context) error {
 
 // GSQLTransact implements the [gsql.ConnectionHandle] interface.
 func (h connHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
-	tx, err := h.inner.Begin(ctx)
+	return withinTransactionOfConn(ctx, h.inner, func(tx *Tx) error {
+		return action(tx)
+	})
+}
+
+func withinTransactionOfConn(ctx context.Context, conn *pgx.Conn, action func(tx *Tx) error) error {
+	tx, err := conn.Begin(ctx)
 	if err != nil {
 		return err
 	}
@@ -93,7 +99,13 @@ func (h poolHandle) GSQLClose(ctx context.Context) error {
 
 // GSQLTransact implements the [gsql.ConnectionHandle] interface.
 func (h poolHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
-	tx, err := h.inner.Begin(ctx)
+	return withinTransactionOfPool(ctx, h.inner, func(tx *Tx) error {
+		return action(tx)
+	})
+}
+
+func withinTransactionOfPool(ctx context.Context, pool *pgxpool.Pool, action func(tx *Tx) error) error {
+	tx, err := pool.Begin(ctx)
 	if err != nil {
 		return err
 	}
@@ -130,7 +142,13 @@ func (h poolConnHandle) GSQLClose(ctx context.Context) error {
 
 // GSQLTransact implements the [gsql.ConnectionHandle] interface.
 func (h poolConnHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
-	tx, err := h.inner.Begin(ctx)
+	return withinTransactionOfPoolConn(ctx, h.inner, func(tx *Tx) error {
+		return action(tx)
+	})
+}
+
+func withinTransactionOfPoolConn(ctx context.Context, conn *pgxpool.Conn, action func(tx *Tx) error) error {
+	tx, err := conn.Begin(ctx)
 	if err != nil {
 		return err
 	}
@@ -167,7 +185,7 @@ func getPreparedStatementName() string {
 	return "oblast_pgx_" + strconv.FormatUint(preparedStatementId.Add(1), 10)
 }
 
-func transact(ctx context.Context, tx pgx.Tx, action func(tx gsql.Handle) error) error {
+func transact(ctx context.Context, tx pgx.Tx, action func(tx *Tx) error) error {
 	err := action(NewTx(tx))
 	if err == nil {
 		return errext.WithCleanup(nil, "tx.Commit", tx.Commit(ctx))
diff --git a/interface.go b/interface.go
index b9d7def..2b9b8cf 100644
--- a/interface.go
+++ b/interface.go
@@ -45,6 +45,15 @@ func (conn *Conn) BeginTx(ctx context.Context, opts pgx.TxOptions) (*Tx, error)
 	return maybeNewTx(tx), err
 }
 
+// WithinTransaction executes an action within a database transaction.
+// The transaction will be committed if the callback returns successfully, or rolled back otherwise.
+//
+// This is equivalent to the GSQLTransact() method of conn's [gsql.ConnectionHandle] implementation,
+// but the callback receives the concrete type [*Tx] instead of a generic [gsql.Handle].
+func (conn *Conn) WithinTransaction(ctx context.Context, action func(*Tx) error) error {
+	return withinTransactionOfConn(ctx, conn.Conn, action)
+}
+
 // Pool wraps [*pgxpool.Pool] into a [gsql.Handle].
 //
 // Because this type has [*pgxpool.Pool] as an embedded field,
@@ -94,6 +103,15 @@ func (pool *Pool) BeginTx(ctx context.Context, opts pgx.TxOptions) (*Tx, error)
 	return maybeNewTx(tx), err
 }
 
+// WithinTransaction executes an action within a database transaction.
+// The transaction will be committed if the callback returns successfully, or rolled back otherwise.
+//
+// This is equivalent to the GSQLTransact() method of pool's [gsql.ConnectionHandle] implementation,
+// but the callback receives the concrete type [*Tx] instead of a generic [gsql.Handle].
+func (pool *Pool) WithinTransaction(ctx context.Context, action func(*Tx) error) error {
+	return withinTransactionOfPool(ctx, pool.Pool, action)
+}
+
 // PoolConn wraps [*pgxpool.Conn] into a [gsql.Handle].
 //
 // Because this type has [*pgxpool.Conn] as an embedded field,
@@ -132,6 +150,15 @@ func (conn *PoolConn) Hijack() *Conn {
 	return maybe(NewConn, conn.Conn.Hijack())
 }
 
+// WithinTransaction executes an action within a database transaction.
+// The transaction will be committed if the callback returns successfully, or rolled back otherwise.
+//
+// This is equivalent to the GSQLTransact() method of conn's [gsql.ConnectionHandle] implementation,
+// but the callback receives the concrete type [*Tx] instead of a generic [gsql.Handle].
+func (conn *PoolConn) WithinTransaction(ctx context.Context, action func(*Tx) error) error {
+	return withinTransactionOfPoolConn(ctx, conn.Conn, action)
+}
+
 // Tx wraps [pgx.Tx] into a [gsql.Handle].
 //
 // Because this type has [pgx.Tx] as an embedded field,