Can't get the correct contextvar set in the request handler in middleware's dispatch method

[gntzh]

Checklist

• The bug is reproducible against the latest release and/or master.
• There are no similar issues or pull requests to fix it yet.

Describe the bug

I try to manage database session with ContextVar and want to close it in the middleware, but I always get none unless the session was set in the middleware (This is valid, but not all handlers require sqlalchemy session).

To reproduce

from contextvars import ContextVar
from fastapi import Request
from sqlalchemy.orm import Session
from sqlalchemy import create_engine
from starlette.middleware.base import BaseHTTPMiddleware
from src.config import settings


engine = create_engine(
    settings.SQLALCHEMY_DATABASE_URI,
)

session_var = ContextVar("session", default=None)


def get_session() -> Session:
    if (rst := session_var.get()) is None:
        rst = Session(engine, future=True)
        print("create a session: ", rst)
        session_var.set(rst)
    print("return a session: ", rst)
    return rst


class DBSessionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # with get_session().begin():
        #     pass
        res = await call_next(request)
        # FIXME Can't get the correct contextvar, unless session is set here.
        print("session in context: ", session_var.get())
        if (session := session_var.get()) is not None:
            session.close()
            print("close the session: ", session)
        return res

Expected behavior

Get the session if it is used in the handler, close it

Actual behavior

Always get none

Debugging material

Environment

• OS: WSL 2/Ubuntu 20.04
• Python version: 3.8.5
• Starlette version: 0.13.6 (via FastAP 0.61.1I)

Additional context

[uriyyo]

Hi @ShoorDay,

The problem not with starlette, there is a problem with how you use contextvars.
Basically, when you set contextvar inside a child task it won't be available at the parent.

An example regarding this problem:

from asyncio import run, create_task
from contextvars import ContextVar
from typing import Optional

VAR: ContextVar[Optional[str]] = ContextVar('VAR', default=None)


async def child():
    print(f'before setting VAR: {VAR.get()}')
    VAR.set('value')
    print(f'after setting VAR: {VAR.get()}')


async def main():
    print(f'before call child: {VAR.get()}')
    await create_task(child())
    print(f'after call child: {VAR.get()}')


if __name__ == '__main__':
    run(main())

The result will be:

before call child: None
before setting: VAR None
after setting: VAR value
after call child: value

In order to make your code working using contextvars you should update code to smth like this:

from contextvars import ContextVar
from dataclasses import dataclass
from typing import Optional

from fastapi import Request
from sqlalchemy.orm import Session
from sqlalchemy import create_engine
from starlette.middleware.base import BaseHTTPMiddleware

from src.config import settings

engine = create_engine(
    settings.SQLALCHEMY_DATABASE_URI,
)

session_var: ContextVar[Optional['SessionHolder']] = ContextVar("session", default=None)


@dataclass
class SessionHolder:
    session: Optional[Session] = None

    @property
    def is_session_inited(self) -> bool:
        return self.session is not None

    def get_session(self) -> Session:
        if self.session is None:
            # self.session = Session(engine, future=True)
            self.session = Session(engine)
            print("create a session: ", self.session)

        print("return a session: ", self.session)
        return self.session


def get_session() -> Session:
    holder = session_var.get()
    return holder.get_session()


class DBSessionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        session_var.set(SessionHolder())

        res = await call_next(request)

        print("session in context: ", session_var.get())
        if (holder := session_var.get()).is_session_inited:
            holder.get_session().close()
            print("close the session: ", holder)

        return res

If you have any questions or maybe I missed smth, feel free to ask)

[gntzh]

@uriyyo thx!

[cglacet]

Hello @uriyyo ! Thanks for the code you provided. How does it compare to using request.state to store some context (like the database session for example). Is there any advantage in doing this the way you showed here?

[uriyyo]

Hi @cglacet,

The main advantage of using contextvars is that you don't need to pass request instance everywhere in order to access some ctx data.

If take into account example above, you can access to session from any place and it will be context-bound:

async def some_route():
    session = get_session()
    # do some staff with session

For more information please follow PEP-567

[cglacet]

@uriyyo oh ok, I see, never heard of this construct in python, that looks interesting. If I get this correctly each asynchronous task will get its own context. I have no idea how Starlette dispatches requests, but I guess each request runs on a single task so it means we get a request based context for free?

Here is a small example that tests contextvars in case someone is interested:

import asyncio 
import random
from contextvars import ContextVar

index_context = ContextVar("index", default=None)

async def test(index):
    index_context.set(index)
    await asyncio.sleep(random.randint(1, 10)/10)
    return index == index_context.get()

success = await asyncio.gather(*[test(i) for i in range(10)])
assert all(success)

[uriyyo]

@cglacet As far as I know, for instance, uvicorn will create a separate task to handle request. So basically, you are totally right, you have request context for free)