From cd84b21e4dca2708b804a990c529ff168113698e Mon Sep 17 00:00:00 2001 From: Simon Kelly Date: Wed, 8 Jul 2026 13:50:57 +0200 Subject: [PATCH] Add queue field support Set automatically from routing key (Celery) and task queue (Procrastinate). Co-Authored-By: Claude Opus 4.8 (1M context) --- taskbadger.yaml | 12 ++++++++ taskbadger/celery.py | 7 ++++- taskbadger/cli/basics.py | 4 +++ .../internal/models/patched_task_request.py | 9 ++++++ taskbadger/internal/models/task.py | 9 ++++++ taskbadger/internal/models/task_request.py | 9 ++++++ taskbadger/procrastinate.py | 10 +++++-- taskbadger/sdk.py | 12 ++++++++ tests/test_celery.py | 30 +++++++++++++++++-- tests/test_celery_system_integration.py | 4 ++- tests/test_procrastinate.py | 18 ++++++++++- tests/test_sdk.py | 25 ++++++++++++++++ 12 files changed, 140 insertions(+), 9 deletions(-) diff --git a/taskbadger.yaml b/taskbadger.yaml index d2a5aa4..c9a56e7 100644 --- a/taskbadger.yaml +++ b/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' diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 2f986fa..b0508d3 100644 --- a/taskbadger/celery.py +++ b/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}) diff --git a/taskbadger/cli/basics.py b/taskbadger/cli/basics.py index bbec648..d803cd8 100644 --- a/taskbadger/cli/basics.py +++ b/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}") diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index 6dbff70..e455f70 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/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, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 0927ce5..632f325 100644 --- a/taskbadger/internal/models/task.py +++ b/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, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index 00c6a51..35198ab 100644 --- a/taskbadger/internal/models/task_request.py +++ b/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, diff --git a/taskbadger/procrastinate.py b/taskbadger/procrastinate.py index fee14c2..7448e27 100644 --- a/taskbadger/procrastinate.py +++ b/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) diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index be1130b..d78b89b 100644 --- a/taskbadger/sdk.py +++ b/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 diff --git a/tests/test_celery.py b/tests/test_celery.py index 88d8c92..cc3ec9a 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -84,7 +84,7 @@ 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 + "new_name", value_max=10, data={"foo": "bar"}, tags={"bar": "baz"}, status=StatusEnum.PENDING, queue="celery" ) @@ -114,7 +114,7 @@ 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) + create.assert_called_once_with("new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery") @pytest.mark.usefixtures("_bind_settings") @@ -147,6 +147,7 @@ def add_with_task_args(self, a, b): value_max=10, data={"foo": "bar", "celery_task_args": [2, 2], "celery_task_kwargs": {}}, status=StatusEnum.PENDING, + queue="celery", ) @@ -184,6 +185,7 @@ def add_with_task_kwargs(self, a, b, c=0): data={"celery_task_args": [2, 2], "celery_task_kwargs": {"c": 3}}, actions=actions, status=StatusEnum.PENDING, + queue="celery", ) @@ -220,6 +222,7 @@ def add_task_custom_serialization(self, a): "tests.test_celery.add_task_custom_serialization", data={"celery_task_args": [{"__type__": "A", "__value__": [2, 2]}], "celery_task_kwargs": {}}, status=StatusEnum.PENDING, + queue="celery", ) @@ -247,7 +250,28 @@ 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) + create.assert_called_once_with(mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery") + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_task_custom_queue(celery_session_app, celery_session_worker): + @celery_session_app.task(bind=True, base=Task) + def add_custom_queue(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() + # The task is routed to a queue the session worker doesn't consume, so we + # only assert on the create call made synchronously during publish. + add_custom_queue.apply_async((2, 2), queue="high_priority") + + assert create.call_args.kwargs["queue"] == "high_priority" @pytest.mark.usefixtures("_bind_settings") diff --git a/tests/test_celery_system_integration.py b/tests/test_celery_system_integration.py index 035d77d..4944a20 100644 --- a/tests/test_celery_system_integration.py +++ b/tests/test_celery_system_integration.py @@ -122,6 +122,7 @@ def add_normal(self, a, b): "tests.test_celery_system_integration.add_normal", status=StatusEnum.PENDING, data={"celery_task_args": [2, 2], "celery_task_kwargs": {}}, + queue="celery", ) assert get_task.call_count == 1 assert update.call_count == 2 @@ -153,7 +154,7 @@ 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 + "tests.test_celery_system_integration.add_normal_with_override", status=StatusEnum.PENDING, queue="celery" ) @@ -187,6 +188,7 @@ def add_with_tags(a, b): name="tests.test_celery_system_integration.add_with_tags", status=StatusEnum.PENDING, tags=TaskRequestTags.from_dict({"tag1": "value1", "tag2": "override"}), + queue="celery", ) create.assert_called_with( client=mock.ANY, diff --git a/tests/test_procrastinate.py b/tests/test_procrastinate.py index 1c110cd..ca9223a 100644 --- a/tests/test_procrastinate.py +++ b/tests/test_procrastinate.py @@ -120,7 +120,7 @@ def add3(a, b): create.assert_called_once() assert create.call_args.args == ("add3",) - assert create.call_args.kwargs == {"status": StatusEnum.PENDING} + assert create.call_args.kwargs == {"status": StatusEnum.PENDING, "queue": "default"} # The injected id should appear in the Procrastinate job's task kwargs. jobs = list(app.connector.jobs.values()) @@ -128,6 +128,21 @@ def add3(a, b): assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id +@pytest.mark.usefixtures("_bind_settings") +def test_defer_sets_queue(app): + @app.task(name="add_queued", queue="high_priority") + def add_queued(a, b): + return 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: + add_queued.defer(a=1, b=2) + + assert create.call_args.kwargs["queue"] == "high_priority" + + def test_defer_no_taskbadger_when_unconfigured(app): @app.task(name="add4") def add4(a, b): @@ -218,6 +233,7 @@ def raw(a): "value_max": 10, "tags": {"env": "test"}, "data": {"k": "v"}, + "queue": "default", } diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 728a55c..343ade0 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -99,6 +99,31 @@ def test_create(settings, patched_create): ) +def test_create_with_queue(settings, patched_create): + api_task = task_for_test() + patched_create.return_value = Response(HTTPStatus.OK, b"", {}, api_task) + + Task.create(name="task name", queue="high_priority") + + request = TaskRequest(name="task name", status=StatusEnum.PENDING, queue="high_priority") + patched_create.assert_called_with( + client=mock.ANY, + organization_slug="org", + project_slug="project", + body=request, + ) + + +def test_update_queue(settings, patched_update): + api_task = task_for_test() + task = Task(api_task) + + patched_update.return_value = Response(HTTPStatus.OK, b"", {}, api_task) + task.update(queue="high_priority") + + _verify_update(settings, patched_update, queue="high_priority") + + def test_before_create_update_task(settings, patched_create): def before_create(task): tags = task.setdefault("tags", {})