Раис Юсупалиев
81 lines
2.4 KiB
Python
81 lines
2.4 KiB
Python
"""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
|