Fixed Celery headers for Beat auto-instrumentation (#2102)

* Fixed celery headers for beat auto instrumentation

---------

Co-authored-by: Ivana Kellyerova <[email protected]>

Author: antonpirker

sentry_sdk/integrations/celery.py

@@ -157,6 +157,13 @@ def apply_async(*args, **kwargs):
                         # tracing tools (dd-trace-py) also employ this exact
                         # workaround and we don't want to break them.
                         kwarg_headers.setdefault("headers", {}).update(headers)
+
+                        # Add the Sentry options potentially added in `sentry_apply_entry`
+                        # to the headers (done when auto-instrumenting Celery Beat tasks)
+                        for key, value in kwarg_headers.items():
+                            if key.startswith("sentry-"):
+                                kwarg_headers["headers"][key] = value
+
                         kwargs["headers"] = kwarg_headers
 
                 return f(*args, **kwargs)
@@ -431,7 +438,9 @@ def sentry_apply_entry(*args, **kwargs):
         )
         headers.update({"sentry-monitor-check-in-id": check_in_id})
 
-        schedule_entry.options.update(headers)
+        # Set the Sentry configuration in the options of the ScheduleEntry.
+        # Those will be picked up in `apply_async` and added to the headers.
+        schedule_entry.options["headers"] = headers
         return original_apply_entry(*args, **kwargs)
 
     Scheduler.apply_entry = sentry_apply_entry

tests/integrations/celery/test_celery.py

@@ -5,11 +5,13 @@
 pytest.importorskip("celery")
 
 from sentry_sdk import Hub, configure_scope, start_transaction
-from sentry_sdk.integrations.celery import CeleryIntegration
+from sentry_sdk.integrations.celery import CeleryIntegration, _get_headers
+
 from sentry_sdk._compat import text_type
 
 from celery import Celery, VERSION
 from celery.bin import worker
+from celery.signals import task_success
 
 try:
     from unittest import mock  # python 3.3 and above
@@ -437,3 +439,29 @@ def dummy_task(x, y):
         celery_invocation(dummy_task, 1, 0)
 
     assert not events
+
+
+def test_task_headers(celery):
+    """
+    Test that the headers set in the Celery Beat auto-instrumentation are passed to the celery signal handlers
+    """
+    sentry_crons_setup = {
+        "sentry-monitor-slug": "some-slug",
+        "sentry-monitor-config": {"some": "config"},
+        "sentry-monitor-check-in-id": "123abc",
+    }
+
+    @celery.task(name="dummy_task")
+    def dummy_task(x, y):
+        return x + y
+
+    def crons_task_success(sender, **kwargs):
+        headers = _get_headers(sender)
+        assert headers == sentry_crons_setup
+
+    task_success.connect(crons_task_success)
+
+    # This is how the Celery Beat auto-instrumentation starts a task
+    # in the monkey patched version of `apply_async`
+    # in `sentry_sdk/integrations/celery.py::_wrap_apply_async()`
+    dummy_task.apply_async(args=(1, 0), headers=sentry_crons_setup)