Merge remote-tracking branch 'origin/main' into sk/task-id

# Conflicts:
#	taskbadger/sdk.py

Author: snopoke

taskbadger.yaml

@@ -699,6 +699,11 @@ components:
           type: string
           description: Queue the task is from
           maxLength: 255
+        external_id:
+          type: string
+          description: Identifier from the originating system (e.g. Celery task ID)
+            for correlating with logs
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'
@@ -799,6 +804,11 @@ components:
           type: string
           description: Queue the task is from
           maxLength: 255
+        external_id:
+          type: string
+          description: Identifier from the originating system (e.g. Celery task ID)
+            for correlating with logs
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'
@@ -902,6 +912,11 @@ components:
           type: string
           description: Queue the task is from
           maxLength: 255
+        external_id:
+          type: string
+          description: Identifier from the originating system (e.g. Celery task ID)
+            for correlating with logs
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'

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/internal/models/patched_task_request.py

@@ -26,6 +26,8 @@ class PatchedTaskRequest:
             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
+            logs
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -53,6 +55,7 @@ class PatchedTaskRequest:
     id: str | Unset = UNSET
     name: str | Unset = UNSET
     queue: str | Unset = UNSET
+    external_id: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -74,6 +77,8 @@ def to_dict(self) -> dict[str, Any]:
 
         queue = self.queue
 
+        external_id = self.external_id
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -135,6 +140,8 @@ def to_dict(self) -> dict[str, Any]:
             field_dict["name"] = name
         if queue is not UNSET:
             field_dict["queue"] = queue
+        if external_id is not UNSET:
+            field_dict["external_id"] = external_id
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -169,6 +176,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
 
         queue = d.pop("queue", UNSET)
 
+        external_id = d.pop("external_id", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -261,6 +270,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             id=id,
             name=name,
             queue=queue,
+            external_id=external_id,
             status=status,
             value=value,
             value_max=value_max,

taskbadger/internal/models/task.py

@@ -33,6 +33,8 @@ class 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
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -67,6 +69,7 @@ class Task:
     public_url: str
     id: str | Unset = UNSET
     queue: str | Unset = UNSET
+    external_id: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -103,6 +106,8 @@ def to_dict(self) -> dict[str, Any]:
 
         queue = self.queue
 
+        external_id = self.external_id
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -173,6 +178,8 @@ def to_dict(self) -> dict[str, Any]:
             field_dict["id"] = id
         if queue is not UNSET:
             field_dict["queue"] = queue
+        if external_id is not UNSET:
+            field_dict["external_id"] = external_id
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -226,6 +233,8 @@ def _parse_value_percent(data: object) -> int | None:
 
         queue = d.pop("queue", UNSET)
 
+        external_id = d.pop("external_id", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -325,6 +334,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             public_url=public_url,
             id=id,
             queue=queue,
+            external_id=external_id,
             status=status,
             value=value,
             value_max=value_max,

taskbadger/internal/models/task_request.py

@@ -26,6 +26,8 @@ class TaskRequest:
         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
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -53,6 +55,7 @@ class TaskRequest:
     name: str
     id: str | Unset = UNSET
     queue: str | Unset = UNSET
+    external_id: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -74,6 +77,8 @@ def to_dict(self) -> dict[str, Any]:
 
         queue = self.queue
 
+        external_id = self.external_id
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -137,6 +142,8 @@ def to_dict(self) -> dict[str, Any]:
             field_dict["id"] = id
         if queue is not UNSET:
             field_dict["queue"] = queue
+        if external_id is not UNSET:
+            field_dict["external_id"] = external_id
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -171,6 +178,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
 
         queue = d.pop("queue", UNSET)
 
+        external_id = d.pop("external_id", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -263,6 +272,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             name=name,
             id=id,
             queue=queue,
+            external_id=external_id,
             status=status,
             value=value,
             value_max=value_max,

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
 

taskbadger/sdk.py

@@ -153,6 +153,7 @@ def create_task(
     tags: dict[str, str] = None,
     queue: str = None,
     task_id: str = None,
+    external_id: str = None,
 ) -> "Task":
     """Create a Task.
 
@@ -170,6 +171,7 @@ def create_task(
         queue: Name of the queue the task is from.
         task_id: ID to assign to the task. May be a UUID or shortened UUID; it must be
             unique and is immutable. If omitted, an ID is generated by the server.
+        external_id: Identifier from the originating system (e.g. Celery task ID) for correlating with logs.
 
     Returns:
         Task: The created Task object.
@@ -182,6 +184,8 @@ def create_task(
         task_dict["id"] = task_id
     if queue is not None:
         task_dict["queue"] = queue
+    if external_id is not None:
+        task_dict["external_id"] = external_id
     if value is not None:
         task_dict["value"] = value
     if value_max is not None:
@@ -227,6 +231,7 @@ def update_task(
     actions: list[Action] = None,
     tags: dict[str, str] = None,
     queue: str = None,
+    external_id: str = None,
 ) -> "Task":
     """Update a task.
     Requires only the task ID and fields to update.
@@ -243,6 +248,7 @@ def update_task(
         actions: Task actions. **Deprecated:** use project-level actions instead.
         tags: Dictionary of namespace -> value tags.
         queue: Name of the queue the task is from.
+        external_id: Identifier from the originating system (e.g. Celery task ID) for correlating with logs.
 
     Returns:
         Task: The updated Task object.
@@ -255,6 +261,7 @@ def update_task(
     max_runtime = _none_to_unset(max_runtime)
     stale_timeout = _none_to_unset(stale_timeout)
     queue = _none_to_unset(queue)
+    external_id = _none_to_unset(external_id)
 
     data = data or UNSET
     body = PatchedTaskRequest(
@@ -266,6 +273,7 @@ def update_task(
         max_runtime=max_runtime,
         stale_timeout=stale_timeout,
         queue=queue,
+        external_id=external_id,
     )
     if actions:
         _warn_actions_deprecated()
@@ -341,6 +349,7 @@ def create(
         tags: dict[str, str] = None,
         queue: str = None,
         task_id: str = None,
+        external_id: str = None,
     ) -> "Task":
         """Create a new task
 
@@ -359,6 +368,7 @@ def create(
             tags=tags,
             queue=queue,
             task_id=task_id,
+            external_id=external_id,
         )
 
     def __init__(self, task):
@@ -444,6 +454,7 @@ def update(
         actions: list[Action] = None,
         tags: dict[str, str] = None,
         queue: str = None,
+        external_id: str = None,
         data_merge_strategy: Any = None,
     ):
         """Generic update method used to update any of the task fields.
@@ -472,6 +483,7 @@ def update(
             actions=actions,
             tags=tags,
             queue=queue,
+            external_id=external_id,
         )
         self._task = task._task