Remove custom_sampling_context (#3747)

- Add support for the sampled flag for start_span and respect it when making sampling decisions. - Rework sampling_context in traces_sampler to work with span attributes instead. Make sure we still have the same data accessible as now. We could go one step further and change the format of sampling_context to just be the actual span attributes without any postprocessing into the current format. I kept the format in line with what we have now to make it easier to update. See #3746 Closes #3739 This is a breaking change since we're removing custom_sampling_context. It'll break multiple integrations until we fix them (see #3746).
getsentry · Nov 12, 2024 · afd8cc5 · afd8cc5
1 parent 3f638f7
commit afd8cc5
Show file tree

Hide file tree

Showing 8 changed files with 42 additions and 36 deletions.
diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md
@@ -19,11 +19,13 @@ Looking to upgrade from Sentry SDK 2.x to 3.x? Here's a comprehensive list of wh
 - Redis integration: In Redis pipeline spans there is no `span["data"]["redis.commands"]` that contains a dict `{"count": 3, "first_ten": ["cmd1", "cmd2", ...]}` but instead `span["data"]["redis.commands.count"]` (containing `3`) and `span["data"]["redis.commands.first_ten"]` (containing `["cmd1", "cmd2", ...]`).
 - clickhouse-driver integration: The query is now available under the `db.query.text` span attribute (only if `send_default_pii` is `True`).
 - `sentry_sdk.init` now returns `None` instead of a context manager.
+- The `sampling_context` argument of `traces_sampler` now additionally contains all span attributes known at span start.
 
 ### Removed
 
 - Spans no longer have a `description`. Use `name` instead.
 - Dropped support for Python 3.6.
+- The `custom_sampling_context` parameter of `start_transaction` has been removed. Use `attributes` instead to set key-value pairs of data that should be accessible in the traces sampler. Note that span attributes need to conform to the [OpenTelemetry specification](https://opentelemetry.io/docs/concepts/signals/traces/#attributes), meaning only certain types can be set as values.
 - The PyMongo integration no longer sets tags. The data is still accessible via span attributes.
 - The PyMongo integration doesn't set `operation_ids` anymore. The individual IDs (`operation_id`, `request_id`, `session_id`) are now accessible as separate span attributes.
 - `sentry_sdk.metrics` and associated metrics APIs have been removed as Sentry no longer accepts metrics data in this form. See https://sentry.zendesk.com/hc/en-us/articles/26369339769883-Upcoming-API-Changes-to-Metrics

diff --git a/sentry_sdk/api.py b/sentry_sdk/api.py
@@ -40,7 +40,6 @@
         ExcInfo,
         MeasurementUnit,
         LogLevelStr,
-        SamplingContext,
     )
     from sentry_sdk.tracing import Span, TransactionKwargs
 
@@ -239,12 +238,8 @@ def flush(
     return get_client().flush(timeout=timeout, callback=callback)
 
 
-def start_span(
-    *,
-    custom_sampling_context=None,
-    **kwargs,  # type: Any
-):
-    # type: (...) -> POTelSpan
+def start_span(**kwargs):
+    # type: (type.Any) -> POTelSpan
     """
     Start and return a span.
 
@@ -257,13 +252,11 @@ def start_span(
     of the `with` block. If not using context managers, call the `finish()`
     method.
     """
-    # TODO: Consider adding type hints to the method signature.
-    return get_current_scope().start_span(custom_sampling_context, **kwargs)
+    return get_current_scope().start_span(**kwargs)
 
 
 def start_transaction(
     transaction=None,  # type: Optional[Transaction]
-    custom_sampling_context=None,  # type: Optional[SamplingContext]
     **kwargs,  # type: Unpack[TransactionKwargs]
 ):
     # type: (...) -> POTelSpan
@@ -295,14 +288,12 @@ def start_transaction(
 
     :param transaction: The transaction to start. If omitted, we create and
         start a new transaction.
-    :param custom_sampling_context: The transaction's custom sampling context.
     :param kwargs: Optional keyword arguments to be passed to the Transaction
         constructor. See :py:class:`sentry_sdk.tracing.Transaction` for
         available arguments.
     """
     return start_span(
         span=transaction,
-        custom_sampling_context=custom_sampling_context,
         **kwargs,
     )
 

diff --git a/sentry_sdk/integrations/opentelemetry/consts.py b/sentry_sdk/integrations/opentelemetry/consts.py
@@ -30,3 +30,4 @@ class SentrySpanAttribute:
     NAME = "sentry.name"
     SOURCE = "sentry.source"
     CONTEXT = "sentry.context"
+    CUSTOM_SAMPLED = "sentry.custom_sampled"  # used for saving start_span(sampled=X)
diff --git a/sentry_sdk/integrations/opentelemetry/sampler.py b/sentry_sdk/integrations/opentelemetry/sampler.py
@@ -2,7 +2,6 @@
 from typing import cast
 
 from opentelemetry import trace
-
 from opentelemetry.sdk.trace.sampling import Sampler, SamplingResult, Decision
 from opentelemetry.trace.span import TraceState
 
@@ -12,6 +11,7 @@
 from sentry_sdk.integrations.opentelemetry.consts import (
     TRACESTATE_SAMPLED_KEY,
     TRACESTATE_SAMPLE_RATE_KEY,
+    SentrySpanAttribute,
 )
 
 from typing import TYPE_CHECKING
@@ -114,28 +114,34 @@ def should_sample(
 
         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)
 
-        sample_rate = None
+        # Explicit sampled value provided at start_span
+        if attributes.get(SentrySpanAttribute.CUSTOM_SAMPLED) is not None:
+            sample_rate = float(attributes[SentrySpanAttribute.CUSTOM_SAMPLED])
+            if sample_rate > 0:
+                return sampled_result(parent_span_context, attributes, sample_rate)
+            else:
+                return dropped_result(parent_span_context, attributes)
 
-        # Check if sampled=True was passed to start_transaction
-        # TODO-anton: Do we want to keep the start_transaction(sampled=True) thing?
+        sample_rate = None
 
         # 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"))
         if has_traces_sampler:
-            # TODO-anton: Make proper sampling_context
-            # TODO-neel-potel: Make proper sampling_context
             sampling_context = {
                 "transaction_context": {
                     "name": name,
+                    "op": attributes.get(SentrySpanAttribute.OP),
                 },
                 "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
             }
-
+            sampling_context.update(attributes)
             sample_rate = client.options["traces_sampler"](sampling_context)
 
         else:

diff --git a/sentry_sdk/integrations/opentelemetry/scope.py b/sentry_sdk/integrations/opentelemetry/scope.py
@@ -24,7 +24,6 @@
     from typing import Tuple, Optional, Generator, Dict, Any
     from typing_extensions import Unpack
 
-    from sentry_sdk._types import SamplingContext
     from sentry_sdk.tracing import TransactionKwargs
 
 
@@ -112,17 +111,17 @@ def _incoming_otel_span_context(self):
 
         return span_context
 
-    def start_transaction(self, custom_sampling_context=None, **kwargs):
-        # type: (Optional[SamplingContext], Unpack[TransactionKwargs]) -> POTelSpan
+    def start_transaction(self, **kwargs):
+        # type: (Unpack[TransactionKwargs]) -> POTelSpan
         """
         .. deprecated:: 3.0.0
             This function is deprecated and will be removed in a future release.
             Use :py:meth:`sentry_sdk.start_span` instead.
         """
-        return self.start_span(custom_sampling_context=custom_sampling_context)
+        return self.start_span(**kwargs)
 
-    def start_span(self, custom_sampling_context=None, **kwargs):
-        # type: (Optional[SamplingContext], Any) -> POTelSpan
+    def start_span(self, **kwargs):
+        # type: (Any) -> POTelSpan
         return POTelSpan(**kwargs, scope=self)
 
 

diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
@@ -946,9 +946,7 @@ def add_breadcrumb(self, crumb=None, hint=None, **kwargs):
         while len(self._breadcrumbs) > max_breadcrumbs:
             self._breadcrumbs.popleft()
 
-    def start_transaction(
-        self, transaction=None, custom_sampling_context=None, **kwargs
-    ):
+    def start_transaction(self, transaction=None, **kwargs):
         # type: (Optional[Transaction], Optional[SamplingContext], Unpack[TransactionKwargs]) -> Union[Transaction, NoOpSpan]
         """
         Start and return a transaction.
@@ -974,7 +972,6 @@ def start_transaction(
 
         :param transaction: The transaction to start. If omitted, we create and
             start a new transaction.
-        :param custom_sampling_context: The transaction's custom sampling context.
         :param kwargs: Optional keyword arguments to be passed to the Transaction
             constructor. See :py:class:`sentry_sdk.tracing.Transaction` for
             available arguments.
@@ -985,8 +982,6 @@ def start_transaction(
 
         try_autostart_continuous_profiler()
 
-        custom_sampling_context = custom_sampling_context or {}
-
         # if we haven't been given a transaction, make one
         transaction = Transaction(**kwargs)
 
@@ -996,7 +991,6 @@ def start_transaction(
             "transaction_context": transaction.to_json(),
             "parent_sampled": transaction.parent_sampled,
         }
-        sampling_context.update(custom_sampling_context)
         transaction._set_initial_sampling_decision(sampling_context=sampling_context)
 
         if transaction.sampled:

diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -43,6 +43,8 @@
 
     from typing_extensions import TypedDict, Unpack
 
+    from opentelemetry.utils import types as OTelSpanAttributes
+
     P = ParamSpec("P")
     R = TypeVar("R")
 
@@ -1202,10 +1204,12 @@ def __init__(
         op=None,  # type: Optional[str]
         description=None,  # type: Optional[str]
         status=None,  # type: Optional[str]
+        sampled=None,  # type: Optional[bool]
         start_timestamp=None,  # type: Optional[Union[datetime, float]]
         origin=None,  # type: Optional[str]
         name=None,  # type: Optional[str]
         source=TRANSACTION_SOURCE_CUSTOM,  # type: str
+        attributes=None,  # type: OTelSpanAttributes
         only_if_parent=False,  # type: bool
         otel_span=None,  # type: Optional[OtelSpan]
         **_,  # type: dict[str, object]
@@ -1230,6 +1234,9 @@ def __init__(
             if skip_span:
                 self._otel_span = INVALID_SPAN
             else:
+                from sentry_sdk.integrations.opentelemetry.consts import (
+                    SentrySpanAttribute,
+                )
                 from sentry_sdk.integrations.opentelemetry.utils import (
                     convert_to_otel_timestamp,
                 )
@@ -1239,12 +1246,18 @@ def __init__(
                     start_timestamp = convert_to_otel_timestamp(start_timestamp)
 
                 span_name = name or description or op or ""
+
+                # Prepopulate some attrs so that they're accessible in traces_sampler
+                attributes = attributes or {}
+                attributes[SentrySpanAttribute.OP] = op
+                if sampled is not None:
+                    attributes[SentrySpanAttribute.CUSTOM_SAMPLED] = sampled
+
                 self._otel_span = tracer.start_span(
-                    span_name, start_time=start_timestamp
+                    span_name, start_time=start_timestamp, attributes=attributes
                 )
 
                 self.origin = origin or DEFAULT_SPAN_ORIGIN
-                self.op = op
                 self.description = description
                 self.name = span_name
                 self.source = source

diff --git a/tests/tracing/test_sampling.py b/tests/tracing/test_sampling.py
@@ -211,13 +211,13 @@ def test_passes_parent_sampling_decision_in_sampling_context(
     assert sampling_context["parent_sampled"]._mock_wraps is parent_sampling_decision
 
 
-def test_passes_custom_samling_context_from_start_transaction_to_traces_sampler(
+def test_passes_attributes_from_start_span_to_traces_sampler(
     sentry_init, DictionaryContaining  # noqa: N803
 ):
     traces_sampler = mock.Mock()
     sentry_init(traces_sampler=traces_sampler)
 
-    start_transaction(custom_sampling_context={"dogs": "yes", "cats": "maybe"})
+    start_transaction(attributes={"dogs": "yes", "cats": "maybe"})
 
     traces_sampler.assert_any_call(
         DictionaryContaining({"dogs": "yes", "cats": "maybe"})