feat: enter_connection

Author: phi-friday

src/typed_diskcache/core/context.py

@@ -3,7 +3,7 @@
 import inspect
 import threading
 from contextlib import contextmanager
-from contextvars import ContextVar, Token, copy_context
+from contextvars import Context, ContextVar, Token, copy_context
 from functools import partial, wraps
 from typing import TYPE_CHECKING, Any, overload
 
@@ -14,13 +14,26 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
 
-__all__ = ["log_context", "context"]
+    from sqlalchemy.engine import Connection
+    from sqlalchemy.ext.asyncio import AsyncConnection
+
+__all__ = [
+    "log_context",
+    "conn_context",
+    "aconn_context",
+    "enter_connection",
+    "context",
+]
 
 _F = TypeVar("_F", bound="Callable[..., Any]", infer_variance=True)
 
 log_context: ContextVar[tuple[str, int]] = ContextVar(
     "log_thread_context", default=(DEFAULT_LOG_CONTEXT, DEFAULT_LOG_THREAD)
 )
+conn_context: ContextVar[Connection | None] = ContextVar("conn_context", default=None)
+aconn_context: ContextVar[AsyncConnection | None] = ContextVar(
+    "aconn_context", default=None
+)
 
 
 @contextmanager
@@ -59,6 +72,33 @@ def context(func_or_context: _F | str) -> _F | Callable[[_F], _F]:
     return partial(_context, name=func_or_context)
 
 
+@contextmanager
+def enter_connection(
+    conn: Connection | AsyncConnection,
+) -> Generator[Context, None, None]:
+    """Enter the connection context.
+
+    Args:
+        conn: The connection to enter.
+
+    Yields:
+        Copy of the current context.
+    """
+    from sqlalchemy.ext.asyncio import AsyncConnection
+
+    if isinstance(conn, AsyncConnection):
+        token = aconn_context.set(conn)
+        reset_context = aconn_context.reset
+    else:
+        token = conn_context.set(conn)
+        reset_context = conn_context.reset
+    context = copy_context()
+    try:
+        yield context
+    finally:
+        reset_context(token)  # pyright: ignore[reportArgumentType]
+
+
 def _context(func: _F, *, name: str) -> _F:
     if inspect.iscoroutinefunction(func):
 

src/typed_diskcache/database/connection.py

@@ -5,9 +5,11 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import anyio.lowlevel
 import sqlalchemy as sa
 
 from typed_diskcache import exception as te
+from typed_diskcache.core.context import aconn_context, conn_context
 from typed_diskcache.core.types import EvictionPolicy
 from typed_diskcache.database import connect as db_connect
 from typed_diskcache.database.model import Cache
@@ -127,12 +129,23 @@ def _async_engine(self) -> AsyncEngine:
     @contextmanager
     def connect(self) -> Generator[SAConnection, None, None]:
         """Connect to the database."""
+        conn = conn_context.get()
+        if conn is not None:
+            yield conn
+            return
+
         with self._sync_engine.connect() as connection:
             yield connection
 
     @asynccontextmanager
     async def aconnect(self) -> AsyncGenerator[AsyncConnection, None]:
         """Connect to the database."""
+        conn = aconn_context.get()
+        if conn is not None:
+            await anyio.lowlevel.checkpoint()
+            yield conn
+            return
+
         async with self._async_engine.connect() as connection:
             yield connection
 

src/typed_diskcache/implement/cache/default/main.py

@@ -14,7 +14,7 @@
 
 from typed_diskcache import exception as te
 from typed_diskcache.core.const import ENOVAL
-from typed_diskcache.core.context import context
+from typed_diskcache.core.context import context, enter_connection
 from typed_diskcache.core.types import (
     Container,
     EvictionPolicy,
@@ -736,9 +736,13 @@ def stats(self, *, enable: bool = True, reset: bool = False) -> Stats:
                 .where(SettingsTable.key == SettingsKey.STATISTICS)
                 .values(value=enable)
             )
-        self.update_settings(self.settings.model_copy(update={"statistics": enable}))
+            with enter_connection(sa_conn) as context:
+                context.run(
+                    self.update_settings,
+                    self.settings.model_copy(update={"statistics": enable}),
+                )
 
-        return stats
+            return stats
 
     @context
     @override
@@ -768,10 +772,12 @@ async def astats(self, *, enable: bool = True, reset: bool = False) -> Stats:
                 .where(SettingsTable.key == SettingsKey.STATISTICS)
                 .values(value=enable)
             )
-        await self.aupdate_settings(
-            self.settings.model_copy(update={"statistics": enable})
-        )
-        return stats
+            with enter_connection(sa_conn) as context:
+                await context.run(
+                    self.aupdate_settings,
+                    self.settings.model_copy(update={"statistics": enable}),
+                )
+            return stats
 
     @override
     def close(self) -> None: