Skip to content

Commit

Permalink
Add fix for using a context with a value in BeforeConnect
Browse files Browse the repository at this point in the history
  • Loading branch information
dharju authored and jackc committed Jan 31, 2023
1 parent 3edc1b5 commit d15cc95
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 3 deletions.
10 changes: 7 additions & 3 deletions pgxpool/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
20 changes: 20 additions & 0 deletions pgxpool/pool_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit d15cc95

Please sign in to comment.