In my previous post I built two decorators, auto_session and wrap_in_transaction, to keep database sessions and transactions out of my service layer. If you did not read it, the short version is: PyMongo's .bind() stores the session in a ContextVar, so every query picks it up automatically and the service code stays clean.
Decorators are a great way to reuse behaviour and avoid repetition. But they have one limitation: a decorator is bound to the whole function call. It always runs before the function and finishes after it. Sometimes that is exactly what you want. Sometimes it is a problem.
Let's look at the problem closely.
The decorators from the previous post
from functools import wraps
from collections.abc import Callable
from pymongo.asynchronous.client_session import _SESSION, AsyncClientSession
def get_client() -> "AsyncMongoClient":
# Function that returns the current Mongo client
...
def auto_session(**session_kwargs) -> Callable:
"""Decorator that sets the session automatically."""
def decorator(f: Callable) -> Callable:
@wraps(f)
async def wrapper(*args, **kwargs):
client = get_client()
async with client.start_session(**session_kwargs) as s:
async with s.bind():
return await f(*args, **kwargs)
return wrapper
return decorator
def wrap_in_transaction(**transaction_kwargs) -> Callable:
"""Decorator that wraps the function in a transaction."""
def decorator(f: Callable) -> Callable:
@wraps(f)
async def wrapper(*args, **kwargs):
db_session: AsyncClientSession | None = _SESSION.get()
if not db_session:
raise Exception("Session not initialized, start a session with .bind() first")
async with await db_session.start_transaction(**transaction_kwargs):
return await f(*args, **kwargs)
return wrapper
return decorator
And we use them on the endpoint level:
router = APIRouter()
@router.post("")
@auto_session() # starts and binds the session
@wrap_in_transaction() # wraps the whole handler in a transaction
async def do_some_service_logic():
return await product_service.run_business_logic()
For a REST endpoint this is perfect: the request is short, and the handler response and the transaction both live under a ~60 second default timeout.
The problem
MongoDB aborts a transaction after 60 seconds by default. For a REST request that is fine. But imagine a long-running task, like importing a big file into the database:
# Parsing a big file can easily take more than 60 seconds.
# After exceeding 60 seconds Mongo will abort the transaction.
@auto_session()
@wrap_in_transaction()
async def long_running_task():
data = load_file("products.json")
categories = parse_categories(data)
await product_service.save_categories(categories)
products = parse_products(data)
await product_service.save_products(products)
Here the decorator wraps the entire function in one transaction, including all the parsing and file I/O. If the task takes longer than 60 seconds, Mongo kills the transaction, and none of the work is committed.
What we actually want is to wrap only the database writes in transactions, and leave everything else outside of them. The classic Python answer is a context manager:
@auto_session()
async def long_running_task():
data = load_file("products.json")
categories = parse_categories(data)
async with atomic_transaction(): # transaction opens here
await product_service.save_categories(categories)
# ...and commits here
products = parse_products(data)
async with atomic_transaction():
await product_service.save_products(products)
Now each write is its own short transaction and we never hit the 60 second limit.
But we still want the decorator form for our REST endpoints. Writing the logic twice (once as a decorator, once as a context manager) is exactly the duplication we were trying to avoid.
Enter AsyncContextDecorator
Python already ships the solution: ContextDecorator and its async twin AsyncContextDecorator.
From the docs:
A base class that enables a context manager to also be used as a decorator.
So all we have to do is turn our transaction logic into a context manager class — __aenter__ / __aexit__ — and inherit from AsyncContextDecorator. That single class can then be used in both ways.
from contextlib import AsyncContextDecorator
from typing import Any
from pymongo.asynchronous.client_session import _SESSION
class Atomic(AsyncContextDecorator):
def __init__(self, **transaction_kwargs: Any) -> None:
self.transaction_kwargs = transaction_kwargs
async def __aenter__(self) -> "Atomic":
db_session = _SESSION.get()
if db_session is None:
raise RuntimeError("No active session. Wrap your entry point with auto_session() first.")
if db_session.in_transaction:
# Already inside a transaction: do not nest.
return self
await db_session.start_transaction(**self.transaction_kwargs)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
db_session = _SESSION.get()
if db_session is None or not db_session.in_transaction:
return
if exc_type is None:
await db_session.commit_transaction()
else:
await db_session.abort_transaction()
Because it inherits from AsyncContextDecorator, this one class now works in every form:
# 1. As a context manager, around a single block
@auto_session()
async def long_running_task():
data = load_file("products.json")
categories = parse_categories(data)
async with Atomic():
await product_service.save_categories(categories)
products = parse_products(data)
# as configurable context manager
async with Atomic(read_concern=ReadConcern("majority")):
await product_service.save_products(products)
# 2. As a bare decorator we can use it in endpoints
@router.post("")
@auto_session()
@Atomic()
async def do_some_service_logic():
return await product_service.run_business_logic()
# 3. As a configured decorator
@router.post("")
@auto_session()
@Atomic(read_concern=ReadConcern("majority"))
async def do_some_service_logic():
return await product_service.run_business_logic()
One class, two ergonomics, zero duplicated logic.
The gotcha: one instance is shared by every call
Here is the part that bites people. When you use the object as a decorator, AsyncContextDecorator does roughly this under the hood:
def __call__(self, func):
@wraps(func)
async def inner(*args, **kwds):
async with self._recreate_cm(): # note: _recreate_cm
return await func(*args, **kwds)
return inner
The decorator object is created once, at import time, and the same instance is reused for every single call — including calls that happen concurrently.
_recreate_cm() defaults to returning self. So if your context manager keeps any per-use state on self (a session, a token, anything mutable), concurrent calls will stomp on each other.
The fix is to override _recreate_cm so each call gets a fresh instance:
class Atomic(AsyncContextDecorator):
def __init__(self, max_retry: int = 2, **transaction_kwargs: Any) -> None:
self.transaction_kwargs = transaction_kwargs
self.max_retry = max_retry
def _recreate_cm(self) -> "Atomic":
# Fresh instance per decorated call -> no shared, mutable state.
return type(self)(max_retry=self.max_retry, **self.transaction_kwargs)
# __aenter__ / __aexit__ as above
...
Supporting both @atomic_transaction and @atomic_transaction(...)
The Atomic class is already usable as @Atomic(). If you also want a paren-less form, add a thin factory so both @atomic_transaction and @atomic_transaction(...) work:
from collections.abc import Awaitable, Callable
from typing import Any, ParamSpec, TypeVar, overload
P = ParamSpec("P")
R = TypeVar("R")
@overload
def atomic_transaction(func: Callable[P, Awaitable[R]], /) -> Callable[P, Awaitable[R]]: ...
@overload
def atomic_transaction(*, max_retry: int = ..., **transaction_kwargs: Any) -> "Atomic": ...
def atomic_transaction(func=None, **transaction_kwargs):
if callable(func):
return Atomic(**transaction_kwargs)(func)
return Atomic(**transaction_kwargs)
Now both styles work:
@atomic_transaction
async def handler():
...
@atomic_transaction(read_concern=ReadConcern("majority"))
async def handler():
...
A small typing note: making func positional-only (/) in the first overload and the options keyword-only (*) in the second is what stops the type checker from reporting the two overloads as overlapping. It is the standard way to type a "decorator with or without arguments".
Summary
- A decorator wraps the whole function call; a context manager wraps just a block.
-
AsyncContextDecoratorlets one class do both, so you never duplicate the logic. - Override
_recreate_cmto return a fresh instance per call, or a shared decorator instance will leak state between concurrent calls. - A thin factory function lets you support both
@decoratorand@decorator(...). - Positional-only / keyword-only overloads keep the type checker happy.
- Core code for working with transactions stays in one place
The full implementation lives in my FastAPI boilerplate under app/libs/beanie_odm_ext/ — the transaction.py and session.py modules show this pattern wired into a working app.
Source code
The full implementation is available in the repository: github.com/Chornuy/fast-api-oauth — see app/libs/beanie_odm_ext/transaction.py and app/libs/beanie_odm_ext/session.py.
If you found this useful, let's connect on LinkedIn
Happy coding!
Top comments (2)
using a ContextVar for session propagation sidesteps explicit passing but adds TLS lookup overhead on every query and masks failures when the var leaks across async tasks, how does your wrapper handle exceptions that abort a transaction without leaving stale sessions
Hi, thanks for the question.
I will split the answer into 3 points:
1) TLS check on each session.
It depends on what you mean.
If we speak about each check of session on a query from _SESSION contextvar, which is a variable from the PyMongo client itself.
PyMongo itself does it already. And even creates a tmp session on each call query.
Here is the source code of pymongo:
github.com/mongodb/mongo-python-dr...
github.com/mongodb/mongo-python-dr...
If we speak about TLS handshake. The heavy lifting is done when socket opens on the pymongo client, the network call to create a session is lightweight, and should not do TLS handshake after the
client.connect()call.Each query will have a session id, but it does not call the server each time to get it.
2) Stale sessions / leaking across async tasks
ContextVars are task-local and thread safe in asyncio: each Task runs with a copy of the context, so a session bound in one task is not visible to a sibling task. It is visible to child tasks spawned inside the bind() block — which is the intended behaviour.
3) Handle exceptions that abort a transaction without leaving stale sessions
It's a good point.
If an error happens, the context manager will call aexit in the pymongo session context manager, and it calls Session reset(), which clears the session.
So an exception that aborts a transaction still unwinds cleanly — the wrapper doesn't need anything extra.
For this post I tried not to overengineer the solution, so it basically wraps the logic behind
with client.start_session(causal_consistency=True) as session:and tried not to add any custom or specific logic to it.But to handle this logic you can write your exception handler around the decorator with try/except, or inside AsyncContextDecorator in .aenter(). Or in the Atomic class