From 148a35b71c8435289df20785ab789deb0d5e8199 Mon Sep 17 00:00:00 2001 From: Walter Oggioni Date: Sat, 19 Sep 2026 00:39:15 +0000 Subject: [PATCH] Add kaya-otel package for OpenTelemetry tracing and metrics Instrument HTTP requests and WebSocket connections via Kaya hooks, covering both ASGI and RSGI. Records handler exceptions, WebSocket close codes, optional header capture, exclusions and lifecycle hooks. Add route-template resolution and exception visibility to kaya-core so trace/metric attributes can use low-cardinality routes and failed spans can record escaped exceptions. --- .gitea/workflows/build.yaml | 12 + README.md | 6 +- packages/kaya-core/src/kaya/core/_app.py | 10 + .../kaya-core/src/kaya/core/_http_context.py | 1 + packages/kaya-core/src/kaya/core/_tree.py | 31 ++ .../kaya-core/src/kaya/core/_websocket.py | 1 + packages/kaya-core/tests/test_asgi.py | 29 ++ packages/kaya-core/tests/test_tree.py | 21 + packages/kaya-otel/README.md | 85 +++ packages/kaya-otel/pyproject.toml | 58 +++ packages/kaya-otel/src/kaya/otel/__init__.py | 10 + packages/kaya-otel/src/kaya/otel/_mixin.py | 491 ++++++++++++++++++ packages/kaya-otel/src/kaya/otel/py.typed | 0 packages/kaya-otel/tests/test_otel.py | 394 ++++++++++++++ requirements-dev.in | 1 + requirements-dev.txt | 33 ++ 16 files changed, 1182 insertions(+), 1 deletion(-) create mode 100644 packages/kaya-otel/README.md create mode 100644 packages/kaya-otel/pyproject.toml create mode 100644 packages/kaya-otel/src/kaya/otel/__init__.py create mode 100644 packages/kaya-otel/src/kaya/otel/_mixin.py create mode 100644 packages/kaya-otel/src/kaya/otel/py.typed create mode 100644 packages/kaya-otel/tests/test_otel.py diff --git a/.gitea/workflows/build.yaml b/.gitea/workflows/build.yaml index 97315bf..44ab9f3 100644 --- a/.gitea/workflows/build.yaml +++ b/.gitea/workflows/build.yaml @@ -60,6 +60,10 @@ jobs: run: | .venv/bin/python -m mypy -p kaya.forwarded .venv/bin/python -m unittest discover -s packages/kaya-forwarded/tests + - name: Check kaya-otel + run: | + .venv/bin/python -m mypy -p kaya.otel + .venv/bin/python -m unittest discover -s packages/kaya-otel/tests - name: Publish kaya-core artifacts env: TWINE_REPOSITORY_URL: ${{ vars.PYPI_REGISTRY_URL }} @@ -132,3 +136,11 @@ jobs: run: | .venv/bin/pyproject-build packages/kaya-forwarded .venv/bin/twine upload --repository gitea packages/kaya-forwarded/dist/*.whl packages/kaya-forwarded/dist/*.tar.gz + - name: Publish kaya-otel artifacts + env: + TWINE_REPOSITORY_URL: ${{ vars.PYPI_REGISTRY_URL }} + TWINE_USERNAME: ${{ vars.PUBLISHER_USERNAME }} + TWINE_PASSWORD: ${{ secrets.PUBLISHER_TOKEN }} + run: | + .venv/bin/pyproject-build packages/kaya-otel + .venv/bin/twine upload --repository gitea packages/kaya-otel/dist/*.whl packages/kaya-otel/dist/*.tar.gz diff --git a/README.md b/README.md index e11608c..3ad7412 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ This repository is a monorepo for the Kaya framework. The code is split into ind - **kaya-openapi** — automatic OpenAPI specification generation (`packages/kaya-openapi/`) - **kaya-cors** — CORS (Cross-Origin Resource Sharing) support (`packages/kaya-cors/`) - **kaya-forwarded** — trusted-proxy `Forwarded`/`X-Forwarded-*` client address resolution (`packages/kaya-forwarded/`) +- **kaya-otel** — OpenTelemetry tracing and metrics (`packages/kaya-otel/`) Additional `kaya-*` packages can be added as new directories under `packages/`. @@ -29,7 +30,7 @@ pip install --index-url https://gitea.woggioni.net/api/packages/woggioni/pypi/si Install the packages in development mode: ```bash -pip install -e packages/kaya-core -e packages/kaya-rsgi -e packages/kaya-session -e packages/kaya-session-redis -e packages/kaya-session-memcache -e packages/kaya-oidc -e packages/kaya-openapi -e packages/kaya-cors -e packages/kaya-forwarded +pip install -e packages/kaya-core -e packages/kaya-rsgi -e packages/kaya-session -e packages/kaya-session-redis -e packages/kaya-session-memcache -e packages/kaya-oidc -e packages/kaya-openapi -e packages/kaya-cors -e packages/kaya-forwarded -e packages/kaya-otel ``` Run the example: @@ -50,6 +51,7 @@ python -m unittest discover -s packages/kaya-oidc/tests python -m unittest discover -s packages/kaya-openapi/tests python -m unittest discover -s packages/kaya-cors/tests python -m unittest discover -s packages/kaya-forwarded/tests +python -m unittest discover -s packages/kaya-otel/tests ``` ## Static analysis @@ -64,6 +66,7 @@ mypy -p kaya.oidc mypy -p kaya.openapi mypy -p kaya.cors mypy -p kaya.forwarded +mypy -p kaya.otel ``` ## Building packages @@ -78,4 +81,5 @@ python -m build packages/kaya-oidc python -m build packages/kaya-openapi python -m build packages/kaya-cors python -m build packages/kaya-forwarded +python -m build packages/kaya-otel ``` diff --git a/packages/kaya-core/src/kaya/core/_app.py b/packages/kaya-core/src/kaya/core/_app.py index 61133f0..baf55b6 100644 --- a/packages/kaya-core/src/kaya/core/_app.py +++ b/packages/kaya-core/src/kaya/core/_app.py @@ -145,6 +145,9 @@ class KayaApp(AbstractKayaApp): await handler.handle_request(ctx, captured) else: await ctx.send_empty(404) + except Exception as exc: + ctx.exception = exc + raise finally: for hook in reversed(self._after_request_hooks): await hook(ctx) @@ -161,6 +164,9 @@ class KayaApp(AbstractKayaApp): await handler.handle_request(ws, captured) else: await ws.close(1000) + except Exception as exc: + ws.exception = exc + raise finally: for hook in reversed(self._after_websocket_hooks): await hook(ws) @@ -173,6 +179,10 @@ class KayaApp(AbstractKayaApp): for mixin in self._mixins: mixin.shutdown(loop) + def route_template(self, url: str, method: HttpMethod = HttpMethod.GET) -> Optional[str]: + """Return the registered route template that would handle ``url``.""" + return self._tree.route_template(url, method) + def route(self, paths: StrOrStrings, methods: Optional[HttpMethod | Sequence[HttpMethod]] = None, diff --git a/packages/kaya-core/src/kaya/core/_http_context.py b/packages/kaya-core/src/kaya/core/_http_context.py index c76aac3..619383d 100644 --- a/packages/kaya-core/src/kaya/core/_http_context.py +++ b/packages/kaya-core/src/kaya/core/_http_context.py @@ -29,6 +29,7 @@ class HttpContext(ABC): server: Optional[Tuple[str, Optional[int]]] request_body: AsyncIterator[bytes] session: Optional[Any] = None + exception: Optional[BaseException] = None @abstractmethod async def stream_body(self, diff --git a/packages/kaya-core/src/kaya/core/_tree.py b/packages/kaya-core/src/kaya/core/_tree.py index 84723cc..8373692 100644 --- a/packages/kaya-core/src/kaya/core/_tree.py +++ b/packages/kaya-core/src/kaya/core/_tree.py @@ -237,6 +237,37 @@ class Tree: # return (handler, unmatched) return None + def route_template(self, url: str, method: HttpMethod = HttpMethod.GET) -> Optional[str]: + """Return the registered route template that would handle ``url``. + + Static segments are returned verbatim, ``${name}``/``${name:int}`` + parameter matchers are reconstructed from the matched nodes, and glob + matchers keep their pattern. Returns ``None`` when no route (including + recursive fallback routes) would handle the request. + """ + result = self.find_node((p for p in PathIterator(urlparse(url).path)), method) + if result is None: + return None + node, captured = result + if len(captured.unmatched_paths) > 0 and not any(handler.recursive for handler in node.handlers): + return None + + parts: List[str] = [] + current: Optional[Node | PathMatcher] = node + while current is not None: + if isinstance(current, Node): + key = current.key + if not isinstance(key, HttpMethod) and key != '/': + parts.append(str(key)) + elif isinstance(current, IntMatcher): + parts.append('${%s:int}' % current.name) + elif isinstance(current, StrMatcher): + parts.append('${%s}' % current.name) + elif isinstance(current, GlobMatcher): + parts.append(current.pattern) + current = current.parent + return '/' + '/'.join(reversed(parts)) if parts else '/' + def parse(self, leaf: str, parent: Optional[Node | PathMatcher]) -> Node | PathMatcher: start = 0 result = index_of_with_escape(leaf, '${', '\\', 0) diff --git a/packages/kaya-core/src/kaya/core/_websocket.py b/packages/kaya-core/src/kaya/core/_websocket.py index a8041e5..2b223a1 100644 --- a/packages/kaya-core/src/kaya/core/_websocket.py +++ b/packages/kaya-core/src/kaya/core/_websocket.py @@ -21,6 +21,7 @@ class WebSocket(ABC): client: Optional[Tuple[str, int]] server: Optional[Tuple[str, Optional[int]]] session: Optional[Any] = None + exception: Optional[BaseException] = None @abstractmethod async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: diff --git a/packages/kaya-core/tests/test_asgi.py b/packages/kaya-core/tests/test_asgi.py index 193441b..fbe3301 100644 --- a/packages/kaya-core/tests/test_asgi.py +++ b/packages/kaya-core/tests/test_asgi.py @@ -250,3 +250,32 @@ class AsgiTest(unittest.TestCase): r = await client.put("/restaurants/42/menu") self.assertEqual(404, r.status_code) + def test_route_template(self): + self.assertEqual('/employee/${employee_id}', + self.app.route_template('/employee/101325', HttpMethod.GET)) + self.assertEqual('/square/${x:int}', self.app.route_template('/square/30', HttpMethod.GET)) + self.assertIsNone(self.app.route_template('/unknown', HttpMethod.GET)) + + @async_test + async def test_exception_exposed_to_after_hooks(self): + app = KayaApp() + seen = [] + + async def after_request(ctx: HttpContext) -> None: + seen.append(ctx.exception) + + app.add_after_request_hook(after_request) + + @app.GET('/raises') + async def raises(ctx: HttpContext) -> None: + raise RuntimeError('boom') + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as client: + with self.assertRaises(RuntimeError): + await client.get('/raises') + + self.assertEqual(1, len(seen)) + self.assertIsInstance(seen[0], RuntimeError) + self.assertEqual('boom', str(seen[0])) + diff --git a/packages/kaya-core/tests/test_tree.py b/packages/kaya-core/tests/test_tree.py index c0397cb..149785f 100644 --- a/packages/kaya-core/tests/test_tree.py +++ b/packages/kaya-core/tests/test_tree.py @@ -80,6 +80,27 @@ class TreeTest(unittest.TestCase): self.assertIs(Maybe.of(handler_num).map(self.handlers.__getitem__).or_none(), Maybe.of_nullable(res).map(lambda it: it[0]).or_none()) + def test_route_template(self): + cases: Tuple[Tuple[str, HttpMethod, Optional[str]], ...] = ( + ('/home/something', HttpMethod.GET, '/home/something'), + ('/home/something_else', HttpMethod.POST, '/home/something_else'), + ('/home/README.md', HttpMethod.GET, '/home/*.md'), + ('/home/something/ciao/blah/README.md', HttpMethod.GET, '/home/something/*/blah/*.md'), + ('/home/bar/ciao/blah/README.md', HttpMethod.GET, '/home/bar/*'), + ('/unknown', HttpMethod.GET, None), + ) + for url, method, expected in cases: + with self.subTest(f'{method} {url}'): + self.assertEqual(expected, self.tree.route_template(url, method)) + + def test_route_template_with_params(self): + tree = Tree() + tree.add((p for p in ('foo', '${id:int}')), HttpMethod.PUT, self.handlers[0]) + tree.add((p for p in ('foo', '${name}')), HttpMethod.GET, self.handlers[1]) + self.assertEqual('/foo/${id:int}', tree.route_template('/foo/42', HttpMethod.PUT)) + self.assertEqual('/foo/${name}', tree.route_template('/foo/bar', HttpMethod.GET)) + self.assertIsNone(tree.route_template('/foo/bar', HttpMethod.DELETE)) + def test_two_method_agnostic_matchers_raise(self): tree = Tree() tree.add((p for p in ('foo', '*')), None, self.handlers[0]) diff --git a/packages/kaya-otel/README.md b/packages/kaya-otel/README.md new file mode 100644 index 0000000..227a63e --- /dev/null +++ b/packages/kaya-otel/README.md @@ -0,0 +1,85 @@ +# kaya-otel + +OpenTelemetry tracing and metrics for the Kaya web framework. + +Provides `OTelMixin`, a `KayaMixin` that instruments HTTP requests and +WebSocket connections with OpenTelemetry spans and exports HTTP server +metrics, shipping everything to an OTLP/HTTP collector. + +## Usage + +```python +from kaya.core import KayaApp, HttpContext +from kaya.otel import OTelMixin + +app = KayaApp(mixins=[ + OTelMixin( + service_name='my-service', + endpoint='http://localhost:4318', + ) +]) + +@app.GET('/') +async def home(ctx: HttpContext): + await ctx.send_str(200, 'Hello World!') +``` + +## Parameters + +- `service_name`: value of the `service.name` resource attribute. +- `endpoint`: base URL of an OTLP/HTTP collector; the signal paths + `/v1/traces` and `/v1/metrics` are appended. When `None`, the exporters + use their own defaults, including the standard + `OTEL_EXPORTER_OTLP_ENDPOINT` environment variable. +- `headers`: extra HTTP headers sent to the collector (e.g. authentication). +- `metric_export_interval_millis`: metric export interval + (default `60000`). +- `tracer_provider` / `meter_provider`: inject custom providers (e.g. with + in-memory exporters for tests) instead of the OTLP defaults. The mixin + only shuts down providers it created itself. +- `resource_attributes`: extra resource attributes merged with + `service.name` when the mixin creates the providers. +- `excluded_paths` / `excluded_path_regexes`: skip tracing and metrics for + matching paths (useful for health checks and metrics endpoints). +- `capture_request_headers` / `capture_response_headers`: opt-in HTTP header + capture as `http.request.header.` / `http.response.header.` + span attributes. Header names are normalized to lowercase with `-` + replaced by `_`; values are captured as string lists. +- `sanitize_headers`: headers captured as `REDACTED` (default: + `authorization`, `proxy-authorization`, `cookie`, `set-cookie`). +- `server_request_hook` / `server_response_hook` / `websocket_connect_hook` / + `websocket_close_hook`: optional synchronous callbacks invoked with the + span and the Kaya context/websocket at the corresponding lifecycle point. + Hook exceptions are logged and do not fail the request. +- `metrics_include_raw_path`: when `True`, metrics use raw `url.path` + attributes (legacy, potentially high cardinality). The default `False` + keeps metrics low-cardinality by using `http.route` when Kaya can resolve + a route template. +- `websocket_error_close_codes`: close codes that mark a websocket span as + failed. Defaults to protocol/application error codes such as `1002`, + `1003` and `1007`-`1011`. + +## Behavior + +- Every HTTP request gets a `SERVER` span named ` ` with the + usual HTTP semantic attributes (`http.request.method`, `url.path`, + `client.address`, `server.address`, ...). When Kaya resolves a route + template, the span name is updated to ` ` and + `http.route` is set. The response status code is recorded as + `http.response.status_code` when the handler sends the response; 5xx + statuses mark the span as failed. Exceptions escaping the handler are + recorded on the span and also mark it as failed. +- Every WebSocket connection gets one span for its whole lifetime. The close + code is recorded as `kaya.websocket.close_code`; exceptions and configured + error close codes mark the span as failed. +- W3C `traceparent`/`tracestate` headers on incoming requests are honored, + so traces propagate from upstream services. +- Metrics: + - `http.server.request.duration` (histogram, seconds), with + `http.request.method`, `http.route` when known and + `http.response.status_code` attributes. + - `http.server.active_requests` (up-down counter), with + `http.request.method` attributes. + +`OTelMixin` is a `KayaMixin`, so the app stays a `KayaApp` and both ASGI and +RSGI keep working. diff --git a/packages/kaya-otel/pyproject.toml b/packages/kaya-otel/pyproject.toml new file mode 100644 index 0000000..a1015a8 --- /dev/null +++ b/packages/kaya-otel/pyproject.toml @@ -0,0 +1,58 @@ +[build-system] +requires = ["setuptools>=61.0", "setuptools-scm>=8"] +build-backend = "setuptools.build_meta" + +[project] +name = "kaya-otel" +dynamic = ["version"] +authors = [ + { name="Walter Oggioni", email="oggioni.walter@gmail.com" }, +] +description = "OpenTelemetry tracing and metrics for the Kaya lightweight ASGI web framework" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +classifiers = [ + 'Development Status :: 3 - Alpha', + 'Topic :: Utilities', + 'Intended Audience :: System Administrators', + 'Intended Audience :: Developers', + 'Environment :: Console', + 'Programming Language :: Python :: 3', +] + +dependencies = [ + "kaya-core", + "opentelemetry-sdk", + "opentelemetry-exporter-otlp-proto-http", +] + +[project.optional-dependencies] +dev = [ + "build", "mypy", "ipdb", "twine", "httpx", "httpx-ws", "kaya-rsgi" +] + +[project.urls] +"Homepage" = "https://github.com/woggioni/kaya" +"Bug Tracker" = "https://github.com/woggioni/kaya/issues" + +[tool.setuptools.packages.find] +where = ["src"] +namespaces = true + +[tool.mypy] +python_version = "3.12" +disallow_untyped_defs = true +show_error_codes = true +no_implicit_optional = true +warn_return_any = true +warn_unused_ignores = true +exclude = ["scripts", "docs", "test"] +strict = true + +[tool.setuptools_scm] +root = "../.." +version_file = "src/kaya/otel/_version.py" + +[tool.setuptools_scm.tag] +prefix = "release/" diff --git a/packages/kaya-otel/src/kaya/otel/__init__.py b/packages/kaya-otel/src/kaya/otel/__init__.py new file mode 100644 index 0000000..23e92c1 --- /dev/null +++ b/packages/kaya-otel/src/kaya/otel/__init__.py @@ -0,0 +1,10 @@ +from pkgutil import extend_path + +__path__ = extend_path(__path__, __name__) + +from ._mixin import OTelMixin + + +__all__ = [ + 'OTelMixin', +] diff --git a/packages/kaya-otel/src/kaya/otel/_mixin.py b/packages/kaya-otel/src/kaya/otel/_mixin.py new file mode 100644 index 0000000..68f240d --- /dev/null +++ b/packages/kaya-otel/src/kaya/otel/_mixin.py @@ -0,0 +1,491 @@ +import re +import time +from asyncio import AbstractEventLoop +from logging import getLogger +from pathlib import Path +from typing import ( + Any, + AsyncGenerator, + Callable, + Dict, + Mapping, + Optional, + Sequence, + Tuple, +) + +from kaya.core import HttpContext, KayaApp, KayaMixin, WebSocket, WebSocketMessage +from kaya.core._types import StrOrStrings +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.metrics import Histogram, UpDownCounter +from opentelemetry.propagators.textmap import Getter +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.trace import Span, SpanKind, Status, StatusCode, Tracer +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +log = getLogger(__name__) + +RequestHook = Callable[[Span, HttpContext], None] +WebSocketHook = Callable[[Span, WebSocket], None] + +_DEFAULT_SANITIZE_HEADERS = ( + 'authorization', + 'proxy-authorization', + 'cookie', + 'set-cookie', +) +_DEFAULT_WEBSOCKET_ERROR_CLOSE_CODES = frozenset((1002, 1003, 1007, 1008, 1009, 1010, 1011)) + + +class _HeadersGetter(Getter[Mapping[str, Sequence[str]]]): + """Extract propagation headers from a kaya request/websocket context. + + Kaya header mappings have lowercase names with one or more values each, + which is all the W3C trace-context propagator needs. + """ + + def get(self, carrier: Mapping[str, Sequence[str]], key: str) -> Optional[list[str]]: + values = carrier.get(key.lower()) + if not values: + return None + return list(values) + + def keys(self, carrier: Mapping[str, Sequence[str]]) -> list[str]: + return list(carrier.keys()) + + +_PROPAGATOR = TraceContextTextMapPropagator() +_GETTER = _HeadersGetter() + + +def _server_attributes(ctx: HttpContext) -> Dict[str, Any]: + attributes: Dict[str, Any] = { + 'http.request.method': str(ctx.method), + 'url.scheme': ctx.scheme, + 'url.path': ctx.path, + } + if ctx.query_string: + attributes['url.query'] = ctx.query_string + if ctx.client is not None: + attributes['client.address'] = ctx.client[0] + attributes['client.port'] = ctx.client[1] + if ctx.server is not None: + attributes['server.address'] = ctx.server[0] + if ctx.server[1] is not None: + attributes['server.port'] = ctx.server[1] + return attributes + + +def _normalize_header_name(name: str) -> str: + return name.lower().replace('-', '_') + + +def _header_values(value: StrOrStrings | Sequence[str]) -> list[str]: + if isinstance(value, str): + return [value] + return list(value) + + +class _SpanHolder: + """Mutable per-request slot shared by the tracing hooks and wrappers.""" + + __slots__ = ('span', 'token', 'status_code', 'start_ns', 'active_attributes', 'close_code') + + def __init__(self, span: Span, token: object, active_attributes: Dict[str, Any]) -> None: + self.span = span + self.token = token + self.status_code: Optional[int] = None + self.start_ns = time.perf_counter_ns() + self.active_attributes = active_attributes + self.close_code: Optional[int] = None + + +class _TracedHttpContext(HttpContext): + """HttpContext wrapper that records the response status code on the span. + + Attributes not explicitly overridden are delegated to the wrapped context + via ``__getattr__``, so it works with any concrete ``HttpContext`` (ASGI + or RSGI) and preserves attributes set by other mixins (e.g. ``session``). + """ + + def __init__(self, ctx: HttpContext, holder: _SpanHolder, mixin: 'OTelMixin') -> None: + object.__setattr__(self, '_ctx', ctx) + object.__setattr__(self, 'session', ctx.session) + object.__setattr__(self, '_holder', holder) + object.__setattr__(self, '_mixin', mixin) + + def __getattr__(self, name: str) -> Any: + if name == '_ctx': + raise AttributeError(name) + return getattr(self._ctx, name) + + def _sent(self, status: int, headers: Optional[Mapping[str, StrOrStrings]]) -> None: + self._holder.status_code = status + self._holder.span.set_attribute('http.response.status_code', status) + self._mixin._capture_response_headers(self._holder.span, headers) + + async def stream_body(self, + status: int, + body_generator: AsyncGenerator[bytes, None], + headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + self._sent(status, headers) + await self._ctx.stream_body(status, body_generator, headers) + + async def send_bytes(self, status: int, body: bytes, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + self._sent(status, headers) + await self._ctx.send_bytes(status, body, headers) + + async def send_str(self, status: int, body: str, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + self._sent(status, headers) + await self._ctx.send_str(status, body, headers) + + async def send_file(self, status: int, path: Path, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + self._sent(status, headers) + await self._ctx.send_file(status, path, headers) + + async def send_empty(self, status: int, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + self._sent(status, headers) + await self._ctx.send_empty(status, headers) + + +class _TracedWebSocket(WebSocket): + """WebSocket wrapper that keeps the connection span reachable. + + Everything is delegated to the wrapped socket via ``__getattr__``, so it + works with any concrete ``WebSocket`` (ASGI or RSGI) and preserves + attributes set by other mixins (e.g. ``session``). + """ + + def __init__(self, ws: WebSocket, holder: _SpanHolder) -> None: + object.__setattr__(self, '_ws', ws) + object.__setattr__(self, 'session', ws.session) + object.__setattr__(self, '_holder', holder) + + def __getattr__(self, name: str) -> Any: + if name == '_ws': + raise AttributeError(name) + return getattr(self._ws, name) + + async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: + await self._ws.accept(headers) + + async def receive(self) -> WebSocketMessage: + return await self._ws.receive() + + async def send_text(self, data: str) -> None: + await self._ws.send_text(data) + + async def send_bytes(self, data: bytes) -> None: + await self._ws.send_bytes(data) + + async def close(self, code: int = 1000) -> None: + self._holder.close_code = code + await self._ws.close(code) + + async def __anext__(self) -> WebSocketMessage: + return await self._ws.__anext__() + + +class OTelMixin(KayaMixin): + """Kaya mixin adding OpenTelemetry tracing and metrics. + + HTTP requests get a ``SERVER`` span named `` `` carrying the + usual HTTP semantic attributes; the response status code is recorded when + the handler sends the response and 5xx statuses mark the span as failed. + Exceptions escaping the handler are recorded and mark the span as failed. + WebSocket connections get one span for the whole connection lifetime. + W3C ``traceparent``/``tracestate`` headers on incoming requests are + honored, so traces propagate from upstream services. + + Metrics exported on the meter provider: + + - ``http.server.request.duration`` (histogram, seconds), with + ``http.request.method``, ``http.route`` when known and + ``http.response.status_code`` attributes; + - ``http.server.active_requests`` (up-down counter), with + ``http.request.method`` attributes. + + Example:: + + app = KayaApp(mixins=[ + OTelMixin( + service_name='my-service', + endpoint='http://localhost:4318', + ) + ]) + + ``endpoint`` is the base URL of an OTLP/HTTP collector (the signal paths + ``/v1/traces`` and ``/v1/metrics`` are appended); when ``None`` the + exporters fall back to their own defaults, including the standard + ``OTEL_EXPORTER_OTLP_ENDPOINT`` environment variable. Custom + ``tracer_provider``/``meter_provider`` instances (e.g. with in-memory + exporters for tests) can be injected instead of the OTLP defaults; the + mixin only shuts down providers it created itself. + """ + + def __init__(self, + service_name: str, + endpoint: Optional[str] = None, + headers: Optional[Mapping[str, str]] = None, + metric_export_interval_millis: float = 60000, + tracer_provider: Optional[TracerProvider] = None, + meter_provider: Optional[MeterProvider] = None, + resource_attributes: Optional[Mapping[str, Any]] = None, + excluded_paths: Sequence[str] = (), + excluded_path_regexes: Sequence[str] = (), + capture_request_headers: Sequence[str] = (), + capture_response_headers: Sequence[str] = (), + sanitize_headers: Sequence[str] = _DEFAULT_SANITIZE_HEADERS, + server_request_hook: Optional[RequestHook] = None, + server_response_hook: Optional[RequestHook] = None, + websocket_connect_hook: Optional[WebSocketHook] = None, + websocket_close_hook: Optional[WebSocketHook] = None, + metrics_include_raw_path: bool = False, + websocket_error_close_codes: Optional[Sequence[int]] = None) -> None: + self._owns_tracer_provider = tracer_provider is None + self._owns_meter_provider = meter_provider is None + if tracer_provider is None: + tracer_provider = TracerProvider(resource=self._create_resource(service_name, resource_attributes)) + span_exporter = OTLPSpanExporter( + endpoint=f'{endpoint}/v1/traces' if endpoint else None, + headers=dict(headers) if headers else None, + ) + tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter)) + if meter_provider is None: + meter_provider = MeterProvider( + resource=self._create_resource(service_name, resource_attributes), + metric_readers=[ + PeriodicExportingMetricReader( + OTLPMetricExporter( + endpoint=f'{endpoint}/v1/metrics' if endpoint else None, + headers=dict(headers) if headers else None, + ), + export_interval_millis=metric_export_interval_millis, + ) + ], + ) + self._tracer_provider = tracer_provider + self._meter_provider = meter_provider + self._tracer: Tracer = tracer_provider.get_tracer('kaya.otel') + meter = meter_provider.get_meter('kaya.otel') + self._request_duration: Histogram = meter.create_histogram( + 'http.server.request.duration', + unit='s', + description='Duration of HTTP server requests', + ) + self._active_requests: UpDownCounter = meter.create_up_down_counter( + 'http.server.active_requests', + unit='{request}', + description='Number of in-flight HTTP server requests', + ) + self._app: Optional[KayaApp] = None + self._excluded_paths = frozenset(excluded_paths) + self._excluded_path_regexes = tuple(re.compile(regex) for regex in excluded_path_regexes) + self._request_headers_to_capture = tuple(capture_request_headers) + self._response_headers_to_capture = tuple(capture_response_headers) + self._sanitize_headers = frozenset(name.lower() for name in sanitize_headers) + self._server_request_hook = server_request_hook + self._server_response_hook = server_response_hook + self._websocket_connect_hook = websocket_connect_hook + self._websocket_close_hook = websocket_close_hook + self._metrics_include_raw_path = metrics_include_raw_path + self._websocket_error_close_codes = ( + frozenset(websocket_error_close_codes) + if websocket_error_close_codes is not None + else _DEFAULT_WEBSOCKET_ERROR_CLOSE_CODES + ) + # Live spans keyed by the id of the (wrapped) context/websocket the + # after hooks receive; entries are removed when the span ends. + self._spans: Dict[int, _SpanHolder] = {} + + @staticmethod + def _create_resource(service_name: str, + resource_attributes: Optional[Mapping[str, Any]]) -> Resource: + attributes: Dict[str, Any] = {'service.name': service_name} + if resource_attributes: + attributes.update(resource_attributes) + return Resource.create(attributes) + + def apply(self, app: KayaApp) -> None: + self._app = app + app.add_before_request_hook(self._before_request) + app.add_after_request_hook(self._after_request) + app.add_before_websocket_hook(self._before_websocket) + app.add_after_websocket_hook(self._after_websocket) + + def _is_excluded(self, path: str) -> bool: + if path in self._excluded_paths: + return True + return any(regex.search(path) is not None for regex in self._excluded_path_regexes) + + @staticmethod + def _call_hook(hook: Optional[Callable[[Span, Any], None]], span: Span, value: Any) -> None: + if hook is None: + return + try: + hook(span, value) + except Exception: + log.exception('OpenTelemetry span hook failed') + + def _capture_headers(self, + span: Span, + headers: Mapping[str, StrOrStrings | Sequence[str]], + configured: Sequence[str], + prefix: str) -> None: + if not configured: + return + lower_headers = {name.lower(): value for name, value in headers.items()} + for name in configured: + key = name.lower() + if key not in lower_headers: + continue + values = _header_values(lower_headers[key]) + if key in self._sanitize_headers: + values = ['REDACTED'] * len(values) + span.set_attribute(f'{prefix}.{_normalize_header_name(name)}', values) + + def _capture_request_headers(self, span: Span, headers: Mapping[str, Sequence[str]]) -> None: + self._capture_headers(span, headers, self._request_headers_to_capture, 'http.request.header') + + def _capture_response_headers(self, + span: Span, + headers: Optional[Mapping[str, StrOrStrings]]) -> None: + if headers is None: + return + self._capture_headers(span, headers, self._response_headers_to_capture, 'http.response.header') + + def _start_span(self, name: str, attributes: Mapping[str, Any], + headers: Mapping[str, Sequence[str]]) -> _SpanHolder: + parent: Context = _PROPAGATOR.extract(carrier=headers, getter=_GETTER) + span = self._tracer.start_span( + name, + context=parent, + kind=SpanKind.SERVER, + attributes=dict(attributes), + ) + self._capture_request_headers(span, headers) + token = otel_context.attach(trace.set_span_in_context(span, parent)) + return _SpanHolder(span, token, {}) + + @staticmethod + def _finish_span(holder: _SpanHolder, error: bool) -> None: + if error: + holder.span.set_status(Status(StatusCode.ERROR)) + holder.span.end() + otel_context.detach(holder.token) # type: ignore[arg-type] + + def _pop_holder(self, obj: Any) -> Optional[_SpanHolder]: + """Find and remove the span holder for a context/websocket. + + The before hooks key holders by the id of the context they received, + but after hooks run on the outermost wrapper (later mixins may have + wrapped the context again), so the id may not match. In that case the + holder is reachable through the wrapper delegation chain as + ``_holder`` and the stale dict entry is reaped by identity. + """ + holder = self._spans.pop(id(obj), None) + if holder is not None: + return holder + candidate = getattr(obj, '_holder', None) + if not isinstance(candidate, _SpanHolder): + return None + for key, value in list(self._spans.items()): + if value is candidate: + del self._spans[key] + return candidate + + def _route_template(self, ctx: HttpContext) -> Optional[str]: + if self._app is None: + return None + return self._app.route_template(ctx.path, ctx.method) + + async def _before_request(self, ctx: HttpContext) -> Optional[HttpContext]: + if self._is_excluded(ctx.path): + return None + holder = self._start_span( + f'{ctx.method} {ctx.path}', + _server_attributes(ctx), + ctx.headers, + ) + holder.active_attributes = {'http.request.method': str(ctx.method)} + if self._metrics_include_raw_path: + holder.active_attributes['url.path'] = ctx.path + self._spans[id(ctx)] = holder + self._active_requests.add(1, holder.active_attributes) + self._call_hook(self._server_request_hook, holder.span, ctx) + return _TracedHttpContext(ctx, holder, self) + + async def _after_request(self, ctx: HttpContext) -> None: + holder = self._pop_holder(ctx) + if holder is None: + return + status_code = holder.status_code + exception = getattr(ctx, 'exception', None) + error = status_code is not None and status_code >= 500 + if isinstance(exception, BaseException): + holder.span.record_exception(exception) + error = True + route_template = self._route_template(ctx) + if route_template is not None: + holder.span.update_name(f'{ctx.method} {route_template}') + holder.span.set_attribute('http.route', route_template) + self._call_hook(self._server_response_hook, holder.span, ctx) + self._finish_span(holder, error) + self._active_requests.add(-1, holder.active_attributes) + metric_attributes: Dict[str, Any] = { + 'http.request.method': str(ctx.method), + } + if self._metrics_include_raw_path: + metric_attributes['url.path'] = ctx.path + elif route_template is not None: + metric_attributes['http.route'] = route_template + if status_code is not None: + metric_attributes['http.response.status_code'] = status_code + duration = (time.perf_counter_ns() - holder.start_ns) / 1e9 + self._request_duration.record(duration, metric_attributes) + + async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]: + if self._is_excluded(ws.path): + return None + holder = self._start_span( + f'WS {ws.path}', + { + 'url.scheme': ws.scheme, + 'url.path': ws.path, + }, + ws.headers, + ) + self._spans[id(ws)] = holder + self._call_hook(self._websocket_connect_hook, holder.span, ws) + return _TracedWebSocket(ws, holder) + + async def _after_websocket(self, ws: WebSocket) -> None: + holder = self._pop_holder(ws) + if holder is None: + return + exception = getattr(ws, 'exception', None) + error = isinstance(exception, BaseException) + if isinstance(exception, BaseException): + holder.span.record_exception(exception) + close_code = holder.close_code + if close_code is not None: + holder.span.set_attribute('kaya.websocket.close_code', close_code) + error = error or close_code in self._websocket_error_close_codes + self._call_hook(self._websocket_close_hook, holder.span, ws) + self._finish_span(holder, error) + + def shutdown(self, loop: AbstractEventLoop) -> None: + # Only providers created by this mixin are shut down; injected ones + # stay under the caller's control. + if self._owns_tracer_provider: + self._tracer_provider.shutdown() + if self._owns_meter_provider: + self._meter_provider.shutdown() diff --git a/packages/kaya-otel/src/kaya/otel/py.typed b/packages/kaya-otel/src/kaya/otel/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/packages/kaya-otel/tests/test_otel.py b/packages/kaya-otel/tests/test_otel.py new file mode 100644 index 0000000..fb7816f --- /dev/null +++ b/packages/kaya-otel/tests/test_otel.py @@ -0,0 +1,394 @@ +import json +import unittest +from typing import Any, Optional, Tuple + +import httpx +from pwo import async_test + +from kaya.core import HttpContext, KayaApp +from kaya.core._asgi import AsgiWebSocket +from kaya.otel import OTelMixin +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +TRACEPARENT = '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01' + + +def make_mixin(**kwargs) -> Tuple[OTelMixin, InMemorySpanExporter, InMemoryMetricReader]: + resource = Resource.create({'service.name': 'test-service'}) + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider(resource=resource) + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader]) + mixin = OTelMixin( + service_name='test-service', + tracer_provider=tracer_provider, + meter_provider=meter_provider, + **kwargs, + ) + return mixin, span_exporter, metric_reader + + +def make_app(mixin: OTelMixin) -> KayaApp: + app = KayaApp(mixins=[mixin]) + + @app.GET('/hello') + async def hello(ctx: HttpContext) -> None: + await ctx.send_str(200, json.dumps({'ok': True})) + + @app.GET('/boom') + async def boom(ctx: HttpContext) -> None: + await ctx.send_str(500, 'boom') + + return app + + +async def request(app: KayaApp, path: str = '/hello', headers: Optional[dict[str, str]] = None) -> httpx.Response: + transport = httpx.ASGITransport(app=app, client=('127.0.0.1', 123)) + async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client: + return await http_client.get(path, headers=headers) + + +class HttpTracingTest(unittest.TestCase): + + @async_test + async def test_request_produces_server_span(self): + mixin, span_exporter, _ = make_mixin() + app = make_app(mixin) + r = await request(app) + self.assertEqual(200, r.status_code) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + span = spans[0] + self.assertEqual('GET /hello', span.name) + assert span.attributes is not None + self.assertEqual('GET', span.attributes['http.request.method']) + self.assertEqual('/hello', span.attributes['url.path']) + self.assertEqual(200, span.attributes['http.response.status_code']) + self.assertEqual(StatusCode.UNSET, span.status.status_code) + + @async_test + async def test_5xx_marks_span_as_error(self): + mixin, span_exporter, _ = make_mixin() + app = make_app(mixin) + r = await request(app, '/boom') + self.assertEqual(500, r.status_code) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(500, spans[0].attributes['http.response.status_code']) # type: ignore[index] + self.assertEqual(StatusCode.ERROR, spans[0].status.status_code) + + @async_test + async def test_traceparent_header_propagates(self): + mixin, span_exporter, _ = make_mixin() + app = make_app(mixin) + await request(app, headers={'traceparent': TRACEPARENT}) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + span = spans[0] + self.assertEqual(0x4bf92f3577b34da6a3ce929d0e0e4736, span.context.trace_id) # type: ignore[union-attr] + assert span.parent is not None + self.assertEqual(0x00f067aa0ba902b7, span.parent.span_id) + + @async_test + async def test_unmatched_route_still_traced(self): + mixin, span_exporter, _ = make_mixin() + app = make_app(mixin) + r = await request(app, '/nowhere') + self.assertEqual(404, r.status_code) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(404, spans[0].attributes['http.response.status_code']) # type: ignore[index] + + @async_test + async def test_no_leaked_spans_after_requests(self): + mixin, span_exporter, _ = make_mixin() + app = make_app(mixin) + await request(app) + await request(app, '/boom') + self.assertEqual(0, len(mixin._spans)) + self.assertEqual(2, len(span_exporter.get_finished_spans())) + + + @async_test + async def test_exception_marks_span_as_error(self): + mixin, span_exporter, _ = make_mixin() + app = KayaApp(mixins=[mixin]) + + @app.GET('/raises') + async def raises(ctx: HttpContext) -> None: + raise RuntimeError('boom') + + transport = httpx.ASGITransport(app=app, client=('127.0.0.1', 123)) + async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client: + with self.assertRaises(RuntimeError): + await http_client.get('/raises') + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + span = spans[0] + self.assertEqual(StatusCode.ERROR, span.status.status_code) + self.assertTrue(any(event.name == 'exception' for event in span.events)) + self.assertEqual(0, len(mixin._spans)) + + @async_test + async def test_excluded_path_skips_tracing_and_metrics(self): + mixin, span_exporter, metric_reader = make_mixin(excluded_paths=('/health',)) + app = KayaApp(mixins=[mixin]) + + @app.GET('/health') + async def health(ctx: HttpContext) -> None: + await ctx.send_str(200, 'ok') + + @app.GET('/hello') + async def hello(ctx: HttpContext) -> None: + await ctx.send_str(200, 'hello') + + await request(app, '/health') + await request(app, '/hello') + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual('GET /hello', spans[0].name) + + metrics = MetricsTest._metric_names(metric_reader) + duration_points = list(metrics['http.server.request.duration'].data.data_points) + self.assertEqual(1, len(duration_points)) + self.assertEqual(1, duration_points[0].count) + + @async_test + async def test_route_template_used_for_span_name_and_metrics(self): + mixin, span_exporter, metric_reader = make_mixin() + app = KayaApp(mixins=[mixin]) + + @app.GET('/items/${item_id:int}') + async def item(ctx: HttpContext, item_id: int) -> None: + await ctx.send_str(200, str(item_id)) + + r = await request(app, '/items/123') + self.assertEqual(200, r.status_code) + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + span = spans[0] + self.assertEqual('GET /items/${item_id:int}', span.name) + assert span.attributes is not None + self.assertEqual('/items/${item_id:int}', span.attributes['http.route']) + self.assertEqual('/items/123', span.attributes['url.path']) + + metrics = MetricsTest._metric_names(metric_reader) + duration_points = list(metrics['http.server.request.duration'].data.data_points) + self.assertEqual(1, len(duration_points)) + duration_attributes = dict(duration_points[0].attributes) + self.assertEqual('GET', duration_attributes['http.request.method']) + self.assertEqual('/items/${item_id:int}', duration_attributes['http.route']) + self.assertEqual(200, duration_attributes['http.response.status_code']) + self.assertNotIn('url.path', duration_attributes) + + active_points = list(metrics['http.server.active_requests'].data.data_points) + self.assertEqual(1, len(active_points)) + self.assertEqual({'http.request.method': 'GET'}, dict(active_points[0].attributes)) + + @async_test + async def test_header_capture_and_sanitization(self): + mixin, span_exporter, _ = make_mixin( + capture_request_headers=('X-Request-Id', 'Authorization'), + capture_response_headers=('X-Response-Id', 'Set-Cookie'), + ) + app = KayaApp(mixins=[mixin]) + + @app.GET('/headers') + async def headers(ctx: HttpContext) -> None: + await ctx.send_str(200, 'ok', { + 'X-Response-Id': 'res-1', + 'Set-Cookie': 'sid=secret', + }) + + r = await request(app, '/headers', headers={ + 'X-Request-Id': 'req-1', + 'Authorization': 'Bearer secret', + }) + self.assertEqual(200, r.status_code) + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + attributes = spans[0].attributes + assert attributes is not None + self.assertEqual(['req-1'], list(attributes['http.request.header.x_request_id'])) + self.assertEqual(['REDACTED'], list(attributes['http.request.header.authorization'])) + self.assertEqual(['res-1'], list(attributes['http.response.header.x_response_id'])) + self.assertEqual(['REDACTED'], list(attributes['http.response.header.set_cookie'])) + + @async_test + async def test_http_hooks_are_called(self): + calls = [] + + def request_hook(span: Any, ctx: HttpContext) -> None: + calls.append(('request', ctx.path)) + span.set_attribute('test.request_hook', True) + + def response_hook(span: Any, ctx: HttpContext) -> None: + calls.append(('response', ctx.path)) + span.set_attribute('test.response_hook', True) + + mixin, span_exporter, _ = make_mixin( + server_request_hook=request_hook, + server_response_hook=response_hook, + ) + app = make_app(mixin) + await request(app) + + self.assertEqual([('request', '/hello'), ('response', '/hello')], calls) + attributes = span_exporter.get_finished_spans()[0].attributes + assert attributes is not None + self.assertTrue(attributes['test.request_hook']) + self.assertTrue(attributes['test.response_hook']) + + +class MetricsTest(unittest.TestCase): + + @staticmethod + def _metric_names(metric_reader: InMemoryMetricReader) -> dict: + data = metric_reader.get_metrics_data() + return {m.name: m for rm in data.resource_metrics for m in rm.scope_metrics[0].metrics + if rm.scope_metrics} + + @async_test + async def test_duration_histogram_and_active_requests(self): + mixin, _, metric_reader = make_mixin() + app = make_app(mixin) + await request(app) + metrics = self._metric_names(metric_reader) + self.assertIn('http.server.request.duration', metrics) + self.assertIn('http.server.active_requests', metrics) + duration_points = list(metrics['http.server.request.duration'].data.data_points) + self.assertEqual(1, len(duration_points)) + self.assertEqual(1, duration_points[0].count) + self.assertGreaterEqual(duration_points[0].sum, 0) + active_points = list(metrics['http.server.active_requests'].data.data_points) + self.assertEqual(1, len(active_points)) + self.assertEqual(0, active_points[0].value) + + +class WebSocketTracingTest(unittest.TestCase): + + @staticmethod + def _make_ws(headers=()): + async def send(message): + pass + + async def receive(): + return {'type': 'websocket.connect'} + + scope = { + 'type': 'websocket', + 'path': '/ws/games/abc', + 'query_string': b'', + 'scheme': 'ws', + 'client': ('127.0.0.1', 12345), + 'server': ('127.0.0.1', 80), + 'headers': headers, + } + return AsgiWebSocket(scope, receive, send) + + @async_test + async def test_websocket_lifecycle_produces_span(self): + mixin, span_exporter, _ = make_mixin() + ws = self._make_ws() + wrapped = await mixin._before_websocket(ws) + self.assertIsNotNone(wrapped) + self.assertEqual(1, len(mixin._spans)) + await mixin._after_websocket(wrapped) # type: ignore[arg-type] + self.assertEqual(0, len(mixin._spans)) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual('WS /ws/games/abc', spans[0].name) + + @async_test + async def test_websocket_traceparent_propagates(self): + mixin, span_exporter, _ = make_mixin() + ws = self._make_ws(headers=[(b'traceparent', TRACEPARENT.encode())]) + wrapped = await mixin._before_websocket(ws) + await mixin._after_websocket(wrapped) # type: ignore[arg-type] + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(0x4bf92f3577b34da6a3ce929d0e0e4736, spans[0].context.trace_id) # type: ignore[union-attr] + + @async_test + async def test_wrapped_websocket_delegates(self): + mixin, _, _ = make_mixin() + ws = self._make_ws() + wrapped = await mixin._before_websocket(ws) + assert wrapped is not None + self.assertEqual('/ws/games/abc', wrapped.path) + self.assertEqual(('127.0.0.1', 12345), wrapped.client) + await mixin._after_websocket(wrapped) + + @async_test + async def test_websocket_error_close_code_marks_span_error(self): + mixin, span_exporter, _ = make_mixin() + ws = self._make_ws() + wrapped = await mixin._before_websocket(ws) + assert wrapped is not None + await wrapped.close(1011) + await mixin._after_websocket(wrapped) + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(StatusCode.ERROR, spans[0].status.status_code) + assert spans[0].attributes is not None + self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code']) + + @async_test + async def test_websocket_exception_marks_span_error(self): + mixin, span_exporter, _ = make_mixin() + ws = self._make_ws() + wrapped = await mixin._before_websocket(ws) + assert wrapped is not None + wrapped.exception = RuntimeError('ws boom') + await mixin._after_websocket(wrapped) + + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(StatusCode.ERROR, spans[0].status.status_code) + self.assertTrue(any(event.name == 'exception' for event in spans[0].events)) + + @async_test + async def test_websocket_hooks_are_called(self): + calls = [] + + def connect_hook(span: Any, ws: Any) -> None: + calls.append(('connect', ws.path)) + span.set_attribute('test.websocket_connect_hook', True) + + def close_hook(span: Any, ws: Any) -> None: + calls.append(('close', ws.path)) + span.set_attribute('test.websocket_close_hook', True) + + mixin, span_exporter, _ = make_mixin( + websocket_connect_hook=connect_hook, + websocket_close_hook=close_hook, + ) + ws = self._make_ws() + wrapped = await mixin._before_websocket(ws) + assert wrapped is not None + await wrapped.close(1000) + await mixin._after_websocket(wrapped) + + self.assertEqual([('connect', '/ws/games/abc'), ('close', '/ws/games/abc')], calls) + spans = span_exporter.get_finished_spans() + self.assertEqual(1, len(spans)) + self.assertEqual(StatusCode.UNSET, spans[0].status.status_code) + assert spans[0].attributes is not None + self.assertTrue(spans[0].attributes['test.websocket_connect_hook']) + self.assertTrue(spans[0].attributes['test.websocket_close_hook']) + + +if __name__ == '__main__': + unittest.main() diff --git a/requirements-dev.in b/requirements-dev.in index 5066501..0b7a468 100644 --- a/requirements-dev.in +++ b/requirements-dev.in @@ -7,6 +7,7 @@ kaya-oidc @ file:./packages/kaya-oidc kaya-openapi @ file:./packages/kaya-openapi kaya-cors @ file:./packages/kaya-cors kaya-forwarded @ file:./packages/kaya-forwarded +kaya-otel @ file:./packages/kaya-otel build fakeredis mypy diff --git a/requirements-dev.txt b/requirements-dev.txt index a652173..1b03e16 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -44,6 +44,8 @@ executing==2.2.1 # via stack-data fakeredis==2.36.2 # via -r requirements-dev.in +googleapis-common-protos==1.75.3 + # via opentelemetry-exporter-otlp-proto-http granian==2.7.9 # via kaya-rsgi h11==0.16.0 @@ -93,6 +95,7 @@ file:./packages/kaya-core # kaya-forwarded # kaya-oidc # kaya-openapi + # kaya-otel # kaya-rsgi # kaya-session file:./packages/kaya-cors @@ -103,6 +106,8 @@ file:./packages/kaya-oidc # via -r requirements-dev.in file:./packages/kaya-openapi # via -r requirements-dev.in +file:./packages/kaya-otel + # via -r requirements-dev.in file:./packages/kaya-rsgi # via -r requirements-dev.in file:./packages/kaya-session @@ -135,6 +140,25 @@ mypy-extensions==1.1.0 # via mypy nh3==0.3.6 # via readme-renderer +opentelemetry-api==1.44.0 + # via + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp-proto-common==1.44.0 + # via opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-http==1.44.0 + # via kaya-otel +opentelemetry-proto==1.44.0 + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-http +opentelemetry-sdk==1.44.0 + # via + # kaya-otel + # opentelemetry-exporter-otlp-proto-http +opentelemetry-semantic-conventions==0.65b0 + # via opentelemetry-sdk packaging==26.2 # via # build @@ -147,6 +171,10 @@ pexpect==4.9.0 # via ipython prompt-toolkit==3.0.52 # via ipython +protobuf==7.36.2 + # via + # googleapis-common-protos + # opentelemetry-proto psutil==7.2.2 # via ipython ptyprocess==0.7.0 @@ -178,6 +206,7 @@ redis==8.0.1 # kaya-session-redis requests==2.34.2 # via + # opentelemetry-exporter-otlp-proto-http # requests-toolbelt # twine requests-toolbelt==1.0.0 @@ -202,6 +231,10 @@ typing-extensions==4.16.0 # via # kaya-core # mypy + # opentelemetry-api + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions # pwo urllib3==2.7.0 # via