Unit of Work pattern in Go

I recently started a new Web project in Go. I'm pretty new to the Go ecosystem, so I'm still discovering the language and what can be done with it.

I was playing around the context object. I really like the context object because it is a very well implemented cancellation token, as well as a holder of many contextual data. I like using it to store data like the user's principal, the request ID and many other goodies that are very useful for logging and traceability.

One of the patterns I have not seen very much in Go is the Unit of Work pattern. Martin Fowler defines this pattern as follows:

Maintains a list of objects affected by a business transaction and coordinates the writing out of changes and the resolution of concurrency problems.

Implementing this pattern in most applications seems like a good idea. With this, you can never find yourself in a half-committed state. Either everything is saved in the database, or nothing is. This is a useful feature to have in a CRUD application. How would we achieve this in Go?

I could implement this pattern manually and ensure I commit or rollback myself when needed. However, I would like to make the pattern as easy to use as possible. What if its usage could be mostly transparent? After all, we know we can commit once we complete our HTTP request successfully, and we should revert when the request errors out.

Let me show you how I implemented it using a middleware and the context API.

Unit of Work Provider

We begin our implementation with the Unit of Work provider.

Let's define some requirements. Some requests do not require a transaction and some do. Most, if not all GET HTTP requests will not involve a transaction, as those mostly perform SELECT statements. We should be able to mark a request as requiring a transaction at the transport level.

The presence or the absence of transaction should be transparent. We should be able to use the same service function with or without transaction to make a query to the database. That's because some isolation levels allow us to read dirty data, allowing us to read our own uncommitted data. This can be required in some cases.

To handle those requirements, we create a Unit of Work provider. The provider is an object that lazily requests a connection for the pool when needed, initializes a transaction (if requested), commits or reverts a transaction (if requested) and frees up the resources at the end of a given scope.

Here is the implementation I went for:

package db

import (
    "context"
    "fmt"
    "github.com/jackc/pgx/v5"
    "github.com/jackc/pgx/v5/pgxpool"
    "sync"
    "sync/atomic"
)

type UnitOfWork struct {
    withTX   bool
    released *atomic.Bool

    initOnce *sync.Once

    conn *pgxpool.Conn
    tx   pgx.Tx


    // querier is the assigned sqlc querier instance
    querier Querier
}

// NewUnitOfWork builds a new UnitOfWork instance.
func NewUnitOfWork() *UnitOfWork {
    return &UnitOfWork{
        withTX:   false,
        released: &atomic.Bool{},

        initOnce: &sync.Once{},

        conn: nil,
        tx:   nil,

        querier: nil,
    }
}

// WithTX sets the UnitOfWork to use a transaction.
// Note: Call this method before calling GetQuerier() for the first time.
func (uow *UnitOfWork) WithTX() {
    uow.withTX = true
}

// TXRequested returns true if the UnitOfWork is configured to use a transaction.
func (uow *UnitOfWork) TXRequested() bool {
    return uow.withTX
}

// GetQuerier returns a Querier instance ready to make database queries.
// Note: This method is lazy-loaded. It will only initialize the database connection when called for the first time.
// All subsequent calls will return the same Querier instance.
func (uow *UnitOfWork) GetQuerier(ctx context.Context) (Querier, error) {
    if uow.released.Load() {
        return nil, fmt.Errorf("UnitOfWork has already been released")
    }

    var err error
    uow.initOnce.Do(func() {
        uow.conn, err = pgx_db.DBPool().Acquire(ctx)
        if err != nil {
            return
        }

        if uow.withTX {
            uow.tx, err = uow.conn.Begin(ctx)
            if err != nil {
                return
            }
            // We create a new instance of the querier object generated by sqlc
            uow.querier = New(uow.tx)
        } else {
            // We create a new instance of the querier object generated by sqlc
            uow.querier = New(uow.conn)
        }
    })

    if err != nil {
        return nil, err
    }

    if uow.querier == nil {
        return nil, fmt.Errorf("error initializing querier")
    }

    return uow.querier, nil
}

// Finalize commits or rolls back the transaction (if necessary) and releases the database connection.
func (uow *UnitOfWork) Finalize(ctx context.Context, isSuccess bool) error {
    if uow.released.Load() {
        return nil
    }

    // UnitOfWork is lazy-loaded. If not used, we don't have anything to do.
    if uow.conn == nil {
        return nil
    }

    defer func() {
        uow.conn.Release()
        uow.released.Store(true)
    }()

    if uow.tx != nil {
        // Just in case...
        defer func(tx pgx.Tx, ctx context.Context) {
            _ = tx.Rollback(ctx)
        }(uow.tx, ctx)

        if isSuccess {
            err := uow.tx.Commit(ctx)
            if err != nil {
                return err
            }
        } else {
            err := uow.tx.Rollback(ctx)
            if err != nil {
                return err
            }
        }
    }

    return nil
}

In my implementation, I decided to return an instance of the Querier in the GetQuerier method. However, feel free to modify this code. You can return any interface you feel makes sense in this case.

Using the Unit of Work Provider

Now that we have a Unit of Work Provider implementation, it's time to look at how we use it. I have a general rule of thumb. If a function creates the Unit of Work (via the NewUnitOfWork factory), it's their responsibility to finalize it (via the Finalize method).

Let's see how to use the Unit of Work Provider through an example.

// gin-handlers.go file
func CreateGizmosV1Handler(c *gin.Context) {
    uow := db.NewUnitOfWork()
    defer uow.Finalize(c, true)

    // Requesting our database queries be wrapped in a Transaction.
    uow.WithTX()

    g, err := gizmos.CreateGizmos(c, uow)
    if err != nil {
        uow.Finalize(c, false)
        _ = c.Error(err)
        return
    }

    c.JSON(http.StatusCreated, g)
}

// gizmos-service.go file
package gizmos

func CreateGizmos(ctx context.Context, uow *UnitOfWork) (*GizmoModel, error) {
    uow, finalize := db.GetSafeUnitOfWork(uow)
    defer finalize(ctx)

    querier, err := uow.GetQuerier(ctx)
    if err != nil {
      return nil, err
    }

    data, err := querier.InsertOneGizmo(ctx)
    if err != nil {
        return nil, err
    }

    return data, nil
}

// unit-of-work.go file
package db

func GetSafeUnitOfWork(uow *UnitOfWork) (*UnitOfWork, func(ctx context.Context)) {
    if uow != nil {
        return uow, func(ctx context.Context) {
            // We do nothing here because we are NOT the owner of the Unit of Work
        }
    }

    uow := db.NewUnitOfWork()
    return uow, func(ctx context.Context) {
        // Passing a success here is irrelevant because we do NOT use transactions.
        _ = uow.Finalize(ctx, true)

        // Log the error here. The risk of an error here is pretty low...
    } 
}

As of you can see above, we create the Unit of Work in the request handler. This means it is its responsibility to finalize it. We use a deferred function to perform this operation since we can call finalize multiple times on a Unit of Work without any problems.

We pass the Unit of Work to the Service layer which performs a GetSafeUnitOfWork function call. This ensures we always have a Unit of Work, even if the caller above passed a nil reference to us. This will allow us to chain multiple service calls using the same Unit of Work in the future. The GetSafeUnitOfWork method returns a finalize method we have to defer. In the case we already have a Unit of Work passed in, we don't do anything, because we are not the owner of the object. If no Unit of Work was passed, we create it and we are thus responsible for its cleanup.

Following that, the GetQuerier call is made and the service performs the database logic.

In the case of an error, the controller finalizes the Unit of Work with the success parameter set to false, which triggers a revert. Otherwise, the Unit of Work commits the data.

Automating the Unit of Work through the Context

There are two ways you can use the Unit of Work pattern we just created. You can create the Unit of Work manually like we did, ensuring you finalize it properly every time.

Another solution we can go for is to initialize the Unit of Work in a middleware and store the Unit of Work in the context.

Some developers don't like to assign any complex objects in the context. That's a fair take. Our Unit of Work implementation can accommodate for both a manual approach or an automated one. I like the context approach because I don't have to "infect" all functions with a Unit of Work. We already pass in a context everywhere, so why not use it for that?

// unit-of-work.middleware.go
package gin

func WithUnitOfWork(c *gin.Context) {
    unitOfWork := db.NewUnitOfWork()
    c.Set(db.UnitOfWorkCtxKey, unitOfWork)

        // We execute our API Handler
    c.Next()

    hasErrors := len(c.Errors) > 0

    err := unitOfWork.Finalize(c, !hasErrors)
    if err != nil {
        _ = c.Error(err)
        return
    }
}

// gin.go

// Using the Unit of Work middleware.
gin.Use(WithUnitOfWork)

// unit-of-work.go
package db

const UnitOfWorkCtxKey = "UnitOfWork"

func GetUnitOfWorkFromCtxOrDefault(ctx context.Context) (*UnitOfWork, func()) {
    unitOfWork := GetUnitOfWorkFromCtx(ctx)
    if unitOfWork == nil {
        uow := NewUnitOfWork()

        return uow, func() {
            _ = uow.Finalize(context.Background(), true)
        }
    }

    return unitOfWork, func() {
        // Do nothing. This is a Unit of Work not owned by this scope, so we don't want to release it.
    }
}

func GetUnitOfWorkFromCtx(ctx context.Context) *UnitOfWork {
    uow, ok := ctx.Value(UnitOfWorkCtxKey).(UnitOfWork)
    if !ok {
        return nil
    }

    return uow
}

The middleware is in charge of creating the Unit of Work, and therefore, it must finalize it. We can now remove the finalize logic from our request handler.

We also created a few functions to make interacting with the Unit of Work in the context a bit easier (GetUnitOfWorkFromCtx and GetUnitOfWorkFromCtxOrDefault).

With those modifications, we can change our service code from earlier to use the context instead.

// gin-handlers.go file
func CreateGizmosV1Handler(c *gin.Context) {
    // Requesting our database queries be wrapped in a Transaction.
    db.GetUnitOfWorkFromCtx(c).WithTX()

    // We no longer have to finalize the Unit of Work, because the middleware does it for us.

    g, err := gizmos.CreateGizmos(c)
    if err != nil {
        _ = c.Error(err)
        return
    }

    c.JSON(http.StatusCreated, g)
}

// gizmos-service.go file
package gizmos

func CreateGizmos(ctx context.Context) (*GizmoModel, error) {
    uow, finalize := db.GetUnitOfWorkFromCtxOrDefault(ctx)
    defer finalize(ctx)

    querier, err := uow.GetQuerier(ctx)
    if err != nil {
      return nil, err
    }

    data, err := querier.InsertOneGizmo(ctx)
    if err != nil {
        return nil, err
    }

    return data, nil
}

There we go. The handling of the Unit of Work is now automated. Remember the Unit of Work uses the database resources lazily. This means if you don't use it, you are not wasting it.

Future possible improvements

Currently, we assume you want the default Isolation Level for your database transaction. This is likely the case, but perhaps not always. It would be worthwhile to have a call in the Unit of Work to specify the Isolation Level that should be used.

We also do not have any pre and post hooks on the Unit of Work finalization. It could be useful if you need to involve another system in your Unit of Work (like a messaging system for instance).

Finally, this implementation allocates memory for each request. For most APIs, I'm convinced this won't make much difference. But for some high throughput APIs, having the Unit of Work being taken from a pool and reused could make a world of difference.

What did you think of this implementation? Is it something you are tempted to bring in your next project? Should I generalize the Unit of Work and expose it in a library? Do you have other improvements you think should be made? Let me know in the comments.