Wrap pgx tx

xataio · Sep 13, 2024 · 4a29918 · 4a29918
1 parent 09321b6
commit 4a29918
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 3 deletions.
diff --git a/internal/postgres/pg_conn.go b/internal/postgres/pg_conn.go
@@ -48,7 +48,7 @@ func (c *Conn) ExecInTx(ctx context.Context, fn func(Tx) error) error {
 		return mapError(err)
 	}
 
-	if err := fn(tx); err != nil {
+	if err := fn(&Txn{Tx: tx}); err != nil {
 		tx.Rollback(ctx)
 		return mapError(err)
 	}

diff --git a/internal/postgres/pg_conn_pool.go b/internal/postgres/pg_conn_pool.go
@@ -49,7 +49,7 @@ func (c *Pool) ExecInTx(ctx context.Context, fn func(Tx) error) error {
 		return mapError(err)
 	}
 
-	if err := fn(tx); err != nil {
+	if err := fn(&Txn{Tx: tx}); err != nil {
 		tx.Rollback(ctx)
 		return mapError(err)
 	}

diff --git a/internal/postgres/pg_querier.go b/internal/postgres/pg_querier.go
@@ -26,7 +26,9 @@ type Rows interface {
 }
 
 type Tx interface {
-	pgx.Tx
+	Query(ctx context.Context, query string, args ...any) (Rows, error)
+	QueryRow(ctx context.Context, query string, args ...any) Row
+	Exec(ctx context.Context, query string, args ...any) (CommandTag, error)
 }
 
 type CommandTag struct {

diff --git a/internal/postgres/pg_tx.go b/internal/postgres/pg_tx.go
@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type Txn struct {
+	pgx.Tx
+}
+
+func (t *Txn) QueryRow(ctx context.Context, query string, args ...any) Row {
+	row := t.Tx.QueryRow(ctx, query, args...)
+	return &mappedRow{inner: row}
+}
+
+func (t *Txn) Query(ctx context.Context, query string, args ...any) (Rows, error) {
+	rows, err := t.Tx.Query(ctx, query, args...)
+	return rows, mapError(err)
+}
+
+func (t *Txn) Exec(ctx context.Context, query string, args ...any) (CommandTag, error) {
+	tag, err := t.Tx.Exec(ctx, query, args...)
+	return CommandTag{tag}, mapError(err)
+}