DEV Community: Gabriel Lemire

Unit of Work pattern in Go

Gabriel Lemire — Thu, 14 Dec 2023 15:42:18 +0000

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.

Type-safe Data Access in Go using Prisma and sqlc

Gabriel Lemire — Tue, 05 Dec 2023 19:32:36 +0000

I recently started a new project in Go where I needed a database. It was an opportunity to select a stack I was happy with. In the past, I had used ORMs like gORM in existing projects.

Now, why not use an ORM? I've seen performance issues too many times with ORMs. I prefer writing my own SQL to avoid surprises. After all, I know the database schema and writing code for a specific purpose very often leads to better performance than generic code. ORMs have to support all kinds of database schemas. I only have to support mine. Having successfully used Knex.js in NodeJS (a popular query builder) in the past, I know writing SQL queries myself is not hard and provides very good performance.

I was browsing awesome-go for ideas on how to setup my data access layer when I stumbled on sqlc. It seemed like a great option. Code generation is a strategy often used in the Go ecosystem and making my queries safe at compile time was an idea I really liked. Knex was great, but it required of me that I test thoroughly my queries at runtime and that I sanitize my query results to ensure type safety within my application.

Here is how you write requests in sqlc. It's plain SQL with a comment that allows you to name your query and give a hint of how many results will be returned.

-- name: GetUserByID :one
SELECT * FROM users uWHERE u.id = sqlc.arg('UserID') AND u.deleted_at IS NULL LIMIT 1;

-- name: InsertUser :one
INSERT INTO users(
   id, name, organization_id,
   created_at, last_updated_at, deleted_at
)
VALUES (
    sqlc.arg('UserID'),
    sqlc.arg('Name'),
    sqlc.arg('OrganizationID'),
    sqlc.arg('CreatedAt'),
    sqlc.arg('LastUpdatedAt'),
    sqlc.arg('DeletedAt')
)
RETURNING *;

In order to provide compile-time safety and Code generation, sqlc requires a Database schema in the form of a single SQL file.

I could have maintained a separate Database schema file for sqlc apart from my migrations, but I felt it would remove some of the type-safety promise sqlc made. I would need to ensure this schema was synced with the real database schema at all times, and I didn't like this idea much.

I also though about appending all migration files together. It would not be too bad at the beginning, but I knew that with time, I would add new columns to tables, drop some of them and have a bunch of changes overall. This could potentially cause issues with sqlc.

I chose another avenue. Enter Prisma.

EDIT: Thank you Andrew for commenting. If you only want to use the migration files without Prisma, sqlc works with a migration folder.

Prisma to the rescue

Prisma is a popular NodeJS ORM. As I said earlier, I'm not a fan of ORMs, even though they can be useful, especially in CRUD scenarios. However, Prisma is not only an ORM, but also a migration tool. To use Prisma, you have to write a schema file. The format used is pretty readable in my opinion.

// This is your Prisma schema file,
// learn more about it in the docs: https://pris.ly/d/prisma-schema
datasource db {
  provider = "postgresql"
  url      = env("MIGRATOR_PG_URL")
}

// We don't care about the client, but it needs to be defined, so we export
// the client to a tmp folder we ignored using a .gitignore file.
generator client {
  provider = "prisma-client-js"
  output = "../tmp/prisma-client"
}

model User {
  id String @id @db.Uuid
  name String

  organization_id String @db.Uuid

  created_at DateTime @db.Timestamptz()
  last_updated_at DateTime? @db.Timestamptz()
  deleted_at DateTime? @db.Timestamptz()

  @@map("users")
}

This format has many advantages. It allows new developers to get familiar with the database schema very quickly because it serves as documentation. It also allows for the creation of new migrations by reconciliating the state of the database with the defined schema.

Prisma allows you to generate SQL migrations through its CLI. You can either provide a database connection so it can generate only what changed, or you can ask it to generate it from scratch. By using the last option, we can generate the schema sqlc requires to do its job.

npx prisma migrate diff --from-empty --to-schema-datamodel=./prisma/schema.prisma --script > ./schema.gen.sql

SQL Schema:

-- CreateTable
CREATE TABLE "users" (
    "id" UUID NOT NULL PRIMARY KEY,
    "name" TEXT NOT NULL,
    "organization_id" UUID NOT NULL,
    "created_at" TIMESTAMPTZ NOT NULL,
    "last_updated_at" TIMESTAMPTZ,
    "deleted_at" TIMESTAMPTZ
);

Then, in sqlc's configuration file, let's take in the generated schema as an input.

version: "2"
sql:
  - engine: "postgresql"
    queries: "sqlc/queries.sql"
    schema: "schema.gen.sql"
    gen:
      go:
        package: "db"
        out: "src/db"
        sql_package: "pgx/v5"
        emit_interface: true
        emit_pointers_for_null_types: true

        output_batch_file_name: "batch.gen.go"
        output_db_file_name: "db.gen.go"
        output_models_file_name: "models.gen.go"
        output_querier_file_name: "querier.gen.go"
        output_copyfrom_file_name: "copyfrom.gen.go"
        output_files_suffix: ".gen.go"

sqlc generate

In the configuration above, I added a .gen.go extension to all generated files. This allows me to ignore them using a .gitignore file. It minimizes the risk of me getting a stale version of those files.

As expected, sqlc generates all the relevant Data Access Objects. In the future, once my schema evolves, it will check my queries for breaking changes, minimizing the risk of discovering something broke in production.

The only problem remaining here is we rely on generated files to ensure this whole chain works. This can be annoying to maintain, but by using a Makefile, we can ensure your dependency chain is always up-to-date.

schema.gen.sql: prisma/schema.prisma ## Generates the SQL schema based on the Prisma schema
 @echo "[INFO] Generating SQL schema"
 @npx prisma migrate diff --from-empty --to-schema-datamodel=./prisma/schema.prisma --script > ./schema.gen.sql

sqlc: ## Generates the SQL code based on the SQLC config and schema
 @echo "[INFO] Installing sqlc"
 @go install github.com/sqlc-dev/sqlc/cmd/sqlc@v1.23.0

 @echo "[INFO] Generating SQL code"
 @sqlc generate
.PHONY: sqlc

Now, in order to generate our Data Access Objects, we can simply call make sqlc . However, we can even go one step further and call it from go generate by embedding the make command in a Go file.

package my_app

//go:generate make schema.gen.sql

As a bonus, sqlc is able to generate code for other languages besides Go. It now supports Kotlin, Python, F#, and there are plans to support C# and Typescript in the future.

EDIT: Typescript support has just launched in alpha. Check it out!

What do you think of this setup? Once in place, it feels great to use. Do you have a solution you like to implement for your Data Access Layer? Do you have a proposal to make this approach better? Let me know.