feat: ✨ transaction can now be placed to the context

BREAKING CHANGE: Querier.Conn now have context as an argument

README.md

@@ -136,6 +136,26 @@ func main() {
   _, err = q.Exec(ctx, "UPDATE table SET something = 1")
   return err
  }, conn.StatementTimeout(time.Second))
+
+tx, err := wrapped.Conn(ctx).Begin(ctx)
+ if err != nil {
+  log.Fatalf("failed to start transaction: %s", err)
+ }
+
+ // Put a transaction in the context, so that all subsequent calls use the transaction
+ txCtx := conn.NewTxContext(ctx, tx)
+ if _, err := wrapped.Exec(txCtx, "UPDATE table SET something = 1"); err != nil {
+  _ = tx.Rollback(ctx)
+  log.Fatalf("failed to exec: %s", err)
+ }
+ if err := wrapped.Get(txCtx, &count, "SELECT COUNT(*) FROM table"); err != nil {
+  _ = tx.Rollback(ctx)
+  log.Fatalf("failed to get: %s", err)
+ }
+
+ if err := tx.Commit(ctx); err != nil {
+  log.Fatalf("failed to commit transaction: %s", err)
+ }
 }
 ```
 

conn/querier.go

@@ -13,7 +13,7 @@ type Querier interface {
 	Get(ctx context.Context, dst interface{}, sql string, args ...interface{}) error
 	Exec(ctx context.Context, sql string, args ...interface{}) (int64, error)
 	Tx(ctx context.Context, f func(q Querier) error, opts ...TxOption) error
-	Conn() PgxConn
+	Conn(ctx context.Context) PgxConn
 }
 
 var _ Querier = &wrappedConn{}
@@ -36,7 +36,7 @@ func WrapConn(conn PgxConn, scanAPI *pgxscan.API) *wrappedConn {
 // Before starting, Select resets the destination slice,
 // so if it's not empty it will overwrite all existing elements.
 func (n *wrappedConn) Select(ctx context.Context, dst interface{}, sql string, args ...interface{}) error {
-	rows, err := n.conn.Query(ctx, sql, args...)
+	rows, err := n.Conn(ctx).Query(ctx, sql, args...)
 	if err != nil {
 		return err
 	}
@@ -48,7 +48,7 @@ func (n *wrappedConn) Select(ctx context.Context, dst interface{}, sql string, a
 // otherwise it returns an error.
 // It scans data from single row into the destination.
 func (n *wrappedConn) Get(ctx context.Context, dst interface{}, sql string, args ...interface{}) error {
-	rows, err := n.conn.Query(ctx, sql, args...)
+	rows, err := n.Conn(ctx).Query(ctx, sql, args...)
 	if err != nil {
 		return err
 	}
@@ -58,7 +58,7 @@ func (n *wrappedConn) Get(ctx context.Context, dst interface{}, sql string, args
 
 // Exec executes a query without returning any rows and return affected rows.
 func (n *wrappedConn) Exec(ctx context.Context, sql string, args ...interface{}) (int64, error) {
-	res, err := n.conn.Exec(ctx, sql, args...)
+	res, err := n.Conn(ctx).Exec(ctx, sql, args...)
 	if err != nil {
 		return 0, err
 	}
@@ -73,7 +73,7 @@ func (n *wrappedConn) Tx(ctx context.Context, f func(q Querier) error, opts ...T
 	for _, o := range opts {
 		o(txOpts)
 	}
-	err := n.conn.BeginFunc(ctx, func(txx pgx.Tx) error {
+	err := n.Conn(ctx).BeginFunc(ctx, func(txx pgx.Tx) error {
 		if txOpts.TransactionTimeout > 0 {
 			if _, err := txx.Exec(ctx, transactionTimeoutQuery, pgx.QuerySimpleProtocol(true), txOpts.TransactionTimeout); err != nil {
 				return fmt.Errorf("set transaction timeout: %w", err)
@@ -92,6 +92,10 @@ func (n *wrappedConn) Tx(ctx context.Context, f func(q Querier) error, opts ...T
 	return err
 }
 
-func (n *wrappedConn) Conn() PgxConn {
+func (n *wrappedConn) Conn(ctx context.Context) PgxConn {
+	if conn, ok := TxFromContext(ctx); ok {
+		return conn
+	}
+
 	return n.conn
 }

conn/querier_test.go

@@ -287,4 +287,27 @@ func TestQuerier(t *testing.T) {
 			is.True(pgErr.Code == "57014")  // 57014 - query_canceled error code
 		})
 	})
+
+	t.Run("Conn", func(t *testing.T) {
+		t.Parallel()
+		querier := WrapConn(db, pgxscan.DefaultAPI)
+
+		t.Run("no transaction in ctx", func(t *testing.T) {
+			is := is.New(t)
+			ctx := context.Background()
+
+			conn := querier.Conn(ctx)
+			is.True(conn == db) // must be original database connection
+		})
+
+		t.Run("transaction in ctx", func(t *testing.T) {
+			is := is.New(t)
+			ctx := context.Background()
+			tx, err := querier.Conn(ctx).Begin(ctx)
+			is.NoErr(err)
+			ctx = NewTxContext(ctx, tx)
+			conn := querier.Conn(ctx)
+			is.True(conn == tx) // must be transaction connection
+		})
+	})
 }

conn/tx_options.go

@@ -1,7 +1,10 @@
 package conn
 
 import (
+	"context"
 	"time"
+
+	"github.com/jackc/pgx/v4"
 )
 
 const (
@@ -31,3 +34,25 @@ func StatementTimeout(d time.Duration) TxOption {
 		o.StatementTimeout = d.Milliseconds()
 	}
 }
+
+// Context options for transaction
+
+type txKeyType uint8
+
+const (
+	txKey txKeyType = 0
+)
+
+// NewTxContext returns a new context carrying the transaction connection
+func NewTxContext(ctx context.Context, tx pgx.Tx) context.Context {
+	if tx == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, txKey, tx)
+}
+
+// TxFromContext extracts the transaction connection if present.
+func TxFromContext(ctx context.Context) (pgx.Tx, bool) {
+	v, ok := ctx.Value(txKey).(pgx.Tx)
+	return v, ok
+}

examples/cluster/cluster.go

@@ -48,7 +48,7 @@ func main() {
 	}
 
 	// If needed, one can access the PgxConn to call pgx methods directly such as SendBatch, CopyFrom ... .
-	conn := db.Primary().Conn()
+	conn := db.Primary().Conn(ctx)
 	_ = conn
 	// If needed, one can access the primary or a replica explicitly.
 	primary, replica := db.Primary(), db.Replica()

examples/conn/conn.go

@@ -48,4 +48,24 @@ func main() {
 		_, err = q.Exec(ctx, "UPDATE table SET something = 1")
 		return err
 	}, conn.StatementTimeout(time.Second))
+
+	tx, err := wrapped.Conn(ctx).Begin(ctx)
+	if err != nil {
+		log.Fatalf("failed to start transaction: %s", err)
+	}
+
+	// Put a transaction in the context, so that all subsequent calls use the transaction
+	txCtx := conn.NewTxContext(ctx, tx)
+	if _, err := wrapped.Exec(txCtx, "UPDATE table SET something = 1"); err != nil {
+		_ = tx.Rollback(ctx)
+		log.Fatalf("failed to exec: %s", err)
+	}
+	if err := wrapped.Get(txCtx, &count, "SELECT COUNT(*) FROM table"); err != nil {
+		_ = tx.Rollback(ctx)
+		log.Fatalf("failed to get: %s", err)
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		log.Fatalf("failed to commit transaction: %s", err)
+	}
 }

rules.mk

@@ -166,7 +166,7 @@ go.bumpdeps:
 go.fmt:
 	@set -e; for dir in $(GOMOD_DIRS); do ( set -e; \
 	  cd $$dir; \
-	  $(GO) run mvdan.cc/gofumpt -extra -w -l -s `go list -f '{{.Dir}}' $(WHAT) | grep -v mocks` \
+	  $(GO) run mvdan.cc/gofumpt -extra -w -l `go list -f '{{.Dir}}' $(WHAT) | grep -v mocks` \
 	); done
 
 VERIFY_STEPS += go.depaware-check

txdb/txdb.go

@@ -132,7 +132,7 @@ func (c *txdbCluster) Replica() conn.Querier {
 
 func (c *txdbCluster) beginOnce(ctx context.Context) (pgx.Tx, error) {
 	if c.tx == nil {
-		tx, err := c.cluster.Primary().Conn().Begin(ctx)
+		tx, err := c.cluster.Primary().Conn(ctx).Begin(ctx)
 		if err != nil {
 			return nil, err
 		}

txdb/txdb_test.go

@@ -102,12 +102,12 @@ func TestIntegration_txdb(t *testing.T) {
 			t.Parallel()
 			is := is.New(t)
 
-			t1, t2 := sendBatch(db.Primary().Conn())
+			t1, t2 := sendBatch(db.Primary().Conn(ctx))
 			is.True(!t1.Equal(t2)) // transaction_timestamp not in transaction must be not equal
 
 			txdb := New(db)
 
-			t1, t2 = sendBatch(txdb.Primary().Conn())
+			t1, t2 = sendBatch(txdb.Primary().Conn(ctx))
 			is.True(t1.Equal(t2)) // transaction_timestamp in transaction must be equal
 		})
 	})