rivertest.Worker and workers that write to SQLite database

[markdboyd]

I'm loving river. It's such a straightforward and robust library.

I'm trying to write some tests for a worker using the rivertest.Worker helpers. In my particular case, my worker logic also needs to write to the database. However, when I try to run my test, it just hangs.

For context, I am using SQLite as the database for my tests. I am also using gorm as an ORM. And I know from experience with SQLite that when a transaction is open and you try to perform other database operations, the database connection will just hang, which explains what I'm seeing.

From what I understand, rivertest.Worker keeps the transaction open indefinitely, as indicated in this warning message in the documentation:

The Work function does not automatically roll back the transaction, nor does it commit it. The transaction is assumed to be owned by the caller, so it is the caller's responsibility to roll back after each individual job or batch of jobs tested.

Thus, my question is: is it possible to use rivertest.Worker to test a worker's Work function that writes to the database, particularly when using SQLite as the database?

I strongly suspect that I am doing something wrong or misunderstanding something, so hopefully someone can point me in the right direction.

Here's some pseudo-code showing what I'm trying to do in case it helps:

type Instance struct {
       Uuid string
}

type CreateArgs struct {
	Instance *Instance
}

func (CreateArgs) Kind() string { return "create" }

type CreateWorker struct {
	river.WorkerDefaults[CreateArgs]
	db                   *gorm.DB
}

func (w *CreateWorker) Work(ctx context.Context, job *river.Job[CreateArgs]) error {
	err := w.db.Save(job.Args.Instance).Error
	return err
}

func TestCreateWorker(t *testing.T) {
        // logic to initialize database as `db` and `riverClient` excluded
	tx := db.Begin()
	if err := tx.Error; err != nil {
		t.Fatal(err)
	}
	defer tx.Rollback()

	sqlTx := tx.Statement.ConnPool.(*sql.Tx)

	testWorker := rivertest.NewWorker(t, riversqlite.New(nil), &river.Config{}, &CreateWorker{
           db: db,
         })

	result, err := testWorker.Work(test.ctx, t, sqlTx, CreateArgs{
		Instance: &Instance{Uuid: "fake-uuid"},
	}, nil)
	if err != nil {
		t.Fatal(err)
	}
}

[brandur]

Hm, that looks right (I think?). It might be work adding some logging in various places just to make sure the transaction is indeed being rolled back as expected (i.e. putting some in the defer too).

Also, try running the test case on its own (i.e. go test . -run TestCreateWorker) just to make sure nothing else is conflicting with it.

[markdboyd]

@brandur What kind of logging were you thinking I should add?

Also, I have confirmed that if I remove the other database operations inside the Work function, then the test completes fine. The problem is probably specific to SQLite, but since a transaction is opened before the Work function and still open during the execution of the Work function, any database operations inside Work are blocked from proceeding, causing the tests to hang.

I also confirmed that the problem persists even if I run the test in isolation.

Unfortunately, given what I understand of SQLite and the behavior of rivertest.NewWorker, I don't think this is resolvable? The only resolution would be if the transaction for inserting the job were committed after inserting the job but before Work is invoked.

My understanding of the code and transaction may be faulty, but in rivertest.Worker.workJob, it seems like the transaction is no longer used after https://github.com/riverqueue/river/blob/master/rivertest/worker.go#L136, so it could probably be committed at that point? However, even if that could work, it seems like it would be a departure from the design philosophy of the testing tools.

My other option is that I can test the worker directly rather than using the rivertest helpers, but that doesn't allow me to test and to verify the actual worker results. I suppose I can also do some kind of mocking for the database operations that occur inside the Work function as well.

Fortunately, I'm only using SQLite for my tests. I think in production with PostgreSQL, this would not occur, because there are multiple database connections (vs just 1 for SQLite) and its handling of transactions is different anyway.

[brandur]

Ah! Okay, I see what's going on here.

Yeah, since it's SQLite, you won't be able to do another DB operation from the worker unless it's sharing the transaction send to the test worker.

Below, I've copy/pasted a full example that shows how you might accomplish this by putting a transaction in your context. So in your test you'd do something like:

	// Inject the transaction into context so the worker can retrieve it.
	ctx = context.WithValue(ctx, txContextKey{}, tx)

Then in a worker:

func (w *RecordingWorker) Work(ctx context.Context, job *river.Job[RecordingArgs]) error {
	tx := ctx.Value(txContextKey{}).(pgx.Tx)

	_, err := tx.Exec(ctx, "INSERT INTO rivertest_example_audit_log (job_kind, message) VALUES ($1, $2)", job.Kind, job.Args.Message)
	return err
}

You could abstract this a little more to have a helper like:

func (w *RecordingWorker) Work(ctx context.Context, job *river.Job[RecordingArgs]) error {
	tx := TxFromContextOrBegin(ctx)

That way, the worker will do reasonably well whether or not it's being run from a test or real life.

I kind of think that we should also maybe have a non-transactional variant of rivertest.Worker.Work for cases like this that'd make use a little easier. I'll float that idea.

Here's the full example:

package rivertest_test

import (
	"context"
	"fmt"
	"testing"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/riverqueue/river"
	"github.com/riverqueue/river/riverdbtest"
	"github.com/riverqueue/river/riverdriver/riverpgxv5"
	"github.com/riverqueue/river/rivershared/riversharedtest"
	"github.com/riverqueue/river/rivershared/util/testutil"
	"github.com/riverqueue/river/rivertest"
)

type txContextKey struct{}

type RecordingArgs struct {
	Message string `json:"message"`
}

func (RecordingArgs) Kind() string { return "recording" }

type RecordingWorker struct {
	river.WorkerDefaults[RecordingArgs]
}

func (w *RecordingWorker) Work(ctx context.Context, job *river.Job[RecordingArgs]) error {
	tx := ctx.Value(txContextKey{}).(pgx.Tx)

	_, err := tx.Exec(ctx, "INSERT INTO rivertest_example_audit_log (job_kind, message) VALUES ($1, $2)", job.Kind, job.Args.Message)
	return err
}

// Example_workerWithTx demonstrates how to inject a transaction into a worker
// via context, allowing the worker to perform additional database operations
// that are visible within the test and automatically rolled back afterward.
func Example_workerWithTx() {
	ctx := context.Background()

	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
	if err != nil {
		panic(err)
	}
	defer dbPool.Close()

	schema := riverdbtest.TestSchema(ctx, testutil.PanicTB(), riverpgxv5.New(dbPool), nil)

	// Start a transaction for the test. All operations—including the worker's
	// own database calls—will happen within this transaction.
	tx, err := dbPool.Begin(ctx)
	if err != nil {
		panic(err)
	}
	defer tx.Rollback(ctx)

	// Set the search path so all queries use the test schema.
	if _, err := tx.Exec(ctx, "SET search_path TO '"+schema+"'"); err != nil {
		panic(err)
	}

	// Create a scratch table that the worker will write to.
	if _, err := tx.Exec(ctx, "CREATE TABLE rivertest_example_audit_log (job_kind text, message text)"); err != nil {
		panic(err)
	}

	config := &river.Config{}

	// Inject the transaction into context so the worker can retrieve it.
	ctx = context.WithValue(ctx, txContextKey{}, tx)

	// Required for purposes of our example here, but in reality t will be the
	// *testing.T that comes from a test's argument.
	t := &testing.T{}

	testWorker := rivertest.NewWorker(t, riverpgxv5.New(nil), config, &RecordingWorker{})
	result, err := testWorker.Work(ctx, t, tx, RecordingArgs{Message: "Hello from worker!"}, nil)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Event kind: %s\n", result.EventKind)

	// Verify the worker's side effect is visible within the same transaction.
	var message string
	if err := tx.QueryRow(ctx, "SELECT message FROM rivertest_example_audit_log WHERE job_kind = 'recording'").Scan(&message); err != nil {
		panic(err)
	}
	fmt.Printf("Audit log message: %s\n", message)

	// Output:
	// Event kind: job_completed
	// Audit log message: Hello from worker!
}

[markdboyd]

@brandur

Thanks for the reply!

Your solution makes sense. But even with using the transaction from the context in my Work function, I'm still observing that the test is hanging.

I think the problem may be due to GORM at this point. My transaction initialization looks like:

			tx := db.Begin()
			if err := tx.Error; err != nil {
				t.Fatal(err)
			}
			defer tx.Rollback()

			sqlTx := tx.Statement.ConnPool.(*sql.Tx)

			ctx := context.WithValue(test.ctx, txContextKey{}, tx)

And inside my Work function, I have something like:

func (w *CreateWorker) Work(ctx context.Context, job *river.Job[CreateArgs]) error {
       tx := ctx.Value(txContextKey{}).(*gorm.DB)
	err := tx.Save(job.Args.Instance).Error
	return err
}

I think the problem may be that GORM's Save method manages a transaction itself, so it's not friendly to transaction re-use. I'll investigate and see if I can figure it out.

[brandur]

Cool, yeah let us know if you're able to resolve, but I think you're right that this is a GORM idiosyncracy.

I don't know GORM very well, but this line looks a little suspicious:

tx := ctx.Value(txContextKey{}).(*gorm.DB)

I would've expected to get a transaction-shaped object out of context rather than a (*gorm.DB).

[markdboyd]

@brandur

So I was wrong, your suggestion does work. I just needed to refactor my code to use the transaction from the context correctly.

I don't know GORM very well, but this line looks a little suspicious:

tx := ctx.Value(txContextKey{}).(*gorm.DB)

Yeah, the GORM method for beginning a transaction actually returns a *gorm.DB interface, not a *sql.Tx interface: https://github.com/go-gorm/gorm/blob/master/finisher_api.go#L665. This is why in GORM you have to do this to access the *sql.Tx interface that river expects in various places:

			tx := db.Begin() // where db is *gorm.DB
			if err := tx.Error; err != nil {
				t.Fatal(err)
			}
			defer tx.Rollback()

			sqlTx := tx.Statement.ConnPool.(*sql.Tx)

Regardless, if that *gorm.DB interface is returned from Begin(), then GORM will perform database operations on that connection as a transaction.

After all that, I am hesitating whether it's worthwhile to access the database connection from the context, given that the driver of the change is my testing environment versus my production environment. But I'm still grateful to know that it is possible to make this work.

Thanks for your help! I'll mark your response as the accepted answer.