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

jararaca/microservice.py

@@ -1,4 +1,10 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import asyncio
 import inspect
+import logging
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from dataclasses import dataclass, field
@@ -9,6 +15,7 @@ from typing import (
     Any,
     AsyncContextManager,
     Callable,
+    Coroutine,
     Generator,
     Literal,
     Protocol,
@@ -18,69 +25,115 @@ from typing import (
     runtime_checkable,
 )
 
-from fastapi import Request, WebSocket
+from fastapi import Request, Response, WebSocket
 
 from jararaca.core.providers import ProviderSpec, T, Token
 from jararaca.messagebus import MessageOf
-from jararaca.messagebus.types import Message
+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:
+    task_name: str
     triggered_at: datetime
     scheduled_to: datetime
     cron_expression: str
-    action: Callable[..., Any]
     context_type: Literal["scheduler"] = "scheduler"
 
 
 @dataclass
-class HttpAppContext:
+class HttpTransactionData:
     request: Request
+    response: Response
     context_type: Literal["http"] = "http"
 
 
 @dataclass
-class MessageBusAppContext:
+class MessageBusTransactionData:
     topic: str
     message: MessageOf[Message]
+    message_type: Type[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
+APP_TYPE = Literal["http", "worker", "beat"]
+
+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_context() -> AppContext:
-    return app_context_ctxvar.get()
+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.
+    """
+
+    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 +159,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 +218,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 +298,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)
 
@@ -226,18 +336,120 @@ def provide_container(container: Container) -> Generator[None, None, None]:
             current_container_ctx.reset(token)
 
 
+class ShutdownState(Protocol):
+
+    def request_shutdown(self) -> None: ...
+
+    def is_shutdown_requested(self) -> bool: ...
+
+    async def wait_for_shutdown(self) -> None: ...
+
+
+shutdown_state_ctx = ContextVar[ShutdownState]("shutdown_state")
+
+
+def is_shutting_down() -> bool:
+    """
+    Check if the application is in the process of shutting down.
+    """
+    return shutdown_state_ctx.get().is_shutdown_requested()
+
+
+def request_shutdown() -> None:
+    """
+    Request the application to shut down.
+    This will set the shutdown event, allowing the application to gracefully shut down.
+    """
+    shutdown_state_ctx.get().request_shutdown()
+
+
+async def wait_for_shutdown() -> None:
+    """
+    Wait for the shutdown event to be set.
+    This function will block until a shutdown is requested.
+    """
+    await shutdown_state_ctx.get().wait_for_shutdown()
+
+
+async def shutdown_race(*concurrent_tasks: Coroutine[Any, Any, Any]) -> bool:
+    """
+    Wait for either a shutdown request or any of the provided tasks to complete.
+    This function will return as soon as a shutdown is requested or any task finishes.
+    Returns True if shutdown was requested, False if a task completed first.
+    """
+
+    tasks = [asyncio.create_task(t) for t in concurrent_tasks + (wait_for_shutdown(),)]
+
+    _, pending = await asyncio.wait(
+        tasks,
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+
+    for task in pending:
+        task.cancel()
+
+    return is_shutting_down()
+
+
+@contextmanager
+def provide_shutdown_state(
+    state: ShutdownState,
+) -> Generator[None, None, None]:
+    """
+    Context manager to provide the shutdown state.
+    This is used to manage the shutdown event for the application.
+    """
+
+    token = shutdown_state_ctx.set(state)
+    try:
+        yield
+    finally:
+        with suppress(ValueError):
+            shutdown_state_ctx.reset(token)
+
+
+app_type_ctx = ContextVar[APP_TYPE]("app_type")
+
+
+def use_app_type() -> APP_TYPE:
+    """
+    Returns the current application type.
+    This function is used to access the application type in the context of an application transaction.
+    If no context is set, it raises a LookupError.
+    """
+    return app_type_ctx.get()
+
+
+@contextmanager
+def providing_app_type(app_type: APP_TYPE) -> Generator[None, None, None]:
+    """
+    Context manager to provide the application type.
+    This is used to set the application type for the current transaction.
+    """
+    token = app_type_ctx.set(app_type)
+    try:
+        yield
+    finally:
+        with suppress(ValueError):
+            app_type_ctx.reset(token)
+
+
 __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/constants.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+TRACEPARENT_KEY = "traceparent"
+
+__all__ = ["TRACEPARENT_KEY"]

jararaca/observability/decorators.py

@@ -1,3 +1,8 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import inspect
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from functools import wraps
@@ -8,29 +13,82 @@ from typing import (
     Callable,
     ContextManager,
     Generator,
-    ParamSpec,
+    Literal,
+    Mapping,
     Protocol,
+    Sequence,
     TypeVar,
+    Union,
 )
 
-from jararaca.microservice import AppContext
+from jararaca.microservice import AppTransactionContext
+
+F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
+
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+
+AttributeMap = Mapping[str, AttributeValue]
 
-P = ParamSpec("P")
-R = TypeVar("R")
+
+class TracingSpan(Protocol): ...
+
+
+class TracingSpanContext(Protocol): ...
 
 
 class TracingContextProvider(Protocol):
 
-    def __call__(
-        self, trace_name: str, context_attributes: dict[str, str]
+    def start_span_context(
+        self, trace_name: str, context_attributes: AttributeMap | None
     ) -> ContextManager[Any]: ...
 
+    def add_event(
+        self,
+        event_name: str,
+        event_attributes: AttributeMap | None = None,
+    ) -> None: ...
+
+    def set_span_status(self, status_code: Literal["OK", "ERROR", "UNSET"]) -> None: ...
+
+    def record_exception(
+        self,
+        exception: Exception,
+        attributes: AttributeMap | None = None,
+        escaped: bool = False,
+    ) -> None: ...
+
+    def set_span_attribute(
+        self,
+        key: str,
+        value: AttributeValue,
+    ) -> None: ...
+
+    def update_span_name(self, new_name: str) -> None: ...
+
+    def add_link(self, span_context: TracingSpanContext) -> None: ...
+
+    def get_current_span(self) -> TracingSpan | None: ...
+    def get_current_span_context(self) -> TracingSpanContext | None: ...
+
 
 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")
@@ -72,22 +130,125 @@ class TracedFunc:
 
     def __call__(
         self,
-        decorated: Callable[P, Awaitable[R]],
-    ) -> Callable[P, Awaitable[R]]:
+        decorated: F,
+    ) -> F:
 
         @wraps(decorated)
         async def wrapper(
-            *args: P.args,
-            **kwargs: P.kwargs,
-        ) -> R:
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
 
             if ctx_provider := get_tracing_ctx_provider():
-                with ctx_provider(
+                with ctx_provider.start_span_context(
                     self.trace_name,
-                    self.trace_mapper(**kwargs),
+                    self.trace_mapper(*args, **kwargs),
                 ):
                     return await decorated(*args, **kwargs)
 
             return await decorated(*args, **kwargs)
 
-        return wrapper
+        return wrapper  # type: ignore[return-value]
+
+
+C = TypeVar("C", bound=type)
+
+
+class TracedClass:
+    """
+    Class decorator that automatically applies tracing to all async methods in a class.
+
+    Usage:
+        @TracedClass()
+        class MyService:
+            async def method1(self) -> str:
+                return "hello"
+
+            async def method2(self, x: int) -> int:
+                return x * 2
+
+            def sync_method(self) -> str:  # Not traced
+                return "sync"
+    """
+
+    def __init__(
+        self,
+        trace_name_prefix: str | None = None,
+        trace_mapper: Callable[..., dict[str, str]] = default_trace_mapper,
+        include_private: bool = False,
+        exclude_methods: set[str] | None = None,
+    ):
+        """
+        Initialize the TracedClass decorator.
+
+        Args:
+            trace_name_prefix: Prefix for trace names. If None, uses class name.
+            trace_mapper: Function to map method arguments to trace attributes.
+            include_private: Whether to trace private methods (starting with _).
+            exclude_methods: Set of method names to exclude from tracing.
+        """
+        self.trace_name_prefix = trace_name_prefix
+        self.trace_mapper = trace_mapper
+        self.include_private = include_private
+        self.exclude_methods = exclude_methods or set()
+
+    def __call__(self, cls: C) -> C:
+        """Apply tracing to all async methods in the class."""
+
+        # Use class name as prefix if not provided
+        trace_prefix = self.trace_name_prefix or cls.__name__
+
+        # Get all methods in the class
+        for name, method in inspect.getmembers_static(
+            cls, predicate=inspect.isfunction
+        ):
+            # Skip if method should be excluded
+            if name in self.exclude_methods:
+                continue
+
+            # Skip private methods unless explicitly included
+            if name.startswith("_") and not self.include_private:
+                continue
+
+            # Only trace async methods
+            if inspect.iscoroutinefunction(method):
+                trace_name = f"{trace_prefix}.{name}"
+                traced_method = TracedFunc(trace_name, self.trace_mapper)(method)
+                setattr(cls, name, traced_method)
+
+        return cls
+
+
+def traced_class(
+    trace_name_prefix: str | None = None,
+    trace_mapper: Callable[..., dict[str, str]] = default_trace_mapper,
+    include_private: bool = False,
+    exclude_methods: set[str] | None = None,
+) -> Callable[[C], C]:
+    """
+    Functional interface for TracedClass decorator.
+
+    Usage:
+        @traced_class(trace_name_prefix="MyService")
+        class MyService:
+            async def method1(self) -> str:
+                return "hello"
+    """
+    return TracedClass(
+        trace_name_prefix=trace_name_prefix,
+        trace_mapper=trace_mapper,
+        include_private=include_private,
+        exclude_methods=exclude_methods,
+    )
+
+
+__all__ = [
+    "TracingContextProvider",
+    "TracingContextProviderFactory",
+    "provide_tracing_ctx_provider",
+    "get_tracing_ctx_provider",
+    "default_trace_mapper",
+    "TracedFunc",
+    "TracedClass",
+    "traced_class",
+]

jararaca/observability/fastapi_exception_handler.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from fastapi import FastAPI, Request, Response
+from fastapi.exception_handlers import (
+    http_exception_handler,
+    request_validation_exception_handler,
+)
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from starlette.exceptions import HTTPException
+
+from jararaca.observability.constants import TRACEPARENT_KEY
+
+
+def setup_fastapi_exception_handler(
+    app: FastAPI, trace_header_name: str = "traceparent"
+) -> None:
+    async def base_http_exception_handler(
+        request: Request, exc: HTTPException | RequestValidationError
+    ) -> JSONResponse | Response:
+
+        if isinstance(exc, RequestValidationError):
+            response = await request_validation_exception_handler(request, exc)
+            response.headers[trace_header_name] = request.scope.get(TRACEPARENT_KEY, "")
+            return response
+        else:
+            err_response = await http_exception_handler(request, exc)
+
+            err_response.headers[trace_header_name] = request.scope.get(
+                TRACEPARENT_KEY, ""
+            )
+            return err_response
+
+    app.exception_handlers[HTTPException] = base_http_exception_handler
+    app.exception_handlers[RequestValidationError] = base_http_exception_handler