jararaca 0.2.37a12__py3-none-any.whl → 0.4.0a5__py3-none-any.whl

jararaca/persistence/interceptors/aiosqa_interceptor.py

@@ -1,23 +1,97 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from contextlib import asynccontextmanager, contextmanager, suppress
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Generator
-
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from typing import Any, AsyncGenerator, Generator, Protocol
+
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    AsyncSessionTransaction,
+    async_sessionmaker,
+    create_async_engine,
+)
 from sqlalchemy.ext.asyncio.engine import AsyncEngine
 
-from jararaca.microservice import AppContext, AppInterceptor
+from jararaca.microservice import AppInterceptor, AppTransactionContext
+from jararaca.persistence.interceptors.constants import DEFAULT_CONNECTION_NAME
+from jararaca.persistence.interceptors.decorators import (
+    INJECT_PERSISTENCE_SESSION_METADATA_TEMPLATE,
+)
+from jararaca.reflect.metadata import get_metadata_value
+
+
+class SessionManager(Protocol):
+    def spawn_session(self, connection_name: str | None = None) -> AsyncSession: ...
 
-ctx_session_map = ContextVar[dict[str, AsyncSession]]("ctx_session_map", default={})
+
+ctx_session_manager: ContextVar[SessionManager | None] = ContextVar(
+    "ctx_session_manager", default=None
+)
 
 
 @contextmanager
-def provide_session(
-    connection_name: str, session: AsyncSession
+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
+)
+
+
+def ensure_name(name: str | None) -> str:
+    return ctx_default_connection_name.get()
+
+
+@dataclass
+class PersistenceCtx:
+    session: AsyncSession
+    tx: AsyncSessionTransaction
+
+
+ctx_session_map = ContextVar[dict[str, PersistenceCtx]]("ctx_session_map", default={})
+
+
+@contextmanager
+def providing_session(
+    session: AsyncSession,
+    tx: AsyncSessionTransaction,
+    connection_name: str | None = None,
+) -> Generator[None, Any, None]:
+    """
+    Context manager to provide a session and transaction for a specific connection name.
+    If no connection name is provided, it uses the default connection name from the context variable.
+    """
+    connection_name = ensure_name(connection_name)
     current_map = ctx_session_map.get({})
 
-    token = ctx_session_map.set({**current_map, connection_name: session})
+    token = ctx_session_map.set(
+        {**current_map, connection_name: PersistenceCtx(session, tx)}
+    )
 
     try:
         yield
@@ -26,21 +100,86 @@ def provide_session(
             ctx_session_map.reset(token)
 
 
-def use_session(connection_name: str = "default") -> AsyncSession:
+provide_session = providing_session
+"""
+Alias for `providing_session` to maintain backward compatibility.
+"""
+
+
+@asynccontextmanager
+async def providing_new_session(
+    connection_name: str | None = None,
+) -> AsyncGenerator[AsyncSession, None]:
+
+    session_manager = use_session_manager()
+    current_session = session_manager.spawn_session(connection_name)
+
+    async with AsyncSession(
+        current_session.bind,
+    ) as new_session, new_session.begin() as new_tx:
+        with providing_session(new_session, new_tx, connection_name):
+            try:
+                yield new_session
+                if new_tx.is_active:
+                    await new_tx.commit()
+            except Exception:
+                if new_tx.is_active:
+                    await new_tx.rollback()
+                raise
+
+
+def use_session(connection_name: str | None = None) -> AsyncSession:
+    connection_name = ensure_name(connection_name)
     current_map = ctx_session_map.get({})
     if connection_name not in current_map:
-        raise ValueError(f"Session not found for connection {connection_name}")
+        raise ValueError(
+            f'Session not found for connection "{connection_name}" in context. Check if your interceptor is correctly set up.'
+        )
 
-    return current_map[connection_name]
+    return current_map[connection_name].session
+
+
+@contextmanager
+def providing_transaction(
+    tx: AsyncSessionTransaction,
+    connection_name: str | None = None,
+) -> Generator[None, Any, None]:
+    connection_name = ensure_name(connection_name)
+
+    current_map = ctx_session_map.get({})
+
+    if connection_name not in current_map:
+        raise ValueError(f"No session found for connection {connection_name}")
+
+    with providing_session(current_map[connection_name].session, tx, connection_name):
+        yield
+
+
+def use_transaction(connection_name: str | None = None) -> AsyncSessionTransaction:
+    current_map = ctx_session_map.get({})
+    if connection_name not in current_map:
+        raise ValueError(f"Transaction not found for connection {connection_name}")
+
+    return current_map[connection_name].tx
 
 
-@dataclass
 class AIOSQAConfig:
-    connection_name: str
     url: str | AsyncEngine
+    connection_name: str
+    inject_default: bool
+
+    def __init__(
+        self,
+        url: str | AsyncEngine,
+        connection_name: str = DEFAULT_CONNECTION_NAME,
+        inject_default: bool = True,
+    ):
+        self.url = url
+        self.connection_name = connection_name
+        self.inject_default = inject_default
 
 
-class AIOSqlAlchemySessionInterceptor(AppInterceptor):
+class AIOSqlAlchemySessionInterceptor(AppInterceptor, SessionManager):
 
     def __init__(self, config: AIOSQAConfig):
         self.config = config
@@ -53,12 +192,37 @@ class AIOSqlAlchemySessionInterceptor(AppInterceptor):
         self.sessionmaker = async_sessionmaker(self.engine)
 
     @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
-        async with self.sessionmaker() as session:
-            with provide_session(self.config.connection_name, session):
-                try:
-                    yield
-                    await session.commit()
-                except Exception as e:
-                    await session.rollback()
-                    raise e
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
+
+        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

jararaca/persistence/interceptors/constants.py

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+DEFAULT_CONNECTION_NAME = "default"

jararaca/persistence/interceptors/decorators.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+
+from jararaca.persistence.interceptors.constants import DEFAULT_CONNECTION_NAME
+from jararaca.reflect.metadata import SetMetadata
+
+INJECT_PERSISTENCE_SESSION_METADATA_TEMPLATE = (
+    "inject_persistence_template_{connection_name}"
+)
+
+
+def set_use_persistence_session(
+    inject: bool, connection_name: str = DEFAULT_CONNECTION_NAME
+) -> SetMetadata:
+    """
+    Set whether to inject the connection metadata for the given connection name.
+    This is useful when you want to control whether the connection metadata
+    should be injected into the context or not.
+    """
+
+    return SetMetadata(
+        INJECT_PERSISTENCE_SESSION_METADATA_TEMPLATE.format(
+            connection_name=connection_name
+        ),
+        inject,
+    )
+
+
+def uses_persistence_session(
+    connection_name: str = DEFAULT_CONNECTION_NAME,
+) -> SetMetadata:
+    """
+    Use connection metadata for the given connection name.
+    This is useful when you want to inject the connection metadata into the context,
+    for example, when you are using a specific connection for a specific operation.
+    """
+    return set_use_persistence_session(True, connection_name=connection_name)
+
+
+def skip_persistence_session(
+    connection_name: str = DEFAULT_CONNECTION_NAME,
+) -> SetMetadata:
+    """
+    Decorator to skip using connection metadata for the given connection name.
+    This is useful when you want to ensure that the connection metadata is not injected
+    into the context, for example, when you are using a different connection for a specific operation.
+    """
+    return set_use_persistence_session(False, connection_name=connection_name)

jararaca/persistence/session.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later

jararaca/persistence/sort_filter.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import re
 from datetime import date, datetime
 from functools import reduce

jararaca/persistence/utilities.py

@@ -1,5 +1,10 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import asyncio
 import logging
+import math
 from datetime import UTC, date, datetime
 from functools import reduce
 from typing import (
@@ -29,6 +34,7 @@ from jararaca.persistence.base import (
     BaseEntity,
     recursive_get_dict,
 )
+from jararaca.persistence.interceptors.aiosqa_interceptor import use_session
 from jararaca.persistence.sort_filter import (
     FilterModel,
     FilterRuleApplier,
@@ -61,7 +67,7 @@ class DatedEntity(BaseEntity):
         DateTime(timezone=True), nullable=False, default=nowutc
     )
     updated_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), nullable=False, default=nowutc
+        DateTime(timezone=True), nullable=False, default=nowutc, onupdate=nowutc
     )
 
 
@@ -213,7 +219,7 @@ class CRUDOperations(Generic[IDENTIFIABLE_T]):
 # region PaginatedFilter
 class PaginatedFilter(BaseModel):
     page: Annotated[int, Field(gt=-1)] = 0
-    page_size: int = 10
+    page_size: Annotated[int, Field(gt=0)] = 10
     sort_models: list[SortModel] = []
     filter_models: list[FilterModel] = []
 
@@ -283,11 +289,13 @@ class QueryOperations(Generic[QUERY_FILTER_T, QUERY_ENTITY_T]):
     def __init__(
         self,
         entity_type: Type[QUERY_ENTITY_T],
-        session_provider: Callable[[], AsyncSession],
-        filters_functions: list[QueryInjector],
+        session_provider: Callable[[], AsyncSession] = use_session,
+        filters_functions: list[QueryInjector] = [],
+        *,
         unique: bool = False,
         sort_rule_applier: SortRuleApplier | None = None,
         filter_rule_applier: FilterRuleApplier | None = None,
+        base_statement: Select[Tuple[QUERY_ENTITY_T]] | None = None,
     ) -> None:
         self.entity_type: type[QUERY_ENTITY_T] = entity_type
         self.session_provider = session_provider
@@ -295,6 +303,7 @@ class QueryOperations(Generic[QUERY_FILTER_T, QUERY_ENTITY_T]):
         self.unique = unique
         self.sort_rule_applier = sort_rule_applier
         self.filter_rule_applier = filter_rule_applier
+        self.base_statement = base_statement
 
     @property
     def session(self) -> AsyncSession:
@@ -303,55 +312,88 @@ class QueryOperations(Generic[QUERY_FILTER_T, QUERY_ENTITY_T]):
     async def query(
         self,
         filter: QUERY_FILTER_T,
+        *,
         interceptors: list[
             Callable[[Select[Tuple[QUERY_ENTITY_T]]], Select[Tuple[QUERY_ENTITY_T]]]
         ] = [],
+        base_statement: (
+            Callable[[Select[Tuple[QUERY_ENTITY_T]]], Select[Tuple[QUERY_ENTITY_T]]]
+            | Select[Tuple[QUERY_ENTITY_T]]
+            | None
+        ) = None,
     ) -> "Paginated[QUERY_ENTITY_T]":
+        """
+        Executes a query with the provided filter and interceptors.
+        Args:
+            filter: The filter to apply to the query.
+            interceptors: A list of functions that can modify the query before execution.
+        Returns:
+            Paginated[QUERY_ENTITY_T]: A paginated result containing the items and metadata.
+        """
+
+        initial_statement = self.base_statement or select(self.entity_type)
+
+        if base_statement and callable(base_statement):
+            initial_statement = base_statement(initial_statement)
+        elif base_statement and isinstance(base_statement, Select):
+            initial_statement = base_statement
+
+        tier_one_filtered_query = self.generate_filtered_query(
+            filter, initial_statement
+        )
 
-        query = reduce(
+        tier_two_filtered_query = reduce(
             lambda query, interceptor: interceptor(query),
             interceptors,
-            select(self.entity_type),
+            tier_one_filtered_query,
         )
 
         if self.sort_rule_applier:
-            query = self.sort_rule_applier.create_query_for_sorting_list(
-                query, filter.sort_models
+            tier_two_filtered_query = (
+                self.sort_rule_applier.create_query_for_sorting_list(
+                    tier_two_filtered_query, filter.sort_models
+                )
             )
 
         if self.filter_rule_applier:
-            query = self.filter_rule_applier.create_query_for_filter_list(
-                query, filter.filter_models
-            )
-
-        unpaginated_total = (
-            await self.session.execute(
-                select(func.count()).select_from(query.subquery())
+            tier_two_filtered_query = (
+                self.filter_rule_applier.create_query_for_filter_list(
+                    tier_two_filtered_query, filter.filter_models
+                )
             )
-        ).scalar_one()
 
-        filtered_query = self.generate_filtered_query(filter, query)
-
-        paginated_query = filtered_query.limit(filter.page_size).offset(
+        paginated_query = tier_two_filtered_query.add_columns(
+            func.count().over().label("total_count")
+        )
+        paginated_query = paginated_query.limit(filter.page_size).offset(
             (filter.page) * filter.page_size
         )
 
-        filtered_total = (
-            await self.session.execute(
-                select(func.count()).select_from(paginated_query.subquery())
-            )
-        ).scalar_one()
+        result = await self.session.execute(paginated_query)
+        result = self.judge_unique(result)
+        rows = result.all()
+
+        if rows:
+            unpaginated_total = rows[0].total_count
+            result_scalars = [row[0] for row in rows]
+        else:
+            result_scalars = []
+            if filter.page == 0:
+                unpaginated_total = 0
+            else:
+                unpaginated_total = (
+                    await self.session.execute(
+                        select(func.count()).select_from(
+                            tier_two_filtered_query.subquery()
+                        )
+                    )
+                ).scalar_one()
 
         return Paginated(
-            items=[
-                e
-                for e in self.judge_unique(
-                    await self.session.execute(paginated_query)
-                ).scalars()
-            ],
-            total=filtered_total,
+            items=result_scalars,
+            total=len(result_scalars),
             unpaginated_total=unpaginated_total,
-            total_pages=int(filtered_total / filter.page_size) + 1,
+            total_pages=math.ceil(unpaginated_total / filter.page_size),
         )
 
     def judge_unique(

jararaca/presentation/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later