Every codebase has that service function - the one drowning in try/except blocks, manual rollbacks, and session management that makes you question your career choices:
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()
You've got rollback calls scattered everywhere, a finally block you'll forget to add next time, and business logic buried under boilerplate. It's fragile, verbose, and one missed rollback() away from a corrupted database state.
Let me show you a pattern that reduces this to:
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)
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
No manual rollback. No finally. No way to forget cleanup. Let's build it.
What is a Unit of Work?
Martin Fowler defined it best:
"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 items, remove items, change quantities - but nothing hits your credit card until you click "Purchase." If you abandon the cart, nothing happens. That's Unit of Work for databases:
- Track changes in memory
- Commit everything, all at once at the end, or
- Rollback everything if something fails
How would you like it if you ordered online groceries to make a pizza, only for your delivery person to come back with just mozzarella, oregano, and raw flour?
Exactly. Your database work needs to be atomic (all-or-nothing). This pattern further provides automatic cleanup and separation between your business logic and persistence mechanics.
The Implementation
Here's the complete, copy-pastable solution:
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()
Here’s how it works:
__aenter__ creates a fresh session when you enter the async with block. Each transaction gets its own isolated session.
async with UnitOfWork() as uow:
...
# the `async with` calls __aenter__, and
# returns our uow object that contains out db session
__aexit__ handles cleanup automatically. The magic is in exc_type - Python passes exception info if something went wrong inside the block. If there's an exception, we rollback. Either way, we close the session.
async with UnitOfWork() as uow:
users: list[User] = uow.find_all_banned_users()
print(users) # When we "leave" that `async with` block, the database
# session is closed!
Explicit commit() keeps you in control. You decide when changes persist. No surprises.
"Why do I have to call commit() manually?"
You might be tempted to put self.commit() inside the __aexit__ method so it saves automatically if no errors occur; this is common, but misguided practice.
Explicit commits are intentional:
- They prevent accidental writes if you
returnearly from a function - They make it obvious to anyone reading your code exactly when the data is saved
- You don’t waste resources automatically committing (or worse, accidentally commit something) on a simple
selectstatement where you only read
As always: explicit is better than implicit
https://peps.python.org/pep-0020/
Using It
First, set up your session factory (do this once at app startup):
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)
Then inject it into your UnitOfWork:
async def main() -> None:
uow: UnitOfWork = 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}")
Here's what happens in different scenarios:
Happy path - code runs, you call commit(), changes persist:
async with uow:
uow.session.add(User(name="Bob"))
await uow.commit() *# Saved to db*
Exception before commit - automatic rollback, nothing persists:
async with uow:
uow.session.add(User(name="Charlie"))
raise ValueError("Something went wrong")
await uow.commit() *# Never reached# __aexit__ catches the exception,
# calls rollback(), Charlie never existed in the database*
Forgot to commit - nothing persists (safe default):
async with uow:
uow.session.add(User(name="Dave"))
*# Oops, forgot commit(), session closes, uncommitted changes discarded*
Real-World Example: Order Processing
Let's look at a more complete, realistic scenario - creating an order with multiple line items where any failure should cancel the entire operation:
from dataclasses import dataclass
@dataclass
class OrderRequest:
user_id: int
items: list[dict]
async def process_order(request: OrderRequest, uow: UnitOfWork) -> Order:
async with uow:
*# Validate user exists*
user = await uow.session.get(User, request.user_id)
if not user:
raise ValueError(f"User {request.user_id} not found")
*# Create order*
order = Order(user_id=request.user_id, status="pending")
uow.session.add(order)
total = 0
for item in request.items:
product = await uow.session.get(Product, item["product_id"])
if not product:
raise ValueError(f"Product {item['product_id']} not found")
if product.stock < item["quantity"]:
raise ValueError(f"Insufficient stock for {product.name}")
line = OrderLine(
order=order,
product_id=product.id,
quantity=item["quantity"],
price=product.price
)
uow.session.add(line)
product.stock -= item["quantity"]
total += product.price * item["quantity"]
order.total = total
await uow.commit()
return order
If the user doesn't exist? Rollback. Product not found? Rollback. Insufficient stock? Rollback. The order, line items, and stock changes all succeed together or fail together.
With proper atomicity (the A in ACID), we always roll back grouped transactions if a single one fails to prevent corrupted or incomplete data.
Important Nuance
"Fancy algorithms are slow when n is small, and n is usually small." — Rob Pike
Not every database operation needs a Unit of Work wrapper. If you're writing a simple script that runs once, or a function that does a single INSERT, the overhead isn't worth it.
Unit of Work shines when:
- Multiple related operations must succeed or fail together
- You're building a layered architecture with services and repositories
- You need testable transaction boundaries
- Different parts of your code need to participate in the same transaction
For a quick one-off query? Just use the session directly.
Conclusion
That's it. You now have a pattern that:
- Guarantees cleanup with context managers
- Rolls back automatically on exceptions
- Makes transaction boundaries explicit and testable
- Separates business logic from persistence boilerplate
The Unit of Work pattern isn't about adding complexity - it's about removing the mental overhead of "did I remember to rollback?" from every line of ORM code you write.
Go find a service function with scattered try/except/finally blocks. Refactor it. Watch the noise disappear.
Questions about Unit of Work or SQLAlchemy patterns? Drop a comment below. :)
Top comments (2)
In my team we use SQLAlchemy… but only as a container for our raw SQL. Keeps us humble.
Did you really loop the order line items? Sorry my SQL optimization heart got an attack after seeing that.