feat(procrastinate): current_task() accessor via ContextVar

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: snopoke

taskbadger/procrastinate.py

@@ -118,6 +118,10 @@ def _update_status(tb_id, status, exception=None):
         return
 
     if exception is not None or status in TERMINAL_STATES:
+        # Bypass the cache for the terminal-state check: the user may have
+        # updated the task to a terminal state via the regular SDK during
+        # the body, which wouldn't be reflected in our local cache.
+        _task_cache.unset(tb_id)
         current = _safe_get_task(tb_id)
         if current is not None and current.status in TERMINAL_STATES:
             return
@@ -260,3 +264,15 @@ def wrap(task):
     if original_task is None:
         return wrap
     return wrap(original_task)
+
+
+def current_task():
+    """Return the TaskBadger Task for the currently-running Procrastinate job.
+
+    Returns ``None`` outside of a tracked task or if the task can't be fetched.
+    Result is cached for the lifetime of the worker process via an LRU.
+    """
+    tb_id = _current_tb_task_id.get()
+    if tb_id is None:
+        return None
+    return _safe_get_task(tb_id)

tests/test_procrastinate.py

@@ -7,7 +7,7 @@
 from procrastinate import testing
 
 from taskbadger import StatusEnum
-from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, track
+from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, _task_cache, current_task, track
 from tests.utils import task_for_test
 
 
@@ -19,6 +19,13 @@ def _check_log_errors(caplog):
         pytest.fail(f"log errors during tests: {errors}")
 
 
+@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()
@@ -248,3 +255,55 @@ def bad():
 
     with pytest.raises(TypeError, match="unexpected keyword"):
         track(name="x", does_not_exist=True)(bad)
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_current_task_inside_body(app):
+    captured = {}
+
+    @track
+    @app.task(name="capture")
+    def capture():
+        captured["task"] = current_task()
+
+    tb = task_for_test()
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.get_task", return_value=tb),
+    ):
+        capture.defer()
+        app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
+
+    assert captured["task"] is not None
+    assert captured["task"].id == tb.id
+
+
+def test_current_task_outside_returns_none():
+    assert current_task() is None
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_user_set_terminal_state_not_overwritten(app):
+    @track
+    @app.task(name="self_complete")
+    def self_complete():
+        pass
+
+    tb_pending = task_for_test(status=StatusEnum.PENDING)
+    tb_done = task_for_test(id=tb_pending.id, status=StatusEnum.SUCCESS)
+
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb_pending),
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+        mock.patch("taskbadger.procrastinate.get_task", return_value=tb_done),
+    ):
+        self_complete.defer()
+        app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
+
+    # The wrapper's post-call SUCCESS update is skipped because the cached
+    # task is already SUCCESS. PROCESSING update is still allowed (early path).
+    statuses = [c.kwargs["status"] for c in update.call_args_list]
+    assert StatusEnum.PROCESSING in statuses
+    # Last attempted SUCCESS call should be suppressed
+    assert statuses.count(StatusEnum.SUCCESS) == 0