Merge pull request jackc#2329 from flimzy/prepareConn

Add PrepareConn hook, which extends BeforeAcquire's behavior to allow…

Author: jackc

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)

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

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)