Transactional methods (FastAPI, SQLAlchemy)

Motivation

When working on FastAPI I had a real discomfort not having @transactionaldecorator like in Spring Boot so I’ve written one myself and wanted to share. The goal is easy, to make a method transactional!

Transactional Method

A transactional method starts a transaction when the method begins, commits it upon successful completion, and rolls it back in the event of an exception to ensure data integrity.

Demo

@transactional
def save_book(title):
    book = save(BookModel(title=title))
    # Return Pydantic model from service (recommended)
    return BookInfo.from_orm(book)

Plan

1. Create a simple database model BookModel and BookInfo schema
2. Implement @transactional decorator
3. Implement @db decorator for injecting database sessions into methods
4. Use @transactional and @dbdecorators in an example
5. Implement sample endpoint for demo (Optional)

Let’s begin by creating a simple BookModel and BookInfo schema

class BookModel(Base):
    id = Column(Integer, primary_key=True, index=True)
    title = Column(String)


class BookInfo(BaseModel):
    id: int
    title: str

    class Config:
        orm_mode = True

Now pay close attention !!! This step is very important, database models are useless outside of the session context (outside of @transactional methods you won’t be able to use them properly), so you should always return schemas (Pydantic models), not database models!

@transactional decorator

db_session_context = contextvars.ContextVar("db_session", default=None)

def transactional(func):
    @wraps(func)
    def wrap_func(*args, **kwargs):
        db_session = db_session_context.get()
        if db_session:
            return func(*args, **kwargs)
        db_session = sessionmaker()
        db_session_context.set(db_session)
        try:
            result = func(*args, **kwargs)
            db_session.commit()
        except Exception as e:
            db_session.rollback()
            raise
        finally:
            db_session.close()
            db_session_context.set(None)
        return result
    return wrap_func

Let’s break it down:

db_session_context = contextvars.ContextVar("db_session", default=None)

Here we utilize a library called ContextVar. It allows us to store the database session and guarantees that any method called from the decorated method will use the same database session.

db_session = db_session_context.get()
if db_session:
    return func(*args, **kwargs)
db_session = sessionmaker()
db_session_context.set(db_session)

Here we grab the session from the context variable, if it already exists then we just continue the execution, if not, then we create one and save it to the context variable.

try:
    result = func(*args, **kwargs) # call the decorated method
    db_session.commit()
except Exception as e:
    db_session.rollback()
    raise # raise exception for outside world to notice
finally:
    db_session.close()
    db_session_context.set(None)
return result

Here we put everything in the try/except/finally clause, if the execution passes without exceptions then we commit the changes, if not we just roll back the transaction and still raise the exception, but either way we close the session and empty the context variable

@db decorator

Now let’s implement a @db decorator to inject database sessions into methods

def db(func):
    @wraps(func)
    def wrap_func(*args, **kwargs):
        db_session = db_session_context.get()
        return func(*args, **kwargs, db=db_session)
    return wrap_func

All it does is grab the session from the session context variable and pass it to the decorated method as a last db parameter.

Now let’s look at an example:

@db
def save(book, db):
    db.add(book)
    db.flush()
    return book


@transactional
def save_book(title):
    book = save(BookModel(title=title))
    # Return Pydantic model from service (recommended)
    return BookInfo.from_orm(book)

Done!

(Optional!) Implement sample endpoint calling save_book method

@book_router.get("/save-book/{title}")
def handle_save_book(title):
    return save_book(title)

Description:

Repository link: https://github.com/kokadva/transactional-methods

This method is quite general and works for Flask too.