Stop Writing try/except Hell: Clean Database Transactions with SQLAlchemy with the Unit Of Work

Source: Dev.to

async def create_order(user_id: int, items_in_basket: list[dict]) -> Order:
    session = get_session()
    try:
        user: User = await session.get(User, user_id)
        if not user:
            await session.rollback()
            raise ValueError("User not found")

        order: Order = Order(user_id=user_id)
        session.add(order)

        for item in items_in_basket:
            line: OrderLineItem = OrderLineItem(**item, order=order)
            session.add(line)

        await session.commit()
        return order
    except DuplicateError:
        await session.rollback()
        raise
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()

The above pattern scatters rollback calls, requires a finally block, and mixes business logic with boilerplate. A cleaner approach uses a Unit of Work (UoW) to manage the session lifecycle automatically.

What is a Unit of Work?

Martin Fowler defined it as:

“A Unit of Work keeps track of everything you do during a business transaction that can affect the database. When you’re done, it figures out everything that needs to be done to alter the database as a result of your work.”

Think of it like an online shopping cart: you add, remove, or change items, but nothing is persisted until you click Purchase. If you abandon the cart, no changes are saved. The same principle applies to database operations:

• Track changes in memory
• Commit everything at once, or
• Rollback everything if something fails

The Implementation

from typing import Self
from types import TracebackType
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

class UnitOfWork:
    """Async context manager for database transactions.

    Commits on success, rolls back on exception, always cleans up.
    """

    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._session_factory = session_factory

    async def __aenter__(self) -> Self:
        self.session = self._session_factory()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if exc_type is not None:
            await self.rollback()
        await self.session.close()

    async def commit(self) -> None:
        await self.session.commit()

    async def rollback(self) -> None:
        await self.session.rollback()

How it works

• __aenter__ creates a fresh AsyncSession for each async with block, giving every transaction its own isolated session.
• __aexit__ automatically rolls back if an exception propagated out of the block, then closes the session.
• Explicit commit() keeps you in control of when changes are persisted.

Explicit is better than implicit – calling commit() manually avoids accidental writes (e.g., when returning early or performing read‑only queries).

Using It

Set up a session factory (once at application start)

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine

engine: AsyncEngine = create_async_engine(
    "postgresql+asyncpg://user:pass@localhost/db"
)

SessionFactory = async_sessionmaker(engine, expire_on_commit=False)

Create a Unit of Work and perform operations

async def main() -> None:
    uow = UnitOfWork(session_factory=SessionFactory)

    async with uow:
        uow.session.add(User(name="Alice", email="alice@example.com"))
        await uow.commit()
        print(f"Created user with ID: {user.id}")

Scenarios

Happy path – commit succeeds

async with uow:
    uow.session.add(User(name="Bob"))
    await uow.commit()   # changes are saved

Exception before commit – automatic rollback

async with uow:
    uow.session.add(User(name="Charlie"))
    raise ValueError("Something went wrong")
    # No commit called; __aexit__ rolls back automatically

Using the Unit of Work in a service function

async def create_order(user_id: int, items_in_basket: list[dict], uow: UnitOfWork) -> Order:
    async with uow:
        user = await uow.session.get(User, user_id)
        if not user:
            raise ValueError("User not found")

        order = Order(user_id=user_id)
        uow.session.add(order)

        for item in items_in_basket:
            uow.session.add(OrderLine(**item, order=order))

        await uow.commit()
        return order

The pattern eliminates scattered rollback calls, removes the need for a finally block, and cleanly separates business logic from persistence concerns.