Skip to content

Unit of Work

Intent

Track everything one business operation changes and make those changes visible together or not at all. The service opens a unit of work, writes through the repositories it hands out and calls commit() once; the unit of work owns BEGIN, COMMIT and ROLLBACK, so domain code never spells them and can never leave a transfer half done.

When to use and when not to

Use it when

  • One operation writes through two or more repositories (debit, credit, ledger) and a failure between them must leave no trace.
  • Your repositories commit nothing themselves, as the Repository page recommends; one small class should own the transaction, not every service.
  • You want a fake as atomic as production: a dict-backed unit of work that really discards uncommitted writes lets a test prove the rollback path.

Leave it out when

  • The operation writes one entity through one repository; its atomic add or update already is the transaction.
  • The writes span services or databases. No single transaction exists there; you need a saga or 2PC, and your commit becomes a compensating action.
  • You want to group writes for throughput. That is buffering; a queue or a bulk insert names it honestly.

Structure

Four roles: the unit-of-work contract, one implementation per storage technology, the repositories each implementation binds to its open transaction, and the service that opens the boundary and commits once.

classDiagram
    class UnitOfWork {
        <<interface>>
        +accounts: AccountRepository
        +ledger: LedgerRepository
        +__enter__() UnitOfWork
        +__exit__(exc_info)
        +commit()
        +rollback()
    }
    class SqliteUnitOfWork {
        -_conn: Connection
    }
    class InMemoryUnitOfWork {
        -_accounts: Dict~str,Account~
        -_entries: Dict~str,LedgerEntry~
        -_working_accounts: Dict~str,Account~
        -_working_entries: Dict~str,LedgerEntry~
        -_lock: Lock
    }
    class AccountRepository {
        <<interface>>
        +get(account_id: str) Account
        +save(account: Account)
    }
    class LedgerRepository {
        <<interface>>
        +append(entry: LedgerEntry)
        +for_account(account_id: str) List~LedgerEntry~
    }
    class SqliteAccountRepository
    class SqliteLedgerRepository
    class InMemoryAccountRepository
    class InMemoryLedgerRepository
    class TransferService {
        -_uow: UnitOfWork
        -_ids: IdGenerator
        -_clock: Clock
        +transfer(source_id: str, target_id: str, amount: Money) LedgerEntry
        +balance(account_id: str) Money
        +history(account_id: str) List~LedgerEntry~
    }
    UnitOfWork <|.. SqliteUnitOfWork
    UnitOfWork <|.. InMemoryUnitOfWork
    AccountRepository <|.. SqliteAccountRepository
    AccountRepository <|.. InMemoryAccountRepository
    LedgerRepository <|.. SqliteLedgerRepository
    LedgerRepository <|.. InMemoryLedgerRepository
    SqliteUnitOfWork *-- SqliteAccountRepository : builds in __enter__
    SqliteUnitOfWork *-- SqliteLedgerRepository
    InMemoryUnitOfWork *-- InMemoryAccountRepository : builds in __enter__
    InMemoryUnitOfWork *-- InMemoryLedgerRepository
    TransferService o-- UnitOfWork : opens, commits once

TransferService holds one UnitOfWork and cannot tell the implementations apart. Each implementation builds its repositories inside __enter__, bound to the transaction it just opened, which is why they hang off the unit of work instead of being injected separately. __exit__ always rolls back; only an explicit commit() lets work survive the block.

Canonical example in Python

Entities, repository contracts and the boundary come first (code/patterns/unit_of_work.py, tested by code/patterns/tests/test_unit_of_work.py):

code/patterns/unit_of_work.py — entities, repository contracts and the boundary
class InsufficientFundsError(InvalidStateError):
    """The source account cannot cover the transfer."""


@dataclass(frozen=True, slots=True)
class Account:
    """A wallet. Frozen: a debit or a credit is a new value handed to ``save``."""

    id: str
    balance: Money

    def debit(self, amount: Money) -> Account:
        if amount > self.balance:
            raise InsufficientFundsError(f"{self.id} holds {self.balance}, cannot debit {amount}")
        return replace(self, balance=self.balance - amount)

    def credit(self, amount: Money) -> Account:
        return replace(self, balance=self.balance + amount)


@dataclass(frozen=True, slots=True)
class LedgerEntry:
    """Append-only record of one transfer; its id doubles as the idempotency key."""

    id: str
    source_id: str
    target_id: str
    amount: Money
    at: float


class AccountRepository(Protocol):
    def get(self, account_id: str) -> Account: ...  # NotFoundError when absent

    def save(self, account: Account) -> None: ...  # insert or replace


class LedgerRepository(Protocol):
    def append(self, entry: LedgerEntry) -> None: ...  # ConflictError on a reused id

    def for_account(self, account_id: str) -> list[LedgerEntry]: ...


class UnitOfWork(Protocol):
    """The boundary: repositories that share one transaction, plus commit and rollback.

    ``commit()`` publishes everything written so far, ``rollback()`` discards it, and
    both leave the block open. ``__exit__`` always rolls back, so leaving without
    ``commit()`` loses the work instead of writing half of it. The repositories exist
    only inside the block. Not reentrant: nesting is a bug here, not a savepoint.
    """

    accounts: AccountRepository
    ledger: LedgerRepository

    def __enter__(self) -> Self: ...

    def __exit__(self, *exc_info: object) -> None: ...

    def commit(self) -> None: ...

    def rollback(self) -> None: ...

Three decisions to say out loud:

  • Explicit commit, implicit rollback. __exit__ rolls back unconditionally, so a block that forgets commit() loses its work loudly in the first test instead of writing half a transfer in production. The opposite convention, commit on a clean exit, is what the @contextmanager version below does; name the one you chose.
  • The repositories belong to the unit of work. They are created in __enter__ against the open transaction and deleted in __exit__, so nobody can hold a repository that writes outside a boundary.
  • The ledger id is the idempotency key. A replayed request reuses it, append raises after both balances were written, and the rollback is all that stands between you and a double debit.

The SQLite implementation is the one place that spells the transaction:

code/patterns/unit_of_work.py — the sqlite3 implementation
SCHEMA = """
CREATE TABLE IF NOT EXISTS accounts (
    id       TEXT PRIMARY KEY,
    cents    INTEGER NOT NULL,
    currency TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS ledger (
    id        TEXT PRIMARY KEY,
    source_id TEXT NOT NULL REFERENCES accounts(id),
    target_id TEXT NOT NULL REFERENCES accounts(id),
    cents     INTEGER NOT NULL,
    currency  TEXT NOT NULL,
    at        REAL NOT NULL
);
"""


def connect(path: str = ":memory:") -> sqlite3.Connection:
    """``autocommit=True`` switches off sqlite3's implicit transactions, so the Unit of
    Work spells BEGIN, COMMIT and ROLLBACK itself and the boundary is visible in one class."""
    conn = sqlite3.connect(path, autocommit=True)
    conn.executescript(SCHEMA)
    return conn


class SqliteAccountRepository:
    def __init__(self, conn: sqlite3.Connection) -> None:
        self._conn = conn

    def get(self, account_id: str) -> Account:
        row = self._conn.execute(
            "SELECT id, cents, currency FROM accounts WHERE id = ?", (account_id,)
        ).fetchone()
        if row is None:
            raise NotFoundError(f"account {account_id} does not exist")
        return Account(id=row[0], balance=Money(row[1], row[2]))

    def save(self, account: Account) -> None:
        self._conn.execute(
            "INSERT OR REPLACE INTO accounts (id, cents, currency) VALUES (?, ?, ?)",
            (account.id, account.balance.cents, account.balance.currency),
        )


class SqliteLedgerRepository:
    def __init__(self, conn: sqlite3.Connection) -> None:
        self._conn = conn

    def append(self, entry: LedgerEntry) -> None:
        amount = entry.amount
        try:
            self._conn.execute(
                "INSERT INTO ledger (id, source_id, target_id, cents, currency, at)"
                " VALUES (?, ?, ?, ?, ?, ?)",
                (entry.id, entry.source_id, entry.target_id, amount.cents, amount.currency, entry.at),
            )
        except sqlite3.IntegrityError as exc:
            raise ConflictError(f"ledger entry {entry.id} already exists") from exc

    def for_account(self, account_id: str) -> list[LedgerEntry]:
        rows = self._conn.execute(
            "SELECT id, source_id, target_id, cents, currency, at FROM ledger"
            " WHERE source_id = ? OR target_id = ? ORDER BY at, id",
            (account_id, account_id),
        ).fetchall()
        return [LedgerEntry(r[0], r[1], r[2], Money(r[3], r[4]), r[5]) for r in rows]


class SqliteUnitOfWork:
    """The only place in the codebase that spells BEGIN, COMMIT and ROLLBACK.

    The repositories are built inside ``__enter__``, bound to the transaction it just
    opened. ``BEGIN IMMEDIATE`` takes SQLite's write lock up front instead of at the
    first write, so a lock upgrade cannot fail halfway through. ``commit`` and
    ``rollback`` end the current transaction and start the next, so the block keeps
    its guarantee until ``__exit__`` discards whatever is left. One connection, one thread.
    """

    accounts: AccountRepository
    ledger: LedgerRepository

    def __init__(self, conn: sqlite3.Connection) -> None:
        self._conn = conn

    def __enter__(self) -> Self:
        self._conn.execute("BEGIN IMMEDIATE")
        self.accounts = SqliteAccountRepository(self._conn)
        self.ledger = SqliteLedgerRepository(self._conn)
        return self

    def __exit__(self, *exc_info: object) -> None:
        del self.accounts, self.ledger  # the repositories live only inside the block
        if self._conn.in_transaction:
            self._conn.execute("ROLLBACK")  # nothing after a commit; everything after an exception

    def commit(self) -> None:
        self._end_transaction("COMMIT")

    def rollback(self) -> None:
        self._end_transaction("ROLLBACK")

    def _end_transaction(self, statement: str) -> None:
        if not self._conn.in_transaction:
            raise InvalidStateError(f"{statement.lower()}() outside a unit of work")
        self._conn.execute(statement)
        self._conn.execute("BEGIN IMMEDIATE")  # the block stays open; later writes need a commit too

connect uses autocommit=True, switching off the driver's implicit transactions so that BEGIN, COMMIT and ROLLBACK appear in exactly one class. BEGIN IMMEDIATE takes the write lock when the block opens, so a lock upgrade cannot fail after the debit. commit() then opens a fresh transaction: a write after a commit still needs its own commit, the rule the fake enforces too.

The in-memory implementation is why the pattern earns its keep in tests. Writes go to a working copy of the committed dicts; commit() replaces the committed dicts with the copy and rollback() rebuilds the copy from them:

code/patterns/unit_of_work.py — the dict-backed implementation
class InMemoryAccountRepository:
    """Writes land in the working copy the unit of work handed over, never in committed state."""

    def __init__(self, accounts: dict[str, Account]) -> None:
        self._accounts = accounts

    def get(self, account_id: str) -> Account:
        try:
            return self._accounts[account_id]
        except KeyError:
            raise NotFoundError(f"account {account_id} does not exist") from None

    def save(self, account: Account) -> None:
        self._accounts[account.id] = account


class InMemoryLedgerRepository:
    def __init__(self, entries: dict[str, LedgerEntry]) -> None:
        self._entries = entries

    def append(self, entry: LedgerEntry) -> None:
        if entry.id in self._entries:
            raise ConflictError(f"ledger entry {entry.id} already exists")
        self._entries[entry.id] = entry

    def for_account(self, account_id: str) -> list[LedgerEntry]:
        mine = (e for e in self._entries.values() if account_id in (e.source_id, e.target_id))
        return sorted(mine, key=lambda e: (e.at, e.id))


class InMemoryUnitOfWork:
    """Change tracking without a database: writes go to a working copy, ``commit`` publishes it.

    ``_lock`` is held from ``__enter__`` to ``__exit__``, so units of work run one at a
    time, the way ``BEGIN IMMEDIATE`` serialises writers in SQLite. It protects the
    committed state (``_accounts``, ``_entries``). Shallow copies are enough because
    the entities are frozen: nothing reachable from the copy can be mutated in place.
    """

    accounts: AccountRepository
    ledger: LedgerRepository

    def __init__(self, accounts: Iterable[Account] = ()) -> None:
        self._accounts: dict[str, Account] = {account.id: account for account in accounts}
        self._entries: dict[str, LedgerEntry] = {}
        self._working_accounts: dict[str, Account] = {}
        self._working_entries: dict[str, LedgerEntry] = {}
        self._active = False
        self._lock = threading.Lock()

    def __enter__(self) -> Self:
        self._lock.acquire()
        self._active = True
        self.accounts = InMemoryAccountRepository(self._working_accounts)
        self.ledger = InMemoryLedgerRepository(self._working_entries)
        self.rollback()  # the working copy starts as the committed state
        return self

    def __exit__(self, *exc_info: object) -> None:
        try:
            self.rollback()  # nothing after a commit; everything after an exception
            del self.accounts, self.ledger
        finally:
            self._active = False
            self._lock.release()

    def commit(self) -> None:
        self._require_open("commit")
        self._accounts = dict(self._working_accounts)  # published under the lock: both or neither
        self._entries = dict(self._working_entries)

    def rollback(self) -> None:
        self._require_open("rollback")
        self._working_accounts.clear()
        self._working_accounts.update(self._accounts)
        self._working_entries.clear()
        self._working_entries.update(self._entries)

    def _require_open(self, operation: str) -> None:
        if not self._active:
            raise InvalidStateError(f"{operation}() outside a unit of work")

The lock is held for the whole block, serialising units of work the way BEGIN IMMEDIATE serialises SQLite writers; every test runs against both classes, so the fake cannot drift into being friendlier than the database.

The service sees neither sqlite3 nor dicts:

code/patterns/unit_of_work.py — the use case
class TransferService:
    """The use case. It sees neither sqlite3 nor dicts: open, write, commit once."""

    def __init__(self, uow: UnitOfWork, ids: IdGenerator, clock: Clock) -> None:
        self._uow = uow
        self._ids = ids
        self._clock = clock

    def transfer(self, source_id: str, target_id: str, amount: Money) -> LedgerEntry:
        if amount.cents <= 0:
            raise ValidationError("amount must be positive")
        if source_id == target_id:
            raise ValidationError("source and target must differ")
        with self._uow as uow:
            source = uow.accounts.get(source_id)
            target = uow.accounts.get(target_id)
            uow.accounts.save(source.debit(amount))
            uow.accounts.save(target.credit(amount))
            entry = LedgerEntry(self._ids.next_id(), source_id, target_id, amount, self._clock.now())
            uow.ledger.append(entry)  # a reused id raises here, after both balances were written
            uow.commit()
        return entry

    def balance(self, account_id: str) -> Money:
        with self._uow as uow:  # a read-only unit of work: leaving without commit is fine
            return uow.accounts.get(account_id).balance

    def history(self, account_id: str) -> list[LedgerEntry]:
        with self._uow as uow:
            return uow.ledger.for_account(account_id)

Validation runs before the block opens, so a bad request never takes the write lock. There is one commit(), on the last line; anything that raises before it, InsufficientFundsError included, leaves through __exit__ and is rolled back.

A transfer step by step: the commit path, and the replayed request that rolls back.

sequenceDiagram
    autonumber
    participant S as TransferService
    participant U as SqliteUnitOfWork
    participant A as SqliteAccountRepository
    participant L as SqliteLedgerRepository
    participant DB as SQLite
    S->>U: __enter__()
    U->>DB: BEGIN IMMEDIATE
    S->>A: get(source), get(target)
    A->>DB: SELECT
    S->>A: save(debited), save(credited)
    A->>DB: INSERT OR REPLACE (x2)
    S->>L: append(entry)
    L->>DB: INSERT INTO ledger
    alt ledger id unused
        S->>U: commit()
        U->>DB: COMMIT, then BEGIN IMMEDIATE
    else ledger id reused (replayed request)
        L-->>S: ConflictError
    end
    S->>U: __exit__()
    U->>DB: ROLLBACK (nothing left after a commit)

Running python -m patterns.unit_of_work prints:

--- transfers over the sqlite unit of work ---
txn-1: alice -> bob 30.00 USD; alice=70.00 USD bob=50.00 USD
rejected before any write: bob holds 50.00 USD, cannot debit 500.00 USD
rolled back after two writes: ledger entry txn-1 already exists
unchanged: alice=70.00 USD bob=50.00 USD
ledger for bob: ['txn-1']
--- transfers over the in-memory unit of work ---
txn-1: alice -> bob 30.00 USD; alice=70.00 USD bob=50.00 USD
rejected before any write: bob holds 50.00 USD, cannot debit 500.00 USD
rolled back after two writes: ledger entry txn-1 already exists
unchanged: alice=70.00 USD bob=50.00 USD
ledger for bob: ['txn-1']
--- @contextmanager variant: the boundary alone ---
block failed: power cut before commit
alice still has 70.00 USD

Pythonic variant

When the boundary is all you need, a generator-based context manager is the whole pattern:

code/patterns/unit_of_work.py — the boundary alone
@contextmanager
def transaction(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
    """The pattern reduced to its boundary: a generator-based context manager."""
    conn.execute("BEGIN IMMEDIATE")
    try:
        yield conn
    except BaseException:
        conn.execute("ROLLBACK")
        raise
    else:
        conn.execute("COMMIT")
  • Commit on success, rollback on exception, decided once by try/except/else; the body mentions neither.
  • except BaseException, not Exception: a KeyboardInterrupt inside the block must roll back too.
  • The stdlib ships one. with conn: on a sqlite3.Connection in its default transaction mode commits on a clean exit and rolls back on an exception; it does not close the connection.
Reach for When
with conn: One connection, one repository, the driver's default transactions
@contextmanager transaction(conn) BEGIN and COMMIT spelled out, or a fixed policy such as BEGIN IMMEDIATE
A UnitOfWork class Several repositories share the transaction, or tests need an atomic fake
A saga The writes cross a service boundary

Real-world usage

  • SQLAlchemy Session uses the name: it tracks loaded objects in an identity map and session.commit() flushes every change as one transaction; with Session(engine) as session, session.begin(): is the context-manager form.
  • Django transaction.atomic() is a context manager and a decorator; nested blocks become savepoints.
  • sqlite3.Connection as a context manager, plus the autocommit attribute added in Python 3.12 that this page relies on.
  • Cosmic Python's AbstractUnitOfWork is the shape used here: repositories as attributes, explicit commit(), rollback in __exit__.
Looks like Unit of Work How to tell them apart
Repository The repository says what is stored; the unit of work says when the writes become visible, together. A repository that commits inside add cannot take part in one.
A bare transaction with conn: is a transaction. A unit of work adds the repositories bound to it and, in ORMs, change tracking, so you never call save at all.
Saga Compensating actions across services, because no shared transaction exists to roll back. A unit of work lives inside one database.
Memento Memento snapshots one object to restore it; a unit of work collects many writes to discard them.
Command with undo Undo replays an inverse operation after the fact; rollback discards writes that were never published. Undo is for users, rollback is for failures.

Where it appears in LLD problems

Interview tips

Interview tip

Name the boundary in one sentence: "the transfer is one unit of work: debit, credit and ledger append commit together or roll back together; the service opens it with with and commits once." Then add the two sentences that mark an SDE2: which test proves the rollback (fail after the second write, assert the first is gone) and what changes when the writes cross a service boundary (a saga, not a bigger transaction).

Common mistake

Repositories that commit inside add, so the second write can no longer be undone. Runner-up: a fake unit of work whose rollback() is pass, which keeps the tests green while the production rollback path stays untested.