Update River test helpers to allow for River drivers

[brandur]

We recently added River database drivers that allow it to be more
agnostic about the database package in use with River rather than being
tightly tied to pgx v5. The test helpers were also quite a recent
addition and when the drivers came in weren't updated to use a similar
system, thereby still being coupled to pgx v5.

The purpose of this change is to get test helpers onto the drivers
system as well. Like with top-level River, it's mostly about getting the
right APIs in place -- we're still tightly coupled to pgx v5 under the
covers, and will have to refactor all of that at a later time.

The old API doesn't quite work anymore, so we bifurcate the helpers into
something more akin to what we see in river.Client, with a transaction
and non-transaction variant. Furthermore, because these are top-level
package helpers, we need to take a driver generic (at least for the
transaction-based variants), making them look like JobCompleteTx:

• RequireInserted(ctx, t, riverpgxv5.New(dbPool), &RequiredArgs{}, nil)
• RequireInsertedTx[*riverpgxv5.Driver](ctx, t, tx, &RequiredArgs{}, nil)
• RequireManyInserted(ctx, t, riverpgxv5.New(dbPool), []rivertest.ExpectedJob{...})
• RequireManyInsertedTx[*riverpgxv5.Driver](ctx, t, tx, []rivertest.ExpectedJob{...})

One approach I experimented with was to do something similar to the
require package by allowing a struct to be created that'd then give
you a shortcut without type parameters necessary (kind of like what
Client does currently):

tRiver := rivertest.New(riverpgxv5.New(dbPool))
tRiver.RequireInsert(ctx, t, &RequiredArgs{}, nil)
tRiver.RequireInsertTx(ctx, t, tx, &RequiredArgs{}, nil)

The problem I ran into there is similar to the top-level insert helpers
we ended up removing from river -- namely that Go has the limitation
that struct-level functions can't have their own type parameters, so
going with that model would no longer allow RequireInsert to return
Job[T] (it'd have to be JobRow instead) which is a non-starter right
now because we want people to be able to easily assert on the return
values of these functions.

I'm still not sure this is the package's perfect form, but is pretty
goo,d and I think this is about as good as we're going to be able to for
now, and is ~identical functionality-wise to what we have before. I'm
still open to continuing to experiment with args equality assertions
being built-in, but we should try to get a reasonably quick fix for the
missing driver problem in so we don't have anyone accidentally start
using an API that's definitively outdated.

[bgentry]

I’m torn on both the verbosity of this and the two mostly redundant helpers (tx and not tx). I do however agree we should fix this now while we continue to think about how we can deliver a better UX on these helpers.

check out my comment on the rollback usage though

[bgentry]

Is this actually safe in pgx? From the docs:

Begin starts a transaction. Unlike database/sql, the context only affects the begin command. i.e. there is no auto-rollback on context cancellation.

Implies to me that if your context is cancelled or times out, the tx will be left open with the code here.

[brandur]

Hm, I'm not sure I totally get the concern in this one — isn't this the same deferred transaction roll back that we're already doing in ten other places?

The rollback might not happen on context cancel, but that might actually make sense. If a context were to be cancelled, it'd cancel other operations ongoing here (or elsewhere), then as the function was returning the defer would invoke roll back. If the roll back cancelled itself on context cancellation, it'd never be issued, because we're already cancelled by that point. So this allows the rollback to succeed normally, even in the event of a cancellation, therefore helping to clean up this resource.

Does that sound right to you?

[brandur]

(Going to merge this since I have a docs page incoming, but will continue discussion on this thread.)

[bgentry]

I'm not sure why I was linking to the Begin() docs while talking about Rollback 🤦‍♂️ However I guess I'm still a bit confused how the contexts work with begin + rollback even though I thought I understood this previously. Begin specifically says the context only affects the begin operation and not the rest of the transaction, whereas AFAICT Rollback does not say anything about how it behaves in response to a cancelled context.

If my context terminates (times out) midway through my transaction, does tx.Rollback(ctx) do anything? How does the transaction get aborted in that case?

[brandur]

Hmm, actually, by my read I think it does respect cancellation ...

You have to trace code all over the place, so I'm not entirely sure, but I think it does.

Rollback:

func (tx *dbTx) Rollback(ctx context.Context) error {
	if tx.closed {
		return ErrTxClosed
	}

	_, err := tx.conn.Exec(ctx, "rollback")
	tx.closed = true
	if err != nil {
		// A rollback failure leaves the connection in an undefined state
		tx.conn.die(fmt.Errorf("rollback failed: %w", err))
		return err
	}

	return nil
}

Exec calls into another exec:

func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) {
	if c.queryTracer != nil {
		ctx = c.queryTracer.TraceQueryStart(ctx, c, TraceQueryStartData{SQL: sql, Args: arguments})
	}

	if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil {
		return pgconn.CommandTag{}, err
	}

	commandTag, err := c.exec(ctx, sql, arguments...)

Which goes through half a doze subfunctions, but eventually ends up somewhere like execExtendedPrefix:

func (pgConn *PgConn) execExtendedPrefix(ctx context.Context, paramValues [][]byte) *ResultReader {
	pgConn.resultReader = ResultReader{
		pgConn: pgConn,
		ctx:    ctx,
	}
	result := &pgConn.resultReader

...

	if ctx != context.Background() {
		select {
		case <-ctx.Done():
			result.concludeCommand(CommandTag{}, newContextAlreadyDoneError(ctx))
			result.closed = true
			pgConn.unlock()
			return result

Which has a context check. You'd honestly have to check with Jack or run tests to be 100% sure, but it looks like it's checked.

I guess the way that the transaction gets rolled back then is implicitly, as the connection is killed the Postgres backend disposes of it automatically.