From d15cc95a8a029fd21eaa4936ab9f994937801405 Mon Sep 17 00:00:00 2001
From: David Harju <dharju@figma.com>
Date: Sun, 29 Jan 2023 16:07:23 -0800
Subject: [PATCH] Add fix for using a context with a value in BeforeConnect

---
 pgxpool/pool.go      | 10 +++++++---
 pgxpool/pool_test.go | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index 8e88ecaf7..e83d8a5df 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -265,7 +265,7 @@ func ConnectConfig(ctx context.Context, config *Config) (*Pool, error) {
 	)
 
 	if !config.LazyConnect {
-		if err := p.checkMinConns(); err != nil {
+		if err := p.checkMinConnsWithContext(ctx); err != nil {
 			// Couldn't create resources for minpool size. Close unhealthy pool.
 			p.Close()
 			return nil, err
@@ -485,17 +485,21 @@ func (p *Pool) checkConnsHealth() bool {
 	return destroyed
 }
 
-func (p *Pool) checkMinConns() error {
+func (p *Pool) checkMinConnsWithContext(ctx context.Context) error {
 	// TotalConns can include ones that are being destroyed but we should have
 	// sleep(500ms) around all of the destroys to help prevent that from throwing
 	// off this check
 	toCreate := p.minConns - p.Stat().TotalConns()
 	if toCreate > 0 {
-		return p.createIdleResources(context.Background(), int(toCreate))
+		return p.createIdleResources(ctx, int(toCreate))
 	}
 	return nil
 }
 
+func (p *Pool) checkMinConns() error {
+	return p.checkMinConnsWithContext(context.Background())
+}
+
 func (p *Pool) createIdleResources(parentCtx context.Context, targetResources int) error {
 	ctx, cancel := context.WithCancel(parentCtx)
 	defer cancel()
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index 4e7120169..4346d6f51 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -73,6 +73,26 @@ func TestLazyConnect(t *testing.T) {
 	assert.Equal(t, context.Canceled, err)
 }
 
+func TestBeforeConnectWithContextWithValueAndOneMinConn(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	assert.NoError(t, err)
+	config.MinConns = 1
+	config.BeforeConnect = func(ctx context.Context, config *pgx.ConnConfig) error {
+		val := ctx.Value("key")
+		if val == nil {
+			return errors.New("no value found with key 'key'")
+		}
+		return nil
+	}
+
+	ctx := context.WithValue(context.Background(), "key", "value")
+
+	_, err = pgxpool.ConnectConfig(ctx, config)
+	assert.NoError(t, err)
+}
+
 func TestConstructorIgnoresContext(t *testing.T) {
 	t.Parallel()