Source code for sqlspec.observability._sampling

"""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",)



[docs]
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.
    """

    __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."""


[docs]
    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