jararaca 0.3.12a15__tar.gz → 0.3.12a16__tar.gz

{jararaca-0.3.12a15 → jararaca-0.3.12a16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jararaca
-Version: 0.3.12a15
+Version: 0.3.12a16
 Summary: A simple and fast API framework for Python
 Home-page: https://github.com/LuscasLeo/jararaca
 Author: Lucas S

{jararaca-0.3.12a15 → jararaca-0.3.12a16}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jararaca"
-version = "0.3.12a15"
+version = "0.3.12a16"
 description = "A simple and fast API framework for Python"
 authors = ["Lucas S <me@luscasleo.dev>"]
 readme = "README.md"

{jararaca-0.3.12a15 → jararaca-0.3.12a16}/src/jararaca/persistence/interceptors/aiosqa_interceptor.py

@@ -1,7 +1,7 @@
 from contextlib import asynccontextmanager, contextmanager, suppress
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Generator
+from typing import Any, AsyncGenerator, Generator, Protocol
 
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
@@ -18,6 +18,42 @@ from jararaca.persistence.interceptors.decorators import (
 )
 from jararaca.reflect.metadata import get_metadata_value
 
+
+class SessionManager(Protocol):
+    def spawn_session(self, connection_name: str | None = None) -> AsyncSession: ...
+
+
+ctx_session_manager: ContextVar[SessionManager | None] = ContextVar(
+    "ctx_session_manager", default=None
+)
+
+
+@contextmanager
+def providing_session_manager(
+    session_manager: SessionManager,
+) -> Generator[None, Any, None]:
+    """
+    Context manager to provide a session manager for the current context.
+    """
+    token = ctx_session_manager.set(session_manager)
+    try:
+        yield
+    finally:
+        with suppress(ValueError):
+            ctx_session_manager.reset(token)
+
+
+def use_session_manager() -> SessionManager:
+    """
+    Retrieve the current session manager from the context variable.
+    Raises ValueError if no session manager is set.
+    """
+    session_manager = ctx_session_manager.get()
+    if session_manager is None:
+        raise ValueError("No session manager set in the context.")
+    return session_manager
+
+
 ctx_default_connection_name: ContextVar[str] = ContextVar(
     "ctx_default_connection_name", default=DEFAULT_CONNECTION_NAME
 )
@@ -71,7 +107,8 @@ async def providing_new_session(
     connection_name: str | None = None,
 ) -> AsyncGenerator[AsyncSession, None]:
 
-    current_session = use_session(connection_name)
+    session_manager = use_session_manager()
+    current_session = session_manager.spawn_session(connection_name)
 
     async with AsyncSession(
         current_session.bind,
@@ -131,7 +168,7 @@ class AIOSQAConfig:
         self.inject_default = inject_default
 
 
-class AIOSqlAlchemySessionInterceptor(AppInterceptor):
+class AIOSqlAlchemySessionInterceptor(AppInterceptor, SessionManager):
 
     def __init__(self, config: AIOSQAConfig):
         self.config = config
@@ -148,27 +185,33 @@ class AIOSqlAlchemySessionInterceptor(AppInterceptor):
         self, app_context: AppTransactionContext
     ) -> AsyncGenerator[None, None]:
 
-        uses_connection_metadata = get_metadata_value(
-            INJECT_PERSISTENCE_SESSION_METADATA_TEMPLATE.format(
-                connection_name=self.config.connection_name
-            ),
-            self.config.inject_default,
-        )
-
-        if not uses_connection_metadata:
-            yield
-            return
-
-        async with self.sessionmaker() as session, session.begin() as tx:
-            token = ctx_default_connection_name.set(self.config.connection_name)
-            with providing_session(session, tx, self.config.connection_name):
-                try:
-                    yield
-                    if tx.is_active:
-                        await tx.commit()
-                except Exception as e:
-                    await tx.rollback()
-                    raise e
-                finally:
-                    with suppress(ValueError):
-                        ctx_default_connection_name.reset(token)
+        with providing_session_manager(self):
+            uses_connection_metadata = get_metadata_value(
+                INJECT_PERSISTENCE_SESSION_METADATA_TEMPLATE.format(
+                    connection_name=self.config.connection_name
+                ),
+                self.config.inject_default,
+            )
+
+            if not uses_connection_metadata:
+                yield
+                return
+
+            async with self.sessionmaker() as session, session.begin() as tx:
+                token = ctx_default_connection_name.set(self.config.connection_name)
+                with providing_session(session, tx, self.config.connection_name):
+                    try:
+                        yield
+                        if tx.is_active:
+                            await tx.commit()
+                    except Exception as e:
+                        await tx.rollback()
+                        raise e
+                    finally:
+                        with suppress(ValueError):
+                            ctx_default_connection_name.reset(token)
+
+    def spawn_session(self, connection_name: str | None = None) -> AsyncSession:
+        connection_name = ensure_name(connection_name)
+        session = self.sessionmaker()
+        return session