feat(procrastinate): ProcrastinateSystemIntegration with auto-track

Author: snopoke

taskbadger/procrastinate.py

@@ -300,3 +300,34 @@ def current_task():
     if tb_id is None:
         return None
     return _safe_get_task(tb_id)
+
+
+def _patch_app_task(app, system):
+    """Replace ``app.task`` with a wrapper that instruments newly-registered
+    tasks under the supplied ``system``. Idempotent — a second call replaces
+    the wrapper but keeps the same original task method."""
+    original = getattr(app, "_taskbadger_original_task", None) or app.task
+    if not getattr(app, "_taskbadger_original_task", None):
+        app._taskbadger_original_task = original
+
+    @functools.wraps(original)
+    def patched(*args, **kwargs):
+        task = original(*args, **kwargs)
+        # ``original`` may return the Task directly or a decorator depending on
+        # call form. Procrastinate's ``app.task`` always returns a decorator
+        # when called with arguments and the Task when called bare.
+        if callable(task) and not hasattr(task, "name"):
+            # decorator form: wrap the returned decorator
+            inner_decorator = task
+
+            @functools.wraps(inner_decorator)
+            def outer(func):
+                t = inner_decorator(func)
+                _instrument_task(t, system=system)
+                return t
+
+            return outer
+        _instrument_task(task, system=system)
+        return task
+
+    app.task = patched

taskbadger/systems/procrastinate.py

@@ -2,8 +2,60 @@
 
 from __future__ import annotations
 
+import re
+
+from taskbadger.procrastinate import _instrument_task, _patch_app_task
 from taskbadger.systems import System
 
 
 class ProcrastinateSystemIntegration(System):
     identifier = "procrastinate"
+
+    def __init__(
+        self,
+        app,
+        auto_track_tasks=True,
+        includes=None,
+        excludes=None,
+        record_task_args=False,
+    ):
+        """
+        Args:
+            app: The ``procrastinate.App`` instance to instrument.
+            auto_track_tasks: Track all tasks regardless of whether they use
+                the ``@taskbadger.procrastinate.track`` decorator.
+            includes: List of task names to include in auto-tracking. Each
+                entry can be a full name or a regex (matched with
+                ``re.fullmatch``).
+            excludes: List of task names to exclude. Same semantics as
+                ``includes``. Exclusions take precedence.
+            record_task_args: Record the task's defer kwargs into the
+                TaskBadger task's ``data`` under ``procrastinate_task_kwargs``.
+        """
+        self.app = app
+        self.auto_track_tasks = auto_track_tasks
+        self.includes = includes
+        self.excludes = excludes
+        self.record_task_args = record_task_args
+
+        for task in list(app.tasks.values()):
+            _instrument_task(task, system=self)
+        _patch_app_task(app, system=self)
+
+    def track_task(self, task_name):
+        if not self.auto_track_tasks:
+            return False
+
+        if self.excludes:
+            for exclude in self.excludes:
+                if re.fullmatch(exclude, task_name):
+                    return False
+
+        if self.includes:
+            for include in self.includes:
+                if re.fullmatch(include, task_name):
+                    break
+            else:
+                return False
+
+        return True

tests/test_procrastinate_system_integration.py

@@ -0,0 +1,88 @@
+from unittest import mock
+
+import procrastinate
+import pytest
+from procrastinate import testing
+
+from taskbadger.procrastinate import _INSTRUMENTED_ATTR, TB_TASK_ID_KWARG, _task_cache
+from taskbadger.systems.procrastinate import ProcrastinateSystemIntegration
+from tests.utils import task_for_test
+
+
+@pytest.fixture(autouse=True)
+def _clear_task_cache():
+    _task_cache.cache.clear()
+    yield
+    _task_cache.cache.clear()
+
+
+@pytest.fixture
+def app():
+    in_memory = testing.InMemoryConnector()
+    app = procrastinate.App(connector=in_memory)
+    with app.open():
+        yield app
+
+
+@pytest.mark.parametrize(
+    ("include", "exclude", "expected"),
+    [
+        (None, None, True),
+        (["myapp.tasks.export_data"], None, True),
+        ([".*export_data"], [], True),
+        ([".*export_da"], [], False),
+        (["myapp.tasks.export_data"], ["myapp.tasks.export_data"], False),
+        ([".*"], ["myapp.tasks.export_data"], False),
+        ([".*"], [".*tasks.*"], False),
+    ],
+)
+def test_task_name_matching(app, include, exclude, expected):
+    integration = ProcrastinateSystemIntegration(app=app, includes=include, excludes=exclude)
+    assert integration.track_task("myapp.tasks.export_data") is expected
+
+
+def test_auto_track_off_returns_false(app):
+    integration = ProcrastinateSystemIntegration(app=app, auto_track_tasks=False)
+    assert integration.track_task("anything") is False
+
+
+def test_wraps_existing_tasks(app):
+    @app.task(name="pre_existing")
+    def pre_existing(a):
+        return a
+
+    assert not getattr(pre_existing, _INSTRUMENTED_ATTR, False)
+    ProcrastinateSystemIntegration(app=app, auto_track_tasks=True)
+    assert getattr(pre_existing, _INSTRUMENTED_ATTR) is True
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_auto_track_creates_pending(app):
+    @app.task(name="auto_target")
+    def auto_target(a):
+        return a
+
+    ProcrastinateSystemIntegration(app=app, auto_track_tasks=True)
+
+    tb = task_for_test()
+    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+        auto_target.defer(a=1)
+
+    create.assert_called_once()
+    # InMemoryConnector.jobs is a dict keyed by int; kwargs under "args"
+    jobs = list(app.connector.jobs.values())
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_auto_track_excludes_skip(app):
+    @app.task(name="myapp.cleanup.flush")
+    def flush():
+        pass
+
+    ProcrastinateSystemIntegration(app=app, auto_track_tasks=True, excludes=[r"myapp\.cleanup\..*"])
+
+    with mock.patch("taskbadger.procrastinate.create_task_safe") as create:
+        flush.defer()
+
+    create.assert_not_called()