Source code for sqlspec.observability._spans

"""Optional OpenTelemetry span helpers."""

from importlib import import_module
from typing import Any

from sqlspec.exceptions import MissingDependencyError
from sqlspec.observability._common import resolve_db_system
from sqlspec.observability._config import TelemetryConfig
from sqlspec.utils.logging import get_logger
from sqlspec.utils.module_loader import ensure_opentelemetry
from sqlspec.utils.type_guards import has_tracer_provider

__all__ = ("SpanManager", "resolve_db_system")


logger = get_logger("sqlspec.observability.spans")




class SpanManager:
    """Lazy OpenTelemetry span manager with graceful degradation."""

    __slots__ = (
        "_enabled",
        "_provider_factory",
        "_resource_attributes",
        "_span_kind",
        "_status_cls",
        "_status_code_cls",
        "_trace_api",
        "_tracer",
    )



    def __init__(self, telemetry: TelemetryConfig | None = None) -> None:
        telemetry = telemetry or TelemetryConfig()
        self._enabled = bool(telemetry.enable_spans)
        self._provider_factory = telemetry.provider_factory
        self._resource_attributes = dict(telemetry.resource_attributes or {})
        self._trace_api: Any | None = None
        self._status_cls: Any | None = None
        self._status_code_cls: Any | None = None
        self._span_kind: Any | None = None
        self._tracer: Any | None = None
        if self._enabled:
            self._resolve_api()

@property def is_enabled(self) -> bool: """Return True once OpenTelemetry spans are available.""" return bool(self._enabled and self._tracer)

def start_query_span( self, *, driver: str, adapter: str, bind_key: str | None, sql: str, operation: str, connection_info: dict[str, Any] | None = None, storage_backend: str | None = None, correlation_id: str | None = None, ) -> Any: """Start a query span with SQLSpec semantic attributes.""" if not self._enabled: return None attributes: dict[str, Any] = { "db.system": resolve_db_system(adapter), "db.operation": operation, "sqlspec.driver": driver, } if sql: attributes["db.statement"] = sql if bind_key: attributes["sqlspec.bind_key"] = bind_key if storage_backend: attributes["sqlspec.storage_backend"] = storage_backend if correlation_id: attributes["sqlspec.correlation_id"] = correlation_id if connection_info: attributes.update(connection_info) attributes.update(self._resource_attributes) return self._start_span("sqlspec.query", attributes)