007 add telemetry

app/observability/tracing.py

@@ -0,0 +1,80 @@
+"""OpenTelemetry tracing setup and instrumentation helpers."""
+
+from typing import cast
+
+from fastapi import FastAPI
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import Tracer, TracerProvider as TraceAPIProvider
+
+_TRACER_PROVIDER: TraceAPIProvider | None = None
+_HTTPX_INSTRUMENTED = False
+
+
+def setup_tracing(
+    *,
+    service_name: str,
+    otlp_endpoint: str,
+    tracer_provider: TraceAPIProvider | None = None,
+) -> TraceAPIProvider:
+    global _TRACER_PROVIDER
+
+    if tracer_provider is None and _TRACER_PROVIDER is not None:
+        return _TRACER_PROVIDER
+
+    if tracer_provider is None:
+        tracer_provider = TracerProvider(
+            resource=Resource.create({"service.name": service_name}),
+        )
+        if otlp_endpoint:
+            span_exporter = OTLPSpanExporter(
+                endpoint=otlp_endpoint,
+                insecure=otlp_endpoint.startswith("http://"),
+            )
+            tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
+
+    _TRACER_PROVIDER = tracer_provider
+    return tracer_provider
+
+
+def instrument_fastapi_app(app: FastAPI) -> None:
+    tracer_provider = _TRACER_PROVIDER
+    if tracer_provider is None:
+        return
+    FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
+
+
+def instrument_httpx_client() -> None:
+    global _HTTPX_INSTRUMENTED
+
+    tracer_provider = _TRACER_PROVIDER
+    if tracer_provider is None:
+        return
+
+    if _HTTPX_INSTRUMENTED:
+        return
+    HTTPXClientInstrumentor().instrument(tracer_provider=tracer_provider)
+    _HTTPX_INSTRUMENTED = True
+
+
+def get_tracer(name: str) -> Tracer:
+    tracer_provider = _TRACER_PROVIDER
+    if tracer_provider is None:
+        return trace.get_tracer(name)
+    return cast(Tracer, tracer_provider.get_tracer(name))
+
+
+def reset_tracing_for_tests() -> None:
+    global _TRACER_PROVIDER, _HTTPX_INSTRUMENTED
+
+    if _HTTPX_INSTRUMENTED:
+        HTTPXClientInstrumentor().uninstrument()
+    if isinstance(_TRACER_PROVIDER, TracerProvider):
+        _TRACER_PROVIDER.shutdown()
+    _TRACER_PROVIDER = None
+    _HTTPX_INSTRUMENTED = False