Support Yugabyte as a Postgres target

client.go

@@ -57,6 +57,8 @@ const (
 	QueueNumWorkersMax       = 10_000
 )
 
+const clientInitialPingTimeout = 10 * time.Second
+
 var (
 	postgresSchemaNameRE = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
 
@@ -1058,9 +1060,26 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		client.testSignals.queueMaintainerLeader = &client.queueMaintainerLeader.TestSignals
 	}
 
+	// Ping asynchronously so that constructing a client remains non-blocking.
+	// PostgreSQL drivers use this ping to select and cache their unique insert
+	// strategy. The first unique insert uses the same resolver, so it can perform
+	// initialization instead if it reaches the database first.
+	if driver.DatabaseName() == riverdriver.DatabaseNamePostgres && driver.PoolIsSet() {
+		go client.initialPing()
+	}
+
 	return client, nil
 }
 
+func (c *Client[TTx]) initialPing() {
+	ctx, cancel := context.WithTimeout(context.Background(), clientInitialPingTimeout)
+	defer cancel()
+
+	if err := c.driver.GetExecutor().Ping(ctx); err != nil {
+		c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": Initial database ping failed", slog.String("err", err.Error()))
+	}
+}
+
 // Start starts the client's job fetching and working loops. Once this is called,
 // the client will run in a background goroutine until stopped. All jobs are
 // run with a context inheriting from the provided context, but with a timeout
@@ -1110,7 +1129,7 @@ func (c *Client[TTx]) Start(ctx context.Context) error {
 		// available, the client appears to have started even though it's completely
 		// non-functional. Here we try to make an initial assessment of health and
 		// return quickly in case of an apparent problem.
-		if err := c.driver.GetExecutor().Exec(fetchCtx, "SELECT 1"); err != nil {
+		if err := c.driver.GetExecutor().Ping(fetchCtx); err != nil {
 			return fmt.Errorf("error making initial connection to database: %w", err)
 		}
 

client_test.go

@@ -49,6 +49,36 @@ import (
 	"github.com/riverqueue/river/rivertype"
 )
 
+type clientInitialPingTestDriver struct {
+	riverdriver.Driver[pgx.Tx]
+
+	executor riverdriver.Executor
+}
+
+func (d *clientInitialPingTestDriver) GetExecutor() riverdriver.Executor { return d.executor }
+
+type clientInitialPingTestExecutor struct {
+	riverdriver.Executor
+
+	pingFunc func(context.Context) error
+}
+
+func (e *clientInitialPingTestExecutor) Ping(ctx context.Context) error { return e.pingFunc(ctx) }
+
+type clientInitialPingTestSignals struct {
+	PingContinue testsignal.TestSignal[struct{}]
+	PingFinished testsignal.TestSignal[error]
+	PingStarted  testsignal.TestSignal[struct{}]
+}
+
+func (s *clientInitialPingTestSignals) Init(tb testing.TB) {
+	tb.Helper()
+
+	s.PingContinue.Init(tb)
+	s.PingFinished.Init(tb)
+	s.PingStarted.Init(tb)
+}
+
 type clientJobCancelTestSignals struct {
 	ContinueWork testsignal.TestSignal[struct{}]
 	JobStarted   testsignal.TestSignal[int64]
@@ -8491,6 +8521,71 @@ func Test_NewClient_Defaults(t *testing.T) {
 	require.IsType(t, &baseservice.UnStubbableTimeGenerator{}, client.baseService.Time)
 }
 
+func Test_NewClient_InitialPing(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	type testBundle struct {
+		client  *Client[pgx.Tx]
+		signals *clientInitialPingTestSignals
+	}
+
+	setup := func(t *testing.T) *testBundle {
+		t.Helper()
+
+		var signals clientInitialPingTestSignals
+		signals.Init(t)
+
+		baseDriver := riverpgxv5.New(riversharedtest.DBPool(ctx, t))
+		schema := riverdbtest.TestSchema(ctx, t, baseDriver, nil)
+		baseExecutor := baseDriver.GetExecutor()
+		driver := &clientInitialPingTestDriver{
+			Driver: baseDriver,
+			executor: &clientInitialPingTestExecutor{
+				Executor: baseExecutor,
+				pingFunc: func(ctx context.Context) error {
+					signals.PingStarted.Signal(struct{}{})
+					signals.PingContinue.WaitOrTimeout()
+
+					err := baseExecutor.Ping(ctx)
+					signals.PingFinished.Signal(err)
+					return err
+				},
+			},
+		}
+
+		client, err := NewClient(driver, &Config{Schema: schema})
+		require.NoError(t, err)
+		t.Cleanup(func() { signals.PingContinue.Signal(struct{}{}) })
+
+		return &testBundle{client: client, signals: &signals}
+	}
+
+	t.Run("InsertCanResolveBeforePing", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+		bundle.signals.PingStarted.WaitOrTimeout()
+
+		_, err := bundle.client.Insert(ctx, noOpArgs{Name: "first"}, &InsertOpts{UniqueOpts: UniqueOpts{ByArgs: true}})
+		require.NoError(t, err)
+
+		bundle.signals.PingContinue.Signal(struct{}{})
+		require.NoError(t, bundle.signals.PingFinished.WaitOrTimeout())
+	})
+
+	t.Run("PingResolvesWithoutInsert", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+		bundle.signals.PingStarted.WaitOrTimeout()
+
+		bundle.signals.PingContinue.Signal(struct{}{})
+		require.NoError(t, bundle.signals.PingFinished.WaitOrTimeout())
+	})
+}
+
 func Test_NewClient_Overrides(t *testing.T) {
 	t.Parallel()
 

internal/rivercommon/river_common.go

@@ -45,11 +45,6 @@ const (
 
 	// MetadataKeyRescueCount records how many times the job has been rescued.
 	MetadataKeyRescueCount = "river:rescue_count"
-
-	// MetadataKeyUniqueNonce is a special metadata key used by the SQLite driver to
-	// determine whether an upsert is was skipped or not because the `(xmax != 0)`
-	// trick we use in Postgres doesn't work in SQLite.
-	MetadataKeyUniqueNonce = "river:unique_nonce"
 )
 
 type ContextKeyClient struct{}

riverdriver/river_driver_interface.go

@@ -289,6 +289,11 @@ type Executor interface {
 	NotificationDeleteBefore(ctx context.Context, params *NotificationDeleteBeforeParams) (int, error)
 
 	NotifyMany(ctx context.Context, params *NotifyManyParams) error
+
+	// Ping checks that the database is reachable and performs any driver-specific
+	// initialization that should happen before the first operation.
+	Ping(ctx context.Context) error
+
 	PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error)
 
 	QueueCreateOrSetUpdatedAt(ctx context.Context, params *QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error)

riverdriver/riverdatabasesql/river_database_sql_driver.go

@@ -16,6 +16,8 @@ import (
 	"io/fs"
 	"math"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -28,6 +30,7 @@ import (
 	"github.com/riverqueue/river/rivershared/uniquestates"
 	"github.com/riverqueue/river/rivershared/util/dbutil"
 	"github.com/riverqueue/river/rivershared/util/ptrutil"
+	"github.com/riverqueue/river/rivershared/util/randutil"
 	"github.com/riverqueue/river/rivershared/util/savepointutil"
 	"github.com/riverqueue/river/rivershared/util/sliceutil"
 	"github.com/riverqueue/river/rivertype"
@@ -38,9 +41,11 @@ var migrationFS embed.FS
 
 // Driver is an implementation of riverdriver.Driver for database/sql.
 type Driver struct {
-	dbPool         *sql.DB
-	listenerDriver *riverpgxv5.Driver
-	replacer       sqlctemplate.Replacer
+	dbPool                 *sql.DB
+	listenerDriver         *riverpgxv5.Driver
+	replacer               sqlctemplate.Replacer
+	uniqueInsertMode       atomic.Uint32
+	uniqueInsertModeInitMu sync.Mutex
 }
 
 // New returns a new database/sql River driver for use with River.
@@ -402,6 +407,16 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt
 }
 
 func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) {
+	uniqueInsertMode, err := e.uniqueInsertMode(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var uniqueNonce string
+	if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce {
+		uniqueNonce = randutil.Hex(8)
+	}
+
 	insertJobsParams := &dbsqlc.JobInsertFastManyParams{
 		ID:           make([]int64, len(params.Jobs)),
 		Args:         make([]string, len(params.Jobs)),
@@ -442,7 +457,16 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo
 		insertJobsParams.CreatedAt[i] = createdAt
 		insertJobsParams.Kind[i] = params.Kind
 		insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec
-		insertJobsParams.Metadata[i] = cmp.Or(string(params.Metadata), "{}")
+		metadata := []byte(cmp.Or(string(params.Metadata), "{}"))
+		if uniqueNonce != "" {
+			var err error
+			metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		insertJobsParams.Metadata[i] = string(metadata)
 		insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec
 		insertJobsParams.Queue[i] = params.Queue
 		insertJobsParams.ScheduledAt[i] = scheduledAt
@@ -452,6 +476,9 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo
 		insertJobsParams.UniqueStates[i] = int32(params.UniqueStates)
 	}
 
+	ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{
+		"unique_skipped_as_duplicate": {Value: uniqueInsertMode.SQL(), Stable: true},
+	}, nil)
 	items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams)
 	if err != nil {
 		return nil, interpretError(err)
@@ -462,7 +489,13 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo
 		if err != nil {
 			return nil, err
 		}
-		return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil
+
+		uniqueSkippedAsDuplicate := row.UniqueSkippedAsDuplicate
+		if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce {
+			uniqueSkippedAsDuplicate = riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce)
+		}
+
+		return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: uniqueSkippedAsDuplicate}, nil
 	})
 }
 
@@ -932,6 +965,11 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan
 	})
 }
 
+func (e *Executor) Ping(ctx context.Context) error {
+	_, err := e.uniqueInsertMode(ctx)
+	return err
+}
+
 func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) {
 	err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key)
 	return &struct{}{}, interpretError(err)
@@ -1090,6 +1128,32 @@ func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableT
 	return interpretError(err)
 }
 
+func (e *Executor) uniqueInsertMode(ctx context.Context) (riverdriver.UniqueInsertMode, error) {
+	if e.driver != nil {
+		if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown {
+			return mode, nil
+		}
+
+		e.driver.uniqueInsertModeInitMu.Lock()
+		defer e.driver.uniqueInsertModeInitMu.Unlock()
+
+		if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown {
+			return mode, nil
+		}
+	}
+
+	productAndVersion, err := dbsqlc.New().PGGetProductAndVersion(ctx, e.dbtx)
+	if err != nil {
+		return riverdriver.UniqueInsertModeUnknown, interpretError(err)
+	}
+
+	mode := riverdriver.UniqueInsertModeFromProductAndVersion(productAndVersion.Product, productAndVersion.VersionNum)
+	if e.driver != nil {
+		e.driver.uniqueInsertMode.Store(uint32(mode))
+	}
+	return mode, nil
+}
+
 type ExecutorTx struct {
 	Executor
 

riverdriver/riverdrivertest/job_insert.go

@@ -88,7 +88,7 @@ func exerciseJobInsert[TTx any](ctx context.Context, t *testing.T,
 				// SQLite needs to set a special metadata key to be able to
 				// check for duplicates. Remove this for purposes of comparing
 				// inserted metadata.
-				job.Metadata, err = sjson.DeleteBytes(job.Metadata, rivercommon.MetadataKeyUniqueNonce)
+				job.Metadata, err = sjson.DeleteBytes(job.Metadata, riverdriver.UniqueInsertMetadataKey)
 				require.NoError(t, err)
 
 				require.Equal(t, idStart+int64(i), job.ID)

riverdriver/riverdrivertest/job_read.go

@@ -55,7 +55,7 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx
 
 			for _, state := range rivertype.JobStates() {
 				require.Contains(t, countsByState, state)
-				switch state { //nolint:exhaustive
+				switch state {
 				case rivertype.JobStateAvailable:
 					require.Equal(t, 2, countsByState[state])
 				case rivertype.JobStateCancelled:
@@ -64,8 +64,10 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx
 					require.Equal(t, 1, countsByState[state])
 				case rivertype.JobStateDiscarded:
 					require.Equal(t, 1, countsByState[state])
-				default:
+				case rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateRunning, rivertype.JobStateScheduled:
 					require.Equal(t, 0, countsByState[state])
+				default:
+					require.FailNow(t, "unknown job state", state)
 				}
 			}
 		})

riverdriver/riverdrivertest/riverdrivertest.go

@@ -54,6 +54,13 @@ func exerciseDriverPool[TTx any](ctx context.Context, t *testing.T,
 ) {
 	t.Helper()
 
+	t.Run("Ping", func(t *testing.T) {
+		t.Parallel()
+
+		exec, _ := executorWithTx(ctx, t)
+		require.NoError(t, exec.Ping(ctx))
+	})
+
 	t.Run("PoolIsSet", func(t *testing.T) {
 		t.Parallel()