How to write a reusable DB transaction wrapper?

Question

I need to write a reusable function where part of it can be run before and after its invoked.

For example:

isNewTnx, rollback, beginErr := DbConnectionManager.Begin(req)

if beginErr != nil {
    return nil, beginErr
}

if isNewTnx {
    defer rollback()
}

Above code starts a transaction and if its a new transaction it controls it, meaning it can either commit or rollback. I'm forced to write this code in a lot of services. Any way to make this more reusable? I know function currying is an option, any other alternatives?

Answer 1

The following code uses types from github.com/jmoiron/sqlx as example. Given a data structure that wraps your database driver (mainly for interface implementation):

type Datastore struct {
    *sqlx.DB
}

You declare a method on it:

func (ds *Datastore) InTransaction(ctx context.Context, fn func(*sqlx.Tx) error) error {
    tx, err := ds.BeginTxx(ctx, nil)
    if err != nil {
        return fmt.Errorf("open tx error: %w", err)
    }
    defer tx.Rollback()

    if err := fn(tx); err != nil {
        return err
    }

    if err := tx.Commit(); err != nil {
        return fmt.Errorf("commit tx error: %w", err)
    }
    return nil
}

Here the trick is to always rollback with defer tx.Rollback(). If the transaction was successfully completed, rolling back does nothing. If an error occurred, it rolls back the transaction.

And then you use is as:

// ds := &Datastore{sqlx.MustConnect(.....)}

err := ds.InTransaction(ctx, func(tx *sqlx.Tx) error {
    // actual application code where you run SQL statements using `tx`
})

---

For functions that need to return something along with a possible error, you can use a generic wrapper around InTransaction function. Unfortunately this must be a top-level function that takes the datastore struct (or whatever appropriate interface) because methods can't have type parameters that weren't already defined on the receiver type:

// InTransaction is an adapter function to call InTransaction with an arbitrarily typed return pair.
func InTransaction[T any](ds *Datastore, ctx context.Context, fn func(tx *sqlx.Tx) (T, error)) (T, error) {
    // this is the possibly empty typed value we want to return
    var t T
    // this err is either the error produced by fn 
    // or the error produced by InTransaction itself, such as begin error, commit error, ...
    err := e.InTransaction(ctx, func(tx *sqlx.Tx) error {
        result, err := fn(tx)
        if err != nil {
            return err
        }
        // propagate the typed value from fn out of InTransaction
        t = result
        return nil
    })
    // return t — possibly empty depending on what fn returned or whether it was called at all — and err
    return t, err
}

And then you use it as:

// ds := &Datastore{sqlx.MustConnect(.....)}

myType, err := InTransaction(ds, ctx, func(tx *sqlx.Tx) (*MyType, error) {
    // actual application code where you run SQL statements using `tx`
    // and may return an arbitrary type and error
})