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

* Fixed celery headers for beat auto instrumentation --------- Co-authored-by: Ivana Kellyerova <ivana.kellyerova@sentry.io>

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)