How to make FastAPI's path operation function to commit on successful return

Baekjun Kim 0
Viewed 50

I'm developing web API server with FastAPI and SQLAlchemy 1.4

What I'm going to do is like

  1. Commit on the end of all path operation function implicitly if there is no error occured.
  2. When HTTPException is occured, rollback the transaction.

In short, how to make FastAPI's path operation function atomic like database transaction.

These code snippets are what I'm using from tiangolo's full-stack-fastapi-postgresql project (https://github.com/tiangolo/full-stack-fastapi-postgresql)

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

engine = create_engine(settings.SQLALCHEMY_DATABASE_URI, pool_pre_ping=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db() -> Generator:
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()

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

app = FastAPI()

@app.put("/users/{id}")
def update_user(
    id: int,
    new_name: str,
    new_login_id: str,
    db: Session = Depends(deps.get_db),
):
    """
    Simple User update API
    """
    user = db.query(User).filter(User.id == id).first()

    if not user:
        raise HTTPExecption(404, "User not found")

    user.name = new_name
    
    if db.query(User).filter(User.login_id == new_login_id).first():
        # if login id is already exist, I want to rollback update name
        raise HTTPExecption(409, "Already used login id")

    user.login_id = new_login_id

    db.add(user)
    # db.commit()
    # I don't want to write db.commit() at the end of every path operation func. 

    return

I've tried in several ways

First, using FastAPI (Starlette) Middleware

I can't find way to pass db (Session) to middleware

Second, Add db.commit() at get_db()'s finally block

db.commit() is called even HTTPException is raised.

So I tried to use FastAPI Error handler for calling db.rollback() when HTTPExecption is raised, but I can't find a way to pass db (Session) to error handler.

How can I acheive it?

sqlalchemy fastapi python
1 Answers

Firstly, i think SQLAlchemy won't commit data if there is some exception occured in a transaction.

Then, if you want to pass db, you can

def get_db(request) -> Generator:
    try:
        db = SessionLocal()
        request.state.db = db
        yield db
    finally:
        db.close()
  1. About exception handler, finally code block will be exceuted before exception handler,so you can't rollback in exception handler.
  2. About middleware, all http path operation exceptions were catched in starlette.ExceptionMiddleware which is fixed in FastAPI app. So if you want to custom it you should create a middleware CustomExceptionMiddleware , you can use db in it: 
    class CustomExceptionMiddleware(starlette.Exception):
        async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
            if scope["type"] != "http":
                await self.app(scope, receive, send)
                return

            response_started = False

            async def sender(message: Message) -> None:
                nonlocal response_started

                if message["type"] == "http.response.start":
                    response_started = True
                await send(message)

            try:
                await self.app(scope, receive, sender)
                a = 0
            except Exception as exc:
                # some operation
                request = Request(scope, receive=receive)
                # request.state.db is usable
                # some operation
    then you must rewrite FastAPI.build_middleware_stack() to replace ExceptionMiddleware by CustomExceptionMiddleware. Now you can execute rollback or other operation in CustomExceptionMiddleware. Also you shoud remove db.close() from get_db() to CustomExceptionMiddleware.
Jedore0
Related