sentry-sdk 2.26.1__py2.py3-none-any.whl → 3.0.0a1__py2.py3-none-any.whl

sentry_sdk/opentelemetry/sampler.py

@@ -0,0 +1,326 @@
+from decimal import Decimal
+from typing import cast
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace.sampling import Sampler, SamplingResult, Decision
+from opentelemetry.trace.span import TraceState
+
+import sentry_sdk
+from sentry_sdk.opentelemetry.consts import (
+    TRACESTATE_SAMPLED_KEY,
+    TRACESTATE_SAMPLE_RAND_KEY,
+    TRACESTATE_SAMPLE_RATE_KEY,
+    SentrySpanAttribute,
+)
+from sentry_sdk.tracing_utils import (
+    _generate_sample_rand,
+    has_tracing_enabled,
+)
+from sentry_sdk.utils import is_valid_sample_rate, logger
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional, Sequence, Union
+    from opentelemetry.context import Context
+    from opentelemetry.trace import Link, SpanKind
+    from opentelemetry.trace.span import SpanContext
+    from opentelemetry.util.types import Attributes
+
+
+def get_parent_sampled(parent_context, trace_id):
+    # type: (Optional[SpanContext], int) -> Optional[bool]
+    if parent_context is None:
+        return None
+
+    is_span_context_valid = parent_context is not None and parent_context.is_valid
+
+    # Only inherit sample rate if `traceId` is the same
+    if is_span_context_valid and parent_context.trace_id == trace_id:
+        # this is getSamplingDecision in JS
+        # if there was no sampling flag, defer the decision
+        dsc_sampled = parent_context.trace_state.get(TRACESTATE_SAMPLED_KEY)
+        if dsc_sampled == "deferred":
+            return None
+
+        if parent_context.trace_flags.sampled is not None:
+            return parent_context.trace_flags.sampled
+
+        if dsc_sampled == "true":
+            return True
+        elif dsc_sampled == "false":
+            return False
+
+    return None
+
+
+def get_parent_sample_rate(parent_context, trace_id):
+    # type: (Optional[SpanContext], int) -> Optional[float]
+    if parent_context is None:
+        return None
+
+    is_span_context_valid = parent_context is not None and parent_context.is_valid
+
+    if is_span_context_valid and parent_context.trace_id == trace_id:
+        parent_sample_rate = parent_context.trace_state.get(TRACESTATE_SAMPLE_RATE_KEY)
+        if parent_sample_rate is None:
+            return None
+
+        try:
+            return float(parent_sample_rate)
+        except Exception:
+            return None
+
+    return None
+
+
+def get_parent_sample_rand(parent_context, trace_id):
+    # type: (Optional[SpanContext], int) -> Optional[Decimal]
+    if parent_context is None:
+        return None
+
+    is_span_context_valid = parent_context is not None and parent_context.is_valid
+
+    if is_span_context_valid and parent_context.trace_id == trace_id:
+        parent_sample_rand = parent_context.trace_state.get(TRACESTATE_SAMPLE_RAND_KEY)
+        if parent_sample_rand is None:
+            return None
+
+        return Decimal(parent_sample_rand)
+
+    return None
+
+
+def dropped_result(span_context, attributes, sample_rate=None, sample_rand=None):
+    # type: (SpanContext, Attributes, Optional[float], Optional[Decimal]) -> SamplingResult
+    """
+    React to a span getting unsampled and return a DROP SamplingResult.
+
+    Update the trace_state with the effective sampled, sample_rate and sample_rand,
+    record that we dropped the event for client report purposes, and return
+    an OTel SamplingResult with Decision.DROP.
+
+    See for more info about OTel sampling:
+    https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.sampling.html
+    """
+    trace_state = _update_trace_state(
+        span_context, sampled=False, sample_rate=sample_rate, sample_rand=sample_rand
+    )
+
+    is_root_span = not (span_context.is_valid and not span_context.is_remote)
+    if is_root_span:
+        # Tell Sentry why we dropped the transaction/root-span
+        client = sentry_sdk.get_client()
+        if client.monitor and client.monitor.downsample_factor > 0:
+            reason = "backpressure"
+        else:
+            reason = "sample_rate"
+
+        if client.transport and has_tracing_enabled(client.options):
+            client.transport.record_lost_event(reason, data_category="transaction")
+
+            # Only one span (the transaction itself) is discarded, since we did not record any spans here.
+            client.transport.record_lost_event(reason, data_category="span")
+
+    return SamplingResult(
+        Decision.DROP,
+        attributes=attributes,
+        trace_state=trace_state,
+    )
+
+
+def sampled_result(span_context, attributes, sample_rate=None, sample_rand=None):
+    # type: (SpanContext, Attributes, Optional[float], Optional[Decimal]) -> SamplingResult
+    """
+    React to a span being sampled and return a sampled SamplingResult.
+
+    Update the trace_state with the effective sampled, sample_rate and sample_rand,
+    and return an OTel SamplingResult with Decision.RECORD_AND_SAMPLE.
+
+    See for more info about OTel sampling:
+    https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.sampling.html
+    """
+    trace_state = _update_trace_state(
+        span_context, sampled=True, sample_rate=sample_rate, sample_rand=sample_rand
+    )
+
+    return SamplingResult(
+        Decision.RECORD_AND_SAMPLE,
+        attributes=attributes,
+        trace_state=trace_state,
+    )
+
+
+def _update_trace_state(span_context, sampled, sample_rate=None, sample_rand=None):
+    # type: (SpanContext, bool, Optional[float], Optional[Decimal]) -> TraceState
+    trace_state = span_context.trace_state
+
+    sampled = "true" if sampled else "false"
+    if TRACESTATE_SAMPLED_KEY not in trace_state:
+        trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, sampled)
+    elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
+        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, sampled)
+
+    if sample_rate is not None:
+        trace_state = trace_state.update(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
+
+    if sample_rand is not None:
+        trace_state = trace_state.update(
+            TRACESTATE_SAMPLE_RAND_KEY, f"{sample_rand:.6f}"  # noqa: E231
+        )
+
+    return trace_state
+
+
+class SentrySampler(Sampler):
+    def should_sample(
+        self,
+        parent_context,  # type: Optional[Context]
+        trace_id,  # type: int
+        name,  # type: str
+        kind=None,  # type: Optional[SpanKind]
+        attributes=None,  # type: Attributes
+        links=None,  # type: Optional[Sequence[Link]]
+        trace_state=None,  # type: Optional[TraceState]
+    ):
+        # type: (...) -> SamplingResult
+        client = sentry_sdk.get_client()
+
+        parent_span_context = trace.get_current_span(parent_context).get_span_context()
+
+        attributes = attributes or {}
+
+        # No tracing enabled, thus no sampling
+        if not has_tracing_enabled(client.options):
+            return dropped_result(parent_span_context, attributes)
+
+        # parent_span_context.is_valid means this span has a parent, remote or local
+        is_root_span = not parent_span_context.is_valid or parent_span_context.is_remote
+
+        sample_rate = None
+
+        parent_sampled = get_parent_sampled(parent_span_context, trace_id)
+        parent_sample_rate = get_parent_sample_rate(parent_span_context, trace_id)
+        parent_sample_rand = get_parent_sample_rand(parent_span_context, trace_id)
+
+        if parent_sample_rand is not None:
+            # We have a sample_rand on the incoming trace or we already backfilled
+            # it in PropagationContext
+            sample_rand = parent_sample_rand
+        else:
+            # We are the head SDK and we need to generate a new sample_rand
+            sample_rand = cast(Decimal, _generate_sample_rand(str(trace_id), (0, 1)))
+
+        # Explicit sampled value provided at start_span
+        custom_sampled = cast(
+            "Optional[bool]", attributes.get(SentrySpanAttribute.CUSTOM_SAMPLED)
+        )
+        if custom_sampled is not None:
+            if is_root_span:
+                sample_rate = float(custom_sampled)
+                if sample_rate > 0:
+                    return sampled_result(
+                        parent_span_context,
+                        attributes,
+                        sample_rate=sample_rate,
+                        sample_rand=sample_rand,
+                    )
+                else:
+                    return dropped_result(
+                        parent_span_context,
+                        attributes,
+                        sample_rate=sample_rate,
+                        sample_rand=sample_rand,
+                    )
+            else:
+                logger.debug(
+                    f"[Tracing.Sampler] Ignoring sampled param for non-root span {name}"
+                )
+
+        # Check if there is a traces_sampler
+        # Traces_sampler is responsible to check parent sampled to have full transactions.
+        has_traces_sampler = callable(client.options.get("traces_sampler"))
+
+        sample_rate_to_propagate = None
+
+        if is_root_span and has_traces_sampler:
+            sampling_context = create_sampling_context(
+                name, attributes, parent_span_context, trace_id
+            )
+            sample_rate = client.options["traces_sampler"](sampling_context)
+            sample_rate_to_propagate = sample_rate
+        else:
+            # Check if there is a parent with a sampling decision
+            if parent_sampled is not None:
+                sample_rate = bool(parent_sampled)
+                sample_rate_to_propagate = (
+                    parent_sample_rate if parent_sample_rate else sample_rate
+                )
+            else:
+                # Check if there is a traces_sample_rate
+                sample_rate = client.options.get("traces_sample_rate")
+                sample_rate_to_propagate = sample_rate
+
+        # If the sample rate is invalid, drop the span
+        if not is_valid_sample_rate(sample_rate, source=self.__class__.__name__):
+            logger.warning(
+                f"[Tracing.Sampler] Discarding {name} because of invalid sample rate."
+            )
+            return dropped_result(parent_span_context, attributes)
+
+        # Down-sample in case of back pressure monitor says so
+        if is_root_span and client.monitor:
+            sample_rate /= 2**client.monitor.downsample_factor
+            if client.monitor.downsample_factor > 0:
+                sample_rate_to_propagate = sample_rate
+
+        # Compare sample_rand to sample_rate to make the final sampling decision
+        sample_rate = float(cast("Union[bool, float, int]", sample_rate))
+        sampled = sample_rand < Decimal.from_float(sample_rate)
+
+        if sampled:
+            if is_root_span:
+                logger.debug(
+                    f"[Tracing.Sampler] Sampled #{name} with sample_rate: {sample_rate} and sample_rand: {sample_rand}"
+                )
+
+            return sampled_result(
+                parent_span_context,
+                attributes,
+                sample_rate=sample_rate_to_propagate,
+                sample_rand=None if sample_rand == parent_sample_rand else sample_rand,
+            )
+        else:
+            if is_root_span:
+                logger.debug(
+                    f"[Tracing.Sampler] Dropped #{name} with sample_rate: {sample_rate} and sample_rand: {sample_rand}"
+                )
+
+            return dropped_result(
+                parent_span_context,
+                attributes,
+                sample_rate=sample_rate_to_propagate,
+                sample_rand=None if sample_rand == parent_sample_rand else sample_rand,
+            )
+
+    def get_description(self) -> str:
+        return self.__class__.__name__
+
+
+def create_sampling_context(name, attributes, parent_span_context, trace_id):
+    # type: (str, Attributes, Optional[SpanContext], int) -> dict[str, Any]
+    sampling_context = {
+        "transaction_context": {
+            "name": name,
+            "op": attributes.get(SentrySpanAttribute.OP) if attributes else None,
+            "source": (
+                attributes.get(SentrySpanAttribute.SOURCE) if attributes else None
+            ),
+        },
+        "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
+    }  # type: dict[str, Any]
+
+    if attributes is not None:
+        sampling_context.update(attributes)
+
+    return sampling_context

sentry_sdk/opentelemetry/scope.py

@@ -0,0 +1,218 @@
+from typing import cast
+from contextlib import contextmanager
+import warnings
+
+from opentelemetry.context import (
+    get_value,
+    set_value,
+    attach,
+    detach,
+    get_current,
+)
+from opentelemetry.trace import (
+    SpanContext,
+    NonRecordingSpan,
+    TraceFlags,
+    TraceState,
+    use_span,
+)
+
+from sentry_sdk.opentelemetry.consts import (
+    SENTRY_SCOPES_KEY,
+    SENTRY_FORK_ISOLATION_SCOPE_KEY,
+    SENTRY_USE_CURRENT_SCOPE_KEY,
+    SENTRY_USE_ISOLATION_SCOPE_KEY,
+    TRACESTATE_SAMPLED_KEY,
+)
+from sentry_sdk.opentelemetry.contextvars_context import (
+    SentryContextVarsRuntimeContext,
+)
+from sentry_sdk.opentelemetry.utils import trace_state_from_baggage
+from sentry_sdk.scope import Scope, ScopeType
+from sentry_sdk.tracing import Span
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Tuple, Optional, Generator, Dict, Any
+
+
+class PotelScope(Scope):
+    @classmethod
+    def _get_scopes(cls):
+        # type: () -> Optional[Tuple[PotelScope, PotelScope]]
+        """
+        Returns the current scopes tuple on the otel context. Internal use only.
+        """
+        return cast(
+            "Optional[Tuple[PotelScope, PotelScope]]", get_value(SENTRY_SCOPES_KEY)
+        )
+
+    @classmethod
+    def get_current_scope(cls):
+        # type: () -> PotelScope
+        """
+        Returns the current scope.
+        """
+        return cls._get_current_scope() or _INITIAL_CURRENT_SCOPE
+
+    @classmethod
+    def _get_current_scope(cls):
+        # type: () -> Optional[PotelScope]
+        """
+        Returns the current scope without creating a new one. Internal use only.
+        """
+        scopes = cls._get_scopes()
+        return scopes[0] if scopes else None
+
+    @classmethod
+    def get_isolation_scope(cls):
+        # type: () -> PotelScope
+        """
+        Returns the isolation scope.
+        """
+        return cls._get_isolation_scope() or _INITIAL_ISOLATION_SCOPE
+
+    @classmethod
+    def _get_isolation_scope(cls):
+        # type: () -> Optional[PotelScope]
+        """
+        Returns the isolation scope without creating a new one. Internal use only.
+        """
+        scopes = cls._get_scopes()
+        return scopes[1] if scopes else None
+
+    @contextmanager
+    def continue_trace(self, environ_or_headers):
+        # type: (Dict[str, Any]) -> Generator[None, None, None]
+        """
+        Sets the propagation context from environment or headers to continue an incoming trace.
+        Any span started within this context manager will use the same trace_id, parent_span_id
+        and inherit the sampling decision from the incoming trace.
+        """
+        self.generate_propagation_context(environ_or_headers)
+
+        span_context = self._incoming_otel_span_context()
+        if span_context is None:
+            yield
+        else:
+            with use_span(NonRecordingSpan(span_context)):
+                yield
+
+    def _incoming_otel_span_context(self):
+        # type: () -> Optional[SpanContext]
+        if self._propagation_context is None:
+            return None
+        # If sentry-trace extraction didn't have a parent_span_id, we don't have an upstream header
+        if self._propagation_context.parent_span_id is None:
+            return None
+
+        trace_flags = TraceFlags(
+            TraceFlags.SAMPLED
+            if self._propagation_context.parent_sampled
+            else TraceFlags.DEFAULT
+        )
+
+        if self._propagation_context.baggage:
+            trace_state = trace_state_from_baggage(self._propagation_context.baggage)
+        else:
+            trace_state = TraceState()
+
+        # for twp to work, we also need to consider deferred sampling when the sampling
+        # flag is not present, so the above TraceFlags are not sufficient
+        if self._propagation_context.parent_sampled is None:
+            trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "deferred")
+
+        span_context = SpanContext(
+            trace_id=int(self._propagation_context.trace_id, 16),
+            span_id=int(self._propagation_context.parent_span_id, 16),
+            is_remote=True,
+            trace_flags=trace_flags,
+            trace_state=trace_state,
+        )
+
+        return span_context
+
+    def start_transaction(self, **kwargs):
+        # type: (Any) -> Span
+        """
+        .. deprecated:: 3.0.0
+            This function is deprecated and will be removed in a future release.
+            Use :py:meth:`sentry_sdk.start_span` instead.
+        """
+        warnings.warn(
+            "The `start_transaction` method is deprecated, please use `sentry_sdk.start_span instead.`",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.start_span(**kwargs)
+
+    def start_span(self, **kwargs):
+        # type: (Any) -> Span
+        return Span(**kwargs)
+
+
+_INITIAL_CURRENT_SCOPE = PotelScope(ty=ScopeType.CURRENT)
+_INITIAL_ISOLATION_SCOPE = PotelScope(ty=ScopeType.ISOLATION)
+
+
+def setup_initial_scopes():
+    # type: () -> None
+    global _INITIAL_CURRENT_SCOPE, _INITIAL_ISOLATION_SCOPE
+    _INITIAL_CURRENT_SCOPE = PotelScope(ty=ScopeType.CURRENT)
+    _INITIAL_ISOLATION_SCOPE = PotelScope(ty=ScopeType.ISOLATION)
+
+    scopes = (_INITIAL_CURRENT_SCOPE, _INITIAL_ISOLATION_SCOPE)
+    attach(set_value(SENTRY_SCOPES_KEY, scopes))
+
+
+def setup_scope_context_management():
+    # type: () -> None
+    import opentelemetry.context
+
+    opentelemetry.context._RUNTIME_CONTEXT = SentryContextVarsRuntimeContext()
+    setup_initial_scopes()
+
+
+@contextmanager
+def isolation_scope():
+    # type: () -> Generator[PotelScope, None, None]
+    context = set_value(SENTRY_FORK_ISOLATION_SCOPE_KEY, True)
+    token = attach(context)
+    try:
+        yield PotelScope.get_isolation_scope()
+    finally:
+        detach(token)
+
+
+@contextmanager
+def new_scope():
+    # type: () -> Generator[PotelScope, None, None]
+    token = attach(get_current())
+    try:
+        yield PotelScope.get_current_scope()
+    finally:
+        detach(token)
+
+
+@contextmanager
+def use_scope(scope):
+    # type: (PotelScope) -> Generator[PotelScope, None, None]
+    context = set_value(SENTRY_USE_CURRENT_SCOPE_KEY, scope)
+    token = attach(context)
+
+    try:
+        yield scope
+    finally:
+        detach(token)
+
+
+@contextmanager
+def use_isolation_scope(isolation_scope):
+    # type: (PotelScope) -> Generator[PotelScope, None, None]
+    context = set_value(SENTRY_USE_ISOLATION_SCOPE_KEY, isolation_scope)
+    token = attach(context)
+
+    try:
+        yield isolation_scope
+    finally:
+        detach(token)