Store acquired connection in context variable

Author: tulir

mautrix/util/async_db/aiosqlite.py

@@ -181,7 +181,7 @@ async def stop(self) -> None:
             self._conns -= 1
             await conn.close()
 
-    def acquire(self) -> AsyncContextManager[LoggingConnection]:
+    def acquire_direct(self) -> AsyncContextManager[LoggingConnection]:
         if self._parent:
             return self._parent.acquire()
         return self._acquire()

mautrix/util/async_db/asyncpg.py

@@ -96,7 +96,7 @@ async def _handle_exception(self, err: Exception) -> None:
             sys.exit(26)
 
     @asynccontextmanager
-    async def acquire(self) -> LoggingConnection:
+    async def acquire_direct(self) -> LoggingConnection:
         async with self.pool.acquire() as conn:
             yield LoggingConnection(
                 self.scheme, conn, self.log, handle_exception=self._handle_exception

mautrix/util/async_db/database.py

@@ -7,6 +7,8 @@
 
 from typing import Any, AsyncContextManager, Type
 from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
 import logging
 
 from yarl import URL
@@ -23,6 +25,8 @@
     from aiosqlite import Cursor
     from asyncpg import Record
 
+conn_var: ContextVar[LoggingConnection | None] = ContextVar("db_connection", default=None)
+
 
 class Database(ABC):
     schemes: dict[str, Type[Database]] = {}
@@ -128,9 +132,22 @@ async def stop(self) -> None:
         pass
 
     @abstractmethod
-    def acquire(self) -> AsyncContextManager[LoggingConnection]:
+    def acquire_direct(self) -> AsyncContextManager[LoggingConnection]:
         pass
 
+    @asynccontextmanager
+    async def acquire(self) -> LoggingConnection:
+        conn = conn_var.get(None)
+        if conn is not None:
+            yield conn
+            return
+        async with self.acquire_direct() as conn:
+            token = conn_var.set(conn)
+            try:
+                yield conn
+            finally:
+                conn_var.reset(token)
+
     async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str | Cursor:
         async with self.acquire() as conn:
             return await conn.execute(query, *args, timeout=timeout)