diff --git a/microbootstrap/bootstrappers/litestar.py b/microbootstrap/bootstrappers/litestar.py index a864c73..736bedc 100644 --- a/microbootstrap/bootstrappers/litestar.py +++ b/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] + ) + ) ] } diff --git a/pyproject.toml b/pyproject.toml index 83125d7..ff7afe3 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/bootstrappers/test_litestar_opentelemetry.py b/tests/bootstrappers/test_litestar_opentelemetry.py index 8893b59..71e1be6 100644 --- a/tests/bootstrappers/test_litestar_opentelemetry.py +++ b/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 diff --git a/tests/instruments/test_opentelemetry.py b/tests/instruments/test_opentelemetry.py index c09aa63..6794749 100644 --- a/tests/instruments/test_opentelemetry.py +++ b/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: