ContextVars in async context

[Marce Coll]

Hi python people, hope this is the correct place to ask this!

For a transactional async decorator I'm building I am using contextvars in order to know when a transaction is open in my current context.

My understanding is that if given the following call stack

A
|- B
|  |- C
|- D
   |- E

If you set a context var at A with value 1, and then override it at B with value 2, then A, D and E will see value 1 and B and C will se value 2. Very similar (although a bit more manual) than dynamic scopes in common lisp.

Now, back to the transactional decorator, from the little documentation there is about this I would imagine that something like this:

@asynccontextmanager
async def transactional(
        endpoint: _base.Endpoint,
        isolation_level: IsolationLevel = IsolationLevel.READ_COMMITTED,
):
    session = TRANSACTION_VAR.get()
    if not session:
        async with endpoint.session() as session, session.begin():
            tok = TRANSACTION_VAR.set(session)
            try:
                await session.execute(text(f"SET TRANSACTION ISOLATION LEVEL {isolation_level.value}"))
                yield session
            finally:
                TRANSACTION_VAR.reset(tok)
    else:
        yield session

should work automatically, and with some preliminary tests it seems it works, but, I don't understand why I need the manual reset. I thought just setting the variable would change it downstream the callstack and reset it when leaving the current context, as seen in the asyncio example in the documentation https://docs.python.org/3/library/contextvars.html#asyncio-support

Now the final question, while the current setup works, when trying to use this as a context manager inside an async generator and then wrapping that with like `asyncio.wait_for`. To illustrate this: 

--
async def gen():
    with transactional():
        ...
        yield data

g = gen()
asyncio.wait_for(g.__anext__(), 1)  # ValueError: <Token ...> was created in a different Context
--

As far as I understand, that moves the awaitable into a new task and when trying to reset tok, it complains that tok was created in a different context. This is unfortunate as sometimes you want to be able to add a timeout using wait_for, particularly in tests. Is there anything I could do to my code to make it more reliable and robust when working with contextvars in an async context?

Sorry for the wall of text, hope it's understandable

Thanks in advance,
-- 
  Marce Coll
  marce at dziban.net