How to write a decorator for path operation?

[gocreating]

First check

• I added a very descriptive title to this issue.
• I used the GitHub search to find a similar issue and didn't find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google "How to X in FastAPI" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.
• After submitting this, I commit to one of:
  • Read open issues with questions until I find 2 issues where I can help someone and add a comment to help there.
  • I already hit the "watch" button in this repository to receive notifications and I commit to help at least 2 people that ask questions in the future.
  • Implement a Pull Request for a confirmed bug.

Example

def with_transaction(func):
    @wraps(func)
    async def wrap_func(request, *args, **kwargs):
        session = request.app.state.database.get_session()
        try:
            result = await func(request, *args, **kwargs)
            session.commit()
        except Exception as e:
            session.rollback()
        finally:
            session.close()
        return result

    return wrap_func

@router.post('/users', response_model=schemas.User)
@with_transaction
def create_user(user: schemas.UserCreate, session: Session = Depends(get_session)):
    user = User(provider=user.provider)
    session.add(user)
    if user.foobar == 'force_error':
        raise Exception('my custom exception')
    return user

Description

My requirements:

1. commit before response validation if path operation succeeds
2. rollback if path operation fails (raises custom exception)
3. such feature is optional for each endpoint

After suffering from the limitation of dependencies, middlewares, and a lot of custom xxx, I finally decided to manage transactions on my own. (sigh, neither dependencies nor middlewares have complete control on request lifecycles)
It seems decorator is the last and optimal practice.
However, the sample decorator above cannot get dependencies injected...

TypeError: wrap_func() missing 1 required positional argument: 'request'

How sad and frustrating such a great framework can't even do this simple and frequently-used feature...

[gocreating]

For guys who want to manage transactions on your own, I've figured out my workaround. Take it if you want.

# dependencies.py

from fastapi import Depends, Request

def get_session(request: Request):
    # here I get sqlalchemy's scoped_session, which should be the same instance per thread
    request.state.session = request.app.state.database.get_session()
    try:
        yield request.state.session
    finally:
        request.state.session.close()

# decorators.py

import functools
import inspect

from fastapi import Depends, HTTPException
from sqlalchemy.orm import Session

from your.dependencies import get_session

def managed_transaction(func):

    @functools.wraps(func)
    def wrap_func(*args, session: Session = Depends(get_session), **kwargs):
        try:
            if inspect.iscoroutinefunction(func):
                result = await func(*args, session=session, **kwargs)
            else:
                result = func(*args, session=session, **kwargs)
            session.commit()
        except HTTPException as e:
            session.rollback()
            raise e
        # don't close session here, or you won't be able to response
        return result

    return wrap_func

from fastapi import APIRouter, Depends, HTTPException, Request
from sqlalchemy.orm import Session

from your.decorators import managed_transaction
from your.dependencies import get_session
from your.sqlalchemy.models.user import User

router = APIRouter()

@router.post('/users', response_model=schemas.User)
@managed_transaction
def create_user(user: schemas.UserCreate, session: Session = Depends(get_session)):
    user = User(foo=user.foo)
    session.add(user)
    if user.foo == 'error':
        raise HTTPException(status_code=403)
    return user

You will get several benefits:

• auto commit after path operation is finished without error
• auto rollback when path operation raises exception
• validate response model and send response after commit
• you can manage the notorious HTTPException before it is caught by fastapi core
• this feature is optional for every endpoints
• you can improve the decorator on your need (e.g. implement a loop to retry path operation function)
• without any hacking, fastapi's dependency injection still works

FYI, none of fastapi's features is possible to abstract transaction management:

• If you use dependencies, then you will end up sending response before commit
• If you use middlewares, then all requests will acquire a session, even when you don't need it at all
• If you use custom API route class, then you don't even get a chance to manage transactions

[artslob]

@gocreating As I understand decorator apply order matters?

In first case router.post will get create_user function unchanged and will analize its signature. This can work.

@managed_transaction
@router.post('/users', response_model=schemas.User)
def create_user(user: schemas.UserCreate, session: Session = Depends(get_session)):

---

In second case managed_transaction will replace create_user with wrap_func(*args, session: Session = Depends(get_session), **kwargs). In this case router.post cant analize signature because it will only see *args and **kwargs. Info about user: schemas.UserCreate would be lost.

@router.post('/users', response_model=schemas.User)
@managed_transaction
def create_user(user: schemas.UserCreate, session: Session = Depends(get_session)):

Am I wrong or right?

[gocreating]

@artslob
I personally think the first case won't work since the signature inference is done through router decorator.
I only used the second case in my side project. It seems not breaking any signature inference.
FYI, here is the project link: https://github.com/gocreating/lation/blob/develop/lation/modules/stock/routers/ptt.py#L69-L76

[gsouveton]

@gocreating Your post really helped me writing custom decorators for my routes.
However, I would like to understand how it works exactly, especially the way you wrote the get_session function.

def get_session(request: Request):
    # here I get sqlalchemy's scoped_session, which should be the same instance per thread
    request.state.session = request.app.state.database.get_session()
    try:
        yield request.state.session
    finally:
        request.state.session.close()

Why do you store the session in request.state.session ? Is there any goal of doing this instead of a random variable name ?

def get_session(request: Request):
    # here I get sqlalchemy's scoped_session, which should be the same instance per thread
    sess = request.app.state.database.get_session()
    try:
        yield sess
    finally:
        sess.close()

And what's the point of using yield request.state.session rather than return request.state.session here ? How does it affect the behaviour ?
Thank you.

[gocreating]

@gsouveton
You can get session by your own way. I store my session in state because that's more convenient for me. No fancy goal. :D
And for yield...actually I forget the reason. Maybe I was reusing get_session as a generator somewhere else in my codebase.

[peterswords]

@gocreating thanks for this, I've been really struggling with the best way to do this. I'm going to try your approach. A few questions:

In your wrapper function you inspect the wrapped function to see if it's an async coroutine:

def managed_transaction(func):

    @functools.wraps(func)
    def wrap_func(*args, session: Session = Depends(get_session), **kwargs):
        try:
            if inspect.iscoroutinefunction(func):
                result = await func(*args, session=session, **kwargs)
            else:
                result = func(*args, session=session, **kwargs)
            session.commit()
        except HTTPException as e:
            session.rollback()
            raise e
        # don't close session here, or you won't be able to response
        return result

    return wrap_func

You inspect the passed function inside the wrapper and therefore take the runtime hit on each call. Since func is available to you before you do the wrapping, why not have two separate wrappers and return the appropriate wrapper at setup time to avoid the runtime hit?

Indeed, I don't see how you can do it any other way because your wrapper is not declared async so you can't await the wrapped function. (In my case I'm using async SQLAlchemy 1.4 so only the async version is needed).

You are passing session as a keyword argument to your wrapped function. This will fail unless the wrapped function uses that specific name for the session parameter. Is there a way of using inspect to avoid this hardcoding?

You are catching HTTPException but your wrapper will be invoked before FastAPI converts path operation exceptions, so shouldn't you be catching a more general type of exception?

[peterswords]

Thanks to @gocreating for most of the idea here. Here's my attempt which I have working. I have made it so that you can use any name for the parameter to a FastAPI path operation which will hold the database session. However, the type of this parameter must be DB and it must be a FastAPI dependency. I have created the DB class as a synonym for AsyncSession to simplify the dependency management -- no parameters need to be passed to the Depends function. The tx decorator is used to indicate that an operation should have automatic transaction management. A path operation will look something like this:

@router.post("/new")
@tx
async def widget_new(widget: WidgetAPI, db: DB = Depends()) -> dict:
    db_obj = Widget(**widget.dict())
    db.add(db_obj)
    return {}

My dependencies.py looks like the following. Note that I am using async I/O with SQLAlchemy 1.4 to a Postgres database, so I only care about async path operations. A different implementation could handle both async and sync.

import functools
from fastapi import Depends
from fastapi.params import Depends as Depends_Type
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
import inspect


engine = create_async_engine(
    "postgresql+asyncpg://scott:tiger@host.docker.internal/test", echo=True, future=True
)

_async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


async def _get_db() -> AsyncSession:
    """
    Dependency function that yields db sessions
    """
    db: AsyncSession = _async_session()
    try:
        yield db
    finally:
        await db.close()


class DB(AsyncSession):
    """Lets you use the following dependency: db: DB = Depends()"""

    def __new__(cls, db: AsyncSession = Depends(_get_db)) -> AsyncSession:
        return db


def tx(func):

    fn = func.__name__
    if not inspect.iscoroutinefunction(func):
        raise Exception(f'Function {fn} must be declared async to use with the @tx decorator')

    for p in inspect.signature(func).parameters.values():
        if (p.annotation is DB and type(p.default) is Depends_Type and (
                p.kind == inspect.Parameter.POSITIONAL_ONLY or
                p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD)):
            db_param_name = p.name
            break
    else:
        raise Exception(f'Function {fn} should have a positional or keyword argument ' +
            'of type DB with a dependency, when using with the @tx decorator')

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            db: DB = kwargs[db_param_name]
            result = await func(*args, **kwargs)
            await db.commit()
        except Exception as e:
            await db.rollback()
            raise e
        return result

    return wrapper

[jeanlst]

@peterswords Great job on the decorator! Quick question, how does one handle the exception raised in

except Exception as e:
    await db.rollback()
    raise e

I have been trying it out and the exception doesn't go upstream, that is, if I put a db.add(whatever) inside a try, except Exception it doesn't get caught by the try, except inside the route function that has the @tx decorator

[peterswords]

@jeanlst , the decorator is wrapped around your route function, not vice versa. Your route function has already exited by the time the raise occurs. Note, the exception is being raised because you got a failure when trying to commit the stuff that you added with db.add. db.add itself doesn't generally throw exceptions. It just adds your object to SQLAlchemy's cache. (That's why you don't have to await it -- it doesn't do any I/O). The failure -- such as a duplicate key exception -- will happen on the commit which, of course, happens inside the tx decorator after your route function completes. If you want it to fail fast, call await db.flush after db.add inside your try ... except block. This will try to write to the database, even though the transaction still won't be committed yet. You'll then get an exception you can catch.

[jeanlst]

Oh I see, thanks for the clarification @peterswords . If I might ask you, what's the flow for the route with the decorator?

Something like the one you used as an example

@router.post("/new")
@tx
async def widget_new(widget: WidgetAPI, db: DB = Depends()) -> dict:
    db_obj = Widget(**widget.dict())
    db.add(db_obj)
    return {}

what's the flow here? The router gets called, it runs db.add and then the return {}, and then the return {} gets pushed forward to the decorator and then the decorator code runs?

[peterswords]

@jeanlst , not quite. Decorator code can run both before and after the route function. I don't want to make assumptions about your understanding of decorators in Python, and you can read better explanations of them elsewhere anyway. But, in short, a decorator is a function that takes a function as input and returns a function as output. The output function can act as a "wrapper" for the input, performing some preprocessing before calling the original input function, and some postprocessing afterward.

That's what FastAPIs @router.post decorator is doing to your route function for instance. It intercepts a web request, converts the request body into Pydantic structures and validates them before passing them to your route function, and then takes the output of your function and converts it into a web response.

Here we've inserted yet another decorator between FastAPI and the route function. The tx wrapper function starts a transaction, calls your route function, and then commits the transaction. It has an exception handler so that any exceptions in your route function result in the transaction being rolled back. Because of the order of the decorators, it is now the tx wrapper function which is further wrapped by the @router.post decorator. So in summary, the normal order of operations would now be:

1. FastAPI intercepts a web request, converts the request data to a Pydantic model, inserts dependencies etc.
2. FastAPI calls the tx wrapper function with Pydantic arguments
3. The tx function starts a transaction and calls the route function
4. The route function does its processsing and returns data
5. The tx function commits the transaction and returns the data
6. FastAPI converts the data into a web response

(Note: you don't have to actually start a SQLAlchemy transaction because it is started implicitly, but I've left it in as part of a conceptual flow).