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.

Author: antonpirker

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

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.

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

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]
     ):

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"