Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 25 additions & 23 deletions microbootstrap/bootstrappers/litestar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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]
)
)
]
}

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
22 changes: 12 additions & 10 deletions tests/bootstrappers/test_litestar_opentelemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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:
Expand All @@ -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
12 changes: 7 additions & 5 deletions tests/instruments/test_opentelemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand All @@ -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:
Expand Down