Support setting task ID on creation

taskbadger.yaml

@@ -685,6 +685,11 @@ components:
     PatchedTaskRequest:
       type: object
       properties:
+        id:
+          type: string
+          minLength: 1
+          description: Task ID. May be set on creation to a UUID or shortened UUID;
+            it must be unique and is immutable thereafter. If omitted, an ID is generated.
         name:
           type: string
           minLength: 1
@@ -783,8 +788,8 @@ components:
       properties:
         id:
           type: string
-          readOnly: true
-          description: Task ID
+          description: Task ID. May be set on creation to a UUID or shortened UUID;
+            it must be unique and is immutable thereafter. If omitted, an ID is generated.
         organization:
           type: string
           readOnly: true
@@ -883,7 +888,6 @@ components:
             to 'value'.
       required:
       - created
-      - id
       - name
       - organization
       - project
@@ -894,6 +898,11 @@ components:
     TaskRequest:
       type: object
       properties:
+        id:
+          type: string
+          minLength: 1
+          description: Task ID. May be set on creation to a UUID or shortened UUID;
+            it must be unique and is immutable thereafter. If omitted, an ID is generated.
         name:
           type: string
           minLength: 1

taskbadger/__init__.py

@@ -3,7 +3,7 @@
 from .internal.models import StatusEnum
 from .mug import Badger, Session
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy, Task, create_task, get_task, init, update_task
+from .sdk import DefaultMergeStrategy, Task, create_task, generate_task_id, get_task, init, update_task
 
 __all__ = [
     "track",
@@ -18,6 +18,7 @@
     "DefaultMergeStrategy",
     "Task",
     "create_task",
+    "generate_task_id",
     "get_task",
     "init",
     "update_task",

taskbadger/internal/models/patched_task_request.py

@@ -22,6 +22,8 @@
 class PatchedTaskRequest:
     """
     Attributes:
+        id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is
+            immutable thereafter. If omitted, an ID is generated.
         name (str | Unset): Name of the task
         queue (str | Unset): Queue the task is from
         external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with
@@ -50,6 +52,7 @@ class PatchedTaskRequest:
         tags (PatchedTaskRequestTags | Unset): Tags for the task represented as a mapping from 'namespace' to 'value'.
     """
 
+    id: str | Unset = UNSET
     name: str | Unset = UNSET
     queue: str | Unset = UNSET
     external_id: str | Unset = UNSET
@@ -68,6 +71,8 @@ class PatchedTaskRequest:
     def to_dict(self) -> dict[str, Any]:
         from ..models.patched_task_request_tags import PatchedTaskRequestTags
 
+        id = self.id
+
         name = self.name
 
         queue = self.queue
@@ -129,6 +134,8 @@ def to_dict(self) -> dict[str, Any]:
         field_dict: dict[str, Any] = {}
         field_dict.update(self.additional_properties)
         field_dict.update({})
+        if id is not UNSET:
+            field_dict["id"] = id
         if name is not UNSET:
             field_dict["name"] = name
         if queue is not UNSET:
@@ -163,6 +170,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         from ..models.patched_task_request_tags import PatchedTaskRequestTags
 
         d = dict(src_dict)
+        id = d.pop("id", UNSET)
+
         name = d.pop("name", UNSET)
 
         queue = d.pop("queue", UNSET)
@@ -258,6 +267,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             tags = PatchedTaskRequestTags.from_dict(_tags)
 
         patched_task_request = cls(
+            id=id,
             name=name,
             queue=queue,
             external_id=external_id,

taskbadger/internal/models/task.py

@@ -22,7 +22,6 @@
 class Task:
     """
     Attributes:
-        id (str): Task ID
         organization (str):
         project (str):
         name (str): Name of the task
@@ -31,6 +30,8 @@ class Task:
         updated (datetime.datetime):
         url (str):
         public_url (str):
+        id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is
+            immutable thereafter. If omitted, an ID is generated.
         queue (str | Unset): Queue the task is from
         external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with
             logs
@@ -58,7 +59,6 @@ class Task:
         tags (TaskTags | Unset): Tags for the task represented as a mapping from 'namespace' to 'value'.
     """
 
-    id: str
     organization: str
     project: str
     name: str
@@ -67,6 +67,7 @@ class Task:
     updated: datetime.datetime
     url: str
     public_url: str
+    id: str | Unset = UNSET
     queue: str | Unset = UNSET
     external_id: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
@@ -84,8 +85,6 @@ class Task:
     def to_dict(self) -> dict[str, Any]:
         from ..models.task_tags import TaskTags
 
-        id = self.id
-
         organization = self.organization
 
         project = self.project
@@ -103,6 +102,8 @@ def to_dict(self) -> dict[str, Any]:
 
         public_url = self.public_url
 
+        id = self.id
+
         queue = self.queue
 
         external_id = self.external_id
@@ -163,7 +164,6 @@ def to_dict(self) -> dict[str, Any]:
         field_dict.update(self.additional_properties)
         field_dict.update(
             {
-                "id": id,
                 "organization": organization,
                 "project": project,
                 "name": name,
@@ -174,6 +174,8 @@ def to_dict(self) -> dict[str, Any]:
                 "public_url": public_url,
             }
         )
+        if id is not UNSET:
+            field_dict["id"] = id
         if queue is not UNSET:
             field_dict["queue"] = queue
         if external_id is not UNSET:
@@ -206,8 +208,6 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         from ..models.task_tags import TaskTags
 
         d = dict(src_dict)
-        id = d.pop("id")
-
         organization = d.pop("organization")
 
         project = d.pop("project")
@@ -229,6 +229,8 @@ def _parse_value_percent(data: object) -> int | None:
 
         public_url = d.pop("public_url")
 
+        id = d.pop("id", UNSET)
+
         queue = d.pop("queue", UNSET)
 
         external_id = d.pop("external_id", UNSET)
@@ -322,7 +324,6 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             tags = TaskTags.from_dict(_tags)
 
         task = cls(
-            id=id,
             organization=organization,
             project=project,
             name=name,
@@ -331,6 +332,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             updated=updated,
             url=url,
             public_url=public_url,
+            id=id,
             queue=queue,
             external_id=external_id,
             status=status,

taskbadger/internal/models/task_request.py

@@ -23,6 +23,8 @@ class TaskRequest:
     """
     Attributes:
         name (str): Name of the task
+        id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is
+            immutable thereafter. If omitted, an ID is generated.
         queue (str | Unset): Queue the task is from
         external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with
             logs
@@ -51,6 +53,7 @@ class TaskRequest:
     """
 
     name: str
+    id: str | Unset = UNSET
     queue: str | Unset = UNSET
     external_id: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
@@ -70,6 +73,8 @@ def to_dict(self) -> dict[str, Any]:
 
         name = self.name
 
+        id = self.id
+
         queue = self.queue
 
         external_id = self.external_id
@@ -133,6 +138,8 @@ def to_dict(self) -> dict[str, Any]:
                 "name": name,
             }
         )
+        if id is not UNSET:
+            field_dict["id"] = id
         if queue is not UNSET:
             field_dict["queue"] = queue
         if external_id is not UNSET:
@@ -167,6 +174,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         d = dict(src_dict)
         name = d.pop("name")
 
+        id = d.pop("id", UNSET)
+
         queue = d.pop("queue", UNSET)
 
         external_id = d.pop("external_id", UNSET)
@@ -261,6 +270,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
 
         task_request = cls(
             name=name,
+            id=id,
             queue=queue,
             external_id=external_id,
             status=status,

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
             )