Implementing the Unit of Work pattern with multiple sources

February 24, 2024

The Unit of Work pattern is a design pattern that helps manage transactions and maintain data consistency in applications. More about the pattern can be found here.

In this article I will explain how the Unit of Work pattern can be implemented in a way that it supports multiple sources. This is useful when you have to manage transactions across different data sources, such as a database and an API.

We will start by explaining the problem that the Unit of Work pattern solves. Then we will provide an example implementation of the pattern.

Example use case

During the process of new account creation for a customer, the following steps are typically undertaken:

A new account record is added to the database using Entity Framework (EF).
An API request is made to an email service provider's API to dispatch an activation email.

However, there's a potential issue. If the database operation is successful but the email operation fails, the customer ends up in a problematic state. The account is not activated, so the customer cannot log in. Moreover, they cannot attempt to create the account again as the email address is already registered in the system.

To avoid this situation, we aim to implement a rollback mechanism. If any operation fails during the account creation process, all executed operations will be rolled back.

Example implementation

The following is a code implementation of the Unit of Work pattern that supports multiple sources. This solution is designed to prevent the situation described in the previous example.

Unit of work and enlistment contracts

public interface IUnitOfWork
{
    /// <summary>
    /// Return the enlistment of the given type
    /// </summary>
    T Enlistment<T>();

    /// <summary>
    /// Commit the enlistments
    /// </summary>
    Task CommitAsync();
}

/// <summary>
/// For each source there should be a different enlistment implementation
/// </summary>
public interface IUnitOfWorkEnlistment
{
    /// <summary>
    /// Commit the changes in a way that it can be rolled back when needed
    /// </summary>
    Task PreCommitAsync();

    /// <summary>
    /// Commit the changes. This method is called when all the pre-commits are successful.
    /// </summary>
    Task CommitAsync();

    /// <summary>
    /// Rollback the changes. This method is called when any of the pre-commits fail.
    /// </summary>
    Task RollbackAsync();
}

Unit of work implementation

public class UnitOfWork(params IUnitOfWorkEnlistment[] enlistments) : IUnitOfWork
{
    public T Enlistment<T>() => enlistments.OfType<T>().SingleOrDefault()
                                ?? throw new InvalidOperationException(
                                    $"Enlistment of type {typeof(T).Name} not found");

    /// <summary>
    /// One-by-one commit the enlistments:
    /// - If any fails to pre-commit, rollback all the already pre-committed enlistments.
    /// - If all pre-commits succeeed, call commit all the enlistments.
    /// </summary>
    public async Task CommitAsync()
    {
        var successfullyPreCommitted = new List<IUnitOfWorkEnlistment>();

        try
        {
            foreach (var enlistment in enlistments)
            {
                await enlistment.PreCommitAsync();
                successfullyPreCommitted.Add(enlistment);
            }
        }
        catch (Exception)
        {
            foreach (var needsRollbackEnlistment in successfullyPreCommitted)
            {
                await needsRollbackEnlistment.RollbackAsync();
            }

            throw;
        }

        foreach (var enlistment in enlistments)
        {
            await enlistment.CommitAsync();
        }
    }
}

EF enlistment

In the EF enlistment we create a new DbContext and start a new database transaction which we potentially can rollback when any of the enlistments pre-commit fails.

public interface IDbContext
{
    DbSet<Account> Accounts { get; }
}

public class DbUnitOfWorkEnlistment
    : IUnitOfWorkEnlistment, IDbContext
{
    private readonly ExampleDbContext _ctx;

    public DbSet<Account> Accounts => _ctx.Accounts;

    public DbUnitOfWorkEnlistment()
    {
        _ctx = new ExampleDbContext();
        _ctx.Database.BeginTransaction();
    }

    public async Task PreCommitAsync()
        => await _ctx.SaveChangesAsync();

    public async Task CommitAsync()
        => await _ctx.Database.CommitTransactionAsync();

    public Task RollbackAsync()
        => _ctx.Database.RollbackTransactionAsync();
}

E-mail API enlistment

For brevity, we are implementing a dummy e-mail API. For every operation we do, we're also describing how we can rollback that operation to go back to the initial state. In terms of the e-mail, rolling back could be cancelling of the pending e-mail.

public interface IEmailApi
{
    Task SendEmailAsync(string to, string body);
}

public class EmailApiUnitOfWorkEnlistment
    : IUnitOfWorkEnlistment, IEmailApi
{
    private readonly List<(Func<Task> preCommitAction, Func<Task> rollbackAction)> _operations = [];
    private readonly List<Func<Task>> _shouldRollback = [];

    public Task SendEmailAsync(string to, string body)
    {
        var commitAction = new Func<Task>(() =>
        {
            Console.WriteLine($"Calling API to send...");
            return Task.CompletedTask;
        });

        var rollbackAction = new Func<Task>(() =>
        {
            Console.WriteLine($"Calling API to cancel...");
            return Task.CompletedTask;
        });

        _operations.Add((commitAction, rollbackAction));

        return Task.CompletedTask;
    }

    public async Task PreCommitAsync()
    {
        foreach (var operation in _operations)
        {
            await operation.preCommitAction();
            _shouldRollback.Add(operation.rollbackAction);
        }
    }

    public async Task RollbackAsync()
    {
        foreach (var rollback in _shouldRollback)
        {
            await rollback();
        }
    }

    public Task CommitAsync() => Task.CompletedTask;
}

Bringing it all together

Below code could be part of your application- or business layer.

var account = new Account()
{
    Email = "[email protected]",
    Password = "123456",
    Code = Guid.NewGuid(),
};

var unitOfWork = new UnitOfWork(
    new DbUnitOfWorkEnlistment(),
    new EmailApiUnitOfWorkEnlistment());

await unitOfWork.Enlistment<IDbContext>()
    .Accounts.AddAsync(account);

await unitOfWork.Enlistment<IEmailApi>()
    .SendEmailAsync(account.Email, $"Code: {account.Code}");

await unitOfWork.CommitAsync();

Summary

I've explained in what way the Unit of Work pattern could prevent an unrepairable state of the application. In the code I gave an example implementation of the Unit of Work pattern which support multiple sources.

You can find the full code example here.