Fix opentelemetry span names (#135)

* fix_opentelemetry_span_names

* fix creating middleware

* test opentelemetry

* fix tests

* fix pyroscope

* Apply suggestion from @vrslev

---------

Co-authored-by: EGOROV Vadim <ruawiih@n-msk-macwiih.raiffeisen.ru>
Co-authored-by: Lev Vereshchagin <levwint@gmail.com>

Author: 3 people

microbootstrap/bootstrappers/litestar.py

@@ -2,19 +2,16 @@
 import typing
 
 import litestar
-import litestar.exceptions
-import litestar.types
 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.middleware import (
-    OpenTelemetryInstrumentationMiddleware,
-)
 from litestar.contrib.prometheus import PrometheusConfig, PrometheusController
+from litestar.middleware import ASGIMiddleware
 from litestar.openapi.plugins import SwaggerRenderPlugin
+from litestar.types.asgi_types import ASGIApp, Scope
 from litestar_offline_docs import generate_static_files_config
 from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
 from opentelemetry.util.http import get_excluded_urls
@@ -43,6 +40,7 @@
 if typing.TYPE_CHECKING:
     from litestar.contrib.opentelemetry import OpenTelemetryConfig
     from litestar.types import ASGIApp, Scope
+    from litestar.types.asgi_types import Receive, Send
 
 
 class LitestarBootstrapper(
@@ -155,34 +153,38 @@ def build_litestar_route_details_from_scope(
     return method, {}
 
 
-class LitestarOpenTelemetryInstrumentationMiddleware(OpenTelemetryInstrumentationMiddleware):
-    def __init__(self, app: ASGIApp, config: OpenTelemetryConfig) -> None:
-        super().__init__(
-            app=app,
-            config=config,
-        )
-        self.open_telemetry_middleware = OpenTelemetryMiddleware(
+class LitestarOpenTelemetryInstrumentationMiddleware(ASGIMiddleware):
+    def __init__(self, config: OpenTelemetryConfig) -> None:
+        self.config = config
+
+    def create_open_telemetry_middleware(self, app: ASGIApp) -> OpenTelemetryMiddleware:
+        return 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]
+            client_request_hook=self.config.client_request_hook_handler,  # type: ignore[arg-type]
+            client_response_hook=self.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,
+            excluded_urls=get_excluded_urls(self.config.exclude_urls_env_key),
+            meter=self.config.meter,
+            meter_provider=self.config.meter_provider,
+            server_request_hook=self.config.server_request_hook_handler,
+            tracer_provider=self.config.tracer_provider,
         )
 
+    async def handle(self, scope: Scope, receive: Receive, send: Send, next_app: ASGIApp) -> None:
+        await self.create_open_telemetry_middleware(next_app)(scope, receive, send)  # type: ignore[arg-type]
+
 
 @LitestarBootstrapper.use_instrument()
 class LitestarOpentelemetryInstrument(OpentelemetryInstrument):
     def bootstrap_before(self) -> dict[str, typing.Any]:
         return {
             "middleware": [
-                LitestarOpentelemetryConfig(
-                    tracer_provider=self.tracer_provider,
-                    middleware_class=LitestarOpenTelemetryInstrumentationMiddleware,
-                ).middleware,
+                LitestarOpenTelemetryInstrumentationMiddleware(
+                    LitestarOpentelemetryConfig(
+                        tracer_provider=self.tracer_provider,
+                        middleware_class=LitestarOpenTelemetryInstrumentationMiddleware,  # type: ignore[arg-type]
+                    )
+                )
             ]
         }
 

pyproject.toml

@@ -41,7 +41,7 @@ dependencies = [
     "rich>=13",
     "sentry-sdk>=2.7",
     "structlog>=24",
-    "pyroscope-io; platform_system != 'Windows'",
+    "pyroscope-io<=1.0.0; platform_system != 'Windows'",
     "opentelemetry-distro[otlp]>=0.54b1",
     "opentelemetry-instrumentation-aio-pika>=0.54b1",
     "opentelemetry-instrumentation-aiohttp-client>=0.54b1",

tests/bootstrappers/test_litestar_opentelemetry.py

@@ -115,21 +115,22 @@ def test_litestar_opentelemetry_instrument_uses_custom_middleware(
     assert "middleware" in bootstrap_result
     assert len(bootstrap_result["middleware"]) == 1
 
-    middleware_config: typing.Final = bootstrap_result["middleware"][0]
-    assert middleware_config.middleware == LitestarOpenTelemetryInstrumentationMiddleware
+    middleware_config: typing.Final = bootstrap_result["middleware"][0].config
+    assert middleware_config.middleware.middleware == LitestarOpenTelemetryInstrumentationMiddleware
 
 
 @pytest.mark.parametrize(
-    ("path", "expected_span_name"),
+    ("path", "expected_span_name", "expected_path_template"),
     [
-        ("/users/123", "GET /users/{user_id}"),
-        ("/users/", "GET /users/"),
-        ("/", "GET /"),
+        ("/users/123", "GET /users/{user_id}", "/users/{user_id}"),
+        ("/users/", "GET /users/", "/users"),
+        ("/", "GET /", "/"),
     ],
 )
 def test_litestar_opentelemetry_integration_with_path_templates(
     path: str,
     expected_span_name: str,
+    expected_path_template: str,
     minimal_opentelemetry_config: OpentelemetryConfig,
 ) -> None:
     @litestar.get("/users/{user_id:int}")
@@ -158,6 +159,7 @@ async def root() -> dict[str, str]:
             response: typing.Final = client.get(path)
         assert response.status_code == HTTP_200_OK
         assert mock_function.called
+        assert mock_function.call_args_list[0].args[0].get("path_template") == expected_path_template
 
 
 def test_litestar_opentelemetry_middleware_initialization() -> None:
@@ -175,8 +177,8 @@ def test_litestar_opentelemetry_middleware_initialization() -> None:
     mock_config.server_request_hook_handler = None
     mock_config.tracer_provider = None
 
-    middleware: typing.Final = LitestarOpenTelemetryInstrumentationMiddleware(app=mock_app, config=mock_config)
+    middleware: typing.Final = LitestarOpenTelemetryInstrumentationMiddleware(config=mock_config)
 
-    assert middleware.app == mock_app
-    assert hasattr(middleware, "open_telemetry_middleware")
-    assert middleware.open_telemetry_middleware is not None
+    assert middleware.config == mock_config
+    otel_middleware = middleware.create_open_telemetry_middleware(mock_app)
+    assert otel_middleware is not None

tests/instruments/test_opentelemetry.py

@@ -7,13 +7,15 @@
 import litestar
 import pytest
 from fastapi.testclient import TestClient as FastAPITestClient
-from litestar.middleware.base import DefineMiddleware
 from litestar.testing import TestClient as LitestarTestClient
 from opentelemetry.instrumentation.dependencies import DependencyConflictError
 
 from microbootstrap import OpentelemetryConfig
 from microbootstrap.bootstrappers.fastapi import FastApiOpentelemetryInstrument
-from microbootstrap.bootstrappers.litestar import LitestarOpentelemetryInstrument
+from microbootstrap.bootstrappers.litestar import (
+    LitestarOpentelemetryInstrument,
+    LitestarOpenTelemetryInstrumentationMiddleware,
+)
 from microbootstrap.instruments import opentelemetry_instrument
 from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryInstrument
 
@@ -61,7 +63,7 @@ def test_litestar_opentelemetry_bootstrap(
     assert "middleware" in opentelemetry_bootstrap_result
     assert isinstance(opentelemetry_bootstrap_result["middleware"], list)
     assert len(opentelemetry_bootstrap_result["middleware"]) == 1
-    assert isinstance(opentelemetry_bootstrap_result["middleware"][0], DefineMiddleware)
+    assert isinstance(opentelemetry_bootstrap_result["middleware"][0], LitestarOpenTelemetryInstrumentationMiddleware)
 
 
 def test_litestar_opentelemetry_teardown(
@@ -83,9 +85,9 @@ def test_litestar_opentelemetry_bootstrap_working(
     opentelemetry_bootstrap_result: typing.Final = test_opentelemetry_instrument.bootstrap_before()
 
     opentelemetry_middleware = opentelemetry_bootstrap_result["middleware"][0]
-    assert isinstance(opentelemetry_middleware, DefineMiddleware)
+    assert isinstance(opentelemetry_middleware, LitestarOpenTelemetryInstrumentationMiddleware)
     async_mock.__name__ = "test-name"
-    opentelemetry_middleware.middleware.__call__ = async_mock  # type: ignore[operator]
+    opentelemetry_middleware.handle = async_mock  # type: ignore[method-assign]
 
     @litestar.get("/test-handler")
     async def test_handler() -> None: