Context as Code
Documentation

Unit of Work Pattern

Transactional consistency across multiple repository operations using the Unit of Work pattern.

Show:

Status: Extracted
Purpose: Transaction management pattern for ensuring consistency across multiple repository operations.

Purpose

The Unit of Work pattern ensures transactional consistency across multiple repository operations, simplifies transaction management in command handlers, and provides automatic commit/rollback on success/failure.

Rationale

  • Transactional Consistency: Ensures all operations in a transaction succeed or fail together
  • Simplified Management: Simplifies transaction management in command handlers
  • Automatic Cleanup: Provides automatic commit/rollback on success/failure
  • Batching: Enables batching multiple database changes into a single transaction
  • Testability: Improves testability by allowing easy transaction rollback in tests

Implementation

Location: infrastructure/database/adapters/sqlalchemy/unit_of_work.py

Pattern:

class SQLAlchemyUnitOfWork:
    """SQLAlchemy-backed UnitOfWork implementation."""
    
    def __init__(
        self,
        session_factory: async_sessionmaker[AsyncSession],
        repositories: dict[str, IRepository],
    ):
        self._session_factory = session_factory
        self._repositories = repositories
        self.session: AsyncSession | None = None
        self._repo_instances: dict[str, IRepository] = {}
    
    @property
    def accounts(self) -> IAccountRepository:
        """Get account repository with session."""
        if self._repo_instances.get('accounts') is None:
            raise RuntimeError("Unit of work not entered")
        return self._repo_instances['accounts']
    
    async def __aenter__(self) -> "SQLAlchemyUnitOfWork":
        """Enter transaction context."""
        self.session = self._session_factory()
        await self.session.__aenter__()
        await self.session.begin()
        
        # Create repository instances with session binding
        for name, repo_class in self._repositories.items():
            self._repo_instances[name] = repo_class(self.session)
        
        return self
    
    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        """Exit transaction context."""
        if exc_type is None:
            await self.session.commit()
        else:
            await self.session.rollback()
        
        await self.session.__aexit__(exc_type, exc, tb)
        self.session = None
        self._repo_instances = {}

Usage Pattern

Command Handler Usage

Pattern:

async def handle_create_account(
    command: CreateAccountCommand,
    uow: IUnitOfWork,
) -> Account:
    """Create a new account."""
    async with uow:
        # All repository operations share the same session
        account = Account.create(...)
        uow.accounts.add(account)
        
        # Transaction commits automatically on success
        # Transaction rolls back automatically on exception
        return account

Benefits

  • Automatic Transaction Management: No manual commit/rollback needed
  • Consistent State: All operations succeed or fail together
  • Error Handling: Automatic rollback on exceptions
  • Clean Code: Simplifies command handler logic

Repository Integration

Session Binding

  • All repositories within a Unit of Work share the same session
  • Repository instances are created with session binding
  • Operations are performed within the same transaction

Repository Access

  • Access repositories via Unit of Work properties
  • Example: uow.accounts, uow.messages, uow.resources
  • Repositories are only available within Unit of Work context

Error Handling

Automatic Rollback

  • On Exception: Transaction automatically rolls back
  • On Success: Transaction automatically commits
  • No Manual Cleanup: Context manager handles cleanup

Exception Propagation

  • Exceptions propagate normally
  • Transaction rollback happens automatically
  • No need to catch exceptions for rollback

Testing

Test Pattern

async def test_create_account():
    """Test account creation with Unit of Work."""
    uow = SQLAlchemyUnitOfWork(session_factory, repositories)
    
    async with uow:
        account = Account.create(...)
        uow.accounts.add(account)
        # Transaction commits in real scenario
    
    # In test, transaction is rolled back automatically
    # Database state is clean for next test

Benefits

  • Isolated Tests: Each test gets a clean transaction
  • Automatic Cleanup: No manual cleanup needed
  • Consistent State: Tests don't affect each other

Port-Adapter Pattern

Interface (Port)

Location: application/unit_of_work.py

Pattern:

class IUnitOfWork(Protocol):
    """Unit of Work interface (tool-agnostic)."""
    
    @property
    def accounts(self) -> IAccountRepository:
        """Get account repository."""
        ...
    
    async def __aenter__(self) -> "IUnitOfWork":
        """Enter transaction context."""
        ...
    
    async def __aexit__(self, exc_type, exc, tb) -> None:
        """Exit transaction context."""
        ...

Implementation (Adapter)

Location: infrastructure/database/adapters/sqlalchemy/unit_of_work.py

  • SQLAlchemy-specific implementation
  • Can be swapped for other adapters (e.g., Django ORM, SQLAlchemy sync)
  • Application layer depends only on interface

Related Documents

  • Session Management - Session factory pattern
  • Repository Implementation - Repository patterns
  • Migration principles (all stacks) — Tool-agnostic migration rules
  • Migration patterns — Alembic — Python migrations
  • Migration patterns — Drizzle Kit — TypeScript migrations

Reference implementation: services/account-api/infrastructure/database/adapters/sqlalchemy/unit_of_work.py

Framework implementations

See also:

Repository Interface (Domain Layer)

Each aggregate has an abstract repository interface in domain/<feature>/repository.py. This is the port — no SQLAlchemy, no database, just Python ABCs.

# domain/account/repository.py
from abc import ABC, abstractmethod
from uuid import UUID
from domain.account.aggregate import Account
 
 
class AccountRepository(ABC):
    @abstractmethod
    async def save(self, session: object, account: Account) -> None: ...
 
    @abstractmethod
    async def find_by_id(self, session: object, account_id: UUID) -> Account | None: ...
 
    @abstractmethod
    async def find_by_email(self, session: object, email: str) -> Account | None: ...
 
    @abstractmethod
    async def delete(self, session: object, account_id: UUID) -> None: ...
 
    @abstractmethod
    async def find_all(
        self, session: object, limit: int = 50, offset: int = 0
    ) -> list[Account]: ...

session is typed as object here — the interface does not know whether the session is a SQLAlchemy AsyncSession or an in-memory dict. Concrete implementations know their session type.

Unit of Work Interface (Application Layer)

# application/unit_of_work.py
from abc import ABC, abstractmethod
from typing import Any
 
 
class UnitOfWork(ABC):
    @abstractmethod
    async def __aenter__(self) -> Any:
        """Returns the session object (type depends on adapter)."""
        ...
 
    @abstractmethod
    async def __aexit__(self, exc_type, exc, tb) -> None:
        """Commits on success, rolls back on exception."""
        ...

SQLAlchemy Implementations (Infrastructure Layer)

ORM Models (infrastructure/database/adapters/sqlalchemy/models/account.py)

from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String, DateTime
from uuid import UUID
import datetime
 
 
class Base(DeclarativeBase):
    pass
 
 
class AccountModel(Base):
    __tablename__ = "accounts"
 
    id: Mapped[UUID] = mapped_column(primary_key=True)
    email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    display_name: Mapped[str] = mapped_column(String(255))
    status: Mapped[str] = mapped_column(String(50), default="active")
    created_at: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))
    updated_at: Mapped[datetime.datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

Mappers (infrastructure/database/adapters/sqlalchemy/mappers/account_mapper.py)

from domain.account.aggregate import Account
from infrastructure.database.adapters.sqlalchemy.models.account import AccountModel
 
 
def account_to_orm(account: Account) -> AccountModel:
    return AccountModel(
        id=account.id,
        email=account.email,
        display_name=account.display_name,
        status=account.status.value,
        created_at=account.created_at,
        updated_at=account.updated_at,
    )
 
 
def account_to_domain(model: AccountModel) -> Account:
    return Account(
        id=model.id,
        email=model.email,
        display_name=model.display_name,
        status=AccountStatus(model.status),
        created_at=model.created_at,
        updated_at=model.updated_at,
    )

Repository Implementation (infrastructure/database/adapters/sqlalchemy/repositories/account_repository.py)

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from domain.account.aggregate import Account
from domain.account.repository import AccountRepository
from .mappers.account_mapper import account_to_orm, account_to_domain
 
 
class SQLAlchemyAccountRepository(AccountRepository):
    async def save(self, session: AsyncSession, account: Account) -> None:
        orm = account_to_orm(account)
        await session.merge(orm)  # insert or update
 
    async def find_by_id(self, session: AsyncSession, account_id) -> Account | None:
        result = await session.execute(
            select(AccountModel).where(AccountModel.id == account_id)
        )
        model = result.scalar_one_or_none()
        return account_to_domain(model) if model else None
 
    async def find_by_email(self, session: AsyncSession, email: str) -> Account | None:
        result = await session.execute(
            select(AccountModel).where(AccountModel.email == email)
        )
        model = result.scalar_one_or_none()
        return account_to_domain(model) if model else None
 
    async def delete(self, session: AsyncSession, account_id) -> None:
        model = await session.get(AccountModel, account_id)
        if model:
            await session.delete(model)
 
    async def find_all(self, session: AsyncSession, limit=50, offset=0) -> list[Account]:
        result = await session.execute(
            select(AccountModel).limit(limit).offset(offset)
        )
        return [account_to_domain(m) for m in result.scalars().all()]

UoW Implementation (infrastructure/database/adapters/sqlalchemy/unit_of_work.py)

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from application.unit_of_work import UnitOfWork
 
 
class SQLAlchemyUnitOfWork(UnitOfWork):
    def __init__(self, session_factory: async_sessionmaker):
        self._session_factory = session_factory
        self._session: AsyncSession | None = None
 
    async def __aenter__(self) -> AsyncSession:
        self._session = self._session_factory()
        return self._session
 
    async def __aexit__(self, exc_type, exc, tb) -> None:
        if exc:
            await self._session.rollback()
        else:
            await self._session.commit()
        await self._session.close()
        self._session = None

In-Memory Implementations (for Testing)

# infrastructure/database/adapters/inmemory/repositories/account_repository.py
from domain.account.repository import AccountRepository
from copy import deepcopy
 
 
class InMemoryAccountRepository(AccountRepository):
    async def save(self, session: dict, account: Account) -> None:
        session[str(account.id)] = deepcopy(account)
 
    async def find_by_id(self, session: dict, account_id) -> Account | None:
        return deepcopy(session.get(str(account_id)))
 
    async def find_by_email(self, session: dict, email: str) -> Account | None:
        return next(
            (deepcopy(a) for a in session.values() if a.email == email), None
        )
 
    async def delete(self, session: dict, account_id) -> None:
        session.pop(str(account_id), None)
 
    async def find_all(self, session: dict, limit=50, offset=0) -> list[Account]:
        return [deepcopy(a) for a in list(session.values())[offset:offset + limit]]
# infrastructure/database/adapters/inmemory/unit_of_work.py
from application.unit_of_work import UnitOfWork
 
 
class InMemoryUnitOfWork(UnitOfWork):
    def __init__(self, store: dict):
        self._store = store
 
    async def __aenter__(self) -> dict:
        return self._store  # no transactions — simple dict
 
    async def __aexit__(self, exc_type, exc, tb) -> None:
        pass  # nothing to commit or rollback

Factory (infrastructure/database/factories.py)

from config.settings import Settings
from application.unit_of_work import UnitOfWork
 
ADAPTER_REGISTRY = {
    "sqlalchemy": _build_sqlalchemy_uow,
    "inmemory": _build_inmemory_uow,
}
 
def build_unit_of_work(settings: Settings) -> UnitOfWork:
    builder = ADAPTER_REGISTRY[settings.database_adapter]
    return builder(settings)

Usage in Application Handlers

# application/account/create.py
class CreateAccountHandler:
    def __init__(self, account_repo: AccountRepository, unit_of_work: UnitOfWork):
        self._repo = account_repo
        self._uow = unit_of_work
 
    async def execute(self, command: CreateAccountCommand) -> Account:
        async with self._uow as session:
            existing = await self._repo.find_by_email(session, command.email)
            if existing:
                raise ConflictError("Email already in use")
 
            account = Account.create(
                email=command.email,
                display_name=command.display_name,
            )
            await self._repo.save(session, account)
            return account
        # UoW commits here on __aexit__ with no exception

Rules

  • Repositories persist and retrieve aggregates only — never sub-entities directly.
  • The session parameter is always obtained from the UoW — never created inside a repository.
  • Mappers (account_to_orm, account_to_domain) are standalone functions, not methods on the model or aggregate.
  • The in-memory adapter must have feature parity with the SQLAlchemy adapter — every test that passes with in-memory must be runnable against SQLAlchemy.
  • Migrations live in infrastructure/database/migrations/ (Alembic). See 06-database-persistence/01-migration-patterns.md.
Exception HierarchyEvent-Driven Patterns