sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/observability/_runtime.py

@@ -0,0 +1,420 @@
+"""Runtime helpers that bundle lifecycle, observer, and span orchestration."""
+
+import re
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.observability._common import compute_sql_hash, get_trace_context, resolve_db_system
+from sqlspec.observability._config import LoggingConfig, ObservabilityConfig
+from sqlspec.observability._dispatcher import LifecycleDispatcher, LifecycleHook
+from sqlspec.observability._observer import StatementObserver, create_event, create_statement_observer
+from sqlspec.observability._spans import SpanManager
+from sqlspec.utils.correlation import CorrelationContext
+from sqlspec.utils.type_guards import has_span_attribute
+
+_LITERAL_PATTERN = re.compile(r"'(?:''|[^'])*'")
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from sqlspec.storage import StorageTelemetry
+
+
+class ObservabilityRuntime:
+    """Aggregates dispatchers, observers, spans, and custom metrics."""
+
+    __slots__ = (
+        "_metrics",
+        "_redaction",
+        "_statement_observers",
+        "bind_key",
+        "config",
+        "config_name",
+        "lifecycle",
+        "span_manager",
+    )
+
+    # Allow test injection with fake span managers (mypyc strict typing workaround)
+    span_manager: "Any"
+
+    def __init__(
+        self, config: ObservabilityConfig | None = None, *, bind_key: str | None = None, config_name: str | None = None
+    ) -> None:
+        config = config.copy() if config else ObservabilityConfig()
+        if config.logging is None:
+            config.logging = LoggingConfig()
+        self.config = config
+        self.bind_key = bind_key
+        self.config_name = config_name or "SQLSpecConfig"
+        lifecycle_config = cast("dict[str, Iterable[LifecycleHook]] | None", config.lifecycle)
+        self.lifecycle = LifecycleDispatcher(lifecycle_config)
+        self.span_manager = SpanManager(config.telemetry)
+        observers: list[StatementObserver] = []
+        if config.statement_observers:
+            observers.extend(config.statement_observers)
+        if config.print_sql:
+            observers.append(create_statement_observer(config.logging))
+        self._statement_observers = tuple(observers)
+        self._redaction = config.redaction.copy() if config.redaction else None
+        self._metrics: dict[str, float] = {}
+
+    @property
+    def has_statement_observers(self) -> bool:
+        """Return True when any observers are registered."""
+
+        return bool(self._statement_observers)
+
+    @property
+    def diagnostics_key(self) -> str:
+        """Derive diagnostics key from bind key or configuration name."""
+
+        if self.bind_key:
+            return self.bind_key
+        return self.config_name
+
+    def base_context(self) -> dict[str, Any]:
+        """Return the base payload for lifecycle events."""
+
+        context = {"config": self.config_name}
+        if self.bind_key:
+            context["bind_key"] = self.bind_key
+        correlation_id = CorrelationContext.get()
+        if correlation_id:
+            context["correlation_id"] = correlation_id
+        return context
+
+    def _build_context(self, **extras: Any) -> dict[str, Any]:
+        context = self.base_context()
+        context.update({key: value for key, value in extras.items() if value is not None})
+        return context
+
+    def lifecycle_snapshot(self) -> dict[str, int]:
+        """Return lifecycle counters keyed under the diagnostics prefix."""
+
+        return self.lifecycle.snapshot(prefix=self.diagnostics_key)
+
+    def metrics_snapshot(self) -> dict[str, float]:
+        """Return accumulated custom metrics with diagnostics prefix."""
+
+        if not self._metrics:
+            return {}
+        prefix = self.diagnostics_key
+        return {f"{prefix}.{name}": value for name, value in self._metrics.items()}
+
+    def increment_metric(self, name: str, amount: float = 1.0) -> None:
+        """Increment a custom metric counter."""
+
+        self._metrics[name] = self._metrics.get(name, 0.0) + amount
+
+    def record_metric(self, name: str, value: float) -> None:
+        """Set a custom metric to an explicit value."""
+
+        self._metrics[name] = value
+
+    def start_migration_span(
+        self, event: str, *, version: "str | None" = None, metadata: "dict[str, Any] | None" = None
+    ) -> Any:
+        """Start a migration span when telemetry is enabled."""
+
+        if not self.span_manager.is_enabled:
+            return None
+        attributes: dict[str, Any] = {"sqlspec.migration.event": event, "sqlspec.config": self.config_name}
+        if self.bind_key:
+            attributes["sqlspec.bind_key"] = self.bind_key
+        correlation_id = CorrelationContext.get()
+        if correlation_id:
+            attributes["sqlspec.correlation_id"] = correlation_id
+        if version:
+            attributes["sqlspec.migration.version"] = version
+        if metadata:
+            for key, value in metadata.items():
+                if value is not None:
+                    attributes[f"sqlspec.migration.{key}"] = value
+        return self.span_manager.start_span(f"sqlspec.migration.{event}", attributes)
+
+    def end_migration_span(
+        self, span: Any, *, duration_ms: "int | None" = None, error: "Exception | None" = None
+    ) -> None:
+        """Finish a migration span, attaching optional duration metadata."""
+
+        if span is None:
+            return
+        if duration_ms is not None and has_span_attribute(span):
+            span.set_attribute("sqlspec.migration.duration_ms", duration_ms)
+        self.span_manager.end_span(span, error=error)
+
+    def emit_pool_create(self, pool: Any) -> None:
+        span = self._start_lifecycle_span("pool.create", subject=pool)
+        try:
+            if self.lifecycle.has_pool_create:
+                self.lifecycle.emit_pool_create(self._build_context(pool=pool))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_pool_destroy(self, pool: Any) -> None:
+        span = self._start_lifecycle_span("pool.destroy", subject=pool)
+        try:
+            if self.lifecycle.has_pool_destroy:
+                self.lifecycle.emit_pool_destroy(self._build_context(pool=pool))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_connection_create(self, connection: Any) -> None:
+        span = self._start_lifecycle_span("connection.create", subject=connection)
+        try:
+            if self.lifecycle.has_connection_create:
+                self.lifecycle.emit_connection_create(self._build_context(connection=connection))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_connection_destroy(self, connection: Any) -> None:
+        span = self._start_lifecycle_span("connection.destroy", subject=connection)
+        try:
+            if self.lifecycle.has_connection_destroy:
+                self.lifecycle.emit_connection_destroy(self._build_context(connection=connection))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_session_start(self, session: Any) -> None:
+        span = self._start_lifecycle_span("session.start", subject=session)
+        try:
+            if self.lifecycle.has_session_start:
+                self.lifecycle.emit_session_start(self._build_context(session=session))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_session_end(self, session: Any) -> None:
+        span = self._start_lifecycle_span("session.end", subject=session)
+        try:
+            if self.lifecycle.has_session_end:
+                self.lifecycle.emit_session_end(self._build_context(session=session))
+        finally:
+            self.span_manager.end_span(span)
+
+    def emit_query_start(self, **extras: Any) -> None:
+        if self.lifecycle.has_query_start:
+            self.lifecycle.emit_query_start(self._build_context(**extras))
+
+    def emit_query_complete(self, **extras: Any) -> None:
+        if self.lifecycle.has_query_complete:
+            self.lifecycle.emit_query_complete(self._build_context(**extras))
+
+    def emit_error(self, exception: Exception, **extras: Any) -> None:
+        if self.lifecycle.has_error:
+            payload = self._build_context(exception=exception)
+            payload.update({key: value for key, value in extras.items() if value is not None})
+            self.lifecycle.emit_error(payload)
+        self.increment_metric("errors", 1.0)
+
+    def emit_statement_event(
+        self,
+        *,
+        sql: str,
+        parameters: Any,
+        driver: str,
+        operation: str,
+        execution_mode: str | None,
+        is_many: bool,
+        is_script: bool,
+        rows_affected: int | None,
+        duration_s: float,
+        storage_backend: str | None,
+        started_at: float | None = None,
+    ) -> None:
+        """Emit a statement event to all registered observers."""
+
+        if not self._statement_observers:
+            return
+        sanitized_sql = self._redact_sql(sql)
+        sanitized_params = self._redact_parameters(parameters)
+        correlation_id = CorrelationContext.get()
+        logging_config = self.config.logging
+        db_system = resolve_db_system(self.config_name)
+        sql_hash = None
+        if logging_config and logging_config.include_sql_hash:
+            sql_hash = compute_sql_hash(sanitized_sql)
+        sql_truncation_length = logging_config.sql_truncation_length if logging_config else 2000
+        sql_original_length = len(sanitized_sql)
+        sql_truncated = sql_original_length > sql_truncation_length
+        trace_id = None
+        span_id = None
+        if logging_config and logging_config.include_trace_context:
+            trace_id, span_id = get_trace_context()
+        event = create_event(
+            sql=sanitized_sql,
+            parameters=sanitized_params,
+            driver=driver,
+            adapter=self.config_name,
+            bind_key=self.bind_key,
+            db_system=db_system,
+            operation=operation,
+            execution_mode=execution_mode,
+            is_many=is_many,
+            is_script=is_script,
+            rows_affected=rows_affected,
+            duration_s=duration_s,
+            correlation_id=correlation_id,
+            storage_backend=storage_backend,
+            started_at=started_at,
+            sql_hash=sql_hash,
+            sql_truncated=sql_truncated,
+            sql_original_length=sql_original_length,
+            trace_id=trace_id,
+            span_id=span_id,
+        )
+        for observer in self._statement_observers:
+            observer(event)
+
+    def start_query_span(self, sql: str, operation: str, driver: str) -> Any:
+        """Start a query span with runtime metadata."""
+
+        sql_hash = compute_sql_hash(sql)
+        connection_info = {"sqlspec.statement.hash": sql_hash, "sqlspec.statement.length": len(sql)}
+        sql_payload = ""
+        if self.config.print_sql:
+            sql_payload = self._redact_sql(sql)
+            sql_payload, truncated = _truncate_text(sql_payload, max_chars=4096)
+            if truncated:
+                connection_info["sqlspec.statement.truncated"] = True
+
+        correlation_id = CorrelationContext.get()
+        return self.span_manager.start_query_span(
+            driver=driver,
+            adapter=self.config_name,
+            bind_key=self.bind_key,
+            sql=sql_payload,
+            operation=operation,
+            connection_info=connection_info,
+            correlation_id=correlation_id,
+        )
+
+    def start_storage_span(
+        self, operation: str, *, destination: str | None = None, format_label: str | None = None
+    ) -> Any:
+        """Start a storage bridge span for read/write operations."""
+
+        if not self.span_manager.is_enabled:
+            return None
+        attributes: dict[str, Any] = {"sqlspec.storage.operation": operation, "sqlspec.config": self.config_name}
+        if self.bind_key:
+            attributes["sqlspec.bind_key"] = self.bind_key
+        correlation_id = CorrelationContext.get()
+        if correlation_id:
+            attributes["sqlspec.correlation_id"] = correlation_id
+        if destination:
+            attributes["sqlspec.storage.destination"] = destination
+        if format_label:
+            attributes["sqlspec.storage.format"] = format_label
+        return self.span_manager.start_span(f"sqlspec.storage.{operation}", attributes)
+
+    def start_span(self, name: str, *, attributes: dict[str, Any] | None = None) -> Any:
+        """Start a custom span enriched with configuration context."""
+
+        if not self.span_manager.is_enabled:
+            return None
+        merged: dict[str, Any] = attributes.copy() if attributes else {}
+        merged.setdefault("sqlspec.config", self.config_name)
+        if self.bind_key:
+            merged.setdefault("sqlspec.bind_key", self.bind_key)
+        correlation_id = CorrelationContext.get()
+        if correlation_id:
+            merged.setdefault("sqlspec.correlation_id", correlation_id)
+        return self.span_manager.start_span(name, merged)
+
+    def end_span(self, span: Any, *, error: Exception | None = None) -> None:
+        """Finish a custom span."""
+
+        self.span_manager.end_span(span, error=error)
+
+    def end_storage_span(
+        self, span: Any, *, telemetry: "StorageTelemetry | None" = None, error: Exception | None = None
+    ) -> None:
+        """Finish a storage span, attaching telemetry metadata when available."""
+
+        if span is None:
+            return
+        if telemetry:
+            telemetry = self.annotate_storage_telemetry(telemetry)
+            self._attach_storage_telemetry(span, telemetry)
+        self.span_manager.end_span(span, error=error)
+
+    def annotate_storage_telemetry(self, telemetry: "StorageTelemetry") -> "StorageTelemetry":
+        """Add bind key / config / correlation metadata to telemetry payloads."""
+
+        annotated = telemetry
+        base = self.base_context()
+        correlation_id = base.get("correlation_id")
+        if correlation_id and not annotated.get("correlation_id"):
+            annotated["correlation_id"] = correlation_id
+        annotated.setdefault("config", self.config_name)
+        if self.bind_key and not annotated.get("bind_key"):
+            annotated["bind_key"] = self.bind_key
+        return annotated
+
+    def _start_lifecycle_span(self, event: str, subject: Any | None = None) -> Any:
+        if not self.span_manager.is_enabled:
+            return None
+        attributes: dict[str, Any] = {"sqlspec.lifecycle.event": event, "sqlspec.config": self.config_name}
+        if self.bind_key:
+            attributes["sqlspec.bind_key"] = self.bind_key
+        correlation_id = CorrelationContext.get()
+        if correlation_id:
+            attributes["sqlspec.correlation_id"] = correlation_id
+        if subject is not None:
+            attributes["sqlspec.lifecycle.subject_type"] = type(subject).__name__
+        return self.span_manager.start_span(f"sqlspec.lifecycle.{event}", attributes)
+
+    def _attach_storage_telemetry(self, span: Any, telemetry: "StorageTelemetry") -> None:
+        if not has_span_attribute(span):
+            return
+        if "backend" in telemetry and telemetry["backend"] is not None:
+            span.set_attribute("sqlspec.storage.backend", telemetry["backend"])
+        if "bytes_processed" in telemetry and telemetry["bytes_processed"] is not None:
+            span.set_attribute("sqlspec.storage.bytes_processed", telemetry["bytes_processed"])
+        if "rows_processed" in telemetry and telemetry["rows_processed"] is not None:
+            span.set_attribute("sqlspec.storage.rows_processed", telemetry["rows_processed"])
+        if "destination" in telemetry and telemetry["destination"] is not None:
+            span.set_attribute("sqlspec.storage.destination", telemetry["destination"])
+        if "format" in telemetry and telemetry["format"] is not None:
+            span.set_attribute("sqlspec.storage.format", telemetry["format"])
+        if "duration_s" in telemetry and telemetry["duration_s"] is not None:
+            span.set_attribute("sqlspec.storage.duration_s", telemetry["duration_s"])
+        if "correlation_id" in telemetry and telemetry["correlation_id"] is not None:
+            span.set_attribute("sqlspec.correlation_id", telemetry["correlation_id"])
+
+    def _redact_sql(self, sql: str) -> str:
+        config = self._redaction
+        if config is None or not config.mask_literals:
+            return sql
+        return _LITERAL_PATTERN.sub("'***'", sql)
+
+    def _redact_parameters(self, parameters: Any) -> Any:
+        config = self._redaction
+        if config is None or not config.mask_parameters:
+            return parameters
+        allow_list = set(config.parameter_allow_list or ())
+        return _mask_parameters(parameters, allow_list)
+
+
+def _mask_parameters(value: Any, allow_list: set[str]) -> Any:
+    if isinstance(value, dict):
+        masked: dict[str, Any] = {}
+        for key, item in value.items():
+            if allow_list and key in allow_list:
+                masked[key] = _mask_parameters(item, allow_list)
+            else:
+                masked[key] = "***"
+        return masked
+    if isinstance(value, list):
+        return [_mask_parameters(item, allow_list) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_mask_parameters(item, allow_list) for item in value)
+    return "***"
+
+
+def _truncate_text(value: str, *, max_chars: int) -> tuple[str, bool]:
+    if len(value) <= max_chars:
+        return value, False
+    return value[:max_chars], True
+
+
+__all__ = ("ObservabilityRuntime",)

sqlspec/observability/_sampling.py

@@ -0,0 +1,188 @@
+"""Sampling configuration for observability data.
+
+This module provides sampling configuration to control the volume of
+logs, metrics, and traces emitted by the observability system.
+Sampling can significantly reduce cloud logging costs while maintaining
+visibility into errors and slow queries.
+"""
+
+import random
+from typing import ClassVar
+
+__all__ = ("SamplingConfig",)
+
+
+class SamplingConfig:
+    """Configuration for log and metric sampling.
+
+    Controls when observability data (logs, spans, metrics) is emitted.
+    Supports both random and deterministic sampling modes, with
+    force-sample conditions for errors and slow queries.
+
+    Attributes:
+        sample_rate: Probability of sampling (0.0 to 1.0). 1.0 means always sample.
+        force_sample_on_error: If True, always sample when an error occurs.
+        force_sample_slow_queries_ms: If set, always sample queries slower than this threshold.
+        deterministic: If True, use hash-based sampling for consistency across distributed systems.
+
+    Example:
+        ```python
+        # Sample 10% of requests, but always sample errors and slow queries
+        config = SamplingConfig(
+            sample_rate=0.1,
+            force_sample_on_error=True,
+            force_sample_slow_queries_ms=100.0,
+            deterministic=True,
+        )
+
+        # Check if a request should be sampled
+        if config.should_sample(
+            correlation_id="abc-123",
+            is_error=False,
+            duration_ms=50.0,
+        ):
+            emit_logs()
+        ```
+    """
+
+    __slots__ = ("deterministic", "force_sample_on_error", "force_sample_slow_queries_ms", "sample_rate")
+
+    DEFAULT_SAMPLE_RATE: ClassVar[float] = 1.0
+    """Default sample rate (100% - sample everything)."""
+
+    DEFAULT_SLOW_QUERY_THRESHOLD_MS: ClassVar[float] = 100.0
+    """Default threshold in milliseconds for slow query detection."""
+
+    HASH_MODULUS: ClassVar[int] = 10000
+    """Modulus for deterministic hash-based sampling."""
+
+    def __init__(
+        self,
+        *,
+        sample_rate: float = 1.0,
+        force_sample_on_error: bool = True,
+        force_sample_slow_queries_ms: float | None = 100.0,
+        deterministic: bool = True,
+    ) -> None:
+        """Initialize sampling configuration.
+
+        Args:
+            sample_rate: Probability of sampling (0.0 to 1.0). Values outside
+                this range are clamped. Defaults to 1.0 (always sample).
+            force_sample_on_error: If True, always sample when an error occurs.
+                Defaults to True.
+            force_sample_slow_queries_ms: If set, always sample queries that take
+                longer than this threshold in milliseconds. Defaults to 100.0ms.
+                Set to None to disable.
+            deterministic: If True, use hash-based sampling that produces consistent
+                results for the same correlation ID across distributed systems.
+                If False, use random sampling. Defaults to True.
+        """
+        self.sample_rate = max(0.0, min(1.0, sample_rate))
+        self.force_sample_on_error = force_sample_on_error
+        self.force_sample_slow_queries_ms = force_sample_slow_queries_ms
+        self.deterministic = deterministic
+
+    def should_sample(
+        self,
+        correlation_id: str | None = None,
+        *,
+        is_error: bool = False,
+        duration_ms: float | None = None,
+        force: bool = False,
+    ) -> bool:
+        """Determine if this request should be sampled.
+
+        Evaluates force-sample conditions first (errors, slow queries, explicit force),
+        then falls back to rate-based sampling.
+
+        Args:
+            correlation_id: The correlation ID for deterministic sampling.
+                If None and deterministic=True, falls back to random sampling.
+            is_error: Whether this request resulted in an error.
+            duration_ms: Query duration in milliseconds, if known.
+            force: Explicit force-sample flag from application code.
+
+        Returns:
+            True if the request should be sampled, False otherwise.
+
+        Example:
+            ```python
+            # Always sampled due to error
+            config.should_sample(is_error=True)  # True
+
+            # Always sampled due to slow query (>100ms default)
+            config.should_sample(duration_ms=150.0)  # True
+
+            # Rate-based sampling
+            config.should_sample(
+                correlation_id="abc-123"
+            )  # depends on rate
+            ```
+        """
+        # Force sample conditions take precedence
+        if force:
+            return True
+
+        if is_error and self.force_sample_on_error:
+            return True
+
+        if (
+            duration_ms is not None
+            and self.force_sample_slow_queries_ms is not None
+            and duration_ms >= self.force_sample_slow_queries_ms
+        ):
+            return True
+
+        # Rate-based sampling
+        if self.sample_rate >= 1.0:
+            return True
+
+        if self.sample_rate <= 0.0:
+            return False
+
+        # Deterministic or random sampling
+        if self.deterministic and correlation_id:
+            # Hash-based sampling for consistency across distributed systems
+            hash_value = hash(correlation_id) % self.HASH_MODULUS
+            threshold = int(self.sample_rate * self.HASH_MODULUS)
+            return hash_value < threshold
+
+        # Fall back to random sampling
+        return random.random() < self.sample_rate  # noqa: S311
+
+    def copy(self) -> "SamplingConfig":
+        """Return a copy of the sampling configuration.
+
+        Returns:
+            A new SamplingConfig instance with the same values.
+        """
+        return SamplingConfig(
+            sample_rate=self.sample_rate,
+            force_sample_on_error=self.force_sample_on_error,
+            force_sample_slow_queries_ms=self.force_sample_slow_queries_ms,
+            deterministic=self.deterministic,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"SamplingConfig(sample_rate={self.sample_rate!r}, "
+            f"force_sample_on_error={self.force_sample_on_error!r}, "
+            f"force_sample_slow_queries_ms={self.force_sample_slow_queries_ms!r}, "
+            f"deterministic={self.deterministic!r})"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SamplingConfig):
+            return NotImplemented
+        return (
+            self.sample_rate == other.sample_rate
+            and self.force_sample_on_error == other.force_sample_on_error
+            and self.force_sample_slow_queries_ms == other.force_sample_slow_queries_ms
+            and self.deterministic == other.deterministic
+        )
+
+    def __hash__(self) -> int:
+        # SamplingConfig is mutable, so it should not be hashable
+        msg = "SamplingConfig objects are mutable and unhashable"
+        raise TypeError(msg)