feat: pre-generate task IDs in the procrastinate integration

Generate the task ID with generate_task_id() before creating the pending task,
so the id is known independently of the create response.

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

Author: snopoke

taskbadger/procrastinate.py

@@ -21,7 +21,7 @@
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy
+from .sdk import DefaultMergeStrategy, generate_task_id
 
 log = logging.getLogger("taskbadger")
 
@@ -164,14 +164,15 @@ async def defer_async(**kwargs):
     task.defer_async = defer_async
 
 
-def _create_pending_task(task, task_kwargs, queue=None):
+def _create_pending_task(task, task_kwargs, queue=None, task_id=None):
     """Create a PENDING TaskBadger task for ``task`` if it should be tracked.
 
     Returns the created TaskBadger task, or ``None`` if Badger isn't
     configured, the task isn't tracked (neither manual nor auto), or the
     create call failed. ``task_kwargs`` is used only for the
     ``record_task_args`` data capture. ``queue`` overrides the queue name
     recorded on the TaskBadger task (defaults to the task's own queue).
+    ``task_id`` assigns the task's id (see ``_maybe_create_pending``).
     """
     if not Badger.is_configured():
         return None
@@ -203,18 +204,22 @@ def _create_pending_task(task, task_kwargs, queue=None):
     if data:
         create_kwargs["data"] = data
 
-    return create_task_safe(name, **create_kwargs)
+    return create_task_safe(name, task_id=task_id, **create_kwargs)
 
 
 def _maybe_create_pending(task, kwargs):
     """Decide whether to track this defer, and if so create the TaskBadger
-    task and inject its id into ``kwargs``. Always returns the kwargs dict."""
-    tb_task = _create_pending_task(task, kwargs)
+    task and inject its id into ``kwargs``. Always returns the kwargs dict.
+
+    The id is generated up front so it's known before the task exists, letting us
+    inject it without depending on the create response."""
+    tb_id = generate_task_id()
+    tb_task = _create_pending_task(task, kwargs, task_id=tb_id)
     if tb_task is None:
         return kwargs
 
     new_kwargs = dict(kwargs)
-    new_kwargs[TB_TASK_ID_KWARG] = tb_task.id
+    new_kwargs[TB_TASK_ID_KWARG] = tb_id
     return new_kwargs
 
 
@@ -314,11 +319,12 @@ 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)
+                candidate_id = generate_task_id()
+                tb_task = _create_pending_task(task, job.task_kwargs, queue=job.queue, task_id=candidate_id)
                 if tb_task is not None:
-                    new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: tb_task.id}
+                    new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: candidate_id}
                     job = job.evolve(task_kwargs=new_kwargs)
-                    tb_id = tb_task.id
+                    tb_id = candidate_id
             job_id = await jm._taskbadger_original_defer_periodic_job(
                 job=job, periodic_id=periodic_id, defer_timestamp=defer_timestamp
             )

tests/test_procrastinate.py

@@ -123,12 +123,13 @@ def add3(a, b):
 
     create.assert_called_once()
     assert create.call_args.args == ("add3",)
-    assert create.call_args.kwargs == {"status": StatusEnum.PENDING, "queue": "default"}
+    task_id = create.call_args.kwargs["task_id"]
+    assert create.call_args.kwargs == {"status": StatusEnum.PENDING, "queue": "default", "task_id": task_id}
 
-    # The injected id should appear in the Procrastinate job's task kwargs.
+    # The generated id is passed to create and injected into the Procrastinate job's task kwargs.
     jobs = list(app.connector.jobs.values())
     assert len(jobs) == 1
-    assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == task_id
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -175,13 +176,13 @@ async def add5(a, b):
 
     tb = task_for_test()
     with (
-        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
         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
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"]
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -194,12 +195,12 @@ def add_ext(a, b):
 
     tb = task_for_test()
     with (
-        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
         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))
+    update.assert_called_once_with(create.call_args.kwargs["task_id"], external_id=str(job_id))
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -212,12 +213,12 @@ async def add_ext_async(a, b):
 
     tb = task_for_test()
     with (
-        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
         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))
+    update.assert_called_once_with(create.call_args.kwargs["task_id"], external_id=str(job_id))
 
 
 def test_defer_no_external_id_when_untracked(app):
@@ -269,15 +270,15 @@ def bare(a):
 
     tb = task_for_test()
     with (
-        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
         mock.patch("taskbadger.procrastinate.update_task_safe"),
     ):
         bare.defer(a=1)
 
     assert getattr(bare, "_taskbadger_manual") is True
     # Inspect the actual Procrastinate job - 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
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"]
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -296,12 +297,14 @@ def raw(a):
 
     create.assert_called_once()
     assert create.call_args.args == ("custom",)
+    task_id = create.call_args.kwargs["task_id"]
     assert create.call_args.kwargs == {
         "status": StatusEnum.PENDING,
         "value_max": 10,
         "tags": {"env": "test"},
         "data": {"k": "v"},
         "queue": "default",
+        "task_id": task_id,
     }
 
 

tests/test_procrastinate_system_integration.py

@@ -65,7 +65,7 @@ def auto_target(a):
     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
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"]
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -121,7 +121,7 @@ def periodic_target(timestamp):
 
     create.assert_called_once()
     jobs_stored = list(app.connector.jobs.values())
-    assert jobs_stored[0]["args"][TB_TASK_ID_KWARG] == tb.id
+    assert jobs_stored[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"]
 
 
 @pytest.mark.usefixtures("_bind_settings")