feat(crons): Make monitor async friendly (#2912)

Author: sentrivana

sentry_sdk/crons/_decorator.py

@@ -0,0 +1,38 @@
+from functools import wraps
+from inspect import iscoroutinefunction
+
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import (
+        Awaitable,
+        Callable,
+        ParamSpec,
+        TypeVar,
+        Union,
+    )
+
+    P = ParamSpec("P")
+    R = TypeVar("R")
+
+
+class MonitorMixin:
+    def __call__(self, fn):
+        # type: (Callable[P, R]) -> Callable[P, Union[R, Awaitable[R]]]
+        if iscoroutinefunction(fn):
+
+            @wraps(fn)
+            async def inner(*args: "P.args", **kwargs: "P.kwargs"):
+                # type: (...) -> R
+                with self:  # type: ignore[attr-defined]
+                    return await fn(*args, **kwargs)
+
+        else:
+
+            @wraps(fn)
+            def inner(*args: "P.args", **kwargs: "P.kwargs"):
+                # type: (...) -> R
+                with self:  # type: ignore[attr-defined]
+                    return fn(*args, **kwargs)
+
+        return inner

sentry_sdk/crons/_decorator_py2.py

@@ -0,0 +1,21 @@
+from functools import wraps
+
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, ParamSpec, TypeVar
+
+    P = ParamSpec("P")
+    R = TypeVar("R")
+
+
+class MonitorMixin:
+    def __call__(self, fn):
+        # type: (Callable[P, R]) -> Callable[P, R]
+        @wraps(fn)
+        def inner(*args, **kwargs):
+            # type: (Any, Any) -> Any
+            with self:  # type: ignore[attr-defined]
+                return fn(*args, **kwargs)
+
+        return inner

sentry_sdk/crons/decorator.py

@@ -1,18 +1,24 @@
-import sys
-
-from sentry_sdk._compat import contextmanager, reraise
+from sentry_sdk._compat import PY2
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.crons import capture_checkin
 from sentry_sdk.crons.consts import MonitorStatus
 from sentry_sdk.utils import now
 
 if TYPE_CHECKING:
-    from typing import Generator, Optional
+    from typing import Optional, Type
+    from types import TracebackType
+
+if PY2:
+    from sentry_sdk.crons._decorator_py2 import MonitorMixin
+else:
+    # This is in its own module so that we don't make Python 2
+    # angery over `async def`s.
+    # Once we drop Python 2, remove the mixin and merge it
+    # into the main monitor class.
+    from sentry_sdk.crons._decorator import MonitorMixin
 
 
-@contextmanager
-def monitor(monitor_slug=None):
-    # type: (Optional[str]) -> Generator[None, None, None]
+class monitor(MonitorMixin):  # noqa: N801
     """
     Decorator/context manager to capture checkin events for a monitor.
 
@@ -39,32 +45,31 @@ def test(arg):
         with sentry_sdk.monitor(monitor_slug='my-fancy-slug'):
             print(arg)
     ```
+    """
 
+    def __init__(self, monitor_slug=None):
+        # type: (Optional[str]) -> None
+        self.monitor_slug = monitor_slug
 
-    """
+    def __enter__(self):
+        # type: () -> None
+        self.start_timestamp = now()
+        self.check_in_id = capture_checkin(
+            monitor_slug=self.monitor_slug, status=MonitorStatus.IN_PROGRESS
+        )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # type: (Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]) -> None
+        duration_s = now() - self.start_timestamp
 
-    start_timestamp = now()
-    check_in_id = capture_checkin(
-        monitor_slug=monitor_slug, status=MonitorStatus.IN_PROGRESS
-    )
+        if exc_type is None and exc_value is None and traceback is None:
+            status = MonitorStatus.OK
+        else:
+            status = MonitorStatus.ERROR
 
-    try:
-        yield
-    except Exception:
-        duration_s = now() - start_timestamp
         capture_checkin(
-            monitor_slug=monitor_slug,
-            check_in_id=check_in_id,
-            status=MonitorStatus.ERROR,
+            monitor_slug=self.monitor_slug,
+            check_in_id=self.check_in_id,
+            status=status,
             duration=duration_s,
         )
-        exc_info = sys.exc_info()
-        reraise(*exc_info)
-
-    duration_s = now() - start_timestamp
-    capture_checkin(
-        monitor_slug=monitor_slug,
-        check_in_id=check_in_id,
-        status=MonitorStatus.OK,
-        duration=duration_s,
-    )

tests/crons/__init__.py

@@ -2,9 +2,8 @@
 import uuid
 
 import sentry_sdk
-from sentry_sdk.crons import capture_checkin
-
 from sentry_sdk import Hub, configure_scope, set_level
+from sentry_sdk.crons import capture_checkin
 
 try:
     from unittest import mock  # python 3.3 and above
@@ -39,95 +38,95 @@ def test_decorator(sentry_init):
 
     with mock.patch(
         "sentry_sdk.crons.decorator.capture_checkin"
-    ) as fake_capture_checking:
+    ) as fake_capture_checkin:
         result = _hello_world("Grace")
         assert result == "Hello, Grace"
 
         # Check for initial checkin
-        fake_capture_checking.assert_has_calls(
+        fake_capture_checkin.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"]
+        assert fake_capture_checkin.call_args[1]["monitor_slug"] == "abc123"
+        assert fake_capture_checkin.call_args[1]["status"] == "ok"
+        assert fake_capture_checkin.call_args[1]["duration"]
+        assert fake_capture_checkin.call_args[1]["check_in_id"]
 
 
 def test_decorator_error(sentry_init):
     sentry_init()
 
     with mock.patch(
         "sentry_sdk.crons.decorator.capture_checkin"
-    ) as fake_capture_checking:
+    ) as fake_capture_checkin:
         with pytest.raises(ZeroDivisionError):
             result = _break_world("Grace")
 
         assert "result" not in locals()
 
         # Check for initial checkin
-        fake_capture_checking.assert_has_calls(
+        fake_capture_checkin.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"]
+        assert fake_capture_checkin.call_args[1]["monitor_slug"] == "def456"
+        assert fake_capture_checkin.call_args[1]["status"] == "error"
+        assert fake_capture_checkin.call_args[1]["duration"]
+        assert fake_capture_checkin.call_args[1]["check_in_id"]
 
 
 def test_contextmanager(sentry_init):
     sentry_init()
 
     with mock.patch(
         "sentry_sdk.crons.decorator.capture_checkin"
-    ) as fake_capture_checking:
+    ) as fake_capture_checkin:
         result = _hello_world_contextmanager("Grace")
         assert result == "Hello, Grace"
 
         # Check for initial checkin
-        fake_capture_checking.assert_has_calls(
+        fake_capture_checkin.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"]
+        assert fake_capture_checkin.call_args[1]["monitor_slug"] == "abc123"
+        assert fake_capture_checkin.call_args[1]["status"] == "ok"
+        assert fake_capture_checkin.call_args[1]["duration"]
+        assert fake_capture_checkin.call_args[1]["check_in_id"]
 
 
 def test_contextmanager_error(sentry_init):
     sentry_init()
 
     with mock.patch(
         "sentry_sdk.crons.decorator.capture_checkin"
-    ) as fake_capture_checking:
+    ) as fake_capture_checkin:
         with pytest.raises(ZeroDivisionError):
             result = _break_world_contextmanager("Grace")
 
         assert "result" not in locals()
 
         # Check for initial checkin
-        fake_capture_checking.assert_has_calls(
+        fake_capture_checkin.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"]
+        assert fake_capture_checkin.call_args[1]["monitor_slug"] == "def456"
+        assert fake_capture_checkin.call_args[1]["status"] == "error"
+        assert fake_capture_checkin.call_args[1]["duration"]
+        assert fake_capture_checkin.call_args[1]["check_in_id"]
 
 
 def test_capture_checkin_simple(sentry_init):