Merge pull request #70 from modern-python/litestar-otel-span-naming

use route template for Litestar OTel span naming

Author: lesnik512

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(),
+            )
         )
 
 

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 == {}