From 657ee90ab53ac9c619aee3e04df81099f1bdc8ec Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Wed, 26 Aug 2026 21:13:29 -0500 Subject: [PATCH] add pgx listeners to database/sql The `database/sql` driver currently operates in poll-only mode because it cannot acquire the dedicated connection required by Postgres `LISTEN`. Cross-process actions like cancelling a running job therefore wait for a poll even when the underlying SQL driver is Pgx. Add `NewWithPgxListener`, which keeps all query and transaction execution on the supplied `*sql.DB` while delegating listener creation to a Pgx driver backed by a separately supplied pool. Preserve `New` as the poll-only constructor and reject a nil listener pool so configuration errors cannot silently disable notifications. Document the database, schema, connection ownership, and pool sizing requirements. Exercise the hybrid driver through the shared listener suite, including notification-driven cancellation and transactional job completion through an ordinary `*sql.Tx`. --- CHANGELOG.md | 4 + riverdriver/riverdatabasesql/go.mod | 2 +- .../river_database_sql_driver.go | 47 ++++++++- .../river_database_sql_driver_test.go | 32 ++++++ .../riverdrivertest/driver_client_test.go | 97 ++++++++++++++++++- riverdriver/riverdrivertest/driver_test.go | 2 +- 6 files changed, 177 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 294a2a24..78d6f8c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added `riverdatabasesql.NewWithPgxListener` for using a dedicated Pgx pool to receive Postgres notifications while continuing to execute jobs and transactions through `database/sql`. [PR #1366](https://github.com/riverqueue/river/pull/1366). + ## [0.45.0] - 2026-08-25 ### Changed diff --git a/riverdriver/riverdatabasesql/go.mod b/riverdriver/riverdatabasesql/go.mod index 8d2b9545..41267f51 100644 --- a/riverdriver/riverdatabasesql/go.mod +++ b/riverdriver/riverdatabasesql/go.mod @@ -9,6 +9,7 @@ require ( github.com/lib/pq v1.12.3 github.com/riverqueue/river v0.45.0 github.com/riverqueue/river/riverdriver v0.45.0 + github.com/riverqueue/river/riverdriver/riverpgxv5 v0.45.0 github.com/riverqueue/river/rivershared v0.45.0 github.com/riverqueue/river/rivertype v0.45.0 github.com/stretchr/testify v1.12.1 @@ -18,7 +19,6 @@ require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/riverqueue/river/riverdriver/riverpgxv5 v0.45.0 // indirect github.com/tidwall/gjson v1.19.0 // indirect github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index e23613cc..1b1864b9 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -18,10 +18,12 @@ import ( "strings" "time" + "github.com/jackc/pgx/v5/pgxpool" "github.com/lib/pq" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riverdatabasesql/internal/dbsqlc" + "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/sqlctemplate" "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivershared/util/dbutil" @@ -36,8 +38,9 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for database/sql. type Driver struct { - dbPool *sql.DB - replacer sqlctemplate.Replacer + dbPool *sql.DB + listenerDriver *riverpgxv5.Driver + replacer sqlctemplate.Replacer } // New returns a new database/sql River driver for use with River. @@ -51,6 +54,38 @@ func New(dbPool *sql.DB) *Driver { } } +// NewWithPgxListener returns a new database/sql River driver with a Pgx-backed +// listener. The database/sql pool continues to be used for all database +// operations other than listening for notifications. The Pgx pool is used only +// to acquire dedicated connections for Postgres LISTEN commands. It panics if +// listenerPool is nil; use New for a poll-only driver. +// +// Both pools are owned by the caller, must connect to the same database, and +// must resolve the same schema. When the River client has no explicit schema, +// both pools' search paths must produce the same current schema. Neither pool +// may be closed while associated River objects are running. +// +// Listener connections are hijacked from listenerPool and never returned. A +// pool dedicated to one River client can generally set MinConns to zero and +// MaxConns to one. Each concurrently running client still needs its own listener +// connection. Because hijacked connections no longer count +// against the pool's maximum, sharing a listener pool between clients may cause +// total live connections to exceed that maximum. Closing listenerPool does not +// close hijacked connections; stopping the associated River clients does. +// +// Applications using PgBouncer must configure the listener pool to use session +// pooling or connect it directly to Postgres. +func NewWithPgxListener(dbPool *sql.DB, listenerPool *pgxpool.Pool) *Driver { + if listenerPool == nil { + panic("riverdatabasesql: listener pool must not be nil") + } + + return &Driver{ + dbPool: dbPool, + listenerDriver: riverpgxv5.New(listenerPool), + } +} + const argPlaceholder = "$" func (d *Driver) ArgPlaceholder() string { return argPlaceholder } @@ -61,7 +96,11 @@ func (d *Driver) GetExecutor() riverdriver.Executor { } func (d *Driver) GetListener(params *riverdriver.GetListenenerParams) riverdriver.Listener { - panic(riverdriver.ErrNotImplemented) + if d.listenerDriver == nil { + panic(riverdriver.ErrNotImplemented) + } + + return d.listenerDriver.GetListener(params) } func (d *Driver) GetMigrationDefaultLines() []string { return []string{riverdriver.MigrationLineMain} } @@ -96,7 +135,7 @@ func (d *Driver) SQLFragmentColumnIn(column string, values any) (string, any, er return fmt.Sprintf("%s = any(@%s)", column, column), pq.Array(values), nil } -func (d *Driver) SupportsListener() bool { return false } +func (d *Driver) SupportsListener() bool { return d.listenerDriver != nil } func (d *Driver) SupportsListenNotify() bool { return true } func (d *Driver) TimePrecision() time.Duration { return time.Microsecond } diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver_test.go b/riverdriver/riverdatabasesql/river_database_sql_driver_test.go index 0d23bb43..b5d2e59e 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver_test.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver_test.go @@ -6,6 +6,7 @@ import ( "errors" "testing" + "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" "github.com/riverqueue/river/riverdriver" @@ -35,6 +36,37 @@ func TestNew(t *testing.T) { }) } +func TestNewWithPgxListener(t *testing.T) { + t.Parallel() + + t.Run("PanicsOnNilListenerPool", func(t *testing.T) { + t.Parallel() + + require.PanicsWithValue(t, "riverdatabasesql: listener pool must not be nil", func() { + NewWithPgxListener(&sql.DB{}, nil) + }) + }) + + t.Run("UsesSeparateListenerPool", func(t *testing.T) { + t.Parallel() + + dbPool := &sql.DB{} + listenerPool := &pgxpool.Pool{} + driver := NewWithPgxListener(dbPool, listenerPool) + + require.Equal(t, dbPool, driver.dbPool) + require.NotNil(t, driver.listenerDriver) + require.True(t, driver.SupportsListener()) + require.Equal(t, dbPool, driver.GetExecutor().(*Executor).dbPool) //nolint:forcetypeassert + + listener1 := driver.GetListener(&riverdriver.GetListenenerParams{Schema: "schema_one"}) + listener2 := driver.GetListener(&riverdriver.GetListenenerParams{Schema: "schema_two"}) + require.NotSame(t, listener1, listener2) + require.Equal(t, "schema_one", listener1.Schema()) + require.Equal(t, "schema_two", listener2.Schema()) + }) +} + func TestBitIntegerToBits(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverdrivertest/driver_client_test.go b/riverdriver/riverdrivertest/driver_client_test.go index 9c9839dc..644cffae 100644 --- a/riverdriver/riverdrivertest/driver_client_test.go +++ b/riverdriver/riverdrivertest/driver_client_test.go @@ -23,6 +23,7 @@ import ( "github.com/riverqueue/river/riverdriver/riversqlite" "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/util/testutil" "github.com/riverqueue/river/rivershared/util/urlutil" "github.com/riverqueue/river/rivertype" @@ -57,7 +58,7 @@ func TestClientWithDriverRiverDatabaseSQLPgx(t *testing.T) { ctx = context.Background() dbPool = riversharedtest.DBPool(ctx, t) stdPool = stdlib.OpenDBFromPool(dbPool) - driver = riverdatabasesql.New(stdPool) + driver = riverdatabasesql.NewWithPgxListener(stdPool, dbPool) ) t.Cleanup(func() { require.NoError(t, stdPool.Close()) }) @@ -70,6 +71,59 @@ func TestClientWithDriverRiverDatabaseSQLPgx(t *testing.T) { ) } +func TestClientWithDriverRiverDatabaseSQLPgxJobCompleteTx(t *testing.T) { + t.Parallel() + + var ( + ctx = context.Background() + dbPool = riversharedtest.DBPool(ctx, t) + stdPool = stdlib.OpenDBFromPool(dbPool) + driver = riverdatabasesql.NewWithPgxListener(stdPool, dbPool) + schema = riverdbtest.TestSchema(ctx, t, driver, nil) + ) + t.Cleanup(func() { require.NoError(t, stdPool.Close()) }) + + var jobCompleted testsignal.TestSignal[int64] + jobCompleted.Init(t) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + config := newTestConfig(t, schema) + config.FetchPollInterval = time.Minute + river.AddWorker(config.Workers, river.WorkFunc(func(ctx context.Context, job *river.Job[JobArgs]) error { + tx, err := stdPool.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() + + updatedJob, err := river.JobCompleteTx[*riverdatabasesql.Driver](ctx, tx, job) + if err != nil { + return err + } + if err := tx.Commit(); err != nil { + return err + } + + jobCompleted.Signal(updatedJob.ID) + return nil + })) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + startClient(ctx, t, client) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + require.Equal(t, insertRes.Job.ID, jobCompleted.WaitOrTimeout()) + + completedJob, err := client.JobGet(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, completedJob.State) +} + func TestClientWithDriverRiverPgxV5(t *testing.T) { t.Parallel() @@ -301,6 +355,47 @@ func ExerciseClient[TTx any](ctx context.Context, t *testing.T, require.Equal(t, insertRes.Job.Kind, event.Job.Kind) }) + t.Run("CancelRunningJobWithListener", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + if bundle.driver.DatabaseName() != riverdriver.DatabaseNamePostgres || !bundle.driver.SupportsListener() { + t.Skip("requires a Postgres listener") + } + config.FetchPollInterval = time.Minute + + client, err := river.NewClient(bundle.driver, config) + require.NoError(t, err) + + var jobStarted testsignal.TestSignal[int64] + jobStarted.Init(t) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + river.AddWorker(bundle.config.Workers, river.WorkFunc(func(ctx context.Context, job *river.Job[JobArgs]) error { + jobStarted.Signal(job.ID) + <-ctx.Done() + return ctx.Err() + })) + + subscribeChan := subscribe(t, client) + startClient(ctx, t, client) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + require.Equal(t, insertRes.Job.ID, jobStarted.WaitOrTimeout()) + + updatedJob, err := client.JobCancel(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.Equal(t, river.EventKindJobCancelled, event.Kind) + require.Equal(t, rivertype.JobStateCancelled, event.Job.State) + }) + t.Run("JobDelete", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverdrivertest/driver_test.go b/riverdriver/riverdrivertest/driver_test.go index bbf9db89..30064d7e 100644 --- a/riverdriver/riverdrivertest/driver_test.go +++ b/riverdriver/riverdrivertest/driver_test.go @@ -68,7 +68,7 @@ func TestDriverRiverDatabaseSQLPgx(t *testing.T) { ctx = context.Background() dbPool = riversharedtest.DBPool(ctx, t) stdPool = stdlib.OpenDBFromPool(dbPool) - driver = riverdatabasesql.New(stdPool) + driver = riverdatabasesql.NewWithPgxListener(stdPool, dbPool) ) t.Cleanup(func() { require.NoError(t, stdPool.Close()) })