diff --git a/microbootstrap/bootstrappers/litestar.py b/microbootstrap/bootstrappers/litestar.py index b071ff8..a864c73 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,56 @@ def bootstrap_before(self) -> dict[str, typing.Any]: LitestarBootstrapper.use_instrument()(PyroscopeInstrument) +def build_span_name(method: str, route: str) -> str: + if not route: + return method + return f"{method} {route}" + + +def build_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. + + """ + path_template: typing.Final = scope.get("path_template") + method: typing.Final = str(scope.get("method", "HTTP")).strip() + if path_template is not None: + path_template_stripped: typing.Final = path_template.strip() + return build_span_name(method, path_template_stripped), {"http.route": path_template_stripped} + + path: typing.Final = scope.get("path") + if path is not None: + path_stripped: typing.Final = path.strip() + return build_span_name(method, path_stripped), {"http.route": path_stripped} + return method, {} + + +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=build_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 +181,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 +209,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..8893b59 --- /dev/null +++ b/tests/bootstrappers/test_litestar_opentelemetry.py @@ -0,0 +1,182 @@ +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, + build_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", + }, + "HTTP /test", + {"http.route": "/test"}, + ), + ( + { + "path": "", + }, + "HTTP", + {"http.route": ""}, + ), + ( + { + "path": " ", + }, + "HTTP", + {"http.route": ""}, + ), + ( + { + "path_template": "", + }, + "HTTP", + {"http.route": ""}, + ), + ( + { + "path_template": " ", + }, + "HTTP", + {"http.route": ""}, + ), + ( + {}, + "HTTP", + {}, + ), + ( + {"method": "GET"}, + "GET", + {}, + ), + ( + { + "path": "/users/123", + "path_template": "/users/{user_id}", + }, + "HTTP /users/{user_id}", + {"http.route": "/users/{user_id}"}, + ), + ], +) +def test_build_litestar_route_details_from_scope( + scope: dict[str, str], + expected_span_name: str, + expected_attributes: dict[str, str], +) -> None: + span_name, attributes = build_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.build_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,