File size: 3,640 Bytes
13fe504 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Optional OpenTelemetry hooks for DataForge backend surfaces."""
from __future__ import annotations
import os
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from importlib import import_module
from typing import Any, Literal
_SENSITIVE_ATTR_FRAGMENTS = ("authorization", "cookie", "token", "key", "secret", "password")
RepairTelemetryStage = Literal[
"detect",
"propose",
"safety_gate",
"smt_verify",
"transaction_create",
"transaction_apply",
"receipt",
"revert",
]
REPAIR_TELEMETRY_SPANS: dict[str, str] = {
"detect": "dataforge.repair.detect",
"propose": "dataforge.repair.propose",
"safety_gate": "dataforge.repair.safety_gate",
"smt_verify": "dataforge.repair.smt_verify",
"transaction_create": "dataforge.repair.transaction_create",
"transaction_apply": "dataforge.repair.transaction_apply",
"receipt": "dataforge.repair.receipt",
"revert": "dataforge.repair.revert",
}
REPAIR_TELEMETRY_STAGES = frozenset(REPAIR_TELEMETRY_SPANS)
def _otel_enabled() -> bool:
"""Return whether optional OpenTelemetry instrumentation is enabled."""
return os.environ.get("DATAFORGE_OTEL_ENABLED", "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
def _safe_attrs(attributes: dict[str, Any]) -> dict[str, str | int | float | bool]:
"""Keep only scalar, non-sensitive telemetry attributes."""
safe: dict[str, str | int | float | bool] = {}
for key, value in attributes.items():
lowered = key.lower()
if any(fragment in lowered for fragment in _SENSITIVE_ATTR_FRAGMENTS):
continue
if lowered in {"row_values", "rows", "payload", "source_bytes", "csv"}:
continue
if isinstance(value, str | int | float | bool):
safe[key] = value
return safe
def configure_fastapi_observability(app: Any, *, service_name: str) -> bool:
"""Instrument a FastAPI app when OpenTelemetry is explicitly enabled."""
if not _otel_enabled():
return False
try:
fastapi_instrumentation = import_module("opentelemetry.instrumentation.fastapi")
trace_module = import_module("opentelemetry.trace")
except ImportError:
return False
app.state.dataforge_service_name = service_name
fastapi_instrumentation.FastAPIInstrumentor.instrument_app(
app,
tracer_provider=trace_module.get_tracer_provider(),
excluded_urls="/api/docs,/docs,/redoc,/openapi.json",
)
return True
def _span_name(stage: str) -> str:
"""Return the span name for a public repair telemetry stage."""
if stage in REPAIR_TELEMETRY_SPANS:
return REPAIR_TELEMETRY_SPANS[stage]
return stage
@contextmanager
def repair_stage_span(stage: RepairTelemetryStage | str, **attributes: Any) -> Iterator[None]:
"""Create a repair-stage span when OpenTelemetry is available."""
if not _otel_enabled():
with nullcontext():
yield
return
try:
trace_module = import_module("opentelemetry.trace")
except ImportError:
with nullcontext():
yield
return
tracer = trace_module.get_tracer("dataforge.repair")
with tracer.start_as_current_span(_span_name(stage)) as span:
if stage in REPAIR_TELEMETRY_SPANS:
span.set_attribute("dataforge.repair.stage", stage)
for key, value in _safe_attrs(attributes).items():
span.set_attribute(key, value)
yield
|