Usage of ContextVar in Fastapi

[Sharleedah]

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.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Structure:
- run.py
- app
  - __init__.py
  - database.py
  - main.py

------------------Content __init__.py------------------
from contextvars import ContextVar
from fastapi import Request

class ContextWrapper:

    def __init__(self, value):
        self.__value: ContextVar = value

    def set(self, value):
        return self.__value.set(value)

    def reset(self, token):
        self.__value.reset(token)
        return

    def __module__(self):
        return self.__value.get()

    @property
    def value(self):
        return self.__value.get()

request: ContextWrapper = ContextWrapper(ContextVar("request", default=None))
db: str = ContextWrapper(ContextVar("db", default=None))

------------------Content database.py------------------
from sqlalchemy.orm import sessionmaker, declarative_base 
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy import Column, String

engine = create_async_engine("sqlite+aiosqlite:///data.db", future=True, echo=False)
async_session = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
Base = declarative_base()

class UserTable(Base):
    __tablename__ = "User"

    mail                = Column(String, nullable=False, primary_key = True) 
    salt                = Column(String, nullable=False)
    hash                = Column(String, nullable=False)
    
    def __repr__(self): #how the object is shown when printed
        return f"User('{self.mail}','{self.count}')"

------------------Content main.py------------------
from fastapi import Depends, FastAPI, Request
from app.database import engine, Base

app = FastAPI(
    title="TEST APP",
    description="",
    version=1.0,
    docs_url= None, 
    redoc_url= None
)

@app.middleware("http")
async def add_middleware(request: Request, call_next):
    from app import db, request as ctx_request 
    from app.database import async_session
    token = ctx_request.set(request)
    async with async_session() as db_session:
        async with db_session.begin():
            db_token = db.set(db_session)
            response = await call_next(request) 
            db.reset(db_token)

    ctx_request.reset(token)
    return response

@app.on_event("startup")
async def startup():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

@app.get("/test", summary="TEST", include_in_schema=False)
async def test():
    from app import db 
    from app.database import UserTable
    user_obj = UserTable(mail="[email protected]",salt="test", hash="hash")
    db.value.add(user_obj)
    await db.value.flush()
    return {"success": True}

------------------Content run.py------------------
from app.main import app 
import uvicorn

if __name__ == '__main__':
    uvicorn.run("run:app", host="0.0.0.0", port=80, reload = True, workers = 1)

Description

I am building an application with FastAPI for 200.000 users, where at some point at 10.000 are accessing the application at the same time. Fastapi will be of course be deployed with gunicorn/uvicorn with auto scaling and a loadbalancer in the aws. But there is one question, which i couldn't solve. I do not get any errors but i cannot rule out that there can be a problem, when people use the application at the same time.

Question:
Is the usage of ContextVar (see defintion in init.py and usage in main.py) secure, when many people access the api at the same time (thread etc. wise)? So that the database session or the request are set correctly in the http middleware, that i can access those contexts from any point in the application (when working on a request) without having to fear that may a function get a request context from another user (authentication etc. wise)
Or do you have a better solution in mind?

Operating System

Linux, Windows, macOS, Other

Operating System Details

No response

FastAPI Version

0.75.0

Python Version

3.10.0

Additional Context

No response

[yinziyan1206]

I guess it is secure but i dont recommand to do this. The problem should be in fastapi but not contextvars.

[Sharleedah]

Thanks a lot for the answer. Yeah i through so aswell, that the problem - if there is one - lies in fastapi, unfortunately i couldn't find anything, which would indicate that it would(n't) cause a problem. Is there a reason why you would not suggest doing doing this?

[yinziyan1206]

contextvars is low level method. It is better to make a higher level method in fastapi framework. I’ve found one project called starlette_context. Maybe it works

[vanishingESCkey]

Or do you have a better solution in mind?

In my project the instance of db session is stored in a request.state so no need to worry about sharing it across users. Got middleware to set it and handy dependency for endpoints. Here is short summary:

# middleware
class DBSessionMiddleware(BaseHTTPMiddleware):  

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        response = Response("Internal server error", status_code=500)

        try:
            SessionFactory = sessionmaker(...)
            request.state.db = SessionFactory()
            response = await call_next(request)
        finally:
            request.state.db.close()

        return response

# depdendency
from sqlalchemy.orm import Session
from fastapi import Request

def get_db(request: Request) -> Session:
    return request.state.db

# endpoint
@app.get("/test")
async def test(db: Session = Depends(get_db)):
  # some db actions ending with db.commit() or db.rollback()

Got also two other middlewares (authentication, logging) that share some data using request.state scope. However I am not an expert of SQL Alchemy so would appreciate comments about disadvantages of this solution if any.

[Sharleedah]

Mhm i mean your solution is better than the default, where you handover each object individually, but you still have to hand over the request for every function... This is what i want to avoid with the ContextVar, where i can access the objects for each request through simple objects without any handover

[Elfpkck]

@Sharleedah I'm wondering what was your final solution. Have you used ContextVar? Any issues?

[Sharleedah]

Hey @Elfpkck
i went with the ContextVar Solution, due to a more lean code in general (i can access the db session, request, session … without having to put the request arg in every function). After tons of tests, mostly stressing the parallelism, there was not a single issue.

[uhLeeshUh]

If others are looking for an answer here, it worked for me to set a ContextVar (which is available to the path operation) as an async dependency to the path operation since FastAPI calls async deps in the same execution context as the path operation. A similar example for setting database state can be found in the FastAPI docs here

[daymien]

Thanks a lot @Sharleedah. Seems your code works with none async session as well

@app.middleware("http")
async def add_middleware(request: Request, call_next):
    from backend.contexts import db
    from backend.contexts import request as ctx_request
    from backend.database import session

    token = ctx_request.set(request)
    with session() as db_session:
        # with db_session.begin(): # this raises sqlalchemy.exc.InvalidRequestError in sync mode
        db_token = db.set(db_session)
        response = await call_next(request)
        db.reset(db_token)

    ctx_request.reset(token)
    return response

I do need the current db session in my modell.py for recursive cte and did not find any other good solution to include this property in my pedantic schema. Your solution just saved my day.

class Item(Model):
    ...
    @property
    def descandants(self):
        top_query = select(Item).where(Item.id == self.id).cte("cte", recursive=True)
        bottom_query = select(Item).join(top_query, Item.ancestor_id == top_query.c.id)
        recursive_query = top_query.union_all(bottom_query)
        query = select(aliased(Item, recursive_query))
        return db.session.execute(query).scalars().all()
    ...

[dlahyani]

@Sharleedah this looks good.
I want to use this approach for similar stuff in our code (e.g. a logger with some per-request state, request TTL).
There's just one thing that I don't understand - what causes the creation of a new Context for every incoming request. I know that when a new Task is pushed on the event-loop you can provide a context, and if None if provided a new Context will get created by calling contextvars.copy_context (which create a new Context that is a copy of the current context). The copied context is great, but it happens when a Task is pushed on the event queue, not for awaits though.

I wonder where in code of FastAPI / uvcorn / asyncio the new context gets created. Did anyone dive to the bottom of this?

[taleinat]

I wonder where in code of FastAPI / uvcorn / asyncio the new context gets created. Did anyone dive to the bottom of this?

It's done by asyncio. See the "asyncio support" section in the Python ContextVar docs:

"Context variables are natively supported in asyncio and are ready to be used without any extra configuration."

[christiansicari]

I am sorry, I don't think this answers the question.
@dlahyani I am trying to have evidence of the same question: "When a new context is generated?".
I didn't come up with a response yet, but to me, the only way to works is that an asyncio.Task is created for each request, I think so because if the requests were managed by a shared Task, this would have compromised the ContextVar, while if for each Request an asyncio.Task is created then, the contextvars are kept isolated.

[dlahyani]

@christiansicari I am a bit out of context (pun-intended) as this was a while back, but in general I think you are right.

[fievelk]

@Sharleedah could you explain why you need to create two contextvars (one for the request and the other for the db session)? Wouldn't just one (db.set(db_session)) be enough?
(Thanks for the solution you posted!)

[tuanvumaihuynh]

Yes you just need to use db context. Sometimes you will need to use request context such as accessing cookies, it may be useful.

[christiansicari]

I am currently experimenting this Dependency to set the Context, and until know I am not facing issues, but any feedback is very welcome

Database Session Geenerator:

@asynccontextmanager
async def get_async_db_session():
    async with SessionLocal() as session:
        try:
            yield session
            await session.commit()
        except:  # noqa
            await session.rollback()
            raise
        finally:
            await session.close()

Dependency

async def context_var_db():
    async with get_async_db_session() as session:
        token = db_session.set(session)
        yield
        db_session.reset(token)

route:

@app.get("/contextvar", dependencies=[Depends(context_var_db)])
async def print_context_var_session_id():
    print("Enjoying my session id", id(db_session.get())

[fievelk]

If you use db_session.commit() in your code, the session.rollback() in the except clause won't have effects. Do you then avoid having session.commit() in your code? Could you give more information about how you deal with this?

[christiansicari]

@fievelk probably missed something about contextmanagers.
A router will get the object yielded and it will apply its business logic. When the router returns the control goes back to get_async_db_session that will continue after the yield instruction, then it will commit.
If the router raises an exception it will terminate, but the exception will be propagated to get_async_db_session that will again continue after the yield instruction. Since it received an exception it will move to the Except block, and it will rollback everything, skipping the commit.
Try this

@app.get("/contextvar-raise", dependencies=[Depends(context_var_db)])
async def print_context_var_session_id():
    """
        I am raising an exception to test the contextvar reset
    """
    raise HTTPException(status_code=404, detail="In the log you should not see a commit message")

@asynccontextmanager
async def get_async_db_session():
    async with SessionLocal() as session:
        try:
            session.session_id = uuid4()
            yield session
            await session.commit()
            print("Committing")
        except:  # noqa
            await session.rollback()
            print("rollbacking")
            raise

your output will look like

INFO:     127.0.0.1:54600 - "GET /openapi.json HTTP/1.1" 200 OK
rollbacking
INFO:     127.0.0.1:54631 - "GET /contextvar-raise HTTP/1.1" 404 Not Found

Not sure if this was your question.
As final remark, yes I NEVER commit in my code, if I want to see changes during a query I can flush, but as pattern, I leave commit and rollback to an outer manager

[fievelk]

@christiansicari thank you, you answered my question in the last remark you made. I inherited a codebase where session.commit() is called within inner functions and, as I mentioned, using session.rollback() doesn't affect data that's already been committed.
I made a few experiments and also came up with the solution of calling session.flush(). However, there are situations in which you don't want to rollback all the operations that have been executed in the request, but only part of them. I suppose in that case using session.rollback() would revert the db state to where it was at the time of the last session.flush()?

[christiansicari]

Ok Sorry, I misunderstood what you asked.
Btw, nope the flush is not like a savepont, if you rollback, you rollback everything, even flushed instructions.
If you want to keep some data you should work with inner transaction, but I don't know much about that, a part that you can handle a father session with multiple son sessions.
Also, I see a request as an atomic operation, therefore in my pattern , committing a part of the changes does not make sense.

[tuanvumaihuynh]

I like your solution, in my IoT application not all the route need session

[fievelk]

@Sharleedah how do you deal with commits and rollbacks? For example, suppose that you call two functions in your code:

@app.get("/some-endpoint")
async def some_endpoint():
    function_a()
    function_b()

Each of these functions needs to store a new item on the database, but if function_b() fails, we need to rollback the changes made by function_a() too.
If we use db.value.commit() internally (in function_a() and function_b()) we won't be able to rollback the changes. How do you handle this situation?

[tuanvumaihuynh]

In your function you should use flush to sync change to db and it not write into db until you commit