Merge pull request #52 from taskbadger/sk/queue

Add queue field support

Author: snopoke

taskbadger.yaml

@@ -690,6 +690,10 @@ components:
           minLength: 1
           description: Name of the task
           maxLength: 255
+        queue:
+          type: string
+          description: Queue the task is from
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'
@@ -786,6 +790,10 @@ components:
           type: string
           description: Name of the task
           maxLength: 255
+        queue:
+          type: string
+          description: Queue the task is from
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'
@@ -881,6 +889,10 @@ components:
           minLength: 1
           description: Name of the task
           maxLength: 255
+        queue:
+          type: string
+          description: Queue the task is from
+          maxLength: 255
         status:
           allOf:
           - $ref: '#/components/schemas/StatusEnum'

taskbadger/celery.py

@@ -122,6 +122,7 @@ def taskbadger_task(self):
 
 @before_task_publish.connect
 def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
+    routing_key = kwargs.get("routing_key")
     headers = headers if "task" in headers else body
     header_kwargs = headers.pop(TB_KWARGS_ARG, {})  # always remove TB headers
     if sender.startswith("celery.") or not Badger.is_configured():
@@ -144,6 +145,8 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     # get kwargs from the task headers (set via apply_async)
     kwargs.update(header_kwargs)
     kwargs["status"] = StatusEnum.PENDING
+    if routing_key and "queue" not in kwargs:
+        kwargs["queue"] = routing_key
     name = kwargs.pop("name", headers["task"])
 
     global_record_task_args = celery_system and celery_system.record_task_args
@@ -242,7 +245,9 @@ def _maybe_create_task(signal_sender):
 
     enter_session()
 
-    task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data)
+    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)
     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/cli/basics.py

@@ -59,6 +59,7 @@ def create(
     ),
     status: StatusEnum = typer.Option(StatusEnum.PROCESSING, help="The initial status of the task."),
     value_max: int = typer.Option(100, help="The maximum value for the task."),
+    queue: str = typer.Option(None, show_default=False, help="The name of the queue the task is from."),
     metadata: list[str] = typer.Option(
         None,
         show_default=False,
@@ -96,6 +97,7 @@ def create(
             actions=actions,
             monitor_id=monitor_id,
             tags=tags,
+            queue=queue,
         )
     except Exception as e:
         err_console.print(f"Error creating task: {e}")
@@ -121,6 +123,7 @@ def update(
     status: StatusEnum = typer.Option(StatusEnum.PROCESSING, help="The status of the task."),
     value: int = typer.Option(None, show_default=False, help="The current task value (progress)."),
     value_max: int = typer.Option(None, show_default=False, help="The maximum value for the task."),
+    queue: str = typer.Option(None, show_default=False, help="The name of the queue the task is from."),
     metadata: list[str] = typer.Option(
         None,
         show_default=False,
@@ -159,6 +162,7 @@ def update(
             data=metadata,
             actions=actions,
             tags=tags,
+            queue=queue,
         )
     except Exception as e:
         err_console.print(f"Error creating task: {e}")

taskbadger/internal/models/patched_task_request.py

@@ -23,6 +23,7 @@ class PatchedTaskRequest:
     """
     Attributes:
         name (str | Unset): Name of the task
+        queue (str | Unset): Queue the task is from
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -48,6 +49,7 @@ class PatchedTaskRequest:
     """
 
     name: str | Unset = UNSET
+    queue: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -65,6 +67,8 @@ def to_dict(self) -> dict[str, Any]:
 
         name = self.name
 
+        queue = self.queue
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -122,6 +126,8 @@ def to_dict(self) -> dict[str, Any]:
         field_dict.update({})
         if name is not UNSET:
             field_dict["name"] = name
+        if queue is not UNSET:
+            field_dict["queue"] = queue
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -152,6 +158,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         d = dict(src_dict)
         name = d.pop("name", UNSET)
 
+        queue = d.pop("queue", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -242,6 +250,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
 
         patched_task_request = cls(
             name=name,
+            queue=queue,
             status=status,
             value=value,
             value_max=value_max,

taskbadger/internal/models/task.py

@@ -31,6 +31,7 @@ class Task:
         updated (datetime.datetime):
         url (str):
         public_url (str):
+        queue (str | Unset): Queue the task is from
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -64,6 +65,7 @@ class Task:
     updated: datetime.datetime
     url: str
     public_url: str
+    queue: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -98,6 +100,8 @@ def to_dict(self) -> dict[str, Any]:
 
         public_url = self.public_url
 
+        queue = self.queue
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -165,6 +169,8 @@ def to_dict(self) -> dict[str, Any]:
                 "public_url": public_url,
             }
         )
+        if queue is not UNSET:
+            field_dict["queue"] = queue
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -216,6 +222,8 @@ def _parse_value_percent(data: object) -> int | None:
 
         public_url = d.pop("public_url")
 
+        queue = d.pop("queue", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -314,6 +322,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
             updated=updated,
             url=url,
             public_url=public_url,
+            queue=queue,
             status=status,
             value=value,
             value_max=value_max,

taskbadger/internal/models/task_request.py

@@ -23,6 +23,7 @@ class TaskRequest:
     """
     Attributes:
         name (str): Name of the task
+        queue (str | Unset): Queue the task is from
         status (StatusEnum | Unset): * `pending` - pending
             * `pre_processing` - pre_processing
             * `processing` - processing
@@ -48,6 +49,7 @@ class TaskRequest:
     """
 
     name: str
+    queue: str | Unset = UNSET
     status: StatusEnum | Unset = StatusEnum.PENDING
     value: int | None | Unset = UNSET
     value_max: int | Unset = UNSET
@@ -65,6 +67,8 @@ def to_dict(self) -> dict[str, Any]:
 
         name = self.name
 
+        queue = self.queue
+
         status: str | Unset = UNSET
         if not isinstance(self.status, Unset):
             status = self.status.value
@@ -124,6 +128,8 @@ def to_dict(self) -> dict[str, Any]:
                 "name": name,
             }
         )
+        if queue is not UNSET:
+            field_dict["queue"] = queue
         if status is not UNSET:
             field_dict["status"] = status
         if value is not UNSET:
@@ -154,6 +160,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         d = dict(src_dict)
         name = d.pop("name")
 
+        queue = d.pop("queue", UNSET)
+
         _status = d.pop("status", UNSET)
         status: StatusEnum | Unset
         if isinstance(_status, Unset):
@@ -244,6 +252,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset:
 
         task_request = cls(
             name=name,
+            queue=queue,
             status=status,
             value=value,
             value_max=value_max,

taskbadger/procrastinate.py

@@ -160,13 +160,14 @@ async def defer_async(**kwargs):
     task.defer_async = defer_async
 
 
-def _create_pending_task(task, task_kwargs):
+def _create_pending_task(task, task_kwargs, queue=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.
+    ``record_task_args`` data capture. ``queue`` overrides the queue name
+    recorded on the TaskBadger task (defaults to the task's own queue).
     """
     if not Badger.is_configured():
         return None
@@ -180,6 +181,9 @@ def _create_pending_task(task, task_kwargs):
     opts = dict(getattr(task, _OPTS_ATTR, {}) or {})
     name = opts.pop("name", None) or task.name
     create_kwargs = {"status": StatusEnum.PENDING}
+    queue = queue or getattr(task, "queue", None)
+    if queue is not None:
+        create_kwargs["queue"] = queue
     for key in ("value_max", "tags"):
         if key in opts and opts[key] is not None:
             create_kwargs[key] = opts[key]
@@ -293,7 +297,7 @@ def _patch_job_manager(app, system):
         async def patched(*, job, periodic_id, defer_timestamp):
             task = app.tasks.get(job.task_name)
             if task is not None:
-                tb_task = _create_pending_task(task, job.task_kwargs)
+                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)

taskbadger/sdk.py

@@ -151,6 +151,7 @@ def create_task(
     actions: list[Action] = None,
     monitor_id: str = None,
     tags: dict[str, str] = None,
+    queue: str = None,
 ) -> "Task":
     """Create a Task.
 
@@ -165,6 +166,7 @@ def create_task(
         actions: Task actions.
         monitor_id: ID of the monitor to associate this task with.
         tags: Dictionary of namespace -> value tags.
+        queue: Name of the queue the task is from.
 
     Returns:
         Task: The created Task object.
@@ -173,6 +175,8 @@ def create_task(
         "name": name,
         "status": status,
     }
+    if queue is not None:
+        task_dict["queue"] = queue
     if value is not None:
         task_dict["value"] = value
     if value_max is not None:
@@ -216,6 +220,7 @@ def update_task(
     stale_timeout: int = None,
     actions: list[Action] = None,
     tags: dict[str, str] = None,
+    queue: str = None,
 ) -> "Task":
     """Update a task.
     Requires only the task ID and fields to update.
@@ -231,6 +236,7 @@ def update_task(
         stale_timeout: Maximum allowed time between updates (seconds).
         actions: Task actions.
         tags: Dictionary of namespace -> value tags.
+        queue: Name of the queue the task is from.
 
     Returns:
         Task: The updated Task object.
@@ -242,6 +248,7 @@ def update_task(
     data = _none_to_unset(data)
     max_runtime = _none_to_unset(max_runtime)
     stale_timeout = _none_to_unset(stale_timeout)
+    queue = _none_to_unset(queue)
 
     data = data or UNSET
     body = PatchedTaskRequest(
@@ -252,6 +259,7 @@ def update_task(
         data=data,
         max_runtime=max_runtime,
         stale_timeout=stale_timeout,
+        queue=queue,
     )
     if actions:
         body.additional_properties = {"actions": [a.to_dict() for a in actions]}
@@ -314,6 +322,7 @@ def create(
         actions: list[Action] = None,
         monitor_id: str = None,
         tags: dict[str, str] = None,
+        queue: str = None,
     ) -> "Task":
         """Create a new task
 
@@ -330,6 +339,7 @@ def create(
             actions=actions,
             monitor_id=monitor_id,
             tags=tags,
+            queue=queue,
         )
 
     def __init__(self, task):
@@ -414,6 +424,7 @@ def update(
         stale_timeout: int = None,
         actions: list[Action] = None,
         tags: dict[str, str] = None,
+        queue: str = None,
         data_merge_strategy: Any = None,
     ):
         """Generic update method used to update any of the task fields.
@@ -441,6 +452,7 @@ def update(
             stale_timeout=stale_timeout,
             actions=actions,
             tags=tags,
+            queue=queue,
         )
         self._task = task._task