feat(procrastinate): worker-side task wrapper with sync/async support

Author: snopoke

taskbadger/procrastinate.py

@@ -11,10 +11,15 @@
 
 from __future__ import annotations
 
+import functools
+import inspect
 import logging
 from contextvars import ContextVar
 
 from .internal.models import StatusEnum
+from .mug import Badger
+from .safe_sdk import update_task_safe
+from .sdk import DefaultMergeStrategy, get_task
 
 log = logging.getLogger("taskbadger")
 
@@ -37,3 +42,93 @@
 _OPTS_ATTR = "_taskbadger_opts"
 
 _current_tb_task_id: ContextVar[str | None] = ContextVar("_current_tb_task_id", default=None)
+
+
+def _instrument_task(task, system=None, manual=False, opts=None):
+    """Wrap a Procrastinate Task's ``func`` so the worker side updates TaskBadger.
+
+    Idempotent: a second call on the same task is a no-op (but ``manual`` and
+    ``opts`` will be merged onto the existing attributes if provided).
+    """
+    if opts is not None:
+        existing_opts = getattr(task, _OPTS_ATTR, {}) or {}
+        merged = {**existing_opts, **opts}
+        setattr(task, _OPTS_ATTR, merged)
+    elif not hasattr(task, _OPTS_ATTR):
+        setattr(task, _OPTS_ATTR, {})
+
+    if manual:
+        setattr(task, _MANUAL_ATTR, True)
+
+    if getattr(task, _INSTRUMENTED_ATTR, False):
+        return
+
+    original_func = task.func
+    is_async = inspect.iscoroutinefunction(original_func)
+
+    if is_async:
+
+        @functools.wraps(original_func)
+        async def wrapped(*args, **kwargs):
+            tb_id = kwargs.pop(TB_TASK_ID_KWARG, None)
+            if tb_id is None:
+                return await original_func(*args, **kwargs)
+            token = _current_tb_task_id.set(tb_id)
+            try:
+                _update_status(tb_id, StatusEnum.PROCESSING)
+                try:
+                    result = await original_func(*args, **kwargs)
+                except Exception as exc:
+                    _update_status(tb_id, StatusEnum.ERROR, exception=exc)
+                    raise
+                _update_status(tb_id, StatusEnum.SUCCESS)
+                return result
+            finally:
+                _current_tb_task_id.reset(token)
+    else:
+
+        @functools.wraps(original_func)
+        def wrapped(*args, **kwargs):
+            tb_id = kwargs.pop(TB_TASK_ID_KWARG, None)
+            if tb_id is None:
+                return original_func(*args, **kwargs)
+            token = _current_tb_task_id.set(tb_id)
+            try:
+                _update_status(tb_id, StatusEnum.PROCESSING)
+                try:
+                    result = original_func(*args, **kwargs)
+                except Exception as exc:
+                    _update_status(tb_id, StatusEnum.ERROR, exception=exc)
+                    raise
+                _update_status(tb_id, StatusEnum.SUCCESS)
+                return result
+            finally:
+                _current_tb_task_id.reset(token)
+
+    task.func = wrapped
+    setattr(task, _INSTRUMENTED_ATTR, True)
+    setattr(task, "_taskbadger_system", system)
+
+
+def _update_status(tb_id, status, exception=None):
+    """Update the TaskBadger task to ``status``. Skips if already terminal."""
+    if not Badger.is_configured():
+        return
+
+    if exception is not None or status in TERMINAL_STATES:
+        try:
+            current = get_task(tb_id)
+        except Exception as e:
+            log.warning("Error fetching task '%s': %s", tb_id, e)
+            current = None
+        if current is not None and current.status in TERMINAL_STATES:
+            return
+        data = None
+        if exception is not None and current is not None:
+            data = DefaultMergeStrategy().merge(current.data, {"exception": str(exception)})
+        if data:
+            update_task_safe(tb_id, status=status, data=data)
+        else:
+            update_task_safe(tb_id, status=status)
+    else:
+        update_task_safe(tb_id, status=status)

tests/test_procrastinate.py

@@ -0,0 +1,105 @@
+import asyncio
+import logging
+from unittest import mock
+
+import procrastinate
+import pytest
+from procrastinate import testing
+
+from taskbadger import StatusEnum
+from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task
+from tests.utils import task_for_test
+
+
+@pytest.fixture(autouse=True)
+def _check_log_errors(caplog):
+    yield
+    errors = [r.getMessage() for r in caplog.get_records("call") if r.levelno == logging.ERROR]
+    if errors:
+        pytest.fail(f"log errors during tests: {errors}")
+
+
+@pytest.fixture
+def app():
+    in_memory = testing.InMemoryConnector()
+    app = procrastinate.App(connector=in_memory)
+    with app.open():
+        yield app
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_worker_updates_task_sync(app):
+    @app.task(name="add")
+    def add(a, b):
+        return a + b
+
+    _instrument_task(add, system=None, manual=True)
+
+    with (
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+        mock.patch("taskbadger.procrastinate.get_task") as get,
+    ):
+        get.return_value = task_for_test(status=StatusEnum.PROCESSING)
+        add.func(a=2, b=3, **{TB_TASK_ID_KWARG: "tb-123"})
+
+    statuses = [call.kwargs["status"] for call in update.call_args_list]
+    assert statuses == [StatusEnum.PROCESSING, StatusEnum.SUCCESS]
+    # The reserved key must not appear in the calls (it's stripped before user fn)
+    assert all(TB_TASK_ID_KWARG not in c.kwargs for c in update.call_args_list)
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_worker_updates_task_async(app):
+    @app.task(name="add_async")
+    async def add_async(a, b):
+        return a + b
+
+    _instrument_task(add_async, system=None, manual=True)
+
+    with (
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+        mock.patch("taskbadger.procrastinate.get_task") as get,
+    ):
+        get.return_value = task_for_test(status=StatusEnum.PROCESSING)
+        result = asyncio.run(add_async.func(a=2, b=3, **{TB_TASK_ID_KWARG: "tb-456"}))
+
+    assert result == 5
+    statuses = [call.kwargs["status"] for call in update.call_args_list]
+    assert statuses == [StatusEnum.PROCESSING, StatusEnum.SUCCESS]
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_worker_marks_error(app):
+    @app.task(name="boom")
+    def boom():
+        raise ValueError("nope")
+
+    _instrument_task(boom, system=None, manual=True)
+
+    with (
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+        mock.patch("taskbadger.procrastinate.get_task") as get,
+    ):
+        get.return_value = task_for_test(status=StatusEnum.PROCESSING, data={"x": 1})
+        with pytest.raises(ValueError, match="nope"):
+            boom.func(**{TB_TASK_ID_KWARG: "tb-789"})
+
+    statuses = [call.kwargs["status"] for call in update.call_args_list]
+    assert statuses == [StatusEnum.PROCESSING, StatusEnum.ERROR]
+    err_call = update.call_args_list[-1]
+    assert err_call.kwargs["data"] == {"x": 1, "exception": "nope"}
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_worker_no_id_runs_clean(app):
+    @app.task(name="add2")
+    def add2(a, b):
+        return a + b
+
+    _instrument_task(add2, system=None, manual=True)
+
+    with mock.patch("taskbadger.procrastinate.update_task_safe") as update:
+        result = add2.func(a=1, b=2)
+
+    assert result == 3
+    update.assert_not_called()