From f7b0684ed31649d7f32e0c3f7b139605806a848d Mon Sep 17 00:00:00 2001 From: Anton Pirker Date: Thu, 16 Mar 2023 16:07:08 +0100 Subject: [PATCH] Add support for Sentry Crons to Celery Beat (#1935) This adds a decorator @sentry.monitor that can be attached to Celery tasks. When the celery tasks are run, a check-in for Sentry Crons is created and also the status of the check-in is set when the tasks fails for finishes. --- sentry_sdk/__init__.py | 1 + sentry_sdk/client.py | 9 ++- sentry_sdk/crons.py | 123 +++++++++++++++++++++++++++++++++++++++++ sentry_sdk/envelope.py | 6 ++ tests/test_crons.py | 88 +++++++++++++++++++++++++++++ 5 files changed, 225 insertions(+), 2 deletions(-) create mode 100644 sentry_sdk/crons.py create mode 100644 tests/test_crons.py diff --git a/sentry_sdk/__init__.py b/sentry_sdk/__init__.py index dc1ba399d1..bb96c97ae6 100644 --- a/sentry_sdk/__init__.py +++ b/sentry_sdk/__init__.py @@ -7,6 +7,7 @@ from sentry_sdk.consts import VERSION # noqa +from sentry_sdk.crons import monitor # noqa from sentry_sdk.tracing import trace # noqa __all__ = [ # noqa diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py index c4be3331fa..22255e80f0 100644 --- a/sentry_sdk/client.py +++ b/sentry_sdk/client.py @@ -440,9 +440,11 @@ def capture_event( .pop("dynamic_sampling_context", {}) ) - # Transactions or events with attachments should go to the /envelope/ + is_checkin = event_opt.get("type") == "check_in" + + # Transactions, events with attachments, and checkins should go to the /envelope/ # endpoint. - if is_transaction or attachments: + if is_transaction or is_checkin or attachments: headers = { "event_id": event_opt["event_id"], @@ -458,11 +460,14 @@ def capture_event( if profile is not None: envelope.add_profile(profile.to_json(event_opt, self.options)) envelope.add_transaction(event_opt) + elif is_checkin: + envelope.add_checkin(event_opt) else: envelope.add_event(event_opt) for attachment in attachments or (): envelope.add_item(attachment.to_envelope_item()) + self.transport.capture_envelope(envelope) else: # All other events go to the /store/ endpoint. diff --git a/sentry_sdk/crons.py b/sentry_sdk/crons.py new file mode 100644 index 0000000000..e652460df4 --- /dev/null +++ b/sentry_sdk/crons.py @@ -0,0 +1,123 @@ +from functools import wraps +import sys +import uuid + +from sentry_sdk import Hub +from sentry_sdk._compat import reraise +from sentry_sdk._types import TYPE_CHECKING +from sentry_sdk.utils import nanosecond_time + + +if TYPE_CHECKING: + from typing import Any, Callable, Dict, Optional + + +class MonitorStatus: + IN_PROGRESS = "in_progress" + OK = "ok" + ERROR = "error" + + +def _create_checkin_event( + monitor_slug=None, check_in_id=None, status=None, duration=None +): + # type: (Optional[str], Optional[str], Optional[str], Optional[float]) -> Dict[str, Any] + options = Hub.current.client.options if Hub.current.client else {} + check_in_id = check_in_id or uuid.uuid4().hex # type: str + # convert nanosecond to millisecond + duration = int(duration * 0.000001) if duration is not None else duration + + checkin = { + "type": "check_in", + "monitor_slug": monitor_slug, + # TODO: Add schedule and schedule_type to monitor config + # "monitor_config": { + # "schedule": "*/10 0 0 0 0", + # "schedule_type": "cron", + # }, + "check_in_id": check_in_id, + "status": status, + "duration": duration, + "environment": options["environment"], + "release": options["release"], + } + + return checkin + + +def capture_checkin(monitor_slug=None, check_in_id=None, status=None, duration=None): + # type: (Optional[str], Optional[str], Optional[str], Optional[float]) -> str + hub = Hub.current + + check_in_id = check_in_id or uuid.uuid4().hex + checkin_event = _create_checkin_event( + monitor_slug=monitor_slug, + check_in_id=check_in_id, + status=status, + duration=duration, + ) + hub.capture_event(checkin_event) + + return checkin_event["check_in_id"] + + +def monitor(monitor_slug=None, app=None): + # type: (Optional[str], Any) -> Callable[..., Any] + """ + Decorator to capture checkin events for a monitor. + + Usage: + ``` + import sentry_sdk + + app = Celery() + + @app.task + @sentry_sdk.monitor(monitor_slug='my-fancy-slug') + def test(arg): + print(arg) + ``` + + This does not have to be used with Celery, but if you do use it with celery, + put the `@sentry_sdk.monitor` decorator below Celery's `@app.task` decorator. + """ + + def decorate(func): + # type: (Callable[..., Any]) -> Callable[..., Any] + if not monitor_slug: + return func + + @wraps(func) + def wrapper(*args, **kwargs): + # type: (*Any, **Any) -> Any + start_timestamp = nanosecond_time() + check_in_id = capture_checkin( + monitor_slug=monitor_slug, status=MonitorStatus.IN_PROGRESS + ) + + try: + result = func(*args, **kwargs) + except Exception: + duration = nanosecond_time() - start_timestamp + capture_checkin( + monitor_slug=monitor_slug, + check_in_id=check_in_id, + status=MonitorStatus.ERROR, + duration=duration, + ) + exc_info = sys.exc_info() + reraise(*exc_info) + + duration = nanosecond_time() - start_timestamp + capture_checkin( + monitor_slug=monitor_slug, + check_in_id=check_in_id, + status=MonitorStatus.OK, + duration=duration, + ) + + return result + + return wrapper + + return decorate diff --git a/sentry_sdk/envelope.py b/sentry_sdk/envelope.py index 2fb1bae387..fed5ed4849 100644 --- a/sentry_sdk/envelope.py +++ b/sentry_sdk/envelope.py @@ -68,6 +68,12 @@ def add_profile( # type: (...) -> None self.add_item(Item(payload=PayloadRef(json=profile), type="profile")) + def add_checkin( + self, checkin # type: Any + ): + # type: (...) -> None + self.add_item(Item(payload=PayloadRef(json=checkin), type="check_in")) + def add_session( self, session # type: Union[Session, Any] ): diff --git a/tests/test_crons.py b/tests/test_crons.py new file mode 100644 index 0000000000..dd632a315a --- /dev/null +++ b/tests/test_crons.py @@ -0,0 +1,88 @@ +import mock +import pytest +import uuid + +import sentry_sdk +from sentry_sdk.crons import capture_checkin + + +@sentry_sdk.monitor(monitor_slug="abc123") +def _hello_world(name): + return "Hello, {}".format(name) + + +@sentry_sdk.monitor(monitor_slug="def456") +def _break_world(name): + 1 / 0 + return "Hello, {}".format(name) + + +def test_decorator(sentry_init): + sentry_init() + + with mock.patch("sentry_sdk.crons.capture_checkin") as fake_capture_checking: + result = _hello_world("Grace") + assert result == "Hello, Grace" + + # Check for initial checkin + fake_capture_checking.assert_has_calls( + [ + mock.call(monitor_slug="abc123", status="in_progress"), + ] + ) + + # Check for final checkin + assert fake_capture_checking.call_args[1]["monitor_slug"] == "abc123" + assert fake_capture_checking.call_args[1]["status"] == "ok" + assert fake_capture_checking.call_args[1]["duration"] + assert fake_capture_checking.call_args[1]["check_in_id"] + + +def test_decorator_error(sentry_init): + sentry_init() + + with mock.patch("sentry_sdk.crons.capture_checkin") as fake_capture_checking: + with pytest.raises(Exception): + result = _break_world("Grace") + + assert "result" not in locals() + + # Check for initial checkin + fake_capture_checking.assert_has_calls( + [ + mock.call(monitor_slug="def456", status="in_progress"), + ] + ) + + # Check for final checkin + assert fake_capture_checking.call_args[1]["monitor_slug"] == "def456" + assert fake_capture_checking.call_args[1]["status"] == "error" + assert fake_capture_checking.call_args[1]["duration"] + assert fake_capture_checking.call_args[1]["check_in_id"] + + +def test_capture_checkin_simple(sentry_init): + sentry_init() + + check_in_id = capture_checkin( + monitor_slug="abc123", + check_in_id="112233", + status=None, + duration=None, + ) + assert check_in_id == "112233" + + +def test_capture_checkin_new_id(sentry_init): + sentry_init() + + with mock.patch("uuid.uuid4") as mock_uuid: + mock_uuid.return_value = uuid.UUID("a8098c1a-f86e-11da-bd1a-00112444be1e") + check_in_id = capture_checkin( + monitor_slug="abc123", + check_in_id=None, + status=None, + duration=None, + ) + + assert check_in_id == "a8098c1af86e11dabd1a00112444be1e"