Repository: jackc/pgx
Pull request #2329 (2025-05-21 15:42:28)

Add PrepareConn hook, which extends BeforeAcquire's behavior to allow…
… canceling the instigating query

- BeforeAcquire now marked as deprecated, and re-implemented in terms of PrepareConn
- PrepareConn now takes precidence over BeforeAcquire if both are provided
- New tests added, so both old and new behavior are tested
- One niggle: AcquireAllIdle does not return an error, so the only recourse that seems reasonable when PrepareConn returns an error in that context, is to destroy the connection.  This more or less retains the spirit of existing functionality, without changing the public API of that method. Although maybe an error-returning variant would be a useful addition as well.

Fixes: #2328


Patch:
diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index 416275f2..666bdf80 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -83,7 +83,7 @@ type Pool struct {
 	config                *Config
 	beforeConnect         func(context.Context, *pgx.ConnConfig) error
 	afterConnect          func(context.Context, *pgx.Conn) error
-	beforeAcquire         func(context.Context, *pgx.Conn) bool
+	prepareConn           func(context.Context, *pgx.Conn) (bool, error)
 	afterRelease          func(*pgx.Conn) bool
 	beforeClose           func(*pgx.Conn)
 	minConns              int32
@@ -118,8 +118,22 @@ type Config struct {
 	// BeforeAcquire is called before a connection is acquired from the pool. It must return true to allow the
 	// acquisition or false to indicate that the connection should be destroyed and a different connection should be
 	// acquired.
+	//
+	// Deprecated: Use PrepareConn instead. If both PrepareConn and BeforeAcquire are set, PrepareConn will take
+	// precedence, ignoring BeforeAcquire.
 	BeforeAcquire func(context.Context, *pgx.Conn) bool
 
+	// PrepareConn is called before a connection is acquired from the pool. If this function returns true, the connection
+	// is considered valid. If the function returns a non-nil error, the instigating query will fail with the returned error.
+	//
+	// Specifically, this means that:
+	//
+	// 	- It must return true and a nil error to allow acquisition and the query to proceed.
+	// 	- If it returns true and an error, the connection will be returned to the pool, and the instigating query will fail with the returned error.
+	// 	- If it returns false, and an error, the query will fail with the returned error, and the connection will be destroyed.
+	// 	- If it returns false and a nil error, the connection will be returned to the pool, and the instigating query will be retried on a new connection.
+	PrepareConn func(context.Context, *pgx.Conn) (bool, error)
+
 	// AfterRelease is called after a connection is released, but before it is returned to the pool. It must return true to
 	// return the connection to the pool or false to destroy the connection.
 	AfterRelease func(*pgx.Conn) bool
@@ -189,11 +203,18 @@ func NewWithConfig(ctx context.Context, config *Config) (*Pool, error) {
 		panic("config must be created by ParseConfig")
 	}
 
+	prepareConn := config.PrepareConn
+	if prepareConn == nil && config.BeforeAcquire != nil {
+		prepareConn = func(ctx context.Context, conn *pgx.Conn) (bool, error) {
+			return config.BeforeAcquire(ctx, conn), nil
+		}
+	}
+
 	p := &Pool{
 		config:                config,
 		beforeConnect:         config.BeforeConnect,
 		afterConnect:          config.AfterConnect,
-		beforeAcquire:         config.BeforeAcquire,
+		prepareConn:           prepareConn,
 		afterRelease:          config.AfterRelease,
 		beforeClose:           config.BeforeClose,
 		minConns:              config.MinConns,
@@ -560,11 +581,23 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
 			}
 		}
 
-		if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) {
-			return cr.getConn(p, res), nil
+		if p.prepareConn != nil {
+			ok, err := p.prepareConn(ctx, cr.conn)
+			if !ok {
+				res.Destroy()
+			}
+			if err != nil {
+				if ok {
+					res.Release()
+				}
+				return nil, err
+			}
+			if !ok {
+				continue
+			}
 		}
 
-		res.Destroy()
+		return cr.getConn(p, res), nil
 	}
 }
 
@@ -588,11 +621,14 @@ func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn {
 	conns := make([]*Conn, 0, len(resources))
 	for _, res := range resources {
 		cr := res.Value()
-		if p.beforeAcquire == nil || p.beforeAcquire(ctx, cr.conn) {
-			conns = append(conns, cr.getConn(p, res))
-		} else {
-			res.Destroy()
+		if p.prepareConn != nil {
+			ok, err := p.prepareConn(ctx, cr.conn)
+			if !ok || err != nil {
+				res.Destroy()
+				continue
+			}
 		}
+		conns = append(conns, cr.getConn(p, res))
 	}
 
 	return conns


Test patch:
diff --git a/pgxpool/common_test.go b/pgxpool/common_test.go
index 009616a1..ef54fb60 100644
--- a/pgxpool/common_test.go
+++ b/pgxpool/common_test.go
@@ -141,6 +141,7 @@ func assertConfigsEqual(t *testing.T, expected, actual *pgxpool.Config, testName
 	// Can't test function equality, so just test that they are set or not.
 	assert.Equalf(t, expected.AfterConnect == nil, actual.AfterConnect == nil, "%s - AfterConnect", testName)
 	assert.Equalf(t, expected.BeforeAcquire == nil, actual.BeforeAcquire == nil, "%s - BeforeAcquire", testName)
+	assert.Equalf(t, expected.PrepareConn == nil, actual.PrepareConn == nil, "%s - PrepareConn", testName)
 	assert.Equalf(t, expected.AfterRelease == nil, actual.AfterRelease == nil, "%s - AfterRelease", testName)
 
 	assert.Equalf(t, expected.MaxConnLifetime, actual.MaxConnLifetime, "%s - MaxConnLifetime", testName)
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index 4cc7f47e..cb8e6810 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -330,6 +330,64 @@ func TestPoolBeforeAcquire(t *testing.T) {
 	assert.EqualValues(t, 12, acquireAttempts)
 }
 
+func TestPoolPrepareConn(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	acquireAttempts := 0
+
+	config.PrepareConn = func(context.Context, *pgx.Conn) (bool, error) {
+		acquireAttempts++
+		var err error
+		if acquireAttempts%3 == 0 {
+			err = errors.New("PrepareConn error")
+		}
+		return acquireAttempts%2 == 0, err
+	}
+
+	db, err := pgxpool.NewWithConfig(ctx, config)
+	require.NoError(t, err)
+	t.Cleanup(db.Close)
+
+	var errorCount int
+	conns := make([]*pgxpool.Conn, 0, 4)
+	for {
+		conn, err := db.Acquire(ctx)
+		if err != nil {
+			errorCount++
+			continue
+		}
+		conns = append(conns, conn)
+		if len(conns) == 4 {
+			break
+		}
+	}
+	const wantErrorCount = 3
+	assert.Equal(t, wantErrorCount, errorCount, "Acquire() should have failed %d times", wantErrorCount)
+
+	for _, c := range conns {
+		c.Release()
+	}
+	waitForReleaseToComplete()
+
+	assert.EqualValues(t, len(conns)*2+wantErrorCount-1, acquireAttempts)
+
+	conns = db.AcquireAllIdle(ctx)
+	assert.Len(t, conns, 1)
+
+	for _, c := range conns {
+		c.Release()
+	}
+	waitForReleaseToComplete()
+
+	assert.EqualValues(t, 14, acquireAttempts)
+}
+
 func TestPoolAfterRelease(t *testing.T) {
 	t.Parallel()
 
@@ -1082,9 +1140,9 @@ func TestConnectEagerlyReachesMinPoolSize(t *testing.T) {
 	acquireAttempts := int64(0)
 	connectAttempts := int64(0)
 
-	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+	config.PrepareConn = func(ctx context.Context, conn *pgx.Conn) (bool, error) {
 		atomic.AddInt64(&acquireAttempts, 1)
-		return true
+		return true, nil
 	}
 	config.BeforeConnect = func(ctx context.Context, cfg *pgx.ConnConfig) error {
 		atomic.AddInt64(&connectAttempts, 1)