Record celery/procrastinate task ID as external_id

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

Author: snopoke

taskbadger/celery.py

@@ -137,7 +137,9 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     ctask = celery.current_app.tasks.get(sender)
 
     # get kwargs from the task class (set via decorator)
-    kwargs = getattr(ctask, TB_KWARGS_ARG, {})
+    # copy: the raw attr is shared across invocations, and we mutate per-publish
+    # values (external_id, status, name) into it below
+    kwargs = dict(getattr(ctask, TB_KWARGS_ARG, {}))
     for attr in dir(ctask):
         if attr.startswith(KWARG_PREFIX) and attr not in IGNORE_ARGS:
             kwargs[attr.removeprefix(KWARG_PREFIX)] = getattr(ctask, attr)
@@ -147,6 +149,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     kwargs["status"] = StatusEnum.PENDING
     if routing_key and "queue" not in kwargs:
         kwargs["queue"] = routing_key
+    kwargs.setdefault("external_id", headers["id"])
     name = kwargs.pop("name", headers["task"])
 
     global_record_task_args = celery_system and celery_system.record_task_args
@@ -247,7 +250,8 @@ def _maybe_create_task(signal_sender):
 
     delivery_info = getattr(signal_sender.request, "delivery_info", None) or {}
     queue = delivery_info.get("routing_key")
-    task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data, queue=queue)
+    external_id = signal_sender.request.id
+    task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data, queue=queue, external_id=external_id)
     if task:
         # Store the task ID in the request so _update_task can find it
         signal_sender.request.update({TB_TASK_ID: task.id})

taskbadger/procrastinate.py

@@ -149,12 +149,16 @@ def _wrap_defer(task):
     @functools.wraps(original_defer)
     def defer(**kwargs):
         kwargs = _maybe_create_pending(task, kwargs)
-        return original_defer(**kwargs)
+        job_id = original_defer(**kwargs)
+        _record_external_id(kwargs, job_id)
+        return job_id
 
     @functools.wraps(original_defer_async)
     async def defer_async(**kwargs):
         kwargs = _maybe_create_pending(task, kwargs)
-        return await original_defer_async(**kwargs)
+        job_id = await original_defer_async(**kwargs)
+        _record_external_id(kwargs, job_id)
+        return job_id
 
     task.defer = defer
     task.defer_async = defer_async
@@ -214,6 +218,18 @@ def _maybe_create_pending(task, kwargs):
     return new_kwargs
 
 
+def _record_external_id(kwargs, job_id):
+    """Record the Procrastinate job id as the TaskBadger task's ``external_id``.
+
+    ``defer`` only returns the DB-assigned job id once the job is enqueued, so this
+    runs as a follow-up update after both ids are known. No-op if the defer wasn't
+    tracked (no injected id) or no job id came back."""
+    tb_id = kwargs.get(TB_TASK_ID_KWARG)
+    if tb_id is None or job_id is None:
+        return
+    update_task_safe(tb_id, external_id=str(job_id))
+
+
 def _serialize_kwargs(kwargs):
     """Return a JSON-roundtrippable copy of the defer kwargs.
 
@@ -296,14 +312,19 @@ def _patch_job_manager(app, system):
         @functools.wraps(original)
         async def patched(*, job, periodic_id, defer_timestamp):
             task = app.tasks.get(job.task_name)
+            tb_id = None
             if task is not None:
                 tb_task = _create_pending_task(task, job.task_kwargs, queue=job.queue)
                 if tb_task is not None:
                     new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: tb_task.id}
                     job = job.evolve(task_kwargs=new_kwargs)
-            return await jm._taskbadger_original_defer_periodic_job(
+                    tb_id = tb_task.id
+            job_id = await jm._taskbadger_original_defer_periodic_job(
                 job=job, periodic_id=periodic_id, defer_timestamp=defer_timestamp
             )
+            if tb_id is not None and job_id is not None:
+                update_task_safe(tb_id, external_id=str(job_id))
+            return job_id
 
         jm.defer_periodic_job = patched
 

tests/test_celery.py

@@ -98,7 +98,13 @@ def add_with_task_args(self, a, b):
         assert result.get(timeout=10, propagate=True) == 4
 
     create.assert_called_once_with(
-        "new_name", value_max=10, data={"foo": "bar"}, tags={"bar": "baz"}, status=StatusEnum.PENDING, queue="celery"
+        "new_name",
+        value_max=10,
+        data={"foo": "bar"},
+        tags={"bar": "baz"},
+        status=StatusEnum.PENDING,
+        queue="celery",
+        external_id=mock.ANY,
     )
 
 
@@ -128,7 +134,9 @@ def add_with_task_args(self, a, b):
         )
         assert result.get(timeout=10, propagate=True) == 4
 
-    create.assert_called_once_with("new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery")
+    create.assert_called_once_with(
+        "new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery", external_id=mock.ANY
+    )
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -162,6 +170,7 @@ def add_with_task_args(self, a, b):
         data={"foo": "bar", "celery_task_args": [2, 2], "celery_task_kwargs": {}},
         status=StatusEnum.PENDING,
         queue="celery",
+        external_id=mock.ANY,
     )
 
 
@@ -200,6 +209,7 @@ def add_with_task_kwargs(self, a, b, c=0):
         actions=actions,
         status=StatusEnum.PENDING,
         queue="celery",
+        external_id=mock.ANY,
     )
 
 
@@ -237,6 +247,7 @@ def add_task_custom_serialization(self, a):
         data={"celery_task_args": [{"__type__": "A", "__value__": [2, 2]}], "celery_task_kwargs": {}},
         status=StatusEnum.PENDING,
         queue="celery",
+        external_id=mock.ANY,
     )
 
 
@@ -264,7 +275,9 @@ def add_with_task_args_in_decorator(self, a, b):
         result = add_with_task_args_in_decorator.delay(2, 2)
         assert result.get(timeout=10, propagate=True) == 4
 
-    create.assert_called_once_with(mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery")
+    create.assert_called_once_with(
+        mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery", external_id=mock.ANY
+    )
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -288,6 +301,49 @@ def add_custom_queue(self, a, b):
     assert create.call_args.kwargs["queue"] == "high_priority"
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_records_task_id_as_external_id(celery_session_app, celery_session_worker):
+    @celery_session_app.task(bind=True, base=Task)
+    def add_external_id(self, a, b):
+        return a + b
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        result = add_external_id.apply_async((2, 2), queue="high_priority")
+
+    assert create.call_args.kwargs["external_id"] == result.id
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_external_id_not_cached_across_invocations(celery_session_app, celery_session_worker):
+    # Tasks defined with taskbadger_kwargs share a class-level dict; external_id
+    # must not leak from one publish to the next.
+    @celery_session_app.task(bind=True, base=Task, taskbadger_kwargs={"value_max": 10})
+    def add_repeat(self, a, b):
+        return a + b
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        result1 = add_repeat.apply_async((2, 2), queue="high_priority")
+        result2 = add_repeat.apply_async((3, 3), queue="high_priority")
+
+    external_ids = [c.kwargs["external_id"] for c in create.call_args_list]
+    assert external_ids == [result1.id, result2.id]
+    assert result1.id != result2.id
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_celery_task_retry(celery_session_app, celery_session_worker):
     """Note: When a task is retried, the celery task ID remains the same but a new TB task

tests/test_celery_system_integration.py

@@ -123,6 +123,7 @@ def add_normal(self, a, b):
         status=StatusEnum.PENDING,
         data={"celery_task_args": [2, 2], "celery_task_kwargs": {}},
         queue="celery",
+        external_id=mock.ANY,
     )
     assert get_task.call_count == 1
     assert update.call_count == 2
@@ -154,7 +155,10 @@ def add_normal_with_override(a, b):
         assert result.get(timeout=10, propagate=True) == 4
 
     create.assert_called_once_with(
-        "tests.test_celery_system_integration.add_normal_with_override", status=StatusEnum.PENDING, queue="celery"
+        "tests.test_celery_system_integration.add_normal_with_override",
+        status=StatusEnum.PENDING,
+        queue="celery",
+        external_id=mock.ANY,
     )
 
 
@@ -189,6 +193,7 @@ def add_with_tags(a, b):
         status=StatusEnum.PENDING,
         tags=TaskRequestTags.from_dict({"tag1": "value1", "tag2": "override"}),
         queue="celery",
+        external_id=result.id,
     )
     create.assert_called_with(
         client=mock.ANY,

tests/test_procrastinate.py

@@ -115,7 +115,10 @@ def add3(a, b):
     _instrument_task(add3, system=None, manual=True)
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         add3.defer(a=1, b=2)
 
     create.assert_called_once()
@@ -137,7 +140,10 @@ def add_queued(a, b):
     _instrument_task(add_queued, system=None, manual=True)
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         add_queued.defer(a=1, b=2)
 
     assert create.call_args.kwargs["queue"] == "high_priority"
@@ -168,13 +174,69 @@ async def add5(a, b):
     _instrument_task(add5, system=None, manual=True)
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb):
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         asyncio.run(add5.defer_async(a=1, b=2))
 
     jobs = list(app.connector.jobs.values())
     assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_defer_records_job_id_as_external_id(app):
+    @app.task(name="add_ext")
+    def add_ext(a, b):
+        return a + b
+
+    _instrument_task(add_ext, system=None, manual=True)
+
+    tb = task_for_test()
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+    ):
+        job_id = add_ext.defer(a=1, b=2)
+
+    update.assert_called_once_with(tb.id, external_id=str(job_id))
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_defer_async_records_job_id_as_external_id(app):
+    @app.task(name="add_ext_async")
+    async def add_ext_async(a, b):
+        return a + b
+
+    _instrument_task(add_ext_async, system=None, manual=True)
+
+    tb = task_for_test()
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+    ):
+        job_id = asyncio.run(add_ext_async.defer_async(a=1, b=2))
+
+    update.assert_called_once_with(tb.id, external_id=str(job_id))
+
+
+def test_defer_no_external_id_when_untracked(app):
+    @app.task(name="add_untracked")
+    def add_untracked(a, b):
+        return a + b
+
+    _instrument_task(add_untracked, system=None, manual=True)
+
+    # Badger is not configured, so no pending task is created and nothing to update.
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=None),
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+    ):
+        add_untracked.defer(a=1, b=2)
+
+    update.assert_not_called()
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_end_to_end_via_worker(app):
     @app.task(name="add6")
@@ -194,7 +256,7 @@ def add6(a, b):
         app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
 
     create.assert_called_once()
-    statuses = [c.kwargs["status"] for c in update.call_args_list]
+    statuses = [c.kwargs["status"] for c in update.call_args_list if "status" in c.kwargs]
     assert statuses == [StatusEnum.PROCESSING, StatusEnum.SUCCESS]
 
 
@@ -206,7 +268,10 @@ def bare(a):
         return a
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb):
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         bare.defer(a=1)
 
     assert getattr(bare, "_taskbadger_manual") is True
@@ -223,7 +288,10 @@ def raw(a):
         return a
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         raw.defer(a=1)
 
     create.assert_called_once()
@@ -248,7 +316,10 @@ def dup(a):
     # Two @track applications must not double-wrap; defer once still creates one
     # PENDING task and injects one id.
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         dup.defer(a=1)
     assert create.call_count == 1
     jobs = list(app.connector.jobs.values())
@@ -312,7 +383,7 @@ def self_complete():
 
     # 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]
+    statuses = [c.kwargs["status"] for c in update.call_args_list if "status" in c.kwargs]
     assert StatusEnum.PROCESSING in statuses
     # Last attempted SUCCESS call should be suppressed
     assert statuses.count(StatusEnum.SUCCESS) == 0
@@ -326,7 +397,10 @@ def recorder(a, b):
         return a + b
 
     tb = task_for_test()
-    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+    ):
         recorder.defer(a=5, b=6)
 
     assert create.call_args.kwargs["data"] == {