Add fix for using a context with a value in BeforeConnect

Author: dharju

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()

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()