Add support for cron jobs in ARQ integration (#2088)

Co-authored-by: Anton Pirker <[email protected]>
Co-authored-by: Ivana Kellyerova <[email protected]>

Author: 3 people

sentry_sdk/integrations/arq.py

@@ -26,10 +26,11 @@
     raise DidNotEnable("Arq is not installed")
 
 if TYPE_CHECKING:
-    from typing import Any, Dict, Optional
+    from typing import Any, Dict, Optional, Union
 
     from sentry_sdk._types import EventProcessor, Event, ExcInfo, Hint
 
+    from arq.cron import CronJob
     from arq.jobs import Job
     from arq.typing import WorkerCoroutine
     from arq.worker import Function
@@ -61,7 +62,7 @@ def setup_once():
 
         patch_enqueue_job()
         patch_run_job()
-        patch_func()
+        patch_create_worker()
 
         ignore_logger("arq.worker")
 
@@ -186,23 +187,40 @@ async def _sentry_coroutine(ctx, *args, **kwargs):
     return _sentry_coroutine
 
 
-def patch_func():
+def patch_create_worker():
     # type: () -> None
-    old_func = arq.worker.func
+    old_create_worker = arq.worker.create_worker
 
-    def _sentry_func(*args, **kwargs):
-        # type: (*Any, **Any) -> Function
+    def _sentry_create_worker(*args, **kwargs):
+        # type: (*Any, **Any) -> Worker
         hub = Hub.current
 
         if hub.get_integration(ArqIntegration) is None:
-            return old_func(*args, **kwargs)
+            return old_create_worker(*args, **kwargs)
 
-        func = old_func(*args, **kwargs)
+        settings_cls = args[0]
 
-        if not getattr(func, "_sentry_is_patched", False):
-            func.coroutine = _wrap_coroutine(func.name, func.coroutine)
-            func._sentry_is_patched = True
+        functions = settings_cls.functions
+        cron_jobs = settings_cls.cron_jobs
 
-        return func
+        settings_cls.functions = [_get_arq_function(func) for func in functions]
+        settings_cls.cron_jobs = [_get_arq_cron_job(cron_job) for cron_job in cron_jobs]
 
-    arq.worker.func = _sentry_func
+        return old_create_worker(*args, **kwargs)
+
+    arq.worker.create_worker = _sentry_create_worker
+
+
+def _get_arq_function(func):
+    # type: (Union[str, Function, WorkerCoroutine]) -> Function
+    arq_func = arq.worker.func(func)
+    arq_func.coroutine = _wrap_coroutine(arq_func.name, arq_func.coroutine)
+
+    return arq_func
+
+
+def _get_arq_cron_job(cron_job):
+    # type: (CronJob) -> CronJob
+    cron_job.coroutine = _wrap_coroutine(cron_job.name, cron_job.coroutine)
+
+    return cron_job

tests/integrations/arq/test_arq.py

@@ -1,16 +1,28 @@
+import asyncio
 import pytest
 
 from sentry_sdk import start_transaction
 from sentry_sdk.integrations.arq import ArqIntegration
 
+import arq.worker
+from arq import cron
 from arq.connections import ArqRedis
 from arq.jobs import Job
 from arq.utils import timestamp_ms
-from arq.worker import Retry, Worker
 
 from fakeredis.aioredis import FakeRedis
 
 
+def async_partial(async_fn, *args, **kwargs):
+    # asyncio.iscoroutinefunction (Used in the integration code) in Python < 3.8
+    # does not detect async functions in functools.partial objects.
+    # This partial implementation returns a coroutine instead.
+    async def wrapped(ctx):
+        return await async_fn(ctx, *args, **kwargs)
+
+    return wrapped
+
+
 @pytest.fixture(autouse=True)
 def patch_fakeredis_info_command():
     from fakeredis._fakesocket import FakeSocket
@@ -28,7 +40,10 @@ def info(self, section):
 
 @pytest.fixture
 def init_arq(sentry_init):
-    def inner(functions, allow_abort_jobs=False):
+    def inner(functions_=None, cron_jobs_=None, allow_abort_jobs_=False):
+        functions_ = functions_ or []
+        cron_jobs_ = cron_jobs_ or []
+
         sentry_init(
             integrations=[ArqIntegration()],
             traces_sample_rate=1.0,
@@ -38,9 +53,16 @@ def inner(functions, allow_abort_jobs=False):
 
         server = FakeRedis()
         pool = ArqRedis(pool_or_conn=server.connection_pool)
-        return pool, Worker(
-            functions, redis_pool=pool, allow_abort_jobs=allow_abort_jobs
-        )
+
+        class WorkerSettings:
+            functions = functions_
+            cron_jobs = cron_jobs_
+            redis_pool = pool
+            allow_abort_jobs = allow_abort_jobs_
+
+        worker = arq.worker.create_worker(WorkerSettings)
+
+        return pool, worker
 
     return inner
 
@@ -70,7 +92,7 @@ async def increase(ctx, num):
 async def test_job_retry(capture_events, init_arq):
     async def retry_job(ctx):
         if ctx["job_try"] < 2:
-            raise Retry
+            raise arq.worker.Retry
 
     retry_job.__qualname__ = retry_job.__name__
 
@@ -105,36 +127,69 @@ async def division(_, a, b=0):
 
     division.__qualname__ = division.__name__
 
-    pool, worker = init_arq([division])
+    cron_func = async_partial(division, a=1, b=int(not job_fails))
+    cron_func.__qualname__ = division.__name__
+
+    cron_job = cron(cron_func, minute=0, run_at_startup=True)
+
+    pool, worker = init_arq(functions_=[division], cron_jobs_=[cron_job])
 
     events = capture_events()
 
     job = await pool.enqueue_job("division", 1, b=int(not job_fails))
     await worker.run_job(job.job_id, timestamp_ms())
 
-    if job_fails:
-        error_event = events.pop(0)
-        assert error_event["exception"]["values"][0]["type"] == "ZeroDivisionError"
-        assert error_event["exception"]["values"][0]["mechanism"]["type"] == "arq"
+    loop = asyncio.get_event_loop()
+    task = loop.create_task(worker.async_run())
+    await asyncio.sleep(1)
 
-    (event,) = events
-    assert event["type"] == "transaction"
-    assert event["transaction"] == "division"
-    assert event["transaction_info"] == {"source": "task"}
+    task.cancel()
+
+    await worker.close()
 
     if job_fails:
-        assert event["contexts"]["trace"]["status"] == "internal_error"
-    else:
-        assert event["contexts"]["trace"]["status"] == "ok"
-
-    assert "arq_task_id" in event["tags"]
-    assert "arq_task_retry" in event["tags"]
-
-    extra = event["extra"]["arq-job"]
-    assert extra["task"] == "division"
-    assert extra["args"] == [1]
-    assert extra["kwargs"] == {"b": int(not job_fails)}
-    assert extra["retry"] == 1
+        error_func_event = events.pop(0)
+        error_cron_event = events.pop(1)
+
+        assert error_func_event["exception"]["values"][0]["type"] == "ZeroDivisionError"
+        assert error_func_event["exception"]["values"][0]["mechanism"]["type"] == "arq"
+
+        func_extra = error_func_event["extra"]["arq-job"]
+        assert func_extra["task"] == "division"
+
+        assert error_cron_event["exception"]["values"][0]["type"] == "ZeroDivisionError"
+        assert error_cron_event["exception"]["values"][0]["mechanism"]["type"] == "arq"
+
+        cron_extra = error_cron_event["extra"]["arq-job"]
+        assert cron_extra["task"] == "cron:division"
+
+    [func_event, cron_event] = events
+
+    assert func_event["type"] == "transaction"
+    assert func_event["transaction"] == "division"
+    assert func_event["transaction_info"] == {"source": "task"}
+
+    assert "arq_task_id" in func_event["tags"]
+    assert "arq_task_retry" in func_event["tags"]
+
+    func_extra = func_event["extra"]["arq-job"]
+
+    assert func_extra["task"] == "division"
+    assert func_extra["kwargs"] == {"b": int(not job_fails)}
+    assert func_extra["retry"] == 1
+
+    assert cron_event["type"] == "transaction"
+    assert cron_event["transaction"] == "cron:division"
+    assert cron_event["transaction_info"] == {"source": "task"}
+
+    assert "arq_task_id" in cron_event["tags"]
+    assert "arq_task_retry" in cron_event["tags"]
+
+    cron_extra = cron_event["extra"]["arq-job"]
+
+    assert cron_extra["task"] == "cron:division"
+    assert cron_extra["kwargs"] == {}
+    assert cron_extra["retry"] == 1
 
 
 @pytest.mark.asyncio