ref: Add support for custom sampling context to span first

Author: sentrivana

sentry_sdk/scope.py

@@ -696,6 +696,13 @@ def get_active_propagation_context(self) -> "PropagationContext":
             isolation_scope._propagation_context = PropagationContext()
         return isolation_scope._propagation_context
 
+    def set_custom_sampling_context(
+        self, custom_sampling_context: "dict[str, Any]"
+    ) -> None:
+        self.get_active_propagation_context()._set_custom_sampling_context(
+            custom_sampling_context
+        )
+
     def clear(self) -> None:
         """Clears the entire scope."""
         self._level: "Optional[LogLevelStr]" = None

sentry_sdk/tracing_utils.py

@@ -417,6 +417,7 @@ class PropagationContext:
         "parent_span_id",
         "parent_sampled",
         "baggage",
+        "custom_sampling_context",
     )
 
     def __init__(
@@ -450,6 +451,8 @@ def __init__(
         if baggage is None and dynamic_sampling_context is not None:
             self.baggage = Baggage(dynamic_sampling_context)
 
+        self.custom_sampling_context: "Optional[dict[str, Any]]" = None
+
     @classmethod
     def from_incoming_data(
         cls, incoming_data: "Dict[str, Any]"
@@ -537,6 +540,11 @@ def update(self, other_dict: "Dict[str, Any]") -> None:
             except AttributeError:
                 pass
 
+    def _set_custom_sampling_context(
+        self, custom_sampling_context: "dict[str, Any]"
+    ) -> None:
+        self.custom_sampling_context = custom_sampling_context
+
     def __repr__(self) -> str:
         return "<PropagationContext _trace_id={} _span_id={} parent_span_id={} parent_sampled={} baggage={}>".format(
             self._trace_id,
@@ -1420,6 +1428,9 @@ def _make_sampling_decision(
             "attributes": dict(attributes) if attributes else {},
         }
 
+        if propagation_context.custom_sampling_context:
+            sampling_context.update(propagation_context.custom_sampling_context)
+
         sample_rate = client.options["traces_sampler"](sampling_context)
     else:
         if propagation_context.parent_sampled is not None:

tests/tracing/test_span_streaming.py

@@ -233,6 +233,62 @@ def traces_sampler(sampling_context):
     assert len(spans) == 1
 
 
+def test_custom_sampling_context(sentry_init):
+    class MyClass: ...
+
+    my_class = MyClass()
+
+    def traces_sampler(sampling_context):
+        assert "class" in sampling_context
+        assert "string" in sampling_context
+        assert sampling_context["class"] == my_class
+        assert sampling_context["string"] == "my string"
+        return 1.0
+
+    sentry_init(
+        traces_sampler=traces_sampler,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    sentry_sdk.get_current_scope().set_custom_sampling_context(
+        {
+            "class": my_class,
+            "string": "my string",
+        }
+    )
+
+    with sentry_sdk.traces.start_span(name="span"):
+        ...
+
+
+def test_new_custom_sampling_context(sentry_init):
+    def traces_sampler(sampling_context):
+        if sampling_context["attributes"]["first"] is True:
+            assert sampling_context["custom_value"] == 1
+        else:
+            assert sampling_context["custom_value"] == 2
+        return 1.0
+
+    sentry_init(
+        traces_sampler=traces_sampler,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    sentry_sdk.traces.new_trace()
+
+    sentry_sdk.get_current_scope().set_custom_sampling_context({"custom_value": 1})
+
+    with sentry_sdk.traces.start_span(name="span", attributes={"first": True}):
+        ...
+
+    sentry_sdk.traces.new_trace()
+
+    sentry_sdk.get_current_scope().set_custom_sampling_context({"custom_value": 2})
+
+    with sentry_sdk.traces.start_span(name="span", attributes={"first": False}):
+        ...
+
+
 def test_span_attributes(sentry_init, capture_envelopes):
     sentry_init(
         traces_sample_rate=1.0,