Implementing Atomic Operations with Unit of Work

1. The Problem: Ensuring Atomic Operations

Imagine a use case like "reallocating an order." This might involve:

1. Loading a Batch aggregate.
2. Calling batch.deallocate(order_line) to remove an allocation.
3. Saving the Batch aggregate.
4. Finding a new Batch to allocate to.
5. Calling new_batch.allocate(order_line).
6. Saving the new_batch aggregate.

If the system crashes between steps 3 and 6, we have a deallocated order that isn't reallocated anywhere—a critical inconsistency. We need a way to ensure that either all these steps succeed, or none of them do. This is the principle of atomicity.

The Unit of Work pattern provides this guarantee. It collects all our changes in memory and commits them to the database in one single transaction at the very end.

2. The Unit of Work and Clean Architecture

In our architecture, where does the responsibility for transactions lie? It's not a domain concern—your users don't talk about "database transactions." It's an application-level concern related to orchestrating a use case.

This confirms our approach: the Application Service will manage the Unit of Work, which in turn handles the transaction. This keeps our domain model pure and our infrastructure concerns neatly separated.

The diagram below shows our target architecture. The Flask/FastAPI entrypoint will start a Unit of Work and pass it to the Application Service. The service then uses the UoW to get repositories and orchestrate the use case, committing everything at the end.

3. Implementing a Unit of Work in Python

While SQLAlchemy's Session object is itself an implementation of the Unit of Work pattern, we will create our own abstraction on top of it. This follows the "Don't mock what you don't own" principle and gives us a simpler, application-specific interface, decoupling us from the full complexity of SQLAlchemy.

We will build our UoW as a context manager. In Python, this is a construct used with the with statement to manage resources like file handles or database connections. It's similar in purpose to Java's try-with-resources, ensuring that cleanup logic (like closing a connection or rolling back a transaction) is always executed.

Let's dive into the implementation, following the excellent guide from the "Cosmic Python" book.

Let's break down the key components from the reading.

The Abstract Unit of Work

This defines our contract using an Abstract Base Class (ABC).

# src/allocation/service_layer/unit_of_work.py

import abc
from . import repository

class AbstractUnitOfWork(abc.ABC):
    batches: repository.AbstractRepository

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.rollback()

    @abc.abstractmethod
    def commit(self):
        raise NotImplementedError

    @abc.abstractmethod
    def rollback(self):
        raise NotImplementedError

• batches: The UoW will provide access to our repositories.
• __enter__: This is called when we enter a with block. It simply returns the UoW instance itself.
• __exit__: This is called when we exit the with block, either normally or due to an exception. Notice that it calls self.rollback() by default. This makes our system safe by default: if we don't explicitly commit, all changes are discarded.

The Concrete SQLAlchemy Unit of Work

This class implements the abstract interface using an SQLAlchemy session.

# src/allocation/service_layer/unit_of_work.py

class SqlAlchemyUnitOfWork(AbstractUnitOfWork):
    def __init__(self, session_factory):
        self.session_factory = session_factory

    def __enter__(self):
        self.session = self.session_factory()  # e.g., Session()
        self.batches = repository.SqlAlchemyRepository(self.session)
        return super().__enter__()

    def __exit__(self, *args):
        super().__exit__(*args)
        self.session.close()

    def commit(self):
        self.session.commit()

    def rollback(self):
        self.session.rollback()

• __init__: It's initialized with a session_factory, which is responsible for creating new database sessions. This allows us to use different databases for production (PostgreSQL) and testing (SQLite).
• __enter__: This is where the magic happens. It creates a new session and instantiates a concrete SqlAlchemyRepository, injecting the session into it. The UoW now acts as a factory for our repositories, ensuring they all share the same session and transaction.
• commit() and rollback(): These methods simply delegate the call to the underlying SQLAlchemy session.

4. Using the Unit of Work in an Application Service

Now we can refactor our Application Service to use our new UoW. The service will no longer depend on a repository directly; instead, it will depend on our AbstractUnitOfWork.

Here is the refactored allocate service:

# src/allocation/service_layer/services.py

def allocate(
    orderid: str, sku: str, qty: int,
    uow: unit_of_work.AbstractUnitOfWork,  # Dependency is now the UoW
) -> str:
    line = OrderLine(orderid, sku, qty)
    with uow:  # Start the transaction boundary
        # 1. Get the repository from the UoW
        batches = uow.batches.list()
        
        # (validation logic...)
        
        # 2. Call the domain model
        batchref = model.allocate(line, batches)
        
        # 3. Explicitly commit at the end
        uow.commit()
    return batchref

This new structure is much more robust:

1. The service depends on an abstraction (AbstractUnitOfWork), upholding the Dependency Inversion Principle.
2. The with uow: block clearly defines the transaction boundary.
3. We get our repository from uow.batches, ensuring it uses the correct session.
4. We must call uow.commit() to make our changes permanent.
5. If any exception occurs before uow.commit() (e.g., during model.allocate), the __exit__ method is triggered, uow.rollback() is called, and the database is left untouched.

This pattern makes handling complex operations that touch multiple aggregates, like the reallocate example from the reading, clean and safe.

Conclusion

Today, we've integrated a crucial pattern for building robust applications. The Unit of Work gives us confidence that our business operations will either complete fully or not at all, preventing data corruption and inconsistencies.

Key Takeaways:

• The Unit of Work (UoW) pattern groups a set of object changes into a single atomic transaction.
• It solves the problem of maintaining data consistency when a use case involves multiple database operations.
• In Python, the UoW is elegantly implemented as a context manager (with statement), providing a clear transaction boundary and ensuring rollbacks happen by default on error.
• The Application Service orchestrates the UoW. It receives a UoW, uses it to access repositories, calls domain logic, and finally commits the UoW.
• This approach creates a clean separation of concerns: the Application Service manages the transaction scope, while the UoW encapsulates the details of session management and commits/rollbacks.

Preview of the Next Lesson

We have now completed the core tactical patterns for building a single, cohesive Bounded Context. We can model our domain, persist it, and orchestrate use cases safely.

But what happens when an action in our domain needs to trigger a process in a different part of the system? For example, after a batch is successfully allocated, perhaps we need to notify the shipping department or send an email to the customer. Tightly coupling an EmailService into our allocate function would violate our architectural principles.

In the next module, we will introduce Domain Events, a powerful pattern for enabling communication between different parts of your system in a loosely coupled way.