Coverage for core/telemetry.py: 87.32%

1import logging 

2from urllib.parse import urljoin 

3 

4from fastapi import FastAPI 

5from opentelemetry import trace 

6from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter 

7from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor 

8from opentelemetry.instrumentation.logging import LoggingInstrumentor 

9from opentelemetry.instrumentation.redis import RedisInstrumentor 

10from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor 

11from opentelemetry.sdk.resources import Resource 

12from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider 

13from opentelemetry.sdk.trace.export import BatchSpanProcessor 

14 

15from core.config import SKIP_METHODS, otel_excluded_urls, settings 

16from core.database import async_engine, engine 

17 

18logger = logging.getLogger("telemetry") 

19 

20EXCLUDED_URLS = otel_excluded_urls() 

21 

22 

23def _span_http_method(span: ReadableSpan) -> str: 

24 attributes = span.attributes or {} 

25 raw = attributes.get("http.request.method") or attributes.get("http.method") or "" 

26 return str(raw).upper() 

27 

28 

29def should_drop_span(span: ReadableSpan) -> bool: 

30 method = _span_http_method(span) 

31 if method in SKIP_METHODS: 

32 return True 

33 name = span.name or "" 

34 return any(name.startswith(f"{skipped} ") for skipped in SKIP_METHODS) 

35 

36 

37class SkipSpanProcessor(SpanProcessor): 

38 """Drop spans FastAPI excluded_urls cannot filter (method-only, e.g. OPTIONS).""" 

39 

40 def __init__(self, wrapped: SpanProcessor) -> None: 

41 self._wrapped = wrapped 

42 

43 def on_start(self, span: Span, parent_context=None) -> None: 

44 self._wrapped.on_start(span, parent_context) 

45 

46 def on_end(self, span: ReadableSpan) -> None: 

47 if should_drop_span(span): 

48 return 

49 self._wrapped.on_end(span) 

50 

51 def shutdown(self) -> None: 

52 self._wrapped.shutdown() 

53 

54 def force_flush(self, timeout_millis: int = 30000) -> bool: 

55 return bool(self._wrapped.force_flush(timeout_millis)) 

56 

57 

58class TraceIdFormatter(logging.Formatter): 

59 def format(self, record: logging.LogRecord) -> str: 

60 if not getattr(record, "otelTraceID", None): 

61 record.otelTraceID = "0" 

62 return super().format(record) 

63 

64 

65def _traces_endpoint(base: str) -> str: 

66 normalized = base.rstrip("/") + "/" 

67 if normalized.endswith("/v1/traces/"): 

68 return normalized.rstrip("/") 

69 return urljoin(normalized, "v1/traces") 

70 

71 

72def _build_resource() -> Resource: 

73 return Resource.create( 

74 { 

75 "service.name": settings.PROJECT_NAME, 

76 "service.version": settings.PROJECT_VERSION, 

77 "deployment.environment": ("development" if settings.DEBUG_MODE else "production"), 

78 } 

79 ) 

80 

81 

82def setup_telemetry(app: FastAPI) -> None: 

83 LoggingInstrumentor().instrument( 

84 inject_trace_context=True, 

85 set_logging_format=False, 

86 enable_log_auto_instrumentation=False, 

87 ) 

88 

89 if not settings.OTEL_ENABLE: 

90 return 

91 

92 resource = _build_resource() 

93 provider = TracerProvider(resource=resource) 

94 exporter = OTLPSpanExporter(endpoint=_traces_endpoint(settings.OTEL_EXPORTER_OTLP_ENDPOINT)) 

95 provider.add_span_processor( 

96 SkipSpanProcessor(BatchSpanProcessor(exporter, schedule_delay_millis=1000)) 

97 ) 

98 trace.set_tracer_provider(provider) 

99 

100 FastAPIInstrumentor.instrument_app(app, excluded_urls=EXCLUDED_URLS) 

101 SQLAlchemyInstrumentor().instrument(engines=[async_engine.sync_engine, engine]) 

102 RedisInstrumentor().instrument() 

103 logger.info("OpenTelemetry tracing enabled") 

104 

105 

106def shutdown_telemetry() -> None: 

107 if not settings.OTEL_ENABLE: 

108 return 

109 provider = trace.get_tracer_provider() 

110 shutdown = getattr(provider, "shutdown", None) 

111 if callable(shutdown): 

112 shutdown()