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
55 changes: 51 additions & 4 deletions lite_bootstrap/bootstrappers/litestar_bootstrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
)
)


Expand Down
50 changes: 50 additions & 0 deletions tests/test_litestar_bootstrap.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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 == {}
Loading