Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions stdlib/sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,22 @@ func init() {
// OptionOpenDB options for configuring the driver when opening a new db pool.
type OptionOpenDB func(*connector)

// ShouldPingParams are passed to OptionShouldPing to decide whether to ping before reusing a connection.
type ShouldPingParams struct {
// Conn is the underlying pgx connection.
Conn *pgx.Conn
// IdleDuration is how long it has been since ResetSession last ran.
IdleDuration time.Duration
}

// OptionShouldPing controls whether stdlib should issue a liveness ping before reusing a connection.
// If the function returns true, stdlib will ping.
// If it returns false, stdlib will skip the ping.
// If not provided, default is ping only when IdleDuration > 1s.
func OptionShouldPing(f func(context.Context, ShouldPingParams) bool) OptionOpenDB {
return func(dc *connector) { dc.ShouldPing = f }
}

// OptionBeforeConnect provides a callback for before connect. It is passed a shallow copy of the ConnConfig that will
// be used to connect, so only its immediate members should be modified. Used only if db is opened with *pgx.ConnConfig.
func OptionBeforeConnect(bc func(context.Context, *pgx.ConnConfig) error) OptionOpenDB {
Expand Down Expand Up @@ -231,6 +247,7 @@ type connector struct {
BeforeConnect func(context.Context, *pgx.ConnConfig) error // function to call before creation of every new connection
AfterConnect func(context.Context, *pgx.Conn) error // function to call after creation of every new connection
ResetSession func(context.Context, *pgx.Conn) error // function is called before a connection is reused
ShouldPing func(context.Context, ShouldPingParams) bool // function to decide if stdlib should ping before reusing a connection
driver *Driver
}

Expand Down Expand Up @@ -282,6 +299,7 @@ func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
driver: c.driver,
connConfig: connConfig,
resetSessionFunc: c.ResetSession,
shouldPing: c.ShouldPing,
psRefCounts: make(map[*pgconn.StatementDescription]int),
}, nil
}
Expand Down Expand Up @@ -389,7 +407,8 @@ type Conn struct {
close func(context.Context) error
driver *Driver
connConfig pgx.ConnConfig
resetSessionFunc func(context.Context, *pgx.Conn) error // Function is called before a connection is reused
resetSessionFunc func(context.Context, *pgx.Conn) error // Function is called before a connection is reused
shouldPing func(context.Context, ShouldPingParams) bool // Function to decide if stdlib should ping before reusing a connection
lastResetSessionTime time.Time

// psRefCounts contains reference counts for prepared statements. Prepare uses the underlying pgx logic to generate
Expand Down Expand Up @@ -537,11 +556,23 @@ func (c *Conn) ResetSession(ctx context.Context) error {
}

now := time.Now()
if now.Sub(c.lastResetSessionTime) > time.Second {
idle := now.Sub(c.lastResetSessionTime)

doPing := idle > time.Second // default behavior: ping only if idle > 1s

if c.shouldPing != nil {
doPing = c.shouldPing(ctx, ShouldPingParams{
Conn: c.conn,
IdleDuration: idle,
})
}

if doPing {
if err := c.conn.PgConn().Ping(ctx); err != nil {
return driver.ErrBadConn
}
}

c.lastResetSessionTime = now

return c.resetSessionFunc(ctx, c.conn)
Expand Down
36 changes: 32 additions & 4 deletions stdlib/sql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,22 @@ import (
"testing"
"time"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/jackc/pgx/v5/stdlib"
"github.com/jackc/pgx/v5/tracelog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func openDB(t testing.TB) *sql.DB {
func openDB(t testing.TB, opts ...stdlib.OptionOpenDB) *sql.DB {
t.Helper()
config, err := pgx.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
require.NoError(t, err)
return stdlib.OpenDB(*config)
return stdlib.OpenDB(*config, opts...)
}

func closeDB(t testing.TB, db *sql.DB) {
Expand Down Expand Up @@ -1374,3 +1376,29 @@ func TestCheckIdleConn(t *testing.T) {

require.NotContains(t, pids, cPID)
}

func TestOptionShouldPing_HookCalledOnReuse(t *testing.T) {
hookCalled := false

db := openDB(t,
stdlib.OptionShouldPing(func(context.Context, stdlib.ShouldPingParams) bool {
hookCalled = true
// Return false to avoid relying on actual ping behavior.
return false
}),
)
defer closeDB(t, db)

// Ensure reuse (so ResetSession runs)
db.SetMaxOpenConns(1)
db.SetMaxIdleConns(1)

// Establish the connection
require.NoError(t, db.Ping())

// Reuse the connection -> should trigger ResetSession -> ShouldPing
_, err := db.Exec("select 1")
require.NoError(t, err)

require.True(t, hookCalled, "hook should be called on reuse")
}