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

jararaca/observability/providers/otel.py

@@ -1,6 +1,10 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import logging
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Protocol
+from typing import Any, AsyncGenerator, Generator, Literal, Protocol
 
 from opentelemetry import metrics, trace
 from opentelemetry._logs import set_logger_provider
@@ -23,32 +27,159 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from jararaca.microservice import AppTransactionContext, Container, Microservice
+from jararaca.messagebus.implicit_headers import (
+    ImplicitHeaders,
+    provide_implicit_headers,
+    use_implicit_headers,
+)
+from jararaca.microservice import (
+    AppTransactionContext,
+    Container,
+    Microservice,
+    use_app_transaction_context,
+)
+from jararaca.observability.constants import TRACEPARENT_KEY
 from jararaca.observability.decorators import (
+    AttributeMap,
+    AttributeValue,
     TracingContextProvider,
     TracingContextProviderFactory,
-    get_tracing_ctx_provider,
+    TracingSpan,
+    TracingSpanContext,
 )
 from jararaca.observability.interceptor import ObservabilityProvider
 
 tracer: trace.Tracer = trace.get_tracer(__name__)
 
 
+def extract_context_attributes(ctx: AppTransactionContext) -> dict[str, Any]:
+    tx_data = ctx.transaction_data
+    extra_attributes: dict[str, Any] = {}
+
+    if tx_data.context_type == "http":
+        extra_attributes = {
+            "http.method": tx_data.request.method,
+            "http.url": str(tx_data.request.url),
+            "http.path": tx_data.request.url.path,
+            "http.route.path": tx_data.request.scope["route"].path,
+            "http.route.endpoint.name": tx_data.request["route"].endpoint.__qualname__,
+            "http.query": tx_data.request.url.query,
+            **{
+                f"http.request.path_param.{k}": v
+                for k, v in tx_data.request.path_params.items()
+            },
+            **{
+                f"http.request.query_param.{k}": v
+                for k, v in tx_data.request.query_params.items()
+            },
+            **{
+                f"http.request.header.{k}": v
+                for k, v in tx_data.request.headers.items()
+            },
+            "http.request.client.host": (
+                tx_data.request.client.host if tx_data.request.client else ""
+            ),
+        }
+    elif tx_data.context_type == "message_bus":
+        extra_attributes = {
+            "bus.message.name": tx_data.message_type.__qualname__,
+            "bus.message.module": tx_data.message_type.__module__,
+            "bus.message.category": tx_data.message_type.MESSAGE_CATEGORY,
+            "bus.message.type": tx_data.message_type.MESSAGE_TYPE,
+            "bus.message.topic": tx_data.message_type.MESSAGE_TOPIC,
+        }
+    elif tx_data.context_type == "websocket":
+        extra_attributes = {
+            "ws.url": str(tx_data.websocket.url),
+        }
+    elif tx_data.context_type == "scheduler":
+        extra_attributes = {
+            "sched.task_name": tx_data.task_name,
+            "sched.scheduled_to": tx_data.scheduled_to.isoformat(),
+            "sched.cron_expression": tx_data.cron_expression,
+            "sched.triggered_at": tx_data.triggered_at.isoformat(),
+        }
+    return {
+        "app.context_type": tx_data.context_type,
+        "controller_member_reflect.rest_controller.class_name": ctx.controller_member_reflect.controller_reflect.controller_class.__qualname__,
+        "controller_member_reflect.rest_controller.module": ctx.controller_member_reflect.controller_reflect.controller_class.__module__,
+        "controller_member_reflect.member_function.name": ctx.controller_member_reflect.member_function.__qualname__,
+        "controller_member_reflect.member_function.module": ctx.controller_member_reflect.member_function.__module__,
+        **extra_attributes,
+    }
+
+
+class OtelTracingSpan(TracingSpan):
+
+    def __init__(self, span: trace.Span) -> None:
+        self.span = span
+
+
+class OtelTracingSpanContext(TracingSpanContext):
+
+    def __init__(self, span_context: trace.SpanContext) -> None:
+        self.span_context = span_context
+
+
 class OtelTracingContextProvider(TracingContextProvider):
 
     def __init__(self, app_context: AppTransactionContext) -> None:
         self.app_context = app_context
 
     @contextmanager
-    def __call__(
+    def start_span_context(
         self,
         trace_name: str,
-        context_attributes: dict[str, str],
+        context_attributes: AttributeMap | None,
     ) -> Generator[None, None, None]:
 
         with tracer.start_as_current_span(trace_name, attributes=context_attributes):
             yield
 
+    def add_event(
+        self, event_name: str, event_attributes: AttributeMap | None = None
+    ) -> None:
+        trace.get_current_span().add_event(name=event_name, attributes=event_attributes)
+
+    def set_span_status(self, status_code: Literal["OK", "ERROR", "UNSET"]) -> None:
+        span = trace.get_current_span()
+        if status_code == "OK":
+            span.set_status(trace.Status(trace.StatusCode.OK))
+        elif status_code == "ERROR":
+            span.set_status(trace.Status(trace.StatusCode.ERROR))
+        else:
+            span.set_status(trace.Status(trace.StatusCode.UNSET))
+
+    def record_exception(
+        self,
+        exception: Exception,
+        attributes: AttributeMap | None = None,
+        escaped: bool = False,
+    ) -> None:
+        span = trace.get_current_span()
+        span.record_exception(exception, attributes=attributes, escaped=escaped)
+
+    def set_span_attribute(self, key: str, value: AttributeValue) -> None:
+        span = trace.get_current_span()
+        span.set_attribute(key, value)
+
+    def update_span_name(self, new_name: str) -> None:
+        span = trace.get_current_span()
+
+        span.update_name(new_name)
+
+    def add_link(self, span_context: TracingSpanContext) -> None:
+        if not isinstance(span_context, OtelTracingSpanContext):
+            return
+        span = trace.get_current_span()
+        span.add_link(span_context.span_context)
+
+    def get_current_span(self) -> TracingSpan | None:
+        return OtelTracingSpan(trace.get_current_span())
+
+    def get_current_span_context(self) -> TracingSpanContext | None:
+        return OtelTracingSpanContext(trace.get_current_span().get_span_context())
+
 
 class OtelTracingContextProviderFactory(TracingContextProviderFactory):
 
@@ -63,15 +194,27 @@ class OtelTracingContextProviderFactory(TracingContextProviderFactory):
     ) -> AsyncGenerator[None, None]:
 
         title: str = "Unmapped App Context Execution"
-        headers = {}
+        headers: dict[str, Any] = {}
         tx_data = app_tx_ctx.transaction_data
-        if tx_data.context_type == "http":
+        extra_attributes = extract_context_attributes(app_tx_ctx)
 
+        if tx_data.context_type == "http":
             headers = dict(tx_data.request.headers)
             title = f"HTTP {tx_data.request.method} {tx_data.request.url}"
+            extra_attributes["http.request.body"] = (await tx_data.request.body())[
+                :5000
+            ].decode(errors="ignore")
 
         elif tx_data.context_type == "message_bus":
             title = f"Message Bus {tx_data.topic}"
+            headers = use_implicit_headers() or {}
+
+        elif tx_data.context_type == "websocket":
+            headers = dict(tx_data.websocket.headers)
+            title = f"WebSocket {tx_data.websocket.url}"
+
+        elif tx_data.context_type == "scheduler":
+            title = f"Scheduler Task {tx_data.task_name}"
 
         carrier = {
             key: value
@@ -90,8 +233,28 @@ class OtelTracingContextProviderFactory(TracingContextProviderFactory):
 
         ctx2 = W3CBaggagePropagator().extract(b2, context=ctx)
 
-        with tracer.start_as_current_span(name=title, context=ctx2):
-            yield
+        with tracer.start_as_current_span(
+            name=title,
+            context=ctx2,
+            attributes={
+                **extra_attributes,
+            },
+        ) as root_span:
+            cx = root_span.get_span_context()
+            span_traceparent_id = hex(cx.trace_id)[2:].rjust(32, "0")
+            if app_tx_ctx.transaction_data.context_type == "http":
+                app_tx_ctx.transaction_data.request.scope[TRACEPARENT_KEY] = (
+                    span_traceparent_id
+                )
+            elif app_tx_ctx.transaction_data.context_type == "websocket":
+                app_tx_ctx.transaction_data.websocket.scope[TRACEPARENT_KEY] = (
+                    span_traceparent_id
+                )
+            tracing_headers: ImplicitHeaders = {}
+            TraceContextTextMapPropagator().inject(tracing_headers)
+            W3CBaggagePropagator().inject(tracing_headers)
+            with provide_implicit_headers(tracing_headers):
+                yield
 
 
 class LoggerHandlerCallback(Protocol):
@@ -99,6 +262,34 @@ class LoggerHandlerCallback(Protocol):
     def __call__(self, logger_handler: logging.Handler) -> None: ...
 
 
+class CustomLoggingHandler(LoggingHandler):
+
+    def _translate(self, record: logging.LogRecord) -> dict[str, Any]:
+        try:
+            ctx = use_app_transaction_context()
+            data = super()._translate(record)
+            extra_attributes = extract_context_attributes(ctx)
+
+            current_span = trace.get_current_span()
+
+            data["attributes"] = {
+                **data.get("attributes", {}),
+                **extra_attributes,
+                **(
+                    {
+                        "span_name": current_span.name,
+                    }
+                    if hasattr(current_span, "name")
+                    and current_span.is_recording() is False
+                    else {}
+                ),
+            }
+
+            return data
+        except LookupError:
+            return super()._translate(record)
+
+
 class OtelObservabilityProvider(ObservabilityProvider):
 
     def __init__(
@@ -143,11 +334,11 @@ class OtelObservabilityProvider(ObservabilityProvider):
             BatchLogRecordProcessor(self.logs_exporter)
         )
 
-        logging_handler = LoggingHandler(
+        logging_handler = CustomLoggingHandler(
             level=logging.DEBUG, logger_provider=logger_provider
         )
 
-        logging_handler.addFilter(lambda _: get_tracing_ctx_provider() is not None)
+        # logging_handler.addFilter(lambda _: get_tracing_ctx_provider() is not None)
 
         self.logging_handler_callback(logging_handler)
 

jararaca/persistence/base.py

@@ -1,4 +1,8 @@
-from typing import Any, Self, Type, TypeVar
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from typing import Any, Callable, Protocol, Self, Type, TypeVar
 
 from pydantic import BaseModel
 from sqlalchemy.ext.asyncio import AsyncAttrs
@@ -21,14 +25,46 @@ def recursive_get_dict(obj: Any) -> Any:
         return obj
 
 
+RESULT_T = TypeVar("RESULT_T", covariant=True)
+ENTITY_T_CONTRA = TypeVar("ENTITY_T_CONTRA", bound="BaseEntity", contravariant=True)
+
+
+class EntityParserType(Protocol[ENTITY_T_CONTRA, RESULT_T]):
+
+    @classmethod
+    def parse_entity(cls, model: ENTITY_T_CONTRA) -> "RESULT_T": ...
+
+
+EntityParserFunc = Callable[[ENTITY_T_CONTRA], RESULT_T]
+
+BASED_BASE_ENTITY_T = TypeVar("BASED_BASE_ENTITY_T", bound="BaseEntity")
+
+
 class BaseEntity(AsyncAttrs, DeclarativeBase):
 
     @classmethod
     def from_basemodel(cls, mutation: T_BASEMODEL) -> "Self":
         intersection = set(cls.__annotations__.keys()) & set(
-            mutation.model_fields.keys()
+            mutation.__class__.model_fields.keys()
         )
         return cls(**{k: getattr(mutation, k) for k in intersection})
 
     def to_basemodel(self, model: Type[T_BASEMODEL]) -> T_BASEMODEL:
         return model.model_validate(recursive_get_dict(self))
+
+    def parse_entity_with_func(
+        self, model_cls: EntityParserFunc["Self", RESULT_T]
+    ) -> RESULT_T:
+        return model_cls(self)
+
+    def parse_entity_with_type(
+        self: BASED_BASE_ENTITY_T,
+        model_cls: Type[EntityParserType[BASED_BASE_ENTITY_T, RESULT_T]],
+    ) -> RESULT_T:
+        return model_cls.parse_entity(self)
+
+    def __rshift__(self, model: EntityParserFunc["Self", RESULT_T]) -> RESULT_T:
+        return self.parse_entity_with_func(model)
+
+    def __and__(self, model: Type[EntityParserType["Self", RESULT_T]]) -> RESULT_T:
+        return self.parse_entity_with_type(model)

jararaca/persistence/exports.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from .base import BaseEntity
 
 __all__ = [

jararaca/persistence/interceptors/__init__.py

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

jararaca/persistence/interceptors/aiosqa_interceptor.py

@@ -1,7 +1,11 @@
+# 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 typing import Any, AsyncGenerator, Generator, Protocol
 
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
@@ -12,9 +16,47 @@ from sqlalchemy.ext.asyncio import (
 from sqlalchemy.ext.asyncio.engine import AsyncEngine
 
 from jararaca.microservice import AppInterceptor, AppTransactionContext
-from jararaca.reflect.metadata import SetMetadata, get_metadata_value
+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_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
 
-DEFAULT_CONNECTION_NAME = "default"
 
 ctx_default_connection_name: ContextVar[str] = ContextVar(
     "ctx_default_connection_name", default=DEFAULT_CONNECTION_NAME
@@ -69,13 +111,21 @@ 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,
     ) as new_session, new_session.begin() as new_tx:
         with providing_session(new_session, new_tx, connection_name):
-            yield new_session
+            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:
@@ -129,50 +179,7 @@ class AIOSQAConfig:
         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):
+class AIOSqlAlchemySessionInterceptor(AppInterceptor, SessionManager):
 
     def __init__(self, config: AIOSQAConfig):
         self.config = config
@@ -189,27 +196,33 @@ class AIOSqlAlchemySessionInterceptor(AppInterceptor):
         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
-                    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

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