Tom
Tom

Reputation: 7992

ContextVar set and reset in the same function fails - created in a different context

I have this function:

async_session = contextvars.ContextVar("async_session")

async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    async with async_session_maker() as session:
        try:
            _token = async_session.set(session)
            yield session
        finally:
            async_session.reset(_token)

This fails with:

ValueError: <Token var=<ContextVar name='async_session' at 0x7e1470e00e00> at 0x7e14706d4680> was created in a different Context

How can this happen? AFAICT the only way for the Context to be changed is for a whole function call. So how can the context change during a yield?

This function is being used as a FastAPI Depends in case that makes a difference - but I can't see how it does. It's running under Python 3.8 and the version of FastAPI is equally ancient - 0.54.

Upvotes: 2

Views: 311

Answers (2)

Tom
Tom

Reputation: 7992

For anyone who comes after me, the actual problem is that my generator function is not async (the example in the question is misleading). FastAPI runs async and sync dependencies in different contexts to avoid the sync dependencies holding up the async loop thread, which is why it is then cleaned up in the wrong context.

Upvotes: 0

jsbueno
jsbueno

Reputation: 110591

If the generator would be used as a regular generator, in a for loop we control, it is really hard to think on how a context for contextvars could change. But the framework can do its own thing: it can store the generator in a variable, and call its __next__ methods from arbitrary contexts.

Actually, ContextVars are not even meant to work with async generators - check the full text of PEP 567 - it was rewritten, simplifying the original proposal at PEP 550, because the iteration with pausing frames and running outer frame code became too complex, and that capability was simply stripped-out.

I have a package to make the use of contextvars more simple - "extracontext" - and I have support for contextvars in async-generators there.

Either way, my proposal is to bring back the "threading.local" namespace, and I did add "context manager" capabilities to "contextvars".

You can just pip install python-extracontext. I didn't write much docs besides docstrings and what is on embedded in the README at https://github.com/jsbueno/extracontext

WIth this package you should be able to use:

from extracontext import ContextLocal

# async_session = contextvars.ContextVar("async_session")
async_session_ns = ContextLocal()

async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    async with (async_session_maker() as session, async_session_ns):
        with async_session_ns:
            async_session_ns.session = session
            # any code wanting this, should be able to use just
            # "async_session_ns.session" in any expression

            yield session

            # no finally blocks needed, as ContextLocal
            # works as a context manager.        

Please tell me if it does not work for you for any reason.

Upvotes: 0

Related Questions