DEV Community: Serhii Chornyi

Share Logic Between Decorators and Context Managers in Python

Serhii Chornyi — Mon, 12 Oct 2026 15:04:00 +0000

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.
AsyncContextDecorator lets one class do both, so you never duplicate the logic.
Override _recreate_cm to 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 @decorator and @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!

Managing MongoDB Session with Beanie 2.0

Serhii Chornyi — Mon, 18 May 2026 16:13:43 +0000

Hi DevTo. It's my first post.

The problem

I have been using Beanie ODM and MongoDB in my work for a long time.

It's a great tool, but managing sessions and transactions across the service layer was problematic for me.

My service layers need to know about the session object to pass it to Beanie ODM model methods to bind it to the session.

It looks like this


class SomeService:


    async def run_business_logic(db_session: AsyncIOMotorClientSession) -> Product:
        chocolate = Category(name="Chocolate", description="A preparation of roasted and ground cacao seeds.")

        tonybar = Product(name="Tony's", price=5.95, category=chocolate)
        await tonybar.save(session=db_session) # set session
        return tonybar


important_service = SomeService()

@router.post("")
async def do_some_service_logic(
    db_session: AsyncIOMotorClientSession = Depends(get_session),
):
    async with await session.start_transaction(): # Start transaction
        return await important_service.run_business_logic(
            db_session=db_session, # We need to pass the session on the service layer to work with Beanie
        )

It's making me feel frustrated, because now I need to pass the session object to my service.

Game changer

But things changed after Beanie released version 2.0, and Beanie moved from Motor in favour of the PyMongo Async client.

So pymongo gives the option to bind all calls to MongoDB with the method .bind()

So basically, pymongo can create a session, and can bind to this session all calls to MongoDB.

Example from the documentation.

async with client.start_session() as s:
    async with s.bind():
        # session=s is passed implicitly
        await client.db.collection.insert_one({"x": 1})

For more details, here is how it works inside the pymongo client:


_SESSION: ContextVar[Optional[AsyncClientSession]] = ContextVar("SESSION", default=None)


class _AsyncBoundSessionContext:
    """Context manager returned by AsyncClientSession.bind() that manages bound state."""

    def __init__(self, session: AsyncClientSession, end_session: bool) -> None:
        self._session = session
        self._session_token: Optional[Token[AsyncClientSession]] = None
        self._end_session = end_session

    async def __aenter__(self) -> AsyncClientSession:
        self._session_token = _SESSION.set(self._session)  # type: ignore[assignment]
        return self._session

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if self._session_token:
            _SESSION.reset(self._session_token)  # type: ignore[arg-type]
            self._session_token = None
        if self._end_session:
            await self._session.end_session()


class AsyncClientSession:

    ...
    def bind(self, end_session: bool = True) -> _AsyncBoundSessionContext:
        return _AsyncBoundSessionContext(self, end_session)

So PyMongo will set the session object in _SESSION ContextVar and on any call to the database will check if there is a session, if there is a session, it will use it.

Because PyMongo sits underneath Beanie, this works transparently. I don't need to change Beanie itself at all. Just create a session and call .bind() method, and everything else will be handled by the pymongo client.

Because _SESSION is ContextVar, this means that the session is task-safe for us.

For a transaction, we can use a session that is set in _SESSION context.

Now we can wrap all this logic in Python decorators to reuse it:


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 session automatically"""

    def decorator(f: Callable) -> Callable:
        @wraps(f)
        async def wrapper(*args, **kwargs):
            client = get_client()  # Return MongoClientAsync
            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 sets transaction automatically with bind session"""

    def decorator(f: Callable) -> Callable:
        @wraps(f)
        async def wrapper(*args, **kwargs):
            # Take session from context directly from pymongo. asynchronous.client_session
            db_session: AsyncClientSession | None = _SESSION.get()
            if not db_session:
                raise Exception("Session not initialized, start session with using .bind() method")

            async with await db_session.start_transaction(**transaction_kwargs):
                return await f(*args, **kwargs)

        return wrapper

    return decorator

And we can use decorators on the endpoint level.

# The Service layer is now clean. No 'session' argument needed!
class SomeService:
    async def run_business_logic(self) -> Product:
        chocolate = Category(name="Chocolate", description="...")
        tonybar = Product(name="Tony's", price=5.95, category=chocolate)
        # Beanie finds the session automatically via PyMongo's ContextVar
        await tonybar.save() 
        return tonybar

# The Controller remains clean and focused on HTTP
@router.post("")
@auto_session() # Automatically starts and binds the session
@wrap_in_transaction() # Wraps the execution in a transaction if the service raise error, Mongo will rollback data
async def do_some_service_logic():
    return await important_service.run_business_logic()

Now our services do not know anything about the session, which means more clean code for us:

Sessions are task-safe, each call of the endpoint will get its own session
We do not bind transactions to a session, we just use it from the context of the pymongo client
Business logic (Our services) stays focused on the domain, not knowing about how to handle database sessions.

DEV Community: Serhii Chornyi

Share Logic Between Decorators and Context Managers in Python

The decorators from the previous post

The problem

Enter AsyncContextDecorator

The gotcha: one instance is shared by every call

Supporting both @atomic_transaction and @atomic_transaction(...)

Summary

Source code

Managing MongoDB Session with Beanie 2.0

The problem

Game changer

Links

Supporting both `@atomic_transaction` and `@atomic_transaction(...)`