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

Question

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

What I’m going to do is like

Commit on the end of all path operation function implicitly if there is no error occured.
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?

Asked By: Baekjun Kim

||

Source

Answer 1

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()

About exception handler, finally code block will be exceuted before exception handler,so you can’t rollback in exception handler.

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.

Answered By: Jedore

Answer 2

I would recommend the following:

sql_engine = ...
SessionLocal = sessionmaker(sql_engine, autoflush=False, autocommit=False)

async def get_session(request: Request) -> Generator:
    try:
        logger.debug("Creating database session.")
        database_session = SessionLocal()
        yield database_session
        database_session.commit()
    finally:
        logger.debug("Closing database session.")
        database_session.close()

But THEN do not forget to flush at the end of your method so the changes are evaluated and the commit is safe:

@router.post(path="/somewhere")
async def create_database(
    your_data: dict = Body()
    session: Session = Depends(database.get_session),
) -> None:
    ...
    logger.debug("Flush changes to database.")
    session.flush()

If data are incorrect or there is BD conflict, flush will fail and your exception handlers will catch the exception before sending the response.

I was looking for a way to raise also an error if commit is called without manual flush, however did not find it (I suppose the best is to rely on testing). Open to suggestions.

Answered By: BorjaEst

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

Question:

Answers: