jararaca 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

jararaca/microservice.py

@@ -1,4 +1,5 @@
 import inspect
+import logging
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from dataclasses import dataclass, field
@@ -23,64 +24,105 @@ from fastapi import Request, WebSocket
 from jararaca.core.providers import ProviderSpec, T, Token
 from jararaca.messagebus import MessageOf
 from jararaca.messagebus.message import Message
+from jararaca.reflect.controller_inspect import ControllerMemberReflect
+
+logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from typing_extensions import TypeIs
 
 
 @dataclass
-class SchedulerAppContext:
+class SchedulerTransactionData:
     triggered_at: datetime
     scheduled_to: datetime
     cron_expression: str
-    action: Callable[..., Any]
     context_type: Literal["scheduler"] = "scheduler"
 
 
 @dataclass
-class HttpAppContext:
+class HttpTransactionData:
     request: Request
     context_type: Literal["http"] = "http"
 
 
 @dataclass
-class MessageBusAppContext:
+class MessageBusTransactionData:
     topic: str
     message: MessageOf[Message]
     context_type: Literal["message_bus"] = "message_bus"
 
 
 @dataclass
-class WebSocketAppContext:
+class WebSocketTransactionData:
     websocket: WebSocket
     context_type: Literal["websocket"] = "websocket"
 
 
-AppContext = (
-    MessageBusAppContext | HttpAppContext | SchedulerAppContext | WebSocketAppContext
+TransactionData = (
+    MessageBusTransactionData
+    | HttpTransactionData
+    | SchedulerTransactionData
+    | WebSocketTransactionData
 )
 
-app_context_ctxvar = ContextVar[AppContext]("app_context")
 
+@dataclass
+class AppTransactionContext:
+    transaction_data: TransactionData
+    controller_member_reflect: ControllerMemberReflect
+
+
+AppContext = AppTransactionContext
+"""
+Alias for AppTransactionContext, used for compatibility with existing code.
+"""
+
+
+app_transaction_context_var = ContextVar[AppTransactionContext]("app_context")
+
+
+def use_app_transaction_context() -> AppTransactionContext:
+    """
+    Returns the current application transaction context.
+    This function is used to access the application transaction context in the context of an application transaction.
+    If no context is set, it raises a LookupError.
+    """
 
-def use_app_context() -> AppContext:
-    return app_context_ctxvar.get()
+    return app_transaction_context_var.get()
+
+
+def use_app_tx_ctx_data() -> TransactionData:
+    """
+    Returns the transaction data from the current app transaction context.
+    This function is used to access the transaction data in the context of an application transaction.
+    """
+
+    return use_app_transaction_context().transaction_data
+
+
+use_app_context = use_app_tx_ctx_data
+"""Alias for use_app_tx_ctx_data, used for compatibility with existing code."""
 
 
 @contextmanager
-def provide_app_context(app_context: AppContext) -> Generator[None, None, None]:
-    token = app_context_ctxvar.set(app_context)
+def provide_app_context(
+    app_context: AppTransactionContext,
+) -> Generator[None, None, None]:
+    token = app_transaction_context_var.set(app_context)
     try:
         yield
     finally:
         with suppress(ValueError):
-            app_context_ctxvar.reset(token)
+            app_transaction_context_var.reset(token)
 
 
 @runtime_checkable
 class AppInterceptor(Protocol):
 
-    def intercept(self, app_context: AppContext) -> AsyncContextManager[None]: ...
+    def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncContextManager[None]: ...
 
 
 class AppInterceptorWithLifecycle(Protocol):
@@ -106,6 +148,49 @@ class Microservice:
     )
 
 
+@dataclass
+class InstantiationNode:
+    property_name: str
+    parent: "InstantiationNode | None" = None
+    source_type: Any | None = None
+    target_type: Any | None = None
+
+
+instantiation_vector_ctxvar = ContextVar[list[InstantiationNode]](
+    "instantiation_vector", default=[]
+)
+
+
+def print_instantiation_vector(
+    instantiation_vector: list[InstantiationNode],
+) -> None:
+    """
+    Prints the instantiation vector for debugging purposes.
+    """
+    for node in instantiation_vector:
+        print(
+            f"Property: {node.property_name}, Source: {node.source_type}, Target: {node.target_type}"
+        )
+
+
+@contextmanager
+def span_instantiation_vector(
+    instantiation_node: InstantiationNode,
+) -> Generator[None, None, None]:
+    """
+    Context manager to track instantiation nodes in a vector.
+    This is useful for debugging and tracing instantiation paths.
+    """
+    current_vector = list(instantiation_vector_ctxvar.get())
+    current_vector.append(instantiation_node)
+    token = instantiation_vector_ctxvar.set(current_vector)
+    try:
+        yield
+    finally:
+        with suppress(ValueError):
+            instantiation_vector_ctxvar.reset(token)
+
+
 class Container:
 
     def __init__(self, app: Microservice) -> None:
@@ -122,40 +207,54 @@ class Container:
                 if provider.use_value:
                     self.instances_map[provider.provide] = provider.use_value
                 elif provider.use_class:
-                    self.get_and_register(provider.use_class, provider.provide)
+                    self._get_and_register(provider.use_class, provider.provide)
                 elif provider.use_factory:
-                    self.get_and_register(provider.use_factory, provider.provide)
+                    self._get_and_register(provider.use_factory, provider.provide)
             else:
-                self.get_and_register(provider, provider)
+                self._get_and_register(provider, provider)
 
-    def instantiate(self, type_: type[Any] | Callable[..., Any]) -> Any:
+    def _instantiate(self, type_: type[Any] | Callable[..., Any]) -> Any:
 
-        dependencies = self.parse_dependencies(type_)
+        dependencies = self._parse_dependencies(type_)
 
-        evaluated_dependencies = {
-            name: self.get_or_register_token_or_type(dependency)
-            for name, dependency in dependencies.items()
-        }
+        evaluated_dependencies: dict[str, Any] = {}
+        for name, dependency in dependencies.items():
+            with span_instantiation_vector(
+                InstantiationNode(
+                    property_name=name,
+                    source_type=type_,
+                    target_type=dependency,
+                )
+            ):
+                evaluated_dependencies[name] = self.get_or_register_token_or_type(
+                    dependency
+                )
 
         instance = type_(**evaluated_dependencies)
 
         return instance
 
-    def parse_dependencies(
+    def _parse_dependencies(
         self, provider: type[Any] | Callable[..., Any]
     ) -> dict[str, type[Any]]:
 
-        signature = inspect.signature(provider)
+        vector = instantiation_vector_ctxvar.get()
+        try:
+            signature = inspect.signature(provider)
+        except ValueError:
+            print("VECTOR:", vector)
+            print_instantiation_vector(vector)
+            raise
 
         parameters = signature.parameters
 
         return {
-            name: self.lookup_parameter_type(parameter)
+            name: self._lookup_parameter_type(parameter)
             for name, parameter in parameters.items()
             if parameter.annotation != inspect.Parameter.empty
         }
 
-    def lookup_parameter_type(self, parameter: inspect.Parameter) -> Any:
+    def _lookup_parameter_type(self, parameter: inspect.Parameter) -> Any:
         if parameter.annotation == inspect.Parameter.empty:
             raise Exception(f"Parameter {parameter.name} has no type annotation")
 
@@ -188,14 +287,14 @@ class Container:
             item_type = bind_to = token_or_type
 
         if token_or_type not in self.instances_map:
-            return self.get_and_register(item_type, bind_to)
+            return self._get_and_register(item_type, bind_to)
 
         return cast(T, self.instances_map[bind_to])
 
-    def get_and_register(
+    def _get_and_register(
         self, item_type: Type[T] | Callable[..., T], bind_to: Any
     ) -> T:
-        instance = self.instantiate(item_type)
+        instance = self._instantiate(item_type)
         self.register(instance, bind_to)
         return cast(T, instance)
 
@@ -227,17 +326,21 @@ def provide_container(container: Container) -> Generator[None, None, None]:
 
 
 __all__ = [
-    "AppContext",
+    "AppTransactionContext",
     "AppInterceptor",
     "AppInterceptorWithLifecycle",
     "Container",
     "Microservice",
-    "SchedulerAppContext",
-    "WebSocketAppContext",
-    "app_context_ctxvar",
+    "SchedulerTransactionData",
+    "WebSocketTransactionData",
+    "app_transaction_context_var",
     "current_container_ctx",
     "provide_app_context",
     "provide_container",
     "use_app_context",
     "use_current_container",
+    "HttpTransactionData",
+    "MessageBusTransactionData",
+    "is_interceptor_with_lifecycle",
+    "AppContext",
 ]

jararaca/observability/decorators.py

@@ -13,7 +13,7 @@ from typing import (
     TypeVar,
 )
 
-from jararaca.microservice import AppContext
+from jararaca.microservice import AppTransactionContext
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -28,9 +28,13 @@ class TracingContextProvider(Protocol):
 
 class TracingContextProviderFactory(Protocol):
 
-    def root_setup(self, app_context: AppContext) -> AsyncContextManager[None]: ...
+    def root_setup(
+        self, app_context: AppTransactionContext
+    ) -> AsyncContextManager[None]: ...
 
-    def provide_provider(self, app_context: AppContext) -> TracingContextProvider: ...
+    def provide_provider(
+        self, app_context: AppTransactionContext
+    ) -> TracingContextProvider: ...
 
 
 tracing_ctx_provider_ctxv = ContextVar[TracingContextProvider]("tracing_ctx_provider")

jararaca/observability/interceptor.py

@@ -2,9 +2,9 @@ from contextlib import asynccontextmanager
 from typing import AsyncContextManager, AsyncGenerator, Protocol
 
 from jararaca.microservice import (
-    AppContext,
     AppInterceptor,
     AppInterceptorWithLifecycle,
+    AppTransactionContext,
     Container,
     Microservice,
 )
@@ -32,7 +32,9 @@ class ObservabilityInterceptor(AppInterceptor, AppInterceptorWithLifecycle):
         self.observability_provider = observability_provider
 
     @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
 
         async with self.observability_provider.tracing_provider.root_setup(app_context):
 

jararaca/observability/providers/otel.py

@@ -23,7 +23,7 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from jararaca.microservice import AppContext, Container, Microservice
+from jararaca.microservice import AppTransactionContext, Container, Microservice
 from jararaca.observability.decorators import (
     TracingContextProvider,
     TracingContextProviderFactory,
@@ -36,7 +36,7 @@ tracer: trace.Tracer = trace.get_tracer(__name__)
 
 class OtelTracingContextProvider(TracingContextProvider):
 
-    def __init__(self, app_context: AppContext) -> None:
+    def __init__(self, app_context: AppTransactionContext) -> None:
         self.app_context = app_context
 
     @contextmanager
@@ -52,22 +52,26 @@ class OtelTracingContextProvider(TracingContextProvider):
 
 class OtelTracingContextProviderFactory(TracingContextProviderFactory):
 
-    def provide_provider(self, app_context: AppContext) -> TracingContextProvider:
+    def provide_provider(
+        self, app_context: AppTransactionContext
+    ) -> TracingContextProvider:
         return OtelTracingContextProvider(app_context)
 
     @asynccontextmanager
-    async def root_setup(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+    async def root_setup(
+        self, app_tx_ctx: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
 
         title: str = "Unmapped App Context Execution"
         headers = {}
+        tx_data = app_tx_ctx.transaction_data
+        if tx_data.context_type == "http":
 
-        if app_context.context_type == "http":
-
-            headers = dict(app_context.request.headers)
-            title = f"HTTP {app_context.request.method} {app_context.request.url}"
+            headers = dict(tx_data.request.headers)
+            title = f"HTTP {tx_data.request.method} {tx_data.request.url}"
 
-        elif app_context.context_type == "message_bus":
-            title = f"Message Bus {app_context.topic}"
+        elif tx_data.context_type == "message_bus":
+            title = f"Message Bus {tx_data.topic}"
 
         carrier = {
             key: value

jararaca/persistence/base.py

@@ -1,6 +1,7 @@
 from typing import Any, Self, Type, TypeVar
 
 from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase
 
 IDENTIFIABLE_SCHEMA_T = TypeVar("IDENTIFIABLE_SCHEMA_T")
@@ -20,7 +21,7 @@ def recursive_get_dict(obj: Any) -> Any:
         return obj
 
 
-class BaseEntity(DeclarativeBase):
+class BaseEntity(AsyncAttrs, DeclarativeBase):
 
     @classmethod
     def from_basemodel(cls, mutation: T_BASEMODEL) -> "Self":

jararaca/persistence/interceptors/aiosqa_interceptor.py

@@ -3,21 +3,53 @@ 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 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.reflect.metadata import SetMetadata, get_metadata_value
 
-ctx_session_map = ContextVar[dict[str, AsyncSession]]("ctx_session_map", default={})
+DEFAULT_CONNECTION_NAME = "default"
+
+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 provide_session(
-    connection_name: str, session: AsyncSession
+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,18 +58,118 @@ 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]:
+
+    current_session = use_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):
+            yield new_session
+
+
+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
+
+
+INJECT_CONNECTION_METADATA = "inject_connection_metadata_{connection_name}"
+
+
+def set_inject_connection(
+    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_CONNECTION_METADATA.format(connection_name=connection_name), inject
+    )
+
+
+def uses_connection(
+    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 SetMetadata(
+        INJECT_CONNECTION_METADATA.format(connection_name=connection_name), True
+    )
+
+
+def dnt_uses_connection(
+    connection_name: str = DEFAULT_CONNECTION_NAME,
+) -> SetMetadata:
+    """
+    Do not use 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 SetMetadata(
+        INJECT_CONNECTION_METADATA.format(connection_name=connection_name), False
+    )
 
 
 class AIOSqlAlchemySessionInterceptor(AppInterceptor):
@@ -53,12 +185,31 @@ 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):
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
+
+        uses_connection_metadata = get_metadata_value(
+            INJECT_CONNECTION_METADATA.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
-                    await session.commit()
+                    if tx.is_active:
+                        await tx.commit()
                 except Exception as e:
-                    await session.rollback()
+                    await tx.rollback()
                     raise e
+                finally:
+                    with suppress(ValueError):
+                        ctx_default_connection_name.reset(token)