diff --git a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py index f48ce18..02defa1 100644 --- a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py @@ -33,12 +33,59 @@ from litestar.static_files import create_static_files_router if import_checker.is_litestar_opentelemetry_installed: - from litestar.contrib.opentelemetry import OpenTelemetryConfig + from litestar.middleware import ASGIMiddleware + from litestar.types.asgi_types import ASGIApp, Receive, Scope, Send + from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware + from opentelemetry.trace import TracerProvider if import_checker.is_opentelemetry_installed: from opentelemetry.trace import get_tracer_provider +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: typing.MutableMapping[str, typing.Any], +) -> tuple[str, dict[str, str]]: + 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, {} + + +if import_checker.is_litestar_opentelemetry_installed: + + class LitestarOpenTelemetryInstrumentationMiddleware(ASGIMiddleware): + def __init__(self, tracer_provider: "TracerProvider", excluded_urls: set[str]) -> None: + self._tracer_provider = tracer_provider + self._excluded_urls = ",".join(excluded_urls) + + async def handle( + self, + scope: "Scope", + receive: "Receive", + send: "Send", + next_app: "ASGIApp", + ) -> None: + await OpenTelemetryMiddleware( + app=next_app, + default_span_details=build_litestar_route_details_from_scope, + excluded_urls=self._excluded_urls, + tracer_provider=self._tracer_provider, + )(scope, receive, send) # ty: ignore + + @dataclasses.dataclass(kw_only=True, slots=True, frozen=True) class LitestarConfig( CorsConfig, @@ -111,10 +158,10 @@ def _build_excluded_urls(self) -> set[str]: def bootstrap(self) -> None: super().bootstrap() self.bootstrap_config.application_config.middleware.append( - OpenTelemetryConfig( + LitestarOpenTelemetryInstrumentationMiddleware( tracer_provider=get_tracer_provider(), - exclude=list(self._build_excluded_urls()), - ).middleware, + excluded_urls=self._build_excluded_urls(), + ) ) diff --git a/tests/test_litestar_bootstrap.py b/tests/test_litestar_bootstrap.py index 7ac07db..be2b2e9 100644 --- a/tests/test_litestar_bootstrap.py +++ b/tests/test_litestar_bootstrap.py @@ -1,9 +1,18 @@ +import dataclasses + +import litestar import pytest import structlog from litestar import status_codes +from litestar.config.app import AppConfig from litestar.testing import TestClient +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import get_tracer_provider from lite_bootstrap import LitestarBootstrapper, LitestarConfig +from lite_bootstrap.bootstrappers.litestar_bootstrapper import build_litestar_route_details_from_scope, build_span_name from tests.conftest import CustomInstrumentor, SentryTestTransport, emulate_package_missing @@ -78,3 +87,44 @@ def test_litestar_bootstrapper_with_missing_instrument_dependency( ) -> None: with emulate_package_missing(package_name), pytest.warns(UserWarning, match=package_name): LitestarBootstrapper(bootstrap_config=litestar_config) + + +def test_litestar_otel_span_naming(litestar_config: LitestarConfig) -> None: + @litestar.get("/items/{item_id:int}") + async def get_item(item_id: int) -> dict[str, int]: + return {"item_id": item_id} + + config = dataclasses.replace(litestar_config, application_config=AppConfig(route_handlers=[get_item])) + bootstrapper = LitestarBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + + tracer_provider = get_tracer_provider() + assert isinstance(tracer_provider, SDKTracerProvider) + exporter = InMemorySpanExporter() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with TestClient(app=application) as client: + response = client.get("/items/42") + assert response.status_code == status_codes.HTTP_200_OK + + spans = exporter.get_finished_spans() + span_names = [s.name for s in spans] + assert any("GET /items/{item_id}" in name for name in span_names) + + +def test_build_span_name_no_route() -> None: + assert build_span_name("GET", "") == "GET" + + +def test_build_litestar_route_details_from_scope_path_fallback() -> None: + scope = {"method": "POST", "path": "/fallback/path"} + name, attrs = build_litestar_route_details_from_scope(scope) + assert name == "POST /fallback/path" + assert attrs == {"http.route": "/fallback/path"} + + +def test_build_litestar_route_details_from_scope_no_path() -> None: + scope = {"type": "lifespan"} + name, attrs = build_litestar_route_details_from_scope(scope) + assert name == "HTTP" + assert attrs == {}