jararaca 0.3.11a16__py3-none-any.whl → 0.4.0a5__py3-none-any.whl

jararaca/microservice.py

@@ -1,3 +1,8 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import asyncio
 import inspect
 import logging
 from contextlib import contextmanager, suppress
@@ -10,6 +15,7 @@ from typing import (
     Any,
     AsyncContextManager,
     Callable,
+    Coroutine,
     Generator,
     Literal,
     Protocol,
@@ -19,7 +25,7 @@ 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
@@ -34,6 +40,7 @@ if TYPE_CHECKING:
 
 @dataclass
 class SchedulerTransactionData:
+    task_name: str
     triggered_at: datetime
     scheduled_to: datetime
     cron_expression: str
@@ -43,6 +50,7 @@ class SchedulerTransactionData:
 @dataclass
 class HttpTransactionData:
     request: Request
+    response: Response
     context_type: Literal["http"] = "http"
 
 
@@ -50,6 +58,7 @@ class HttpTransactionData:
 class MessageBusTransactionData:
     topic: str
     message: MessageOf[Message]
+    message_type: Type[Message]
     context_type: Literal["message_bus"] = "message_bus"
 
 
@@ -59,6 +68,8 @@ class WebSocketTransactionData:
     context_type: Literal["websocket"] = "websocket"
 
 
+APP_TYPE = Literal["http", "worker", "beat"]
+
 TransactionData = (
     MessageBusTransactionData
     | HttpTransactionData
@@ -325,6 +336,104 @@ 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__ = [
     "AppTransactionContext",
     "AppInterceptor",

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,23 +13,72 @@ from typing import (
     Callable,
     ContextManager,
     Generator,
-    ParamSpec,
+    Literal,
+    Mapping,
     Protocol,
+    Sequence,
     TypeVar,
+    Union,
 )
 
 from jararaca.microservice import AppTransactionContext
 
-P = ParamSpec("P")
-R = TypeVar("R")
+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]
+
+
+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):
 
@@ -76,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

jararaca/observability/hooks.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import logging
+import typing
+from contextlib import contextmanager
+from typing import Any, Generator, Literal
+
+from jararaca.observability.decorators import (
+    AttributeMap,
+    AttributeValue,
+    TracingContextProvider,
+    TracingSpan,
+    TracingSpanContext,
+    get_tracing_ctx_provider,
+)
+
+
+@contextmanager
+def start_span(
+    name: str,
+    attributes: AttributeMap | None = None,
+) -> Generator[None, Any, None]:
+    if trace_context_provider := get_tracing_ctx_provider():
+        with trace_context_provider.start_span_context(
+            trace_name=name, context_attributes=attributes
+        ):
+            yield
+    else:
+        yield
+
+
+def spawn_trace(
+    name: str,
+    attributes: AttributeMap | None = None,
+) -> typing.ContextManager[None]:
+    logging.warning(
+        "spawn_trace is deprecated, use start_span as context manager instead."
+    )
+    return start_span(name=name, attributes=attributes)
+
+
+def add_event(
+    name: str,
+    attributes: AttributeMap | None = None,
+) -> None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        trace_context_provider.add_event(
+            event_name=name,
+            event_attributes=attributes,
+        )
+
+
+def set_span_status(status_code: Literal["OK", "ERROR", "UNSET"]) -> None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        trace_context_provider.set_span_status(status_code=status_code)
+
+
+def record_exception(
+    exception: Exception,
+    attributes: AttributeMap | None = None,
+    escaped: bool = False,
+) -> None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        trace_context_provider.record_exception(
+            exception=exception,
+            attributes=attributes,
+            escaped=escaped,
+        )
+
+
+def set_span_attribute(
+    key: str,
+    value: AttributeValue,
+) -> None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        trace_context_provider.set_span_attribute(
+            key=key,
+            value=value,
+        )
+
+
+def get_tracing_provider() -> TracingContextProvider | None:
+    return get_tracing_ctx_provider()
+
+
+def get_current_span_context() -> TracingSpanContext | None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        return trace_context_provider.get_current_span_context()
+    return None
+
+
+def get_current_span() -> TracingSpan | None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        return trace_context_provider.get_current_span()
+    return None
+
+
+def add_span_link(span_context: TracingSpanContext) -> None:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        trace_context_provider.add_link(span_context=span_context)

jararaca/observability/interceptor.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from contextlib import asynccontextmanager
 from typing import AsyncContextManager, AsyncGenerator, Protocol
 

jararaca/observability/providers/__init__.py

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