Public way to supply the client to `JobCompleteTx` when completing in a different client’s transaction type

[pckilgore]

We run two clients against one database, and JobCompleteTx currently can’t be used cleanly in that topology without importing rivertest into production code.

Setup. Our application transactions come from a database/sql-based ORM pool, so transactional job inserts and JobCompleteTx must run on a Client[*sql.Tx] (riverprodatabasesql — we’re a Pro customer).

But a database/sql driver can’t LISTEN, which means remote cancellation of running jobs never reaches workers. The fix was to split: keep the *sql.Tx client for inserts/completion, and run workers on a riverpropgxv5 client so producers get a notifier. This seems like it will work well. The cancel’s pg_notify is emitted by the SQL itself, so any client can send and the pgx workers receive.

Problem. JobCompleteTx[TDriver] resolves its client with ClientFromContext[TTx](ctx). Inside a worker running under the pgx client, the context carries a *Client[pgx.Tx], so the *sql.Tx-typed lookup fails, even though we have a perfectly good *sql.Tx client to complete through.

The only public way to fix the context is rivertest.WorkContext.

// production code today:
ctx = rivertest.WorkContext(ctx, sqlTxClient)
_, err := river.JobCompleteTx[*riverprodatabasesql.Driver](ctx, tx, job)

Ask

Either of these would remove the rivertest import from production code:

1. Export the context setter (e.g. river.WithClient[TTx](ctx, client)), or
2. A JobCompleteTx variant that accepts the client explicitly, e.g. client.JobCompleteTx(ctx, tx, job) or river.JobCompleteTxClient(ctx, client, tx, job).

But very flexible here if you advise something else.

Observed on v0.39.0; the relevant code is unchanged on master.

[CyJaySong]

Here is my solution. My English is not good, so I used AI to read my project and sorted out the following content

---

Running River's JobCompleteTx Inside a Business Transaction While Keeping PostgreSQL LISTEN/NOTIFY

One-sentence summary

Use database/sql + the pgx stdlib driver for business transactions, use the native riverpgxv5 driver for the consumer to keep LISTEN/NOTIFY alive, and bridge the gap by extracting pgx.Tx from *sql.Tx.

---

1. The problem

Many Go projects access the database through database/sql (or an ORM on top of it). River's consumer side, however, is best used with riverpgxv5, which is built on pgxpool.Pool, because only the pgx driver supports PostgreSQL's LISTEN/NOTIFY real-time notifications.

Inside a worker, you have to do two things at once:

1. Business data writes (inside a *sql.Tx transaction).
2. Mark the job as completed (river.JobCompleteTx requires a pgx.Tx).

Both must happen in the same transaction. Otherwise, you can end up with:

• business data committed but the job not completed, or
• the job completed but business data rolled back.

But river.JobCompleteTx pulls the client driver from the worker context. Since the consumer is riverpgxv5, it expects a pgx.Tx, not *sql.Tx.

So the core question becomes:

How do you turn a *sql.Tx into a pgx.Tx?

---

2. Splitting the architecture: two drivers

┌─────────────────────────────────────────────────────────────────┐
│  Producer: riverdatabasesql (based on *sql.DB / *sql.Tx)       │
│  - Inserts jobs                                                 │
│  - Can send pg_notify inside a transaction                      │
│  - Cannot LISTEN (cannot be the consumer)                       │
└──────────────────┬──────────────────────────────────────────────┘
                   │ same transaction: business writes + inserts + pg_notify
                   ▼
┌─────────────────────────────────────────────────────────────────┐
│  Consumer: riverpgxv5 (based on pgxpool.Pool)                  │
│  - LISTENs on NOTIFY                                            │
│  - Fetches new jobs after notification                          │
│  - JobCompleteTx needs pgx.Tx                                   │
└─────────────────────────────────────────────────────────────────┘

Role	Driver	Transaction type	LISTEN support
Producer	riverdatabasesql	*sql.Tx / *sql.DB	Cannot listen, but can pg_notify
Consumer	riverpgxv5	pgx.Tx / pgxpool.Pool	Supports LISTEN/NOTIFY

riverdatabasesql.SupportsListener() = false, SupportsListenNotify() = true: it cannot listen, but it can call pg_notify(...).

riverpgxv5.SupportsListener() = true: it holds a dedicated pgx connection, runs LISTEN "river.river_insert", and calls pgx.Conn.WaitForNotification to receive events.

---

3. Bridging *sql.Tx to pgx.Tx

The key assumption is: database/sql must use the pgx stdlib driver (sql.Open("pgx", dsn)).

When you begin a transaction through database/sql with the pgx stdlib driver, the underlying driver.Tx is a wrapTx that looks like this:

type wrapTx struct {
    ctx context.Context
    tx  pgx.Tx
}

func (wtx wrapTx) Commit() error   { return wtx.tx.Commit(wtx.ctx) }
func (wtx wrapTx) Rollback() error { return wtx.tx.Rollback(wtx.ctx) }

So wrapTx already stores the pgx.Tx. GetTxFromDriverTx is just a type assertion:

func GetTxFromDriverTx(driverTx driver.Tx) (pgx.Tx, bool) {
    if wtx, ok := driverTx.(wrapTx); ok {
        return wtx.tx, true
    }
    return nil, false
}

Full extraction pipeline

sql.Tx has an unexported field txi of type driver.Tx. The full pipeline is:

1. Use reflection to find the offset of txi inside sql.Tx.
2. Use unsafe to read the driver.Tx value at that offset.
3. Assert the driver.Tx as wrapTx and return the inner pgx.Tx.

package pgsql

import (
    "database/sql"
    "database/sql/driver"
    "fmt"
    "log"
    "reflect"
    "unsafe"

    "github.com/jackc/pgx/v5"
    "github.com/jackc/pgx/v5/stdlib"
)

var sqlTxTxiOffset uintptr

func init() {
    txType := reflect.TypeOf(sql.Tx{})
    txiField, found := txType.FieldByName("txi")
    if !found {
        log.Fatal("pgsql: sql.Tx.txi field not found, Go version may be incompatible")
    }
    sqlTxTxiOffset = txiField.Offset
}

func getDriverTx(sqlTx *sql.Tx) driver.Tx {
    if sqlTx == nil {
        return nil
    }

    txPtr := unsafe.Pointer(sqlTx)
    txiPtr := unsafe.Pointer(uintptr(txPtr) + sqlTxTxiOffset)
    return *(*driver.Tx)(txiPtr)
}

func GetPgxTx(sqlTx *sql.Tx) pgx.Tx {
    driverTx := getDriverTx(sqlTx)
    if driverTx == nil {
        return nil
    }

    pgxTx, ok := stdlib.GetTxFromDriverTx(driverTx)
    if !ok {
        return nil
    }
    return pgxTx
}

Then call JobCompleteTx like this:

pgxTx := GetPgxTx(sqlTx)
if pgxTx == nil {
    return errors.New("failed to extract pgx.Tx from sql.Tx")
}

_, err := river.JobCompleteTx[*riverpgxv5.Driver](ctx, pgxTx, job)

Note on GetTxFromDriverTx availability

wrapTx is an unexported type inside pgx stdlib. Different versions or forks may expose GetTxFromDriverTx directly, or they may not. If your version does not expose it, you can add the type assertion yourself. The core idea remains the same: find the driver.Tx implementation that wraps pgx.Tx and pull it out.

---

4. Typical worker transaction flow

func (w *MyWorker) Work(ctx context.Context, job *river.Job[MyJobArgs]) error {
    tx, err := db.BeginTx(ctx, nil) // *sql.Tx, business transaction
    if err != nil {
        return err
    }
    defer tx.Rollback()

    // 1. Business data write
    if err := businessRepo.DoSomething(ctx, tx); err != nil {
        return err
    }

    // 2. Optional: insert downstream jobs in the same transaction
    if err := producerClient.InsertManyTx(ctx, tx, downstreamJobs); err != nil {
        return err
    }

    // 3. Convert *sql.Tx to pgx.Tx and complete the current job
    pgxTx, ok := GetPgxTx(tx)
    if !ok {
        return errors.New("failed to get pgx.Tx")
    }

    if _, err := river.JobCompleteTx[*riverpgxv5.Driver](ctx, pgxTx, job); err != nil {
        return err
    }

    return tx.Commit(ctx)
}

This gives you:

• business data writes, downstream job inserts, and current job completion all in one *sql.Tx;
• all-or-nothing semantics: any failure rolls back everything;
• JobCompleteTx uses the riverpgxv5 driver but operates on the same underlying pgx transaction, so no second transaction is opened.

---

5. How LISTEN/NOTIFY stays intact

Sending side

When the producer inserts a job inside a transaction, it executes:

SELECT pg_notify('river.river_insert', payload);

Because pg_notify is called inside a transaction, PostgreSQL delays the actual delivery until the transaction commits. The notification and the transaction commit are atomically coupled.

Receiving side

The consumer uses the riverpgxv5 driver, which holds a dedicated pgx connection that runs LISTEN. When it receives a NOTIFY, it immediately fetches new jobs.

Relationship to JobCompleteTx

JobCompleteTx does not itself send a NOTIFY. Its role is to make sure that "job completion" and "business changes / downstream job inserts" commit together.

If the worker also inserts downstream jobs in the same transaction, the pg_notify calls for those new jobs are also delayed until COMMIT. After the commit, the consumer receives the notifications and processes the new jobs. If the transaction rolls back, no notifications are sent and no downstream jobs are triggered.

---

6. What this solves

1. Atomicity: business data, job completion, and downstream job inserts all commit or roll back together.
2. Real-time delivery: the consumer still uses native pgx LISTEN/NOTIFY instead of falling back to polling.
3. Compatibility: the business layer keeps using database/sql + ORM; you don't have to migrate the whole data layer to raw pgx.
4. No dual writes: you don't have to commit the business transaction first and then separately call job.Complete(), which avoids inconsistency if the process crashes in between.

---

7. Why not just use riverdatabasesql for the consumer?

A natural question: since the business transaction is *sql.Tx, why not make the consumer riverdatabasesql and pass *sql.Tx directly to JobCompleteTx?

The answer is LISTEN/NOTIFY:

• riverdatabasesql.SupportsListener() = false, so the consumer would have to poll;
• riverpgxv5.SupportsListener() = true, so the consumer can receive real-time notifications.

Therefore the consumer must stay on riverpgxv5, and riverpgxv5.UnwrapExecutor(tx) requires pgx.Tx. That is the fundamental reason for converting *sql.Tx back into pgx.Tx.

---

8. Risks and caveats

sql.Tx.txi is unexported

The Go standard library does not guarantee the name, order, or type of sql.Tx fields. You should:

• locate the txi offset dynamically at runtime with reflection;
• validate the extraction at startup with a real transaction.

func ValidatePgxTxExtractionWithTx(sqlTx *sql.Tx) error {
    if sqlTx == nil {
        return fmt.Errorf("sqlTx is nil")
    }

    pgxTx := GetPgxTx(sqlTx)
    if pgxTx == nil {
        return fmt.Errorf("failed to extract pgx.Tx from sql.Tx")
    }

    if reflect.TypeOf(pgxTx) == nil {
        return fmt.Errorf("extracted pgx.Tx type is nil")
    }

    return nil
}

The driver must be pgx stdlib

If database/sql uses lib/pq or another driver, the driver.Tx implementation will not be wrapTx, and GetTxFromDriverTx will fail.

wrapTx is unexported

Different versions or forks of pgx may change wrapTx. Some expose GetTxFromDriverTx; some do not. Verify the pgx stdlib implementation in your version before relying on it.

Avoid over-relying on unsafe

Reading sql.Tx.txi with unsafe is lightweight but exploits a standard-library gap. Use it only within a controlled Go version range, and validate it in CI with a startup check.

---

9. Summary

The core idea of this approach is:

Use the same underlying pgx connection/transaction, but expose it through two different abstractions.

• The business ORM sees database/sql.
• The River consumer sees pgx.
• By extracting driver.Tx from sql.Tx via unsafe, and then asserting it as pgx stdlib's wrapTx to retrieve the inner pgx.Tx, both sides operate on the same transaction.

This keeps LISTEN/NOTIFY real-time delivery alive while making business data changes and job state changes atomically commit inside the same transaction.

---

[bgentry]

Hi @pckilgore, thanks for the detailed explanation and sorry for the trouble. We're looking into a few options here, one of which is making the database/sql driver support LISTEN via an embedded pgx pool that's only used for listeners. Do you think this would potentially make it so you don't need to balance multiple types of clients with different drivers?

[pckilgore]

Big fan of #1366, that the right solution IMO.

[pckilgore]

Or give me an interface there I can back with pgx