Auto monitoring beat update (#1989)

- Small update to support Celery 4 and 5
- Changed the name of the schedule shelf file that we patch to have the suffix `-patched-by-sentry-sdk` instead of `.new` so in case there is an error with this new shelf file somewhere the users know that it is patched by the sentry sdk.
- Additionally some minor tweaks to make code more readable

sentry_sdk/integrations/celery.py

@@ -3,6 +3,7 @@
 import sys
 import shutil
 import functools
+import tempfile
 
 from sentry_sdk.consts import OP
 from sentry_sdk._compat import reraise
@@ -320,6 +321,11 @@ def sentry_workloop(*args, **kwargs):
 def _get_headers(task):
  # type: (Task) -> Dict[str, Any]
  headers = task.request.get("headers") or {}
+
+ if "headers" in headers:
+ headers.update(headers["headers"])
+ del headers["headers"]
+
  return headers
 
 
@@ -392,9 +398,11 @@ def _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks):
  add_updated_periodic_task()
 
  # Start Celery Beat (with new (cloned) schedule, because old one is still in use)
- new_schedule_filename = sender.schedule_filename + ".new"
- shutil.copy2(sender.schedule_filename, new_schedule_filename)
- app.Beat(schedule=new_schedule_filename).run()
+ cloned_schedule = tempfile.NamedTemporaryFile(suffix="-patched-by-sentry-sdk")
+ with open(sender.schedule_filename, "rb") as original_schedule:
+ shutil.copyfileobj(original_schedule, cloned_schedule)
+
+ app.Beat(schedule=cloned_schedule.name).run()
 
 
 # Nested functions do not work as Celery hook receiver,
@@ -480,9 +488,7 @@ def crons_task_before_run(sender, **kwargs):
  if "sentry-monitor-slug" not in headers:
  return
 
- monitor_config = (
- headers["sentry-monitor-config"] if "sentry-monitor-config" in headers else {}
- )
+ monitor_config = headers.get("sentry-monitor-config", {})
 
  start_timestamp_s = now()
 
@@ -506,9 +512,7 @@ def crons_task_success(sender, **kwargs):
  if "sentry-monitor-slug" not in headers:
  return
 
- monitor_config = (
- headers["sentry-monitor-config"] if "sentry-monitor-config" in headers else {}
- )
+ monitor_config = headers.get("sentry-monitor-config", {})
 
  start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
 
@@ -529,9 +533,7 @@ def crons_task_failure(sender, **kwargs):
  if "sentry-monitor-slug" not in headers:
  return
 
- monitor_config = (
- headers["sentry-monitor-config"] if "sentry-monitor-config" in headers else {}
- )
+ monitor_config = headers.get("sentry-monitor-config", {})
 
  start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
 
@@ -552,9 +554,7 @@ def crons_task_retry(sender, **kwargs):
  if "sentry-monitor-slug" not in headers:
  return
 
- monitor_config = (
- headers["sentry-monitor-config"] if "sentry-monitor-config" in headers else {}
- )
+ monitor_config = headers.get("sentry-monitor-config", {})
 
  start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
 

tests/integrations/celery/test_celery_beat_crons.py

@@ -1,3 +1,4 @@
+import tempfile
 import mock
 
 import pytest
@@ -37,6 +38,20 @@ def test_get_headers():
 
  assert _get_headers(fake_task) == {"bla": "blub"}
 
+ fake_task.request.update(
+ {
+ "headers": {
+ "headers": {
+ "tri": "blub",
+ "bar": "baz",
+ },
+ "bla": "blub",
+ },
+ }
+ )
+
+ assert _get_headers(fake_task) == {"bla": "blub", "tri": "blub", "bar": "baz"}
+
 
 @pytest.mark.parametrize(
  "seconds, expected_tuple",
@@ -273,16 +288,20 @@ def test_reinstall_patched_tasks():
 
  add_updated_periodic_tasks = [mock.MagicMock(), mock.MagicMock(), mock.MagicMock()]
 
- with mock.patch("sentry_sdk.integrations.celery.shutil.copy2") as mock_copy2:
- _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks)
+ mock_open = mock.Mock(return_value=tempfile.NamedTemporaryFile())
 
- sender.stop.assert_called_once_with()
+ with mock.patch("sentry_sdk.integrations.celery.open", mock_open):
+ with mock.patch(
+ "sentry_sdk.integrations.celery.shutil.copyfileobj"
+ ) as mock_copyfileobj:
+ _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks)
 
- add_updated_periodic_tasks[0].assert_called_once_with()
- add_updated_periodic_tasks[1].assert_called_once_with()
- add_updated_periodic_tasks[2].assert_called_once_with()
+ sender.stop.assert_called_once_with()
 
- mock_copy2.assert_called_once_with(
- "test_schedule_filename", "test_schedule_filename.new"
- )
- fake_beat.run.assert_called_once_with()
+ add_updated_periodic_tasks[0].assert_called_once_with()
+ add_updated_periodic_tasks[1].assert_called_once_with()
+ add_updated_periodic_tasks[2].assert_called_once_with()
+
+ mock_copyfileobj.assert_called_once()
+
+ fake_beat.run.assert_called_once_with()