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