From b742c45ce893d96864ec9d907141223a0ea728f1 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Thu, 28 Mar 2024 12:29:34 +0100 Subject: [PATCH] feat(crons): Make `monitor` async friendly (#2912) --- sentry_sdk/crons/_decorator.py | 38 ++++++++ sentry_sdk/crons/_decorator_py2.py | 21 +++++ sentry_sdk/crons/decorator.py | 63 +++++++------ tests/crons/__init__.py | 0 tests/{ => crons}/test_crons.py | 51 +++++------ tests/crons/test_crons_async_py3.py | 136 ++++++++++++++++++++++++++++ 6 files changed, 254 insertions(+), 55 deletions(-) create mode 100644 sentry_sdk/crons/_decorator.py create mode 100644 sentry_sdk/crons/_decorator_py2.py create mode 100644 tests/crons/__init__.py rename tests/{ => crons}/test_crons.py (82%) create mode 100644 tests/crons/test_crons_async_py3.py diff --git a/sentry_sdk/crons/_decorator.py b/sentry_sdk/crons/_decorator.py new file mode 100644 index 0000000000..5a15000a48 --- /dev/null +++ b/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 diff --git a/sentry_sdk/crons/_decorator_py2.py b/sentry_sdk/crons/_decorator_py2.py new file mode 100644 index 0000000000..9e1da797e2 --- /dev/null +++ b/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 diff --git a/sentry_sdk/crons/decorator.py b/sentry_sdk/crons/decorator.py index 34f4d0ac95..38653ca161 100644 --- a/sentry_sdk/crons/decorator.py +++ b/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, - ) diff --git a/tests/crons/__init__.py b/tests/crons/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_crons.py b/tests/crons/test_crons.py similarity index 82% rename from tests/test_crons.py rename to tests/crons/test_crons.py index 39d02a5d47..0b31494acf 100644 --- a/tests/test_crons.py +++ b/tests/crons/test_crons.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,22 +38,22 @@ 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): @@ -62,24 +61,24 @@ def test_decorator_error(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): @@ -87,22 +86,22 @@ def test_contextmanager(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): @@ -110,24 +109,24 @@ def test_contextmanager_error(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): diff --git a/tests/crons/test_crons_async_py3.py b/tests/crons/test_crons_async_py3.py new file mode 100644 index 0000000000..6e00b594bd --- /dev/null +++ b/tests/crons/test_crons_async_py3.py @@ -0,0 +1,136 @@ +import pytest + +import sentry_sdk + +try: + from unittest import mock # python 3.3 and above +except ImportError: + import mock # python < 3.3 + + +@sentry_sdk.monitor(monitor_slug="abc123") +async def _hello_world(name): + return "Hello, {}".format(name) + + +@sentry_sdk.monitor(monitor_slug="def456") +async def _break_world(name): + 1 / 0 + return "Hello, {}".format(name) + + +async def my_coroutine(): + return + + +async def _hello_world_contextmanager(name): + with sentry_sdk.monitor(monitor_slug="abc123"): + await my_coroutine() + return "Hello, {}".format(name) + + +async def _break_world_contextmanager(name): + with sentry_sdk.monitor(monitor_slug="def456"): + await my_coroutine() + 1 / 0 + return "Hello, {}".format(name) + + +@pytest.mark.asyncio +async def test_decorator(sentry_init): + sentry_init() + + with mock.patch( + "sentry_sdk.crons.decorator.capture_checkin" + ) as fake_capture_checkin: + result = await _hello_world("Grace") + assert result == "Hello, Grace" + + # Check for initial checkin + fake_capture_checkin.assert_has_calls( + [ + mock.call(monitor_slug="abc123", status="in_progress"), + ] + ) + + # Check for final checkin + 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"] + + +@pytest.mark.asyncio +async def test_decorator_error(sentry_init): + sentry_init() + + with mock.patch( + "sentry_sdk.crons.decorator.capture_checkin" + ) as fake_capture_checkin: + with pytest.raises(ZeroDivisionError): + result = await _break_world("Grace") + + assert "result" not in locals() + + # Check for initial checkin + fake_capture_checkin.assert_has_calls( + [ + mock.call(monitor_slug="def456", status="in_progress"), + ] + ) + + # Check for final checkin + 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"] + + +@pytest.mark.asyncio +async def test_contextmanager(sentry_init): + sentry_init() + + with mock.patch( + "sentry_sdk.crons.decorator.capture_checkin" + ) as fake_capture_checkin: + result = await _hello_world_contextmanager("Grace") + assert result == "Hello, Grace" + + # Check for initial checkin + fake_capture_checkin.assert_has_calls( + [ + mock.call(monitor_slug="abc123", status="in_progress"), + ] + ) + + # Check for final checkin + 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"] + + +@pytest.mark.asyncio +async def test_contextmanager_error(sentry_init): + sentry_init() + + with mock.patch( + "sentry_sdk.crons.decorator.capture_checkin" + ) as fake_capture_checkin: + with pytest.raises(ZeroDivisionError): + result = await _break_world_contextmanager("Grace") + + assert "result" not in locals() + + # Check for initial checkin + fake_capture_checkin.assert_has_calls( + [ + mock.call(monitor_slug="def456", status="in_progress"), + ] + ) + + # Check for final checkin + 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"]