Turn custom_sampling_context into span attrs in Celery integration (#3813)

Author: sentrivana

MIGRATION_GUIDE.md

@@ -21,6 +21,16 @@ Looking to upgrade from Sentry SDK 2.x to 3.x? Here's a comprehensive list of wh
 - 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.
+- If you're using the Celery integration, the `sampling_context` argument of `traces_sampler` doesn't contain the `celery_job` dictionary anymore. Instead, the individual keys are now available as:
+
+  | Dictionary keys        | Sampling context key |
+  | ---------------------- | -------------------- |
+  | `celery_job["args"]`   | `celery.job.args`    |
+  | `celery_job["kwargs"]` | `celery.job.kwargs`  |
+  | `celery_job["task"]`   | `celery.job.task`    |
+
+  Note that all of these are serialized, i.e., not the original `args` and `kwargs` but rather OpenTelemetry-friendly span attributes.
+
 - If you're using the AIOHTTP integration, the `sampling_context` argument of `traces_sampler` doesn't contain the `aiohttp_request` object anymore. Instead, some of the individual properties of the request are accessible, if available, as follows:
 
   | Request property | Sampling context key(s)         |
@@ -71,15 +81,15 @@ Looking to upgrade from Sentry SDK 2.x to 3.x? Here's a comprehensive list of wh
   | `client`       | `client.address`, `client.port` |
   | full URL       | `url.full`                      |
 
-- If you're using the RQ integration, the `sampling_context` argument of `traces_sampler` doesn't contain the `rq_job` object anymore. Instead, the individual properties of the scope, if available, are accessible as follows:
+- If you're using the RQ integration, the `sampling_context` argument of `traces_sampler` doesn't contain the `rq_job` object anymore. Instead, the individual properties of the job and the queue, if available, are accessible as follows:
 
   | RQ property     | Sampling context key(s)      |
   | --------------- | ---------------------------- |
   | `rq_job.args`   | `rq.job.args`                |
   | `rq_job.kwargs` | `rq.job.kwargs`              |
   | `rq_job.func`   | `rq.job.func`                |
   | `queue.name`    | `messaging.destination.name` |
-  | `job.id`        | `messaging.message.id`       |
+  | `rq_job.id`     | `messaging.message.id`       |
 
   Note that `rq.job.args`, `rq.job.kwargs`, and `rq.job.func` are serialized and not the actual objects on the job.
 

sentry_sdk/integrations/celery/__init__.py

@@ -20,6 +20,7 @@
     ensure_integration_enabled,
     event_from_exception,
     reraise,
+    _serialize_span_attribute,
 )
 
 from typing import TYPE_CHECKING
@@ -318,15 +319,9 @@ def _inner(*args, **kwargs):
                     name=task.name,
                     source=TRANSACTION_SOURCE_TASK,
                     origin=CeleryIntegration.origin,
-                    custom_sampling_context={
-                        "celery_job": {
-                            "task": task.name,
-                            # for some reason, args[1] is a list if non-empty but a
-                            # tuple if empty
-                            "args": list(args[1]),
-                            "kwargs": args[2],
-                        }
-                    },
+                    # for some reason, args[1] is a list if non-empty but a
+                    # tuple if empty
+                    attributes=_prepopulate_attributes(task, list(args[1]), args[2]),
                 ) as transaction:
                     transaction.set_status(SPANSTATUS.OK)
                     return f(*args, **kwargs)
@@ -516,3 +511,12 @@ def sentry_publish(self, *args, **kwargs):
             return original_publish(self, *args, **kwargs)
 
     Producer.publish = sentry_publish
+
+
+def _prepopulate_attributes(task, args, kwargs):
+    attributes = {
+        "celery.job.task": task.name,
+        "celery.job.args": _serialize_span_attribute(args),
+        "celery.job.kwargs": _serialize_span_attribute(kwargs),
+    }
+    return attributes

tests/integrations/celery/test_celery.py

@@ -13,6 +13,7 @@
     _wrap_task_run,
 )
 from sentry_sdk.integrations.celery.beat import _get_headers
+from sentry_sdk.utils import _serialize_span_attribute
 from tests.conftest import ApproxDict
 
 
@@ -430,7 +431,7 @@ def dummy_task(self, x, y):
 
 
 def test_traces_sampler_gets_task_info_in_sampling_context(
-    init_celery, celery_invocation, DictionaryContaining  # noqa:N803
+    init_celery, celery_invocation
 ):
     traces_sampler = mock.Mock()
     celery = init_celery(traces_sampler=traces_sampler)
@@ -445,10 +446,13 @@ def walk_dogs(x, y):
         walk_dogs, [["Maisey", "Charlie", "Bodhi", "Cory"], "Dog park round trip"], 1
     )
 
-    traces_sampler.assert_any_call(
-        # depending on the iteration of celery_invocation, the data might be
-        # passed as args or as kwargs, so make this generic
-        DictionaryContaining({"celery_job": dict(task="dog_walk", **args_kwargs)})
+    sampling_context = traces_sampler.call_args_list[1][0][0]
+    assert sampling_context["celery.job.task"] == "dog_walk"
+    assert sampling_context["celery.job.args"] == _serialize_span_attribute(
+        args_kwargs["args"]
+    )
+    assert sampling_context["celery.job.kwargs"] == _serialize_span_attribute(
+        args_kwargs["kwargs"]
     )