CODEWITHSUNDEEP
pythonsqlalchemyfastapipython-contextvars
Альберт Александров
Альберт Александров

Reputation: 857

What happens with ContextVar if don't reset it?

With code below what will happen to SqlAlchemy session that is set in ContextVar if not to reset it?

from contextvars import ContextVar

from fastapi import FastAPI, BackgroundTasks
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession

engine = create_async_engine("postgresql+asyncpg://postgres:postgres@localhost:5433/postgres", isolation_level="AUTOCOMMIT")
session_context = ContextVar("session")

app = FastAPI()


class Repository:

    async def get_one(self):
        return await self.execute(text("select 1"))

    async def get_two(self):
        return await self.execute(text("select 2"))

    async def execute(self, statement):
        try:
            session = session_context.get()
        except LookupError:
            session = AsyncSession(engine)
            session_context.set(session)
        print(session)
        result = (await session.execute(statement)).scalar()
        await session.close()  # for some reason I need to close session every time
        return result


async def check_connections_statuses():
    print(engine.pool.status())
    print(session_context.get())


@app.get("/")
async def main(background_tasks: BackgroundTasks):
    repo = Repository()
    print(await repo.get_one())
    print(await repo.get_two())
    background_tasks.add_task(check_connections_statuses)

The code reproduces output:

<sqlalchemy.ext.asyncio.session.AsyncSession object at 0x121bb31a0>
1
<sqlalchemy.ext.asyncio.session.AsyncSession object at 0x121bb31a0>
2
INFO:     127.0.0.1:59836 - "GET / HTTP/1.1" 200 OK
Pool size: 5  Connections in pool: 1 Current Overflow: -4 Current Checked out connections: 0
<sqlalchemy.ext.asyncio.session.AsyncSession object at 0x121bb31a0>

It seems that connections does not leak. But what about session object in ContextVar? As you could see it is still in ContextVar. Is it save not to reset it? Or Python somehow resets it?

Upvotes: 1

Views: 63

Answers (1)

jsbueno
jsbueno

Reputation: 110591

Yes, it will be reset as soon as the current HTTP view is done. In the example above, it happens when the co-routine main returns.

Contextvars where added to the language so that concurrent Async tasks can "see" independent context information, even if the tasks run the same functions (think of several HTTP requests in parallel running your main view). The act of creating a concurrent task in asyncio, which FastAPI does for your before calling main, automatically forks the context, and that task will run in the new context. When the task is done, the forked context is disposed (unless explicitly saved with something like a contextvars.copy_context() call while the task is active)

The snippet below demonstrates how it works. Note that it is a stand alone program setting up its own asyncio loop, while in a FastAPI app, that is done by the framework + application server (asgi or uvicorn)

import asyncio
import contextvars


myctx = contextvars.ContextVar("myctx", default=None)

async def task(n):
    await asyncio.sleep(0.1)
    myctx.set(n)
    await asyncio.sleep(n)
    print(f"Value in contextvar: {myctx.get()}")

async def main():
    t1 = task(1)
    t2 = task(2)
    await asyncio.gather(t1, t2)
    print(f"Value in contextvar, after running tasks: {myctx.get()}")


asyncio.run(main())

Its output being:

Value in contextvar: 1
Value in contextvar: 2
Value in contextvar, after running tasks: None

Upvotes: 2

Related Questions