Context manager monitor (getsentry#2290)

* Commented a confusing line of code

* monitor can now also be used as a contextmanager

* fixed so this also works as contextmanager

* added unit tests

* added type hints

* contextmanager docstring

* Combine import into one line

* Minor changes to docstring

Author: szokeasaurusrex

sentry_sdk/_compat.py

@@ -1,4 +1,6 @@
 import sys
+import contextlib
+from functools import wraps
 
 from sentry_sdk._types import TYPE_CHECKING
 
@@ -8,6 +10,7 @@
     from typing import Any
     from typing import Type
     from typing import TypeVar
+    from typing import Callable
 
     T = TypeVar("T")
 
@@ -35,8 +38,44 @@ def implements_str(cls):
         cls.__str__ = lambda x: unicode(x).encode("utf-8")  # noqa
         return cls
 
+    # The line below is written as an "exec" because it triggers a syntax error in Python 3
     exec("def reraise(tp, value, tb=None):\n raise tp, value, tb")
 
+    def contextmanager(func):
+        # type: (Callable) -> Callable
+        """
+        Decorator which creates a contextmanager that can also be used as a
+        decorator, similar to how the built-in contextlib.contextmanager
+        function works in Python 3.2+.
+        """
+        contextmanager_func = contextlib.contextmanager(func)
+
+        @wraps(func)
+        class DecoratorContextManager:
+            def __init__(self, *args, **kwargs):
+                # type: (...) -> None
+                self.the_contextmanager = contextmanager_func(*args, **kwargs)
+
+            def __enter__(self):
+                # type: () -> None
+                self.the_contextmanager.__enter__()
+
+            def __exit__(self, *args, **kwargs):
+                # type: (...) -> None
+                self.the_contextmanager.__exit__(*args, **kwargs)
+
+            def __call__(self, decorated_func):
+                # type: (Callable) -> Callable[...]
+                @wraps(decorated_func)
+                def when_called(*args, **kwargs):
+                    # type: (...) -> Any
+                    with self.the_contextmanager:
+                        return_val = decorated_func(*args, **kwargs)
+                    return return_val
+
+                return when_called
+
+        return DecoratorContextManager
 
 else:
     import urllib.parse as urlparse  # noqa
@@ -59,6 +98,9 @@ def reraise(tp, value, tb=None):
             raise value.with_traceback(tb)
         raise value
 
+    # contextlib.contextmanager already can be used as decorator in Python 3.2+
+    contextmanager = contextlib.contextmanager
+
 
 def with_metaclass(meta, *bases):
     # type: (Any, *Any) -> Any

sentry_sdk/crons/decorator.py

@@ -1,23 +1,22 @@
-from functools import wraps
 import sys
 
-from sentry_sdk._compat import reraise
+from sentry_sdk._compat import contextmanager, reraise
 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 Any, Callable, Optional
+    from typing import Generator, Optional
 
 
+@contextmanager
 def monitor(monitor_slug=None):
-    # type: (Optional[str]) -> Callable[..., Any]
+    # type: (Optional[str]) -> Generator[None, None, None]
     """
-    Decorator to capture checkin events for a monitor.
+    Decorator/context manager to capture checkin events for a monitor.
 
-    Usage:
+    Usage (as decorator):
     ```
     import sentry_sdk
 
@@ -31,44 +30,41 @@ def test(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 = now()
-            check_in_id = capture_checkin(
-                monitor_slug=monitor_slug, status=MonitorStatus.IN_PROGRESS
-            )
-
-            try:
-                result = func(*args, **kwargs)
-            except Exception:
-                duration_s = now() - start_timestamp
-                capture_checkin(
-                    monitor_slug=monitor_slug,
-                    check_in_id=check_in_id,
-                    status=MonitorStatus.ERROR,
-                    duration=duration_s,
-                )
-                exc_info = sys.exc_info()
-                reraise(*exc_info)
+    Usage (as context manager):
+    ```
+    import sentry_sdk
 
-            duration_s = now() - start_timestamp
-            capture_checkin(
-                monitor_slug=monitor_slug,
-                check_in_id=check_in_id,
-                status=MonitorStatus.OK,
-                duration=duration_s,
-            )
+    def test(arg):
+        with sentry_sdk.monitor(monitor_slug='my-fancy-slug'):
+            print(arg)
+    ```
 
-            return result
 
-        return wrapper
+    """
 
-    return decorate
+    start_timestamp = now()
+    check_in_id = capture_checkin(
+        monitor_slug=monitor_slug, status=MonitorStatus.IN_PROGRESS
+    )
+
+    try:
+        yield
+    except Exception:
+        duration_s = now() - start_timestamp
+        capture_checkin(
+            monitor_slug=monitor_slug,
+            check_in_id=check_in_id,
+            status=MonitorStatus.ERROR,
+            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/test_crons.py

@@ -21,6 +21,17 @@ def _break_world(name):
     return "Hello, {}".format(name)
 
 
+def _hello_world_contextmanager(name):
+    with sentry_sdk.monitor(monitor_slug="abc123"):
+        return "Hello, {}".format(name)
+
+
+def _break_world_contextmanager(name):
+    with sentry_sdk.monitor(monitor_slug="def456"):
+        1 / 0
+        return "Hello, {}".format(name)
+
+
 def test_decorator(sentry_init):
     sentry_init()
 
@@ -69,6 +80,54 @@ def test_decorator_error(sentry_init):
         assert fake_capture_checking.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:
+        result = _hello_world_contextmanager("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_contextmanager_error(sentry_init):
+    sentry_init()
+
+    with mock.patch(
+        "sentry_sdk.crons.decorator.capture_checkin"
+    ) as fake_capture_checking:
+        with pytest.raises(Exception):
+            result = _break_world_contextmanager("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()