diff --git a/microbootstrap/bootstrappers/litestar.py b/microbootstrap/bootstrappers/litestar.py index b071ff8..d350107 100644 --- a/microbootstrap/bootstrappers/litestar.py +++ b/microbootstrap/bootstrappers/litestar.py @@ -7,19 +7,32 @@ import typing_extensions from litestar import openapi from litestar.config.cors import CORSConfig as LitestarCorsConfig -from litestar.contrib.opentelemetry.config import OpenTelemetryConfig as LitestarOpentelemetryConfig +from litestar.contrib.opentelemetry.config import ( + OpenTelemetryConfig as LitestarOpentelemetryConfig, +) +from litestar.contrib.opentelemetry.middleware import ( + OpenTelemetryInstrumentationMiddleware, +) from litestar.contrib.prometheus import PrometheusConfig, PrometheusController from litestar.openapi.plugins import SwaggerRenderPlugin from litestar_offline_docs import generate_static_files_config +from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware +from opentelemetry.util.http import get_excluded_urls from sentry_sdk.integrations.litestar import LitestarIntegration from microbootstrap.bootstrappers.base import ApplicationBootstrapper from microbootstrap.config.litestar import LitestarConfig from microbootstrap.instruments.cors_instrument import CorsInstrument -from microbootstrap.instruments.health_checks_instrument import HealthChecksInstrument, HealthCheckTypedDict +from microbootstrap.instruments.health_checks_instrument import ( + HealthChecksInstrument, + HealthCheckTypedDict, +) from microbootstrap.instruments.logging_instrument import LoggingInstrument from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryInstrument -from microbootstrap.instruments.prometheus_instrument import LitestarPrometheusConfig, PrometheusInstrument +from microbootstrap.instruments.prometheus_instrument import ( + LitestarPrometheusConfig, + PrometheusInstrument, +) from microbootstrap.instruments.pyroscope_instrument import PyroscopeInstrument from microbootstrap.instruments.sentry_instrument import SentryInstrument from microbootstrap.instruments.swagger_instrument import SwaggerInstrument @@ -27,6 +40,11 @@ from microbootstrap.settings import LitestarSettings +if typing.TYPE_CHECKING: + from litestar.contrib.opentelemetry import OpenTelemetryConfig + from litestar.types import ASGIApp, Scope + + class LitestarBootstrapper( ApplicationBootstrapper[LitestarSettings, litestar.Litestar, LitestarConfig], ): @@ -106,6 +124,53 @@ def bootstrap_before(self) -> dict[str, typing.Any]: LitestarBootstrapper.use_instrument()(PyroscopeInstrument) +def get_litestar_route_details_from_scope( + scope: Scope, +) -> tuple[str, dict[str, str]]: + """Retrieve the span name and attributes from the ASGI scope for Litestar routes. + + Args: + scope: The ASGI scope instance. + + Returns: + A tuple of the span name and a dict of attrs. + + """ + # Try to get the route pattern from Litestar + path_template = scope.get("path_template") + if path_template: + method = str(scope.get("method", "")).strip() + if method and path_template: + return f"{method} {path_template}", {"http.route": path_template} + + # Fallback to default behavior + path = scope.get("path", "").strip() + method = str(scope.get("method", "")).strip() + if method and path: + return f"{method} {path}", {"http.route": path} + + return path, {"http.route": path} + + +class LitestarOpenTelemetryInstrumentationMiddleware(OpenTelemetryInstrumentationMiddleware): + def __init__(self, app: ASGIApp, config: OpenTelemetryConfig) -> None: + super().__init__( + app=app, + config=config, + ) + self.open_telemetry_middleware = OpenTelemetryMiddleware( + app=app, + client_request_hook=config.client_request_hook_handler, # type: ignore[arg-type] + client_response_hook=config.client_response_hook_handler, # type: ignore[arg-type] + default_span_details=get_litestar_route_details_from_scope, + excluded_urls=get_excluded_urls(config.exclude_urls_env_key), + meter=config.meter, + meter_provider=config.meter_provider, + server_request_hook=config.server_request_hook_handler, + tracer_provider=config.tracer_provider, + ) + + @LitestarBootstrapper.use_instrument() class LitestarOpentelemetryInstrument(OpentelemetryInstrument): def bootstrap_before(self) -> dict[str, typing.Any]: @@ -113,9 +178,9 @@ def bootstrap_before(self) -> dict[str, typing.Any]: "middleware": [ LitestarOpentelemetryConfig( tracer_provider=self.tracer_provider, - exclude=self.define_exclude_urls(), + middleware_class=LitestarOpenTelemetryInstrumentationMiddleware, ).middleware, - ], + ] } @@ -141,7 +206,10 @@ class LitestarPrometheusController(PrometheusController): **self.instrument_config.prometheus_additional_params, ) - return {"route_handlers": [LitestarPrometheusController], "middleware": [litestar_prometheus_config.middleware]} + return { + "route_handlers": [LitestarPrometheusController], + "middleware": [litestar_prometheus_config.middleware], + } @classmethod def get_config_type(cls) -> type[LitestarPrometheusConfig]: diff --git a/tests/bootstrappers/test_litestar_opentelemetry.py b/tests/bootstrappers/test_litestar_opentelemetry.py new file mode 100644 index 0000000..a83abb0 --- /dev/null +++ b/tests/bootstrappers/test_litestar_opentelemetry.py @@ -0,0 +1,136 @@ +import typing +from unittest.mock import Mock, patch + +import litestar +import pytest +from litestar.contrib.opentelemetry.config import OpenTelemetryConfig as LitestarOpentelemetryConfig +from litestar.status_codes import HTTP_200_OK +from litestar.testing import TestClient + +from microbootstrap import LitestarSettings +from microbootstrap.bootstrappers.litestar import ( + LitestarBootstrapper, + LitestarOpentelemetryInstrument, + LitestarOpenTelemetryInstrumentationMiddleware, + get_litestar_route_details_from_scope, +) +from microbootstrap.config.litestar import LitestarConfig +from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig + + +@pytest.mark.parametrize( + ("scope", "expected_span_name", "expected_attributes"), + [ + ( + { + "path": "/users/123", + "path_template": "/users/{user_id}", + "method": "GET", + }, + "GET /users/{user_id}", + {"http.route": "/users/{user_id}"}, + ), + ( + { + "path": "/users/123", + "method": "POST", + }, + "POST /users/123", + {"http.route": "/users/123"}, + ), + ( + { + "path": "/test", + }, + "/test", + {"http.route": "/test"}, + ), + ], +) +def test_get_litestar_route_details_from_scope( + scope: dict[str, str], + expected_span_name: str, + expected_attributes: dict[str, str], +) -> None: + span_name, attributes = get_litestar_route_details_from_scope(scope) # type: ignore[arg-type] + + assert span_name == expected_span_name + assert attributes == expected_attributes + + +def test_litestar_opentelemetry_instrument_uses_custom_middleware( + minimal_opentelemetry_config: OpentelemetryConfig, +) -> None: + opentelemetry_instrument: typing.Final = LitestarOpentelemetryInstrument(minimal_opentelemetry_config) + opentelemetry_instrument.bootstrap() + + bootstrap_result: typing.Final = opentelemetry_instrument.bootstrap_before() + + assert "middleware" in bootstrap_result + assert len(bootstrap_result["middleware"]) == 1 + + middleware_config: typing.Final = bootstrap_result["middleware"][0] + assert middleware_config.middleware == LitestarOpenTelemetryInstrumentationMiddleware + + +@pytest.mark.parametrize( + ("path", "expected_span_name"), + [ + ("/users/123", "GET /users/{user_id}"), + ("/users/", "GET /users/"), + ("/", "GET /"), + ], +) +def test_litestar_opentelemetry_integration_with_path_templates( + path: str, + expected_span_name: str, + minimal_opentelemetry_config: OpentelemetryConfig, +) -> None: + @litestar.get("/users/{user_id:int}") + async def get_user(user_id: int) -> dict[str, int]: + return {"user_id": user_id} + + @litestar.get("/users/") + async def list_users() -> dict[str, str]: + return {"message": "list of users"} + + @litestar.get("/") + async def root() -> dict[str, str]: + return {"message": "root"} + + with patch("microbootstrap.bootstrappers.litestar.get_litestar_route_details_from_scope") as mock_function: + mock_function.return_value = (expected_span_name, {"http.route": path}) + + application: typing.Final = ( + LitestarBootstrapper(LitestarSettings()) + .configure_instrument(minimal_opentelemetry_config) + .configure_application(LitestarConfig(route_handlers=[get_user, list_users, root])) + .bootstrap() + ) + + with TestClient(app=application) as client: + response: typing.Final = client.get(path) + assert response.status_code == HTTP_200_OK + assert mock_function.called + + +def test_litestar_opentelemetry_middleware_initialization() -> None: + mock_app: typing.Final = Mock() + + mock_config: typing.Final = Mock(spec=LitestarOpentelemetryConfig) + mock_config.scopes = ["http"] + mock_config.exclude = [] + mock_config.exclude_opt_key = None + mock_config.client_request_hook_handler = None + mock_config.client_response_hook_handler = None + mock_config.exclude_urls_env_key = None + mock_config.meter = None + mock_config.meter_provider = None + mock_config.server_request_hook_handler = None + mock_config.tracer_provider = None + + middleware: typing.Final = LitestarOpenTelemetryInstrumentationMiddleware(app=mock_app, config=mock_config) + + assert middleware.app == mock_app + assert hasattr(middleware, "open_telemetry_middleware") + assert middleware.open_telemetry_middleware is not None diff --git a/tests/conftest.py b/tests/conftest.py index 79aeb97..dc8aeca 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -99,7 +99,7 @@ def minimal_health_checks_config() -> HealthChecksConfig: @pytest.fixture def minimal_opentelemetry_config() -> OpentelemetryConfig: return OpentelemetryConfig( - opentelemetry_endpoint="/my-engdpoint", + opentelemetry_endpoint="/my-endpoint", opentelemetry_namespace="namespace", opentelemetry_container_name="container-name", opentelemetry_generate_health_check_spans=False,