diff --git a/CHANGELOG.md b/CHANGELOG.md index 60d4b1f0..97bc01b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,30 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +* `JobsClient.cancel_where(filter, reason)` — bulk cancel all jobs matching a + `JobFilter` in a single set-based operation. Pending/scheduled jobs go straight + to terminal `cancelled`; running jobs get cooperative cancel (`cancel_phase=1`). + Returns `BulkCancelResult` with counts and affected IDs. Empty filters are + rejected with `EmptyFilterError` unless `allow_empty_filter=True` is passed. +* `SubJobEnqueuer.enqueue()` now accepts `tags`, `inherit_tags`, + `schedule_to_close`, `start_to_close`, and `heartbeat_timeout` parameters. + Sub-jobs inherit the parent job's tags by default (`inherit_tags=True`); pass + `inherit_tags=False` to suppress inheritance for a specific sub-job. +* `BulkCancelResult` and `EmptyFilterError` exported from `taskq` top-level. + +### Changed + +* **Sub-jobs now inherit parent tags by default.** Every `ctx.jobs.enqueue()` + call inside an actor body now propagates the parent job's tags to the sub-job, + making sub-jobs findable by `JobFilter(tags=...)` and cancellable by + `cancel_where`. Pass `inherit_tags=False` per-call to opt out. This is a + behavior change for any code that relied on sub-job tags being empty — + inherited tags make sub-jobs visible to tag-based filters and bulk cancels. + ## [0.2.2](https://github.com/AZX-PBC-OSS/TaskQ/compare/v0.2.1...v0.2.2) (2026-07-22) diff --git a/README.md b/README.md index dbce89ae..97425f6c 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,9 @@ Async-native, Postgres-backed background job library for Python 3.12+. isn't inside an actor should instead poll `BatchHandle.status(db_connection)`. See [Jobs & Clients](docs/guides/jobs-clients.md#enqueue_batch). - **Cancellation** — cooperative cancellation with grace periods and - force-cancel sweeps; `ctx.check_cancelled()` inside actor bodies. + force-cancel sweeps; `ctx.check_cancelled()` inside actor bodies; bulk + `cancel_where(filter)` for set-based cancellation by tag, queue, actor, + or batch ID. - **Progress tracking** — `ctx.progress(...)` events buffered and published to subscribers and the admin UI. - **Workgroups** — multi-worker process supervision with a shared heartbeat diff --git a/docs/architecture.md b/docs/architecture.md index b3522e50..2e17b4a3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -102,6 +102,7 @@ class Backend(Protocol): # Cancel signals async def write_cancel_request(self, job_id, reason) -> bool: ... + async def cancel_where(self, filter: JobFilter, reason: str | None) -> BulkCancelResult: ... async def poll_cancel_flags(self, worker_id) -> list[CancelFlag]: ... # Admin operations diff --git a/docs/guides/actors.md b/docs/guides/actors.md index 2d1705e5..653fb203 100644 --- a/docs/guides/actors.md +++ b/docs/guides/actors.md @@ -689,6 +689,13 @@ INSERTs are part of the parent's database transaction: This is the correct default for fan-out patterns where sub-jobs should only exist if the parent completes successfully. +**Tag inheritance.** Sub-jobs enqueued via `ctx.jobs.enqueue()` inherit the parent +job's tags by default. This makes sub-jobs findable by `JobFilter(tags=...)` and +cancellable by `cancel_where`. Pass `inherit_tags=False` to suppress inheritance +for a specific sub-job, or pass explicit `tags=[...]` to merge with inherited +tags. See [Jobs & Clients — Sub-job tag inheritance](jobs-clients.md#tag-inheritance) +for the full semantics table. + **Autonomous fallback.** If no LOOP-scope `asyncpg.Connection` is registered in the DI container, `ctx.jobs.enqueue()` falls back to the worker pool and commits each INSERT independently. In this mode, sub-jobs are persisted even if the parent subsequently raises diff --git a/docs/guides/cancellation.md b/docs/guides/cancellation.md index 48ca5557..d7e53e7e 100644 --- a/docs/guides/cancellation.md +++ b/docs/guides/cancellation.md @@ -30,6 +30,22 @@ result = await handle.cancel(reason="deadline exceeded") `JobHandle.cancel()` delegates directly to `JobsClient.cancel(handle.job_id, reason)`. The handle must have been constructed with a `JobsClient` (i.e. via `client.enqueue()` or `client.get()`); handles obtained from inside an actor body via `ctx.jobs.enqueue()` do not have a client and will raise `RuntimeError`. +### Via `JobsClient.cancel_where()` + +```python +from taskq import JobFilter + +result = await client.cancel_where( + JobFilter(tags=("tenant-acme",), active=True), + reason="tenant offboarded", +) +``` + +`cancel_where()` cancels all jobs matching a `JobFilter` in a single set-based SQL +operation. Pending/scheduled jobs go straight to terminal `cancelled`; running jobs get +`cancel_phase=1` (cooperative cancel). Returns a `BulkCancelResult` with counts and +affected IDs. See [jobs-clients.md](jobs-clients.md#cancel_where) for the full API. + ### Effect by prior status | Prior status | Effect | diff --git a/docs/guides/jobs-clients.md b/docs/guides/jobs-clients.md index 30ddf59e..0caf9b9b 100644 --- a/docs/guides/jobs-clients.md +++ b/docs/guides/jobs-clients.md @@ -52,13 +52,14 @@ Terminal statuses (`succeeded`, `failed`, `cancelled`, `crashed`, `abandoned`) h 6. [`JobHandle[R]`](#jobhandler) 7. [`get()`](#get) 8. [`cancel()`](#cancel) -9. [`list()`](#list) -10. [`SubJobEnqueuer`](#subjobenqueuer) -11. [Error handling](#error-handling) -12. [Full enqueue-and-wait example](#full-enqueue-and-wait-example) -13. [Idempotency example](#idempotency-example) -14. [Batch enqueue example](#batch-enqueue-example) -15. [Tags](#tags) +9. [`cancel_where()`](#cancel_where) +10. [`list()`](#list) +11. [`SubJobEnqueuer`](#subjobenqueuer) +12. [Error handling](#error-handling) +13. [Full enqueue-and-wait example](#full-enqueue-and-wait-example) +14. [Idempotency example](#idempotency-example) +15. [Batch enqueue example](#batch-enqueue-example) +16. [Tags](#tags) --- @@ -636,6 +637,63 @@ else: --- +## `cancel_where()` + +```python +async def cancel_where( + self, + filter: JobFilter, + reason: str | None = None, + *, + allow_empty_filter: bool = False, +) -> BulkCancelResult: ... +``` + +Cancel all jobs matching `filter` in a single set-based operation. Pending/scheduled +jobs go straight to terminal `cancelled`; running jobs get `cancel_phase=1` (cooperative +cancel) — the worker's heartbeat observes the phase change and sets the in-process +`cancel_event` at the next tick. + +**Guardrail:** a filter with no predicates (no `queue`, `status`, `actor`, +`identity_key`, `batch_id`, `tags`, or `active`) is rejected with `EmptyFilterError` +unless `allow_empty_filter=True` is passed. This prevents accidental full-table cancels. + +**Filter fields used:** `queue`, `status`, `actor`, `identity_key`, `batch_id`, `tags`, +`active`. The `limit`, `cursor`, and `order_by` fields are ignored — a bulk cancel is +not paginated. + +**Snapshot boundary:** jobs matching the filter that are enqueued *after* the +statement's snapshot escape this call. Stop producers first, or issue a follow-up call; +the returned counts make non-convergence detectable. + +```python +result = await client.cancel_where( + JobFilter(tags=("tenant-acme",), active=True), + reason="tenant offboarded", +) +print(f"cancelled {result.cancelled_directly} pending, " + f"requested cancel for {result.cancel_requested} running") +``` + +### `BulkCancelResult` + +Frozen Pydantic model returned by `cancel_where()`. + +| Field | Type | Description | +|---|---|---| +| `cancelled_directly` | `int` | Pending/scheduled jobs moved straight to terminal `cancelled`. | +| `cancel_requested` | `int` | Running jobs with `cancel_phase=1` set (cooperative cancel). | +| `cancelled_ids` | `list[UUID]` | IDs of jobs cancelled directly. | +| `cancel_requested_ids` | `list[UUID]` | IDs of running jobs with cancel requested. | +| `total_affected` | `int` (property) | `cancelled_directly + cancel_requested`. | + +For tenant-scale cancels (10^5+ matching rows), partition via filter (e.g. +`JobFilter(queue=..., tags=...)` to split by queue). A single transaction covering 10^6 +rows would hold locks too long. The counts let the caller verify completeness and issue +follow-up calls for remaining partitions. + +--- + ## `list()` ```python @@ -781,17 +839,52 @@ async def enqueue( unique_for: timedelta | None = None, unique_states: tuple[JobStatus, ...] | None = None, max_pending: int | None = None, + tags: list[str] | None = None, + inherit_tags: bool = True, + schedule_to_close: datetime | None = None, + start_to_close: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, ) -> JobHandle[R]: ... ``` Enqueues a single sub-job. Accepts the same options as `JobsClient.enqueue()` except: - No `queue` override (sub-jobs use the actor's declared queue). -- No `schedule_to_close`, `start_to_close`, or `heartbeat_timeout` (set on the actor declaration). - No explicit `trace_id` / `span_id` (extracted from the active OTel span). - `connection` may be passed to use a specific `asyncpg.Connection` rather than the LOOP-scope connection. +#### Tag inheritance + +Sub-jobs inherit the parent job's tags by default (`inherit_tags=True`). When no +explicit `tags` are passed, the sub-job carries the parent's tags. When explicit +`tags` are passed, they merge with the parent's tags (parent first, deduplicated). + +| `inherit_tags` | `tags` | Resulting job tags | +|---|---|---| +| `True` (default) | `None` | Parent job's tags (or `()` if parent has none) | +| `True` (default) | `["new-tag"]` | Parent tags + explicit tags, merged (union, parent-first, deduped) | +| `False` | `None` | `()` (no inheritance) | +| `False` | `["new-tag"]` | `("new-tag",)` (explicit only) | + +Pass `inherit_tags=False` to suppress inheritance for a specific sub-job. + +**Blast radius:** inherited tags make sub-jobs visible to `cancel_where` filters +matching those tags. A shared/utility sub-job enqueued by a tenant-tagged parent +will be swept up in that tenant's `cancel_where`. + +`enqueue_batch()` does **not** inherit parent tags — batch items carry their own +`EnqueueItem.tags`. This asymmetry is deliberate: batch fan-out callers typically +set per-item tags explicitly. + +#### `schedule_to_close` / `start_to_close` / `heartbeat_timeout` + +`schedule_to_close` and `start_to_close` override the actor's declared defaults for +this specific sub-job. `heartbeat_timeout` has no actor-level declaration — the +per-call value is the only source. Note that `schedule_to_close` bounds total +wall-clock time *including* time snoozed on `wait_for_batch` — finalizer-style +sub-jobs that snooze for long periods should set it generously or not at all. + The per-call `max_pending=` argument is resolved against the operator-owned stored cap and the `@actor(...)` literal, not in place of them: against a non-NULL stored `actor_config.max_pending` the tighter of the two wins (`min(stored, per_call)`) — explicit @@ -1014,7 +1107,7 @@ Tags are user-defined keyword labels stored in `jobs.tags text[]`. They have no handle = await client.enqueue( send_email, EmailPayload(to="user@example.com"), - tags=["notification", "priority:high", "tenant:acme"], + tags=["notification", "priority-high", "tenant-acme"], ) ``` @@ -1046,7 +1139,7 @@ Use `JobFilter.tags` with array-overlap semantics (matches jobs that have **any* page = await client.list(JobFilter( actor="send_email", status="failed", - tags=["priority:high", "tenant:acme"], + tags=["priority-high", "tenant-acme"], limit=50, )) ``` diff --git a/pyproject.toml b/pyproject.toml index db397dc6..7d61d1ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -202,6 +202,11 @@ ignore = [ "src/taskq/backend/_terminal.py" = ["S608", "SIM117"] "src/taskq/backend/_enqueue.py" = ["S608", "SIM117"] "src/taskq/backend/_reads.py" = ["S608"] +# S608: schema name validated against _IDENT_RE before interpolation; same +# rationale as _reads.py. All user-supplied values use $N parameter binding. +# SIM117: pool.acquire() and conn.transaction() cannot be flattened. +# S311: random module used for jitter in deadlock retry backoff, not crypto. +"src/taskq/backend/_cancel_bulk.py" = ["S608", "SIM117", "S311"] "src/taskq/backend/_dispatch.py" = ["S608", "SIM117"] "src/taskq/backend/_sql_templates.py" = ["S608"] # SIM117: pool.acquire() and conn.transaction() cannot be flattened into a @@ -257,6 +262,7 @@ ignore = [ # Same rationale — schema name validated via PostgresBackend constructor; # all user-supplied values use $N parameter binding. "tests/test_backend_equivalence.py" = ["S608"] +"tests/test_cancel_where_pg.py" = ["S608"] # Same rationale — schema name validated against _IDENT_RE; all # user-supplied values use $N parameter binding. "tests/test_heartbeat_integration.py" = ["S608"] diff --git a/src/taskq/__init__.py b/src/taskq/__init__.py index e18c8f1f..35fbdcb8 100644 --- a/src/taskq/__init__.py +++ b/src/taskq/__init__.py @@ -37,7 +37,7 @@ ScheduleRecord, ) from taskq.batch import BatchCompletionStatus, BatchHandle, EnqueueItem, wait_for_batch -from taskq.client import CancelResult, JobEvent, JobHandle, JobsClient, TaskQ +from taskq.client import BulkCancelResult, CancelResult, JobEvent, JobHandle, JobsClient, TaskQ from taskq.client._enqueuer import SubJobEnqueuer from taskq.connections import ConnFactory, PoolFactory, RedisFactory, WorkerConnections from taskq.context import JobContext @@ -48,6 +48,7 @@ BackpressureError, DependencyCycle, DIError, + EmptyFilterError, IllegalStateTransition, JobFailed, MaxPendingExceededError, @@ -93,6 +94,7 @@ "BackpressureError", "BatchCompletionStatus", "BatchHandle", + "BulkCancelResult", "CancelPhase", "CancelResult", "ConnFactory", @@ -100,6 +102,7 @@ "DIError", "DependencyCycle", "DstStrategy", + "EmptyFilterError", "EnqueueItem", "ErrorReporter", "EventRow", diff --git a/src/taskq/backend/__init__.py b/src/taskq/backend/__init__.py index 9598d32f..e436a030 100644 --- a/src/taskq/backend/__init__.py +++ b/src/taskq/backend/__init__.py @@ -20,6 +20,7 @@ AttemptRow, Backend, BackendDeps, + BulkCancelResult, CancelFlag, DstStrategy, EnqueueArgs, @@ -57,6 +58,7 @@ def __getattr__(name: str) -> object: "AttemptRow", "Backend", "BackendDeps", + "BulkCancelResult", "CancelFlag", "DstStrategy", "EnqueueArgs", diff --git a/src/taskq/backend/_cancel_bulk.py b/src/taskq/backend/_cancel_bulk.py new file mode 100644 index 00000000..29f59839 --- /dev/null +++ b/src/taskq/backend/_cancel_bulk.py @@ -0,0 +1,192 @@ +"""Bulk cancel SQL implementation for PostgresBackend. + +Two-statement pattern within a single transaction, mirroring the +single-job ``write_cancel_request`` path: + +1. ``cancel_pending_scheduled`` — UPDATE pending/scheduled rows to + terminal ``cancelled`` with EPQ-safe predicates on the target table. +2. ``cancel_running`` — UPDATE running rows with ``cancel_phase=0`` to + ``cancel_phase=1`` (cooperative cancel), using a fresh snapshot that + catches jobs dispatched between statements. + +The two-statement approach eliminates the race where a job transitioning +``pending→running`` mid-statement escapes both CTEs in a single-shot +design: statement 1's EPQ guard rejects the now-running row, and +statement 2's fresh snapshot sees it as running and sets +``cancel_phase=1``. + +Events are inserted via ``executemany`` within the same transaction. +NOTIFY is sent by the caller (``PostgresBackend.cancel_where``) after +commit because the ``taskq.cancel.notify_sent`` counter lives in +``postgres.py``. +""" + +import asyncio +import random +from typing import NamedTuple +from uuid import UUID + +import asyncpg + +from taskq.backend._filter_sql import build_filter_conditions +from taskq.backend._protocol import BulkCancelResult, JobFilter +from taskq.backend._records import jsonb_param +from taskq.backend._sql_templates import SqlTemplates + +__all__ = ["_cancel_where"] + + +class NotifyTarget(NamedTuple): + """A running job that needs a post-commit NOTIFY.""" + + job_id: UUID + worker_id: UUID + + +async def _cancel_where( + pool: asyncpg.Pool, + schema: str, + sql: SqlTemplates, + filter: JobFilter, + reason: str | None, +) -> tuple[BulkCancelResult, list[NotifyTarget]]: + filter_sql = build_filter_conditions(filter) + conditions_str = " AND ".join(filter_sql.conditions) if filter_sql.conditions else "TRUE" + params = list(filter_sql.params) + + # Statement 1: cancel pending/scheduled → terminal 'cancelled' + # EPQ-safe: predicates on the target table (j.status) are re-evaluated + # for concurrently-modified rows. + cancel_ps_sql = f""" + WITH matching AS ( + SELECT id, status + FROM "{schema}".jobs + WHERE {conditions_str} + AND status IN ('pending', 'scheduled') + ORDER BY id + ), + cancelled AS ( + UPDATE "{schema}".jobs AS j + SET status = 'cancelled', finished_at = clock_timestamp() + FROM ( + SELECT id, status AS prev_status + FROM matching + ) AS prev + WHERE j.id = prev.id + AND j.status IN ('pending', 'scheduled') + RETURNING j.id, prev.prev_status + ) + SELECT + (SELECT count(*)::int FROM cancelled) AS cancelled_directly, + (SELECT array_agg(id ORDER BY id) FROM cancelled) AS cancelled_ids, + (SELECT array_agg(prev_status ORDER BY id) FROM cancelled) AS cancelled_prev_statuses + """ + + # Statement 2: cooperative cancel for running jobs with cancel_phase=0 + # Fresh snapshot — catches jobs dispatched between statements 1 and 2. + cancel_running_sql = f""" + WITH matching AS ( + SELECT id, locked_by_worker + FROM "{schema}".jobs + WHERE {conditions_str} + AND status = 'running' + AND cancel_phase = 0 + ORDER BY id + ), + cancel_requested AS ( + UPDATE "{schema}".jobs AS j + SET cancel_requested_at = now(), cancel_phase = 1 + FROM ( + SELECT id, locked_by_worker + FROM matching + ) AS prev + WHERE j.id = prev.id + AND j.status = 'running' + AND j.cancel_phase = 0 + RETURNING j.id, prev.locked_by_worker + ) + SELECT + (SELECT count(*)::int FROM cancel_requested) AS cancel_requested, + (SELECT array_agg(id ORDER BY id) FROM cancel_requested) AS cancel_requested_ids, + (SELECT array_agg(locked_by_worker ORDER BY id) FROM cancel_requested) AS cancel_requested_workers + """ + + cr_detail = jsonb_param({"reason": reason} if reason is not None else {}) + + cancelled_ids: list[UUID] = [] + cancel_requested_ids: list[UUID] = [] + notify_targets: list[NotifyTarget] = [] + + for attempt in range(3): + try: + async with pool.acquire() as conn: + async with conn.transaction(): + # Statement 1: pending/scheduled → cancelled + ps_row = await conn.fetchrow(cancel_ps_sql, *params) + if ps_row is not None: + cancelled_ids = list(ps_row["cancelled_ids"] or []) + prev_statuses: dict[UUID, str] = dict( + zip( + cancelled_ids, + ps_row["cancelled_prev_statuses"] or [], + strict=True, + ) + ) + + if cancelled_ids: + await conn.executemany( + sql.insert_event, + [ + ( + jid, + "state_change", + jsonb_param( + { + "from_state": prev_statuses[jid], + "to_state": "cancelled", + } + ), + ) + for jid in cancelled_ids + ], + ) + await conn.executemany( + sql.insert_event, + [(jid, "cancel_request", cr_detail) for jid in cancelled_ids], + ) + + # Statement 2: running → cooperative cancel (fresh snapshot) + running_row = await conn.fetchrow(cancel_running_sql, *params) + if running_row is not None: + cancel_requested_ids = list(running_row["cancel_requested_ids"] or []) + notify_targets = [ + NotifyTarget(job_id=jid, worker_id=wid) + for jid, wid in zip( + cancel_requested_ids, + running_row["cancel_requested_workers"] or [], + strict=True, + ) + if wid is not None + ] + + if cancel_requested_ids: + await conn.executemany( + sql.insert_event, + [ + (jid, "cancel_request", cr_detail) + for jid in cancel_requested_ids + ], + ) + break + except asyncpg.DeadlockDetectedError: + if attempt == 2: + raise + await asyncio.sleep(0.1 * (2**attempt) + random.random() * 0.05) + + result = BulkCancelResult( + cancelled_directly=len(cancelled_ids), + cancel_requested=len(cancel_requested_ids), + cancelled_ids=tuple(cancelled_ids), + cancel_requested_ids=tuple(cancel_requested_ids), + ) + return result, notify_targets diff --git a/src/taskq/backend/_filter_sql.py b/src/taskq/backend/_filter_sql.py new file mode 100644 index 00000000..d9dc0d37 --- /dev/null +++ b/src/taskq/backend/_filter_sql.py @@ -0,0 +1,102 @@ +"""Shared filter→SQL WHERE condition builder. + +Extracted from ``_reads._list_jobs`` so that ``cancel_where`` (bulk +cancel) reuses the exact same filter logic. Only predicate fields +(queue, status, actor, identity_key, batch_id, tags, active) are +translated to conditions. The ``cursor``, ``limit``, and ``order_by`` +fields are NOT handled here — callers apply them separately. + +This module is SQL-only; the in-memory backend filters via its own +implementation in ``testing/_reads.py``. Filter semantics between the +two backends are verified by ``test_job_filter.py`` and the per-backend +cancel_where test suites. +""" + +from dataclasses import dataclass + +from taskq._json import dumps_str +from taskq.backend._protocol import JobFilter +from taskq.backend.statemachine import ACTIVE_STATUSES, TERMINAL_STATUSES + +__all__ = ["FilterSQL", "build_filter_conditions"] + + +@dataclass(frozen=True, slots=True) +class FilterSQL: + """Built SQL fragments and parameters from a JobFilter. + + ``conditions`` and ``params`` are stored as tuples so the container + itself is immutable (prevents reassignment of the field). Individual + param elements (e.g. ``list[str]`` for ``ANY()``) are inherently + mutable — callers consume them immediately without mutation. + """ + + conditions: tuple[str, ...] = () + params: tuple[object, ...] = () + + +def build_filter_conditions(filter: JobFilter) -> FilterSQL: + """Build WHERE clause conditions and parameters from a JobFilter. + + Shared between ``_list_jobs`` (reads) and ``cancel_where`` (writes) + so the filter semantics are identical for query and mutation. + + Only predicate fields (queue, status, actor, identity_key, batch_id, + tags, active) are translated to conditions. The ``cursor``, ``limit``, + and ``order_by`` fields are NOT handled here — callers apply them + separately: + + - ``_list_jobs`` appends the cursor keyset condition and LIMIT/OFFSET + after calling this helper (preserving the existing behavior). + - ``cancel_where`` ignores cursor/limit/order_by entirely (bulk writes + are not paginated). + """ + conditions: list[str] = [] + params: list[object] = [] + n = 0 + + def _next_param(expr: str) -> str: + nonlocal n + n += 1 + return f"{expr} = ${n}" + + def _next_any_param(expr: str) -> str: + nonlocal n + n += 1 + return f"{expr} = ANY(${n})" + + if filter.queue is not None: + conditions.append(_next_param("queue")) + params.append(filter.queue) + + if filter.status is not None: + if isinstance(filter.status, str): + conditions.append(_next_param("status")) + params.append(filter.status) + else: + conditions.append(_next_any_param("status")) + params.append(list(filter.status)) + elif filter.active is not None: + statuses = list(ACTIVE_STATUSES) if filter.active else list(TERMINAL_STATUSES) + conditions.append(_next_any_param("status")) + params.append(statuses) + + if filter.actor is not None: + conditions.append(_next_param("actor")) + params.append(filter.actor) + + if filter.identity_key is not None: + conditions.append(_next_param("identity_key")) + params.append(filter.identity_key) + + if filter.batch_id is not None: + n += 1 + conditions.append(f"metadata @> ${n}::jsonb") + params.append(dumps_str({"batch_id": str(filter.batch_id)})) + + if filter.tags is not None and len(filter.tags) > 0: + n += 1 + conditions.append(f"tags && ${n}::text[]") + params.append(list(filter.tags)) + + return FilterSQL(conditions=tuple(conditions), params=tuple(params)) diff --git a/src/taskq/backend/_protocol.py b/src/taskq/backend/_protocol.py index dc96005f..623efa49 100644 --- a/src/taskq/backend/_protocol.py +++ b/src/taskq/backend/_protocol.py @@ -49,6 +49,7 @@ "AttemptRow", "Backend", "BackendDeps", + "BulkCancelResult", "CancelFlag", "CancelPhase", "DstStrategy", @@ -423,7 +424,13 @@ class CancelFlag: @dataclass(frozen=True, slots=True) class JobFilter: - """Filter parameters for :meth:`Backend.list_jobs`. + """Filter parameters for :meth:`Backend.list_jobs` and + :meth:`Backend.cancel_where`. + + For ``cancel_where``, the ``limit``, ``cursor``, and ``order_by`` + fields are ignored — a bulk cancel is not paginated. Use + :meth:`has_predicates` to check whether the filter has at least one + predicate before passing it to ``cancel_where``. Heads-up: ``active=True`` is **not** Celery's 'active' — Celery's means 'currently executing' (``running`` only), TaskQ's means 'not @@ -523,6 +530,26 @@ def __post_init__(self) -> None: "terminal/non-terminal meta-filter" ) + def has_predicates(self) -> bool: + """Return True if at least one filter predicate is set. + + Used by ``JobsClient.cancel_where`` to reject empty filters that + would match the entire table. New predicate fields added to + ``JobFilter`` MUST be added here and in + ``build_filter_conditions`` — the two are kept in sync manually. + Non-predicate fields (``limit``, ``cursor``, ``order_by``) are + excluded by design. + """ + return ( + self.queue is not None + or self.status is not None + or self.actor is not None + or self.identity_key is not None + or self.batch_id is not None + or (self.tags is not None and len(self.tags) > 0) + or self.active is not None + ) + @dataclass(frozen=True, slots=True) class ScheduleCreateArgs: @@ -616,6 +643,35 @@ class ScheduleRecord(BaseModel): metadata: dict[str, object] +class BulkCancelResult(BaseModel): + """Structured outcome of a bulk cancellation request. + + Returned by ``JobsClient.cancel_where()`` so callers can inspect + how many jobs were cancelled directly (pending/scheduled → terminal + 'cancelled') vs how many had cooperative cancel requested (running → + cancel_phase=1). + """ + + model_config = ConfigDict(frozen=True) + + cancelled_directly: int + """Count of pending/scheduled jobs moved straight to terminal 'cancelled'.""" + + cancel_requested: int + """Count of running jobs with cancel_phase=1 set (cooperative cancel).""" + + cancelled_ids: tuple[UUID, ...] + """IDs of jobs cancelled directly (pending/scheduled → cancelled).""" + + cancel_requested_ids: tuple[UUID, ...] + """IDs of running jobs with cancel requested.""" + + @property + def total_affected(self) -> int: + """Total jobs affected by the bulk cancel.""" + return self.cancelled_directly + self.cancel_requested + + @dataclass(frozen=True, slots=True) class ErrorInfo: """Structured error information for terminal writes.""" @@ -690,8 +746,8 @@ def dispatcher_pool(self) -> "asyncpg.Pool | None": class Backend(Protocol): """Contract that both PostgresBackend and InMemoryBackend satisfy. - 31 async methods plus two sync methods (``subscribe_wake`` and - ``subscribe_cancel_wake``) (33 methods total) covering enqueue, + 34 async methods plus two sync methods (``subscribe_wake`` and + ``subscribe_cancel_wake``) (36 methods total) covering enqueue, dispatch, heartbeat, terminal writes, attempt history, cancel signals, scheduling / sweeps, read, NOTIFY hook, and schedule CRUD. Method order grouped for review-grep ergonomics. @@ -946,6 +1002,31 @@ async def write_cancel_request( reason: str | None, ) -> bool: ... + async def cancel_where( + self, + filter: JobFilter, + reason: str | None, + ) -> BulkCancelResult: + """Cancel all jobs matching *filter* in a set-based operation. + + Pending/scheduled jobs → terminal 'cancelled'. + Running jobs → cancel_phase=1 (cooperative cancel + NOTIFY). + + The filter's ``limit``, ``cursor``, and ``order_by`` fields are + ignored — this is a bulk write, not a paginated read. + + **Guardrail:** the client layer (:meth:`JobsClient.cancel_where`) + rejects empty filters (no predicates) with + :class:`EmptyFilterError`. Backend implementations receive a + filter that has already been validated. A direct backend call + with ``JobFilter()`` renders ``WHERE TRUE`` and cancels the + entire table — callers using the backend directly are + responsible for validating the filter. + + Returns a :class:`BulkCancelResult` with counts and affected IDs. + """ + ... + async def poll_cancel_flags( self, worker_id: UUID, diff --git a/src/taskq/backend/_reads.py b/src/taskq/backend/_reads.py index f613de69..24029fbe 100644 --- a/src/taskq/backend/_reads.py +++ b/src/taskq/backend/_reads.py @@ -8,8 +8,8 @@ from datetime import timedelta from typing import TYPE_CHECKING -from taskq._json import dumps_str from taskq.backend._cursor import decode_cursor +from taskq.backend._filter_sql import build_filter_conditions from taskq.backend._protocol import ( AttemptRow, EventRow, @@ -24,7 +24,6 @@ jsonb_to_dict, ) from taskq.backend._sql_templates import SqlTemplates -from taskq.backend.statemachine import ACTIVE_STATUSES, TERMINAL_STATUSES from taskq.constants import RECLAIM_EVENT_VISIBILITY_DELAY if TYPE_CHECKING: @@ -55,49 +54,10 @@ async def _list_jobs( schema: str, filters: JobFilter, ) -> list[JobRow]: - conditions: list[str] = [] - params: list[object] = [] - n = 0 - - def _next_param(expr: str) -> str: - nonlocal n - n += 1 - return f"{expr} = ${n}" - - def _next_any_param(expr: str) -> str: - nonlocal n - n += 1 - return f"{expr} = ANY(${n})" - - if filters.queue is not None: - conditions.append(_next_param("queue")) - params.append(filters.queue) - if filters.status is not None: - if isinstance(filters.status, str): - conditions.append(_next_param("status")) - params.append(filters.status) - else: - conditions.append(_next_any_param("status")) - params.append(list(filters.status)) - elif filters.active is not None: - statuses = list(ACTIVE_STATUSES) if filters.active else list(TERMINAL_STATUSES) - conditions.append(_next_any_param("status")) - params.append(statuses) - if filters.actor is not None: - conditions.append(_next_param("actor")) - params.append(filters.actor) - if filters.identity_key is not None: - conditions.append(_next_param("identity_key")) - params.append(filters.identity_key) - if filters.batch_id is not None: - n += 1 - conditions.append(f"metadata @> ${n}::jsonb") - params.append(dumps_str({"batch_id": str(filters.batch_id)})) - - if filters.tags is not None and len(filters.tags) > 0: - n += 1 - conditions.append(f"tags && ${n}::text[]") - params.append(list(filters.tags)) + filter_sql = build_filter_conditions(filters) + conditions: list[str] = list(filter_sql.conditions) + params: list[object] = list(filter_sql.params) + n = len(params) if filters.cursor is not None: cursor_priority, cursor_scheduled_at, cursor_id = decode_cursor(filters.cursor) diff --git a/src/taskq/backend/postgres.py b/src/taskq/backend/postgres.py index 4d0f85e2..f50d0783 100644 --- a/src/taskq/backend/postgres.py +++ b/src/taskq/backend/postgres.py @@ -9,9 +9,11 @@ (:mod:`taskq.backend._schedules`), terminal writes (:mod:`taskq.backend._terminal`), enqueue (:mod:`taskq.backend._enqueue`), reads (:mod:`taskq.backend._reads`), -and dispatch (:mod:`taskq.backend._dispatch`) live in companion -submodules; this module holds the cohesive core: ``__init__``, heartbeat, -cancel signals, NOTIFY, and schedule CRUD wiring. +bulk cancel (:mod:`taskq.backend._cancel_bulk`), filter SQL builder +(:mod:`taskq.backend._filter_sql`), and dispatch +(:mod:`taskq.backend._dispatch`) live in companion submodules; this +module holds the cohesive core: ``__init__``, heartbeat, cancel +signals, NOTIFY, and schedule CRUD wiring. """ import asyncio @@ -23,6 +25,7 @@ import structlog from taskq._json import dumps_str +from taskq.backend._cancel_bulk import _cancel_where from taskq.backend._dispatch import ( _dispatch_batch as _dispatch, ) @@ -41,6 +44,7 @@ AttemptOutcome, AttemptRow, BackendDeps, + BulkCancelResult, CancelFlag, ConnLike, EnqueueArgs, @@ -157,6 +161,14 @@ ) +def _cancel_notify_payload(job_id: UUID, worker_id: UUID) -> str: + return dumps_str({"type": "cancel", "job_id": str(job_id), "worker_id": str(worker_id)}) + + +def _cancel_notify_channels(schema: str, worker_id: UUID) -> list[str]: + return [events_channel(schema), worker_channel(schema, str(worker_id))] + + class PostgresBackend: """Production backend backed by Postgres. @@ -583,15 +595,8 @@ async def write_cancel_request( job_id=str(job_id), ) if _locked_by_worker is not None: - payload = dumps_str( - { - "type": "cancel", - "job_id": str(job_id), - "worker_id": str(_locked_by_worker), - } - ) - fleet_ch = events_channel(self._schema_name) - worker_ch = worker_channel(self._schema_name, str(_locked_by_worker)) + payload = _cancel_notify_payload(job_id, _locked_by_worker) + fleet_ch, worker_ch = _cancel_notify_channels(self._schema_name, _locked_by_worker) async with self._worker_pool.acquire() as notify_conn: await notify_conn.execute( "SELECT pg_notify($1, $2), pg_notify($3, $4)", @@ -619,6 +624,39 @@ async def poll_cancel_flags( for rec in recs ] + async def cancel_where( + self, + filter: JobFilter, + reason: str | None, + ) -> BulkCancelResult: + result, notify_targets = await _cancel_where( + self._worker_pool, self._schema_name, self._sql, filter, reason + ) + if notify_targets: + channels: list[str] = [] + payloads: list[str] = [] + for target in notify_targets: + payload = _cancel_notify_payload(target.job_id, target.worker_id) + ch = _cancel_notify_channels(self._schema_name, target.worker_id) + channels.extend(ch) + payloads.extend([payload, payload]) + try: + async with self._worker_pool.acquire() as notify_conn: + await notify_conn.execute( + "SELECT pg_notify(channel, payload) " + "FROM unnest($1::text[], $2::text[]) AS t(channel, payload)", + channels, + payloads, + ) + _cancel_notify_sent_counter.add(len(notify_targets), {"schema": self._schema_name}) + except Exception: + logger.warning( + "cancel_where_notify_failed", + notify_count=len(notify_targets), + exc_info=True, + ) + return result + # ── Admin operations ────────────────────────────────────────────── async def retry_job(self, job_id: JobId) -> bool: diff --git a/src/taskq/client/__init__.py b/src/taskq/client/__init__.py index 8418f1f5..4fac140a 100644 --- a/src/taskq/client/__init__.py +++ b/src/taskq/client/__init__.py @@ -13,9 +13,10 @@ from taskq.client._handle import JobHandle from taskq.client._jobs import JobsClient from taskq.client._taskq import JobEvent, TaskQ -from taskq.types import CancelResult +from taskq.types import BulkCancelResult, CancelResult __all__ = [ + "BulkCancelResult", "CancelResult", "JobEvent", "JobHandle", diff --git a/src/taskq/client/_enqueuer.py b/src/taskq/client/_enqueuer.py index e1e9c32b..b7a72a57 100644 --- a/src/taskq/client/_enqueuer.py +++ b/src/taskq/client/_enqueuer.py @@ -5,11 +5,19 @@ One instance per loop — survives across dispatches so the per-100-enqueue re-warning fires on the loop-level counter, not per-job. +Parent-tag propagation: the consumer sets the parent job's tags via +``set_parent_tags()`` before actor invocation and resets them after +(via ``parent_tags()`` context manager or manual token reset). The +``contextvars.ContextVar`` ensures concurrent consumers in the same +event loop each see their own parent tags — asyncio Tasks copy the +context at creation time. """ from __future__ import annotations -from collections.abc import Mapping, Sequence +import contextlib +import contextvars +from collections.abc import Generator, Mapping, Sequence from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, cast from uuid import UUID @@ -39,10 +47,45 @@ from taskq.actor import ActorRef -__all__ = ["SubJobEnqueuer"] +__all__ = ["SubJobEnqueuer", "parent_tags", "set_parent_tags"] _log: structlog.stdlib.BoundLogger = structlog.get_logger(__name__) +_parent_tags_var: contextvars.ContextVar[tuple[str, ...]] = contextvars.ContextVar( + "taskq_parent_tags", + default=(), +) + + +def set_parent_tags(tags: tuple[str, ...]) -> contextvars.Token[tuple[str, ...]]: + """Set the parent job's tags for sub-job tag inheritance. + + Called by the consumer before actor invocation. The returned token + must be used to reset the context after the actor completes — use + ``_parent_tags_var.reset(token)`` or the ``parent_tags()`` context + manager. + """ + return _parent_tags_var.set(tags) + + +@contextlib.contextmanager +def parent_tags(tags: tuple[str, ...]) -> Generator[None, None, None]: + """Context manager that sets parent tags for the duration of the block. + + Ensures the ContextVar is reset on all exit paths (success, exception, + cancellation). Use this at worker entry points instead of manual + set/reset:: + + with parent_tags(tuple(job.tags)): + # actor invocation, sub-job enqueues, etc. + ... + """ + token = _parent_tags_var.set(tags) + try: + yield + finally: + _parent_tags_var.reset(token) + class SubJobEnqueuer: """Enqueue sub-jobs from within an actor body. @@ -90,6 +133,11 @@ async def enqueue[P: BaseModel, R: BaseModel | None]( unique_for: timedelta | None = None, unique_states: tuple[JobStatus, ...] | None = None, max_pending: int | None = None, + tags: list[str] | None = None, + inherit_tags: bool = True, + schedule_to_close: datetime | None = None, + start_to_close: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, ) -> JobHandle[R]: """Enqueue a sub-job. ``max_pending`` is a per-call limit resolved against the operator-owned stored cap and the ``@actor(...)`` @@ -113,6 +161,7 @@ async def enqueue[P: BaseModel, R: BaseModel | None]( actor_ref.max_pending, per_call=max_pending, ) + resolved_tags = self._resolve_tags(tags, inherit_tags) args = build_enqueue_args( actor_ref, payload, @@ -125,6 +174,10 @@ async def enqueue[P: BaseModel, R: BaseModel | None]( idempotency_scope=idempotency_scope, trace_id=extracted_trace_id, span_id=extracted_span_id, + tags=resolved_tags, + schedule_to_close=schedule_to_close, + start_to_close=start_to_close, + heartbeat_timeout=heartbeat_timeout, unique_for=unique_for, unique_states=unique_states, max_pending=effective_max_pending, @@ -140,6 +193,33 @@ async def enqueue[P: BaseModel, R: BaseModel | None]( client=None, ) + def _resolve_tags( + self, + tags: list[str] | None, + inherit_tags: bool, + ) -> list[str] | None: + """Resolve tags with parent inheritance. + + Returns a list suitable for build_enqueue_args, or None for empty. + Deduplication is order-preserving (parent first); the downstream + ``_validate_and_dedup_tags`` in ``build_enqueue_args`` also + deduplicates, but we do it here so the merge result is clean. + """ + parent_tags = _parent_tags_var.get() if inherit_tags else () + + if tags is None: + if parent_tags: + return list(parent_tags) + return None + + if not inherit_tags or not parent_tags: + return tags + + if not tags: + return list(parent_tags) + + return list(dict.fromkeys((*parent_tags, *tags))) + def _resolve_connection( self, connection: asyncpg.Connection | None, @@ -274,6 +354,9 @@ async def enqueue_batch( idempotency_key=item.idempotency_key, idempotency_scope=item.idempotency_scope, identity_key=item.identity_key, + tags=list(item.tags) if item.tags else None, + inherit_tags=False, + start_to_close=item.start_to_close, ) handles.append(handle) except Exception as exc: diff --git a/src/taskq/client/_jobs.py b/src/taskq/client/_jobs.py index 59661185..7f6ee27e 100644 --- a/src/taskq/client/_jobs.py +++ b/src/taskq/client/_jobs.py @@ -45,8 +45,8 @@ from taskq.client._args import build_batch_args, build_enqueue_args, enqueue_span from taskq.client._capacity import DEFAULT_CAPACITY_CACHE_TTL, ActorCapacityCache from taskq.client._handle import JobHandle -from taskq.exceptions import PayloadValidationError, SchemaNotMigratedError -from taskq.types import CancelResult +from taskq.exceptions import EmptyFilterError, PayloadValidationError, SchemaNotMigratedError +from taskq.types import BulkCancelResult, CancelResult if TYPE_CHECKING: import asyncpg @@ -733,6 +733,42 @@ async def cancel( ) return result + async def cancel_where( + self, + filter: JobFilter, + reason: str | None = None, + *, + allow_empty_filter: bool = False, + ) -> BulkCancelResult: + """Cancel all jobs matching *filter* in a single set-based operation. + + Pending/scheduled jobs are moved straight to terminal 'cancelled' + (no running actor to cooperate with). Running jobs get + ``cancel_phase=1`` set (cooperative cancel) — the worker's + heartbeat-driven cancel controller observes the phase change and + sets the in-process ``cancel_event``. + + **Guardrail:** a filter with no predicates (no queue, status, + actor, identity_key, batch_id, tags, or active) is rejected with + :class:`EmptyFilterError` unless ``allow_empty_filter=True`` is + passed. + + **Filter fields used:** ``queue``, ``status``, ``actor``, + ``identity_key``, ``batch_id``, ``tags``, ``active``. The + ``limit``, ``cursor``, and ``order_by`` fields are ignored. + + Returns a :class:`BulkCancelResult` with counts and affected IDs. + """ + from taskq.obs import record_cancel_requested + + if not allow_empty_filter and not filter.has_predicates(): + raise EmptyFilterError() + + record_cancel_requested() + + with self._translate_schema_errors(): + return await self._backend.cancel_where(filter, reason) + # ── Schedule CRUD ──────────────────────────────────────────────────── async def create_schedule[P: BaseModel, R: BaseModel | None]( diff --git a/src/taskq/client/_taskq.py b/src/taskq/client/_taskq.py index fc21d32a..15308820 100644 --- a/src/taskq/client/_taskq.py +++ b/src/taskq/client/_taskq.py @@ -72,7 +72,7 @@ async def create_task(payload: MyPayload): from taskq.constants import RECLAIM_EVENT_VISIBILITY_DELAY, progress_channel, wake_channel from taskq.cron import ScheduleHandle from taskq.progress._events import ProgressEvent -from taskq.types import CancelResult +from taskq.types import BulkCancelResult, CancelResult __all__ = ["EventRow", "JobEvent", "TaskQ"] @@ -413,6 +413,18 @@ async def cancel( """Request cancellation of a job. Raises :class:`KeyError` if not found.""" return await self._require_open().cancel(job_id, reason) + async def cancel_where( + self, + filter: JobFilter, + reason: str | None = None, + *, + allow_empty_filter: bool = False, + ) -> BulkCancelResult: + """Cancel all jobs matching *filter*. See :meth:`JobsClient.cancel_where`.""" + return await self._require_open().cancel_where( + filter, reason, allow_empty_filter=allow_empty_filter + ) + # ── Schedule operations ───────────────────────────────────────────────── async def create_schedule[P: BaseModel, R: BaseModel | None]( diff --git a/src/taskq/exceptions.py b/src/taskq/exceptions.py index 6b0b107e..6b470580 100644 --- a/src/taskq/exceptions.py +++ b/src/taskq/exceptions.py @@ -387,6 +387,24 @@ def __init__(self, schema: str) -> None: ) +class EmptyFilterError(TaskQError): + """Raised when cancel_where is called with a filter that has no predicates. + + A filter with no queue, status, actor, identity_key, batch_id, tags, or + active predicate would match every job in the table — almost certainly + a bug. The guardrail is intentionally loud: the caller must add at least + one predicate or explicitly bypass with ``allow_empty_filter=True``. + """ + + def __init__(self) -> None: + super().__init__( + "cancel_where requires at least one filter predicate " + "(queue, status, actor, identity_key, batch_id, tags, or active); " + "an empty filter would cancel the entire table. " + "Pass allow_empty_filter=True to override this guardrail." + ) + + class ScopedIdempotencyMigrationPendingError(TaskQError): """``idempotency_scope`` was used, but the schema has not yet had ``01.00.03_01_post_idempotency_scope_drop_old_index.sql`` applied. diff --git a/src/taskq/testing/_cancel_bulk.py b/src/taskq/testing/_cancel_bulk.py new file mode 100644 index 00000000..0458616a --- /dev/null +++ b/src/taskq/testing/_cancel_bulk.py @@ -0,0 +1,71 @@ +"""Bulk cancel implementation for InMemoryBackend. + +Module-level function following the same pattern as ``testing/_reads.py``, +``testing/_terminal.py``, etc. +""" + +from dataclasses import replace as dc_replace +from typing import TYPE_CHECKING +from uuid import UUID + +from taskq.backend._protocol import BulkCancelResult, CancelPhase, JobFilter + +if TYPE_CHECKING: + from taskq.testing.in_memory import InMemoryBackend + +__all__ = ["_cancel_where"] + + +async def _cancel_where( + self: "InMemoryBackend", + filter: JobFilter, + reason: str | None, +) -> BulkCancelResult: + from taskq.testing._reads import _list_jobs + + # Sanitize the filter: cancel_where ignores limit, cursor, and order_by. + # Use a very large limit instead of None because JobFilter.limit is typed + # as int (not int | None) with a __post_init__ guard against negatives. + # cursor=None disables keyset slicing. order_by=None selects the default + # priority/scheduled_at/id sort, which is harmless for cancel. + sanitized = dc_replace(filter, limit=2**31, cursor=None, order_by=None) + rows = await _list_jobs(self, sanitized) + + cancelled_ids: list[UUID] = [] + cancel_requested_ids: list[UUID] = [] + + for row in rows: + if row.status in ("pending", "scheduled"): + now = self._clock.now() + self._jobs[row.id] = dc_replace( + row, + status="cancelled", + finished_at=now, + ) + self._append_state_change_event( + job_id=row.id, + from_state=row.status, + to_state="cancelled", + now=now, + ) + self._append_cancel_request_event(row.id, now, reason) + cancelled_ids.append(row.id) + + elif row.status == "running" and row.cancel_phase == CancelPhase.NONE: + now = self._clock.now() + self._jobs[row.id] = dc_replace( + row, + cancel_requested_at=now, + cancel_phase=CancelPhase.COOPERATIVE, + ) + self._append_cancel_request_event(row.id, now, reason) + for event in self._cancel_wake_subscribers: + event.set() + cancel_requested_ids.append(row.id) + + return BulkCancelResult( + cancelled_directly=len(cancelled_ids), + cancel_requested=len(cancel_requested_ids), + cancelled_ids=tuple(cancelled_ids), + cancel_requested_ids=tuple(cancel_requested_ids), + ) diff --git a/src/taskq/testing/in_memory.py b/src/taskq/testing/in_memory.py index 56d663fc..c5b3d58a 100644 --- a/src/taskq/testing/in_memory.py +++ b/src/taskq/testing/in_memory.py @@ -36,6 +36,7 @@ BACKEND_PROTOCOL_VERSION, AttemptOutcome, AttemptRow, + BulkCancelResult, CancelFlag, CancelPhase, EnqueueArgs, @@ -51,6 +52,7 @@ ) from taskq.backend.clock import Clock from taskq.retry import OnRetryExhausted, OnSuccess, RetryClassifierHook, RetryPolicy +from taskq.testing._cancel_bulk import _cancel_where from taskq.testing._dispatch import _dispatch_batch, _set_queue_mode from taskq.testing._enqueue import ( _enqueue, @@ -622,6 +624,13 @@ async def poll_cancel_flags( and row.locked_by_worker == worker_id ] + async def cancel_where( + self, + filter: JobFilter, + reason: str | None, + ) -> BulkCancelResult: + return await _cancel_where(self, filter, reason) + # ── Admin operations ────────────────────────────────────────────── async def retry_job(self, job_id: JobId) -> bool: diff --git a/src/taskq/types.py b/src/taskq/types.py index 9673db1b..c1378d38 100644 --- a/src/taskq/types.py +++ b/src/taskq/types.py @@ -1,12 +1,15 @@ """Client-facing result and event-detail types. -``CancelResult`` is the structured return value of ``JobsClient.cancel()`` - ``StateChangeEvent`` is the JSON payload stored in +``CancelResult`` is the structured return value of ``JobsClient.cancel()``. +``BulkCancelResult`` is the structured return value of +``JobsClient.cancel_where()`` (defined in ``backend._protocol`` to avoid +a circular import — see its docstring). +``StateChangeEvent`` is the JSON payload stored in ``job_events.detail`` for rows with ``kind='state_change'``. -These types live here — not in ``taskq.backend`` — so the Backend protocol -remains pydantic-free and the layering contract is enforceable by import -inspection. +``BulkCancelResult`` is re-exported here (not defined) because +``types.py`` imports from ``backend._protocol`` — defining it here would +create a circular import (``_protocol → types → _protocol``). """ from dataclasses import dataclass @@ -14,9 +17,9 @@ from pydantic import BaseModel, ConfigDict -from taskq.backend._protocol import JobId, JobStatus +from taskq.backend._protocol import BulkCancelResult, JobId, JobStatus -__all__ = ["CancelResult", "StateChangeEvent"] +__all__ = ["BulkCancelResult", "CancelResult", "StateChangeEvent"] class CancelResult(BaseModel): diff --git a/src/taskq/worker/_consumer.py b/src/taskq/worker/_consumer.py index a851bf87..b6218c1c 100644 --- a/src/taskq/worker/_consumer.py +++ b/src/taskq/worker/_consumer.py @@ -29,7 +29,7 @@ JobRow, ) from taskq.backend.clock import Clock -from taskq.client._enqueuer import SubJobEnqueuer +from taskq.client._enqueuer import SubJobEnqueuer, _parent_tags_var from taskq.constants import MAX_RESULT_BYTES from taskq.context import JobContext from taskq.exceptions import ( @@ -353,6 +353,8 @@ async def consume_one_job( _buf = _ProgressBuffer(job_id=job.id, base_seq=job.progress_seq) _progress_buffers[job.id] = _buf + _parent_tags_token = _parent_tags_var.set(tuple(job.tags)) + try: validated_payload = ( validated_payload @@ -593,6 +595,7 @@ async def consume_one_job( await active_jobs.deregister(job.id) finally: + _parent_tags_var.reset(_parent_tags_token) if acquired and rate_limit_registry is not None: try: await asyncio.shield( diff --git a/src/taskq/worker/run.py b/src/taskq/worker/run.py index fbfaa1fb..fb80c00e 100644 --- a/src/taskq/worker/run.py +++ b/src/taskq/worker/run.py @@ -46,7 +46,7 @@ from taskq.backend._protocol import Backend, JobRow from taskq.backend._records import jsonb_param from taskq.backend.clock import Clock -from taskq.client._enqueuer import SubJobEnqueuer +from taskq.client._enqueuer import SubJobEnqueuer, parent_tags from taskq.constants import ( _IDENT_RE, # pyright: ignore[reportPrivateUsage] # Why: canonical identifier regex; copying would drift the validation pattern. ) @@ -346,59 +346,60 @@ async def consumer_loop_stub( if current_task is None: raise RuntimeError("consumer_loop_stub must run inside a TaskGroup") - ctx: JobContext[_StubPayload] = JobContext( - job_id=job.id, - actor=job.actor, - queue=job.queue, - attempt=job.attempt, - worker_id=worker_id, - payload=_StubPayload(), - jobs=SubJobEnqueuer( - loop_scope_resolved=None, - worker_pool=None, - backend=backend, - ), - log=bind_job_context( - _consumer_log, + with parent_tags(tuple(job.tags)): + ctx: JobContext[_StubPayload] = JobContext( job_id=job.id, actor=job.actor, queue=job.queue, attempt=job.attempt, - identity_key=job.identity_key, - trace_id=job.trace_id or "", - ), - ) + worker_id=worker_id, + payload=_StubPayload(), + jobs=SubJobEnqueuer( + loop_scope_resolved=None, + worker_pool=None, + backend=backend, + ), + log=bind_job_context( + _consumer_log, + job_id=job.id, + actor=job.actor, + queue=job.queue, + attempt=job.attempt, + identity_key=job.identity_key, + trace_id=job.trace_id or "", + ), + ) - await deps.active_jobs.register(job.id, current_task, ctx) # type: ignore[arg-type] # Why: JobContext[_StubPayload] is a JobContext[BaseModel]; pyright cannot widen Generic[TChild] to Generic[TParent] without explicit covariance. + await deps.active_jobs.register(job.id, current_task, ctx) # type: ignore[arg-type] # Why: JobContext[_StubPayload] is a JobContext[BaseModel]; pyright cannot widen Generic[TChild] to Generic[TParent] without explicit covariance. - try: try: - await asyncio.wait_for( - ctx.cancel_event.wait(), - timeout=stub_work_timeout, - ) + try: + await asyncio.wait_for( + ctx.cancel_event.wait(), + timeout=stub_work_timeout, + ) + except asyncio.CancelledError: + with contextlib.suppress(asyncio.CancelledError): + await asyncio.shield(backend.mark_cancelled(job.id, worker_id)) + raise + except TimeoutError: + pass + + if ctx.cancellation_requested: + await asyncio.shield(backend.mark_cancelled(job.id, worker_id)) + else: + await asyncio.shield(backend.mark_succeeded(job.id, worker_id, None)) + except asyncio.CancelledError: with contextlib.suppress(asyncio.CancelledError): await asyncio.shield(backend.mark_cancelled(job.id, worker_id)) raise - except TimeoutError: - pass - - if ctx.cancellation_requested: - await asyncio.shield(backend.mark_cancelled(job.id, worker_id)) - else: - await asyncio.shield(backend.mark_succeeded(job.id, worker_id, None)) - except asyncio.CancelledError: - with contextlib.suppress(asyncio.CancelledError): - await asyncio.shield(backend.mark_cancelled(job.id, worker_id)) - raise - - except Exception: - _consumer_log.exception("consumer-stub-error", job_id=str(job.id)) + except Exception: + _consumer_log.exception("consumer-stub-error", job_id=str(job.id)) - finally: - await deps.active_jobs.deregister(job.id) + finally: + await deps.active_jobs.deregister(job.id) async def di_consumer_loop( diff --git a/tests/e2e/actors.py b/tests/e2e/actors.py index 7f21e709..6c02ab27 100644 --- a/tests/e2e/actors.py +++ b/tests/e2e/actors.py @@ -705,3 +705,73 @@ async def concurrent_tracked_worker( "ct_finished", {"run_id": payload.run_id, "job_index": payload.job_index}, ) + + +# ── Tagged pipeline actors (sub-job tag inheritance / merge) ────────────── + + +class PipelineStagePayload(BaseModel): + run_id: str + stage: int + total_stages: int = 3 + + +@actor(name="pipeline_stage", queue="e2e") +async def pipeline_stage( + payload: PipelineStagePayload, + ctx: JobContext[PipelineStagePayload], + *, + pool: asyncpg.Pool, +) -> None: + """Linear pipeline: each stage enqueues the next via ctx.jobs.enqueue().""" + await _record_effect( + pool, + ctx, + "stage", + {"run_id": payload.run_id, "stage": payload.stage}, + ) + await asyncio.sleep(0.05) + + if payload.stage < payload.total_stages: + await ctx.jobs.enqueue( + pipeline_stage, + PipelineStagePayload( + run_id=payload.run_id, + stage=payload.stage + 1, + total_stages=payload.total_stages, + ), + ) + + +class TaggedPipelineStagePayload(BaseModel): + run_id: str + stage: int + total_stages: int = 3 + + +@actor(name="tagged_pipeline_stage", queue="e2e") +async def tagged_pipeline_stage( + payload: TaggedPipelineStagePayload, + ctx: JobContext[TaggedPipelineStagePayload], + *, + pool: asyncpg.Pool, +) -> None: + """Pipeline that passes explicit tags to sub-job enqueue.""" + await _record_effect( + pool, + ctx, + "stage", + {"run_id": payload.run_id, "stage": payload.stage}, + ) + await asyncio.sleep(0.05) + + if payload.stage < payload.total_stages: + await ctx.jobs.enqueue( + tagged_pipeline_stage, + TaggedPipelineStagePayload( + run_id=payload.run_id, + stage=payload.stage + 1, + total_stages=payload.total_stages, + ), + tags=[f"stage-{payload.stage + 1}"], + ) diff --git a/tests/e2e/test_cancellation.py b/tests/e2e/test_cancellation.py index 21ecc593..69f07da9 100644 --- a/tests/e2e/test_cancellation.py +++ b/tests/e2e/test_cancellation.py @@ -1,27 +1,17 @@ -"""Cancellation e2e — cooperative cancel mid-run + clean pre-dispatch cancel. - -Scenario: -long report → ``wait_for_handle_status(handle, "running")`` → -``handle.cancel()`` → terminal ``cancelled``; effects show only early stages. - -Real cancel semantics, verified against the library (not guessed): +"""Cancellation e2e — single-job cancel + bulk cancel_by_filter. +Single-job path: - ``JobsClient.cancel`` → ``PostgresBackend.write_cancel_request`` - (backend/postgres.py) branches on the row's current status - (backend/_sql_templates.py): - - - ``pending``/``scheduled`` → ``cancel_pending_scheduled`` UPDATEs the row - straight to ``status='cancelled'``. The cancel API does NOT reject - pre-dispatch cancels: ``CancelResult.cancellation_initiated`` is True, - ``new_status`` is ``'cancelled'``, and the worker never dispatches the - job — no actor code runs, so no effects exist. - - ``running`` → ``cancel_running`` sets ``cancel_phase=1``; the worker's - heartbeat-driven CancelController sets the in-process ``cancel_event``, - the actor's ``ctx.check_cancelled()`` raises ``asyncio.CancelledError`` - at the next stage boundary, and the consumer writes ``mark_cancelled``. - -- ``handle.wait()`` on any non-success terminal status raises ``JobFailed`` - carrying the row (``client/_handle.py._extract_result``). + branches on the row's current status: + - ``pending``/``scheduled`` → straight to ``status='cancelled'`` + - ``running`` → ``cancel_phase=1`` (cooperative), actor sees + ``ctx.check_cancelled()`` at next stage boundary + +Bulk path (``cancel_where``): +- pending/scheduled → terminal ``cancelled`` (set-based, single SQL) +- running → ``cancel_phase=1`` (cooperative, batched NOTIFY) +- empty filter rejected with ``EmptyFilterError`` +- batch_id filter cancels all jobs in a batch Every test requests ``e2e_worker`` explicitly: the worker container fixture is not autouse, so no worker (and no dispatch) exists unless a test pulls it in. @@ -31,13 +21,20 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING +from uuid import uuid4 import pytest -from taskq import JobFailed +from taskq import JobFailed, JobFilter +from taskq.batch import EnqueueItem from ._assertions import fetch_effects, wait_for_effects, wait_for_handle_status -from .actors import GenerateReportPayload, generate_report +from .actors import ( + GenerateReportPayload, + ImportContactsChunkPayload, + generate_report, + import_contacts_chunk, +) if TYPE_CHECKING: import asyncpg @@ -139,3 +136,137 @@ async def test_cancel_before_dispatch_is_clean( # The job never dispatched, so no actor code ran: no effects of any kind. assert await fetch_effects(e2e_pg_pool, e2e_schema.schema_name, run_id) == [] + + +# ── Bulk cancel (cancel_where) ───────────────────────────────────────── + + +async def test_cancel_where_pending_jobs_by_tag( + e2e_client: TaskQ, + e2e_worker: E2EWorker, + e2e_pg_pool: asyncpg.Pool, + e2e_schema: E2ESchema, + run_id: str, +) -> None: + """cancel_where cancels all scheduled jobs matching a tag filter in one + set-based write. Untagged jobs are unaffected.""" + tag = f"tenant-{run_id[:8]}" + future = datetime.now(UTC) + timedelta(seconds=120) + + for i in range(5): + await e2e_client.enqueue( + generate_report, + GenerateReportPayload(run_id=run_id, report_id=f"r-{i}"), + scheduled_at=future, + tags=[tag], + ) + for i in range(2): + await e2e_client.enqueue( + generate_report, + GenerateReportPayload(run_id=f"other-{run_id}", report_id=f"o-{i}"), + scheduled_at=future, + ) + + result = await e2e_client.cancel_where( + JobFilter(tags=(tag,)), + reason="tenant offboarded", + ) + + assert result.cancelled_directly == 5 + assert result.cancel_requested == 0 + assert result.total_affected == 5 + + tagged = await e2e_client.list(JobFilter(tags=(tag,), status="cancelled")) + assert len(tagged.jobs) == 5 + + untagged = await e2e_client.list(JobFilter(status="scheduled", queue="e2e")) + assert len(untagged.jobs) == 2 + + +async def test_cancel_where_running_cooperative( + e2e_client: TaskQ, + e2e_worker: E2EWorker, + e2e_pg_pool: asyncpg.Pool, + e2e_schema: E2ESchema, + run_id: str, +) -> None: + """cancel_where on a running job sets cooperative cancel (cancel_phase=1). + The actor observes the cancel signal at the next stage boundary and + reaches terminal 'cancelled' via the normal cooperative path.""" + tag = f"run-{run_id[:8]}" + + handle = await e2e_client.enqueue( + generate_report, + GenerateReportPayload( + run_id=run_id, + report_id=f"r-{run_id[:8]}", + stages=4, + stage_latency_ms=2000, + ), + tags=[tag], + ) + await wait_for_handle_status(handle, "running", timeout=30) + await wait_for_effects( + e2e_pg_pool, + e2e_schema.schema_name, + run_id, + kind="stage", + min_count=1, + timeout=30, + ) + + result = await e2e_client.cancel_where( + JobFilter(tags=(tag,), status="running"), + reason="abort run", + ) + + assert result.total_affected >= 1, "cancel_where should affect the running job" + assert result.cancelled_directly == 0 + + await wait_for_handle_status(handle, "cancelled", timeout=30) + + +async def test_cancel_where_empty_filter_raises( + e2e_client: TaskQ, + e2e_worker: E2EWorker, +) -> None: + """Empty filter is rejected even in e2e.""" + from taskq.exceptions import EmptyFilterError + + with pytest.raises(EmptyFilterError): + await e2e_client.cancel_where(JobFilter(), reason="oops") + + +async def test_cancel_where_batch_id_filter( + e2e_client: TaskQ, + e2e_worker: E2EWorker, + e2e_pg_pool: asyncpg.Pool, + e2e_schema: E2ESchema, + run_id: str, +) -> None: + """cancel_where with batch_id cancels all jobs in a batch.""" + future = datetime.now(UTC) + timedelta(seconds=120) + bid = uuid4() + + items = [ + EnqueueItem( + actor_ref=import_contacts_chunk, + payload=ImportContactsChunkPayload( + run_id=run_id, + upload_id=str(bid), + chunk_id=i, + start_row=i * 100, + end_row=(i + 1) * 100, + ), + scheduled_at=future, + ) + for i in range(5) + ] + await e2e_client.enqueue_batch(items, batch_id=bid) + + result = await e2e_client.cancel_where( + JobFilter(batch_id=bid), + reason="batch abort", + ) + + assert result.cancelled_directly == 5 diff --git a/tests/e2e/test_sub_job_tags.py b/tests/e2e/test_sub_job_tags.py new file mode 100644 index 00000000..8b5c15ba --- /dev/null +++ b/tests/e2e/test_sub_job_tags.py @@ -0,0 +1,112 @@ +"""Sub-job tag inheritance and merge e2e — real pipeline in a worker container. + +Scenario: +parent job enqueued with a tag enqueues sub-jobs via ``ctx.jobs.enqueue()``. +Sub-jobs with no explicit tags inherit the parent tag; sub-jobs with explicit +tags carry both the inherited parent tag and the explicit tag. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from taskq import JobFilter + +from ._assertions import wait_for_effects +from .actors import ( + PipelineStagePayload, + TaggedPipelineStagePayload, + pipeline_stage, + tagged_pipeline_stage, +) + +if TYPE_CHECKING: + import asyncpg + + from taskq import TaskQ + + from .conftest import E2ESchema, E2EWorker + +pytestmark = [pytest.mark.e2e, pytest.mark.timeout(900)] + + +async def test_sub_job_inherits_tags_in_pipeline( + e2e_client: TaskQ, + e2e_worker: E2EWorker, + e2e_pg_pool: asyncpg.Pool, + e2e_schema: E2ESchema, + run_id: str, +) -> None: + """Sub-jobs enqueued via ctx.jobs.enqueue() inherit parent tags. + + A pipeline of 3 stages is enqueued with tag "run-{run_id}". Each stage + enqueues the next via ctx.jobs.enqueue() with no explicit tags. All 3 + jobs should be findable by the tag filter. + """ + tag = f"run-{run_id[:8]}" + await e2e_client.enqueue( + pipeline_stage, + PipelineStagePayload(run_id=run_id, stage=1, total_stages=3), + tags=[tag], + ) + + await wait_for_effects( + e2e_pg_pool, + e2e_schema.schema_name, + run_id, + kind="stage", + min_count=3, + timeout=30, + ) + + page = await e2e_client.list(JobFilter(tags=(tag,))) + assert len(page.jobs) == 3, ( + f"Expected 3 jobs with tag {tag!r}, found {len(page.jobs)}: {[j.id for j in page.jobs]}" + ) + + +async def test_sub_job_explicit_tags_merge_with_parent( + e2e_client: TaskQ, + e2e_worker: E2EWorker, + e2e_pg_pool: asyncpg.Pool, + e2e_schema: E2ESchema, + run_id: str, +) -> None: + """Explicit tags on sub-job merge with inherited parent tags. + + Stage 1 is enqueued with parent tag "run-{run_id}". Each stage enqueues + the next with an explicit stage tag (e.g. "stage-2"). The sub-job should + carry both the inherited run tag and the explicit stage tag. + """ + parent_tag = f"run-{run_id[:8]}" + + await e2e_client.enqueue( + tagged_pipeline_stage, + TaggedPipelineStagePayload(run_id=run_id, stage=1, total_stages=3), + tags=[parent_tag], + ) + + await wait_for_effects( + e2e_pg_pool, + e2e_schema.schema_name, + run_id, + kind="stage", + min_count=3, + timeout=30, + ) + + page = await e2e_client.list(JobFilter(tags=(parent_tag,))) + assert len(page.jobs) == 3, ( + f"Expected 3 jobs with parent tag {parent_tag!r}, found {len(page.jobs)}" + ) + + stage2 = await e2e_client.list(JobFilter(tags=("stage-2",))) + assert len(stage2.jobs) == 2, ( + f"Expected 2 jobs with stage-2 tag (stage 2 + stage 3 inherited it), found {len(stage2.jobs)}" + ) + stage2_only = [j for j in stage2.jobs if "stage-3" not in j.tags] + assert len(stage2_only) == 1, "Expected exactly one job with stage-2 but not stage-3" + assert parent_tag in stage2_only[0].tags + assert "stage-2" in stage2_only[0].tags diff --git a/tests/e2e/worker_entry.py b/tests/e2e/worker_entry.py index 71b8bce3..8b4855a8 100644 --- a/tests/e2e/worker_entry.py +++ b/tests/e2e/worker_entry.py @@ -22,12 +22,14 @@ import_contacts_csv, long_running_job, loop_blocker_job, + pipeline_stage, quick_result, rebuild_search_index, send_welcome_email, short_lived_job, slow_deliver_webhook, sync_user_profile, + tagged_pipeline_stage, ) from e2e.di import build_registry from taskq import ActorRef @@ -53,6 +55,8 @@ "cron_heartbeat": cron_heartbeat, "short_lived_job": short_lived_job, "concurrent_tracked_worker": concurrent_tracked_worker, + "pipeline_stage": pipeline_stage, + "tagged_pipeline_stage": tagged_pipeline_stage, } diff --git a/tests/test_backend_protocol.py b/tests/test_backend_protocol.py index 3090dbd4..c8d5b188 100644 --- a/tests/test_backend_protocol.py +++ b/tests/test_backend_protocol.py @@ -215,9 +215,9 @@ def test_isinstance_dict_is_false(self) -> None: class TestMethodCount: - def test_exactly_thirty_six_public_members(self) -> None: + def test_exactly_thirty_seven_public_members(self) -> None: public = [m for m in dir(Backend) if not m.startswith("_")] - assert len(public) == 36, f"Expected 36 public members, got {len(public)}: {public}" + assert len(public) == 37, f"Expected 37 public members, got {len(public)}: {public}" def test_all_member_names_present(self) -> None: expected = { @@ -243,6 +243,7 @@ def test_all_member_names_present(self) -> None: "get_events", "poll_reclaim_events", "write_cancel_request", + "cancel_where", "poll_cancel_flags", "scheduled_to_pending", "deadline_sweep", @@ -262,6 +263,13 @@ def test_all_member_names_present(self) -> None: assert actual == expected +async def test_protocol_has_cancel_where() -> None: + """Backend protocol declares cancel_where.""" + from taskq.backend._protocol import Backend + + assert hasattr(Backend, "cancel_where") + + # ── Return-type annotations ──────────────────────────────────────────── diff --git a/tests/test_bulk_cancel_types.py b/tests/test_bulk_cancel_types.py new file mode 100644 index 00000000..6bd38d10 --- /dev/null +++ b/tests/test_bulk_cancel_types.py @@ -0,0 +1,75 @@ +from uuid import uuid4 + +import pytest +from pydantic import ValidationError + +from taskq.exceptions import EmptyFilterError, TaskQError +from taskq.types import BulkCancelResult + + +class TestBulkCancelResult: + def test_construction(self) -> None: + ids = [uuid4() for _ in range(3)] + result = BulkCancelResult( + cancelled_directly=2, + cancel_requested=1, + cancelled_ids=ids[:2], + cancel_requested_ids=ids[2:], + ) + assert result.cancelled_directly == 2 + assert result.cancel_requested == 1 + assert result.total_affected == 3 + assert len(result.cancelled_ids) == 2 + assert len(result.cancel_requested_ids) == 1 + + def test_frozen(self) -> None: + result = BulkCancelResult( + cancelled_directly=0, + cancel_requested=0, + cancelled_ids=[], + cancel_requested_ids=[], + ) + with pytest.raises(ValidationError): + result.cancelled_directly = 1 # type: ignore[misc] + + def test_zero_counts(self) -> None: + result = BulkCancelResult( + cancelled_directly=0, + cancel_requested=0, + cancelled_ids=[], + cancel_requested_ids=[], + ) + assert result.total_affected == 0 + + def test_ids_are_tuples(self) -> None: + """ID fields are tuples, not lists — frozen immutability.""" + result = BulkCancelResult( + cancelled_directly=1, + cancel_requested=0, + cancelled_ids=[uuid4()], + cancel_requested_ids=[], + ) + assert isinstance(result.cancelled_ids, tuple) + assert isinstance(result.cancel_requested_ids, tuple) + + def test_list_input_coerced_to_tuple(self) -> None: + """Pydantic v2 coerces list inputs to tuples.""" + ids = [uuid4() for _ in range(2)] + result = BulkCancelResult( + cancelled_directly=2, + cancel_requested=0, + cancelled_ids=ids, + cancel_requested_ids=[], + ) + assert isinstance(result.cancelled_ids, tuple) + assert len(result.cancelled_ids) == 2 + + +class TestEmptyFilterError: + def test_is_taskq_error(self) -> None: + assert issubclass(EmptyFilterError, TaskQError) + + def test_message_mentions_guardrail(self) -> None: + err = EmptyFilterError() + assert "allow_empty_filter" in str(err) + assert "filter predicate" in str(err) diff --git a/tests/test_cancel_where.py b/tests/test_cancel_where.py new file mode 100644 index 00000000..0865579a --- /dev/null +++ b/tests/test_cancel_where.py @@ -0,0 +1,290 @@ +"""Unit tests for InMemoryBackend.cancel_where (bulk cancel).""" + +from dataclasses import replace +from datetime import UTC, datetime +from uuid import uuid4 + +from taskq.backend._protocol import CancelPhase, JobFilter +from taskq.testing.clock import FakeClock +from taskq.testing.in_memory import InMemoryBackend +from taskq.testing.jobs import make_enqueue_args +from taskq.types import BulkCancelResult + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +async def test_cancel_where_pending_jobs() -> None: + """cancel_where moves pending jobs straight to 'cancelled'.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + for _ in range(3): + await backend.enqueue(make_enqueue_args(tags=("tenant-acme", "run-001"), scheduled_at=_NOW)) + for _ in range(2): + await backend.enqueue(make_enqueue_args(tags=("tenant-other",), scheduled_at=_NOW)) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancelled_directly == 3 + assert result.cancel_requested == 0 + assert result.total_affected == 3 + assert len(result.cancelled_ids) == 3 + + remaining = await backend.list_jobs(JobFilter(tags=("tenant-other",))) + assert len(remaining) == 2 + assert all(r.status == "pending" for r in remaining) + + +async def test_cancel_where_running_jobs_cooperative() -> None: + """cancel_where sets cancel_phase=1 for running jobs (cooperative).""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + backend._jobs[row.id] = replace( + backend._jobs[row.id], status="running", locked_by_worker=uuid4() + ) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancelled_directly == 0 + assert result.cancel_requested == 1 + assert len(result.cancel_requested_ids) == 1 + + updated = await backend.get(row.id) + assert updated is not None + assert updated.status == "running" + assert updated.cancel_phase == CancelPhase.COOPERATIVE + + +async def test_cancel_where_mixed_statuses() -> None: + """cancel_where handles both pending and running in one call.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + for _ in range(2): + await backend.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + + args3 = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row3 = await backend.enqueue(args3) + backend._jobs[row3.id] = replace( + backend._jobs[row3.id], status="running", locked_by_worker=uuid4() + ) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancelled_directly == 2 + assert result.cancel_requested == 1 + assert result.total_affected == 3 + + +async def test_cancel_where_no_matches_returns_zero() -> None: + """cancel_where with a filter matching nothing returns zero counts.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + await backend.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + + result = await backend.cancel_where( + JobFilter(tags=("nonexistent",)), + reason="offboard", + ) + + assert result.cancelled_directly == 0 + assert result.cancel_requested == 0 + assert result.total_affected == 0 + + +async def test_cancel_where_already_cancelled_not_affected() -> None: + """Already-terminal jobs are not re-cancelled.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + backend._jobs[row.id] = replace(backend._jobs[row.id], status="cancelled", finished_at=_NOW) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.total_affected == 0 + + +async def test_cancel_where_filter_by_batch_id() -> None: + """cancel_where works with batch_id filter.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + bid = uuid4() + for _ in range(3): + args = make_enqueue_args( + tags=("tenant-acme",), + scheduled_at=_NOW, + metadata={"batch_id": str(bid)}, + ) + await backend.enqueue(args) + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW, metadata={"batch_id": str(uuid4())})) + + result = await backend.cancel_where( + JobFilter(batch_id=bid), + reason="batch abort", + ) + + assert result.cancelled_directly == 3 + + +async def test_cancel_where_filter_by_queue_and_actor() -> None: + """cancel_where works with queue and actor filters.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + await backend.enqueue(make_enqueue_args(queue="default", actor="worker-a", scheduled_at=_NOW)) + await backend.enqueue(make_enqueue_args(queue="default", actor="worker-b", scheduled_at=_NOW)) + await backend.enqueue(make_enqueue_args(queue="priority", actor="worker-a", scheduled_at=_NOW)) + + result = await backend.cancel_where( + JobFilter(queue="default", actor="worker-a"), + reason="abort", + ) + + assert result.cancelled_directly == 1 + + +async def test_cancel_where_active_filter() -> None: + """cancel_where with active=True targets only non-terminal jobs.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + args3 = make_enqueue_args(scheduled_at=_NOW) + row3 = await backend.enqueue(args3) + backend._jobs[row3.id] = replace(backend._jobs[row3.id], status="succeeded", finished_at=_NOW) + + result = await backend.cancel_where( + JobFilter(active=True), + reason="drain", + ) + + assert result.cancelled_directly == 2 + + +async def test_cancel_where_ignores_filter_limit() -> None: + """cancel_where cancels ALL matching jobs even when filter.limit is small. + + Guards against the H3 bug: _list_jobs applies filters.limit (default + 100). If _cancel_where reuses _list_jobs without sanitizing the filter, + a caller passing JobFilter(limit=5, tags=...) would cancel only 5 jobs. + """ + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + for _ in range(11): + await backend.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",), limit=5), + reason="offboard", + ) + + assert result.cancelled_directly == 11 + assert result.total_affected == 11 + + +async def test_cancel_where_already_cooperative_cancel_not_recounted() -> None: + """A running job already in cooperative cancel (cancel_phase=1) is + NOT double-counted by a subsequent cancel_where call.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + worker_id = uuid4() + backend._jobs[row.id] = replace( + backend._jobs[row.id], + status="running", + locked_by_worker=worker_id, + cancel_phase=CancelPhase.COOPERATIVE, + cancel_requested_at=_NOW, + ) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancel_requested == 0 + assert result.total_affected == 0 + + +async def test_cancel_where_events_inserted() -> None: + """cancel_where inserts state_change + cancel_request events for + pending jobs, matching single-job write_cancel_request semantics.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + + await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="test", + ) + + events = await backend.get_events(row.id) + kinds = [e.kind for e in events] + assert "state_change" in kinds + assert "cancel_request" in kinds + sc = [e for e in events if e.kind == "state_change"] + assert sc[0].detail["from_state"] in ("pending", "scheduled") + assert sc[0].detail["to_state"] == "cancelled" + + +async def test_cancel_where_running_event_only_cancel_request() -> None: + """Running jobs get only cancel_request (no state_change).""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + backend._jobs[row.id] = replace( + backend._jobs[row.id], status="running", locked_by_worker=uuid4() + ) + + await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="test", + ) + + events = await backend.get_events(row.id) + kinds = [e.kind for e in events] + assert "cancel_request" in kinds + assert "state_change" not in kinds + + +async def test_cancel_where_wakes_cancel_subscribers() -> None: + """cancel_where sets cancel wake events for running jobs.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + backend._jobs[row.id] = replace( + backend._jobs[row.id], status="running", locked_by_worker=uuid4() + ) + + async with backend.subscribe_cancel_wake() as event: + assert not event.is_set() + await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="test", + ) + assert event.is_set() + + +async def test_cancel_where_returns_bulk_cancel_result_type() -> None: + """cancel_where returns a BulkCancelResult instance.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + await backend.enqueue(make_enqueue_args(tags=("x",), scheduled_at=_NOW)) + + result = await backend.cancel_where(JobFilter(tags=("x",)), reason=None) + + assert isinstance(result, BulkCancelResult) diff --git a/tests/test_cancel_where_client.py b/tests/test_cancel_where_client.py new file mode 100644 index 00000000..f6e50176 --- /dev/null +++ b/tests/test_cancel_where_client.py @@ -0,0 +1,181 @@ +"""Client-layer unit tests for JobsClient.cancel_where and TaskQ.cancel_where.""" + +from datetime import UTC, datetime + +import pytest +from opentelemetry.sdk.metrics._internal.point import NumberDataPoint +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Metric + +from taskq.backend._protocol import JobFilter +from taskq.client._jobs import JobsClient +from taskq.exceptions import EmptyFilterError +from taskq.testing.clock import FakeClock +from taskq.testing.in_memory import InMemoryBackend +from taskq.testing.jobs import make_enqueue_args +from taskq.types import BulkCancelResult + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +@pytest.fixture +def otel_requested_reader(monkeypatch: pytest.MonkeyPatch) -> InMemoryMetricReader: + """Per-test OTel meter isolation for the cancel-requested counter.""" + + from opentelemetry.sdk.metrics import MeterProvider + + import taskq.obs as obs_mod + import taskq.obs._otel as otel_mod + + reader = InMemoryMetricReader() + new_provider = MeterProvider(metric_readers=[reader]) + new_meter = new_provider.get_meter(obs_mod.INSTRUMENTATION_NAME, otel_mod._version()) + + monkeypatch.setattr( + otel_mod, + "_cancellation_requested", + new_meter.create_counter("taskq.cancellation.requested"), + ) + + return reader + + +def _collect_metrics(reader: InMemoryMetricReader) -> list[Metric]: + md = reader.get_metrics_data() + assert md is not None + results: list[Metric] = [] + for rm in md.resource_metrics: + for sm in rm.scope_metrics: + results.extend(sm.metrics) + return results + + +def _data_points(reader: InMemoryMetricReader, metric_name: str) -> list[NumberDataPoint]: + for m in _collect_metrics(reader): + if m.name == metric_name: + return list(m.data.data_points) # type: ignore[return-value] # Why: counter metrics always produce NumberDataPoint instances. + return [] + + +async def test_client_cancel_where_with_tags() -> None: + """JobsClient.cancel_where cancels jobs by tag filter.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + for _ in range(3): + await backend.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + + result = await client.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert isinstance(result, BulkCancelResult) + assert result.cancelled_directly == 3 + + +async def test_client_cancel_where_empty_filter_raises() -> None: + """Empty filter (no predicates) raises EmptyFilterError.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + with pytest.raises(EmptyFilterError, match="filter predicate"): + await client.cancel_where(JobFilter(), reason="oops") + + +async def test_client_cancel_where_empty_tags_tuple_raises() -> None: + """JobFilter(tags=()) is an empty filter — must raise EmptyFilterError.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + with pytest.raises(EmptyFilterError): + await client.cancel_where(JobFilter(tags=()), reason="oops") + + +async def test_client_cancel_where_empty_filter_override() -> None: + """allow_empty_filter=True bypasses the guardrail.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + + result = await client.cancel_where( + JobFilter(), + reason="drain all", + allow_empty_filter=True, + ) + + assert result.cancelled_directly == 2 + + +async def test_client_cancel_where_with_status_filter() -> None: + """Status filter alone is a valid predicate (not empty).""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + + result = await client.cancel_where( + JobFilter(status="pending"), + reason="drain pending", + ) + + assert result.cancelled_directly == 1 + + +async def test_client_cancel_where_with_active_filter() -> None: + """Active filter alone is a valid predicate (not empty).""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + await backend.enqueue(make_enqueue_args(scheduled_at=_NOW)) + + result = await client.cancel_where( + JobFilter(active=True), + reason="drain active", + ) + + assert result.cancelled_directly == 1 + + +async def test_client_cancel_where_increments_counter( + otel_requested_reader: InMemoryMetricReader, +) -> None: + """cancel_where increments taskq.cancellation.requested once.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend) + + await backend.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + + await client.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="test", + ) + + dps = _data_points(otel_requested_reader, "taskq.cancellation.requested") + assert len(dps) >= 1 + assert dps[0].value == 1 + + +async def test_client_cancel_where_translates_schema_errors() -> None: + """cancel_where wraps UndefinedTableError in SchemaNotMigratedError.""" + from taskq.exceptions import SchemaNotMigratedError + + backend = InMemoryBackend(clock=FakeClock(_NOW)) + client = JobsClient(backend, settings=type("S", (), {"schema_name": "test_schema"})()) + + await backend.enqueue(make_enqueue_args(tags=("x",), scheduled_at=_NOW)) + + import asyncpg + + original = backend.cancel_where + + async def raise_undefined(*args: object, **kwargs: object) -> None: + raise asyncpg.exceptions.UndefinedTableError("relation does not exist") + + backend.cancel_where = raise_undefined # type: ignore[method-assign] + try: + with pytest.raises(SchemaNotMigratedError): + await client.cancel_where(JobFilter(tags=("x",)), reason="test") + finally: + backend.cancel_where = original # type: ignore[method-assign] diff --git a/tests/test_cancel_where_pg.py b/tests/test_cancel_where_pg.py new file mode 100644 index 00000000..36ac446d --- /dev/null +++ b/tests/test_cancel_where_pg.py @@ -0,0 +1,367 @@ +"""Integration tests for PostgresBackend.cancel_where.""" + +import asyncio +from datetime import UTC, datetime +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import asyncpg +import pytest + +from taskq.backend._cancel_bulk import _cancel_where +from taskq.backend._protocol import Backend, JobFilter +from taskq.testing.fixtures import JobsApp +from taskq.testing.jobs import make_enqueue_args + +_NOW = datetime(2026, 1, 1, tzinfo=UTC) + + +@pytest.mark.integration +class TestCancelWherePostgres: + async def test_pg_cancel_where_pending(self, backend_pair: Backend) -> None: + """PostgresBackend.cancel_where cancels pending jobs by tag.""" + for _ in range(3): + await backend_pair.enqueue(make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW)) + await backend_pair.enqueue(make_enqueue_args(tags=("other",), scheduled_at=_NOW)) + + result = await backend_pair.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancelled_directly == 3 + assert result.cancel_requested == 0 + + remaining = await backend_pair.list_jobs(JobFilter(tags=("other",))) + assert len(remaining) == 1 + assert remaining[0].status in ("pending", "scheduled") + + async def test_pg_cancel_where_events_inserted(self, backend_pair: Backend) -> None: + """cancel_where inserts job_events for cancelled jobs.""" + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend_pair.enqueue(args) + + await backend_pair.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="test", + ) + + events = await backend_pair.get_events(row.id) + kinds = [e.kind for e in events] + assert "state_change" in kinds + assert "cancel_request" in kinds + sc = [e for e in events if e.kind == "state_change"] + assert sc[0].detail.get("from_state") in ("pending", "scheduled") + assert sc[0].detail.get("to_state") == "cancelled" + + async def test_pg_cancel_where_reason_with_quotes(self, backend_pair: Backend) -> None: + """Reason containing double-quotes does not cause DataError.""" + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend_pair.enqueue(args) + + result = await backend_pair.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason='offboard "tenant-acme" \\ done', + ) + + assert result.cancelled_directly == 1 + events = await backend_pair.get_events(row.id) + cr = [e for e in events if e.kind == "cancel_request"] + assert cr[0].detail.get("reason") == 'offboard "tenant-acme" \\ done' + + async def test_pg_cancel_where_batch_id_filter(self, backend_pair: Backend) -> None: + """cancel_where works with batch_id filter on Postgres.""" + bid = uuid4() + for _ in range(3): + await backend_pair.enqueue( + make_enqueue_args( + scheduled_at=_NOW, + metadata={"batch_id": str(bid)}, + ) + ) + + result = await backend_pair.cancel_where( + JobFilter(batch_id=bid), + reason="batch abort", + ) + + assert result.cancelled_directly == 3 + + async def test_pg_cancel_where_running_cooperative( + self, + clean_jobs_app: JobsApp, + ) -> None: + """cancel_where sets cancel_phase=1 for running jobs on Postgres.""" + backend = clean_jobs_app.backend + deps = clean_jobs_app.deps + schema = deps.settings.schema_name + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + + worker_id = uuid4() + async with deps.worker_pool.acquire() as conn: + await conn.execute( + f'UPDATE "{schema}".jobs ' + f"SET status = 'running', locked_by_worker = $1 WHERE id = $2", + worker_id, + row.id, + ) + + result = await backend.cancel_where( + JobFilter(tags=("tenant-acme",)), + reason="offboard", + ) + + assert result.cancelled_directly == 0 + assert result.cancel_requested == 1 + + updated = await backend.get(row.id) + assert updated is not None + assert updated.status == "running" + assert updated.cancel_phase == 1 + + events = await backend.get_events(row.id) + kinds = [e.kind for e in events] + assert "cancel_request" in kinds + assert "state_change" not in kinds + + async def test_pg_cancel_where_does_not_clobber_concurrent_claim( + self, + clean_jobs_app: JobsApp, + ) -> None: + """EPQ regression: a job claimed while cancel_where executes must + NOT be overwritten to terminal 'cancelled'.""" + backend = clean_jobs_app.backend + deps = clean_jobs_app.deps + schema = deps.settings.schema_name + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + worker_id = uuid4() + + claim_conn = await deps.worker_pool.acquire() + try: + claim_tx = claim_conn.transaction() + await claim_tx.start() + await claim_conn.execute( + f'UPDATE "{schema}".jobs ' + f"SET status = 'running', locked_by_worker = $1 WHERE id = $2", + worker_id, + row.id, + ) + + cancel_task = asyncio.create_task( + backend.cancel_where(JobFilter(tags=("tenant-acme",)), reason="offboard") + ) + await asyncio.sleep(0.2) + await claim_tx.commit() + result = await cancel_task + + assert result.cancelled_directly == 0 + updated = await backend.get(row.id) + assert updated is not None + assert updated.status == "running" + finally: + await deps.worker_pool.release(claim_conn) + + async def test_pg_cancel_where_notify_sent_for_running( + self, + clean_jobs_app: JobsApp, + pg_dsn: str, + ) -> None: + """Batched NOTIFY fires on the fleet and per-worker channels.""" + import asyncpg + + from taskq.constants import events_channel, worker_channel + + backend = clean_jobs_app.backend + deps = clean_jobs_app.deps + schema = deps.settings.schema_name + + args = make_enqueue_args(tags=("tenant-acme",), scheduled_at=_NOW) + row = await backend.enqueue(args) + worker_id = uuid4() + async with deps.worker_pool.acquire() as conn: + await conn.execute( + f'UPDATE "{schema}".jobs ' + f"SET status = 'running', locked_by_worker = $1 WHERE id = $2", + worker_id, + row.id, + ) + + received: list[str] = [] + listen_conn = await asyncpg.connect(pg_dsn) + try: + await listen_conn.add_listener( + events_channel(schema), + lambda _c, _p, _ch, payload: received.append(payload), + ) + await listen_conn.add_listener( + worker_channel(schema, str(worker_id)), + lambda _c, _p, _ch, payload: received.append(payload), + ) + result = await backend.cancel_where(JobFilter(tags=("tenant-acme",)), reason="offboard") + assert result.cancel_requested == 1 + await asyncio.sleep(0.3) + finally: + await listen_conn.close() + + assert len(received) == 2 + + +class TestDeadlockRetry: + """Unit tests for the deadlock retry loop in _cancel_where. + + These don't need a real Postgres — they mock the pool to simulate + deadlock-then-success and deadlock-then-exhaustion scenarios. + """ + + @staticmethod + def _mock_pool_and_conn( + fetch_rows: list[dict[str, object] | None] | None = None, + ) -> tuple[MagicMock, MagicMock]: + conn = MagicMock() + + if fetch_rows is not None: + conn.fetchrow = AsyncMock(side_effect=fetch_rows) + else: + conn.fetchrow = AsyncMock(return_value=None) + + conn.executemany = AsyncMock(return_value=None) + + tx = AsyncMock() + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=False) + conn.transaction = MagicMock(return_value=tx) + + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=False) + + pool = MagicMock() + pool.acquire = MagicMock(return_value=conn) + return pool, conn + + @staticmethod + def _ps_success_row() -> dict[str, object]: + return { + "cancelled_directly": 1, + "cancelled_ids": [uuid4()], + "cancelled_prev_statuses": ["pending"], + } + + @staticmethod + def _running_empty_row() -> dict[str, object]: + return { + "cancel_requested": 0, + "cancel_requested_ids": [], + "cancel_requested_workers": [], + } + + @staticmethod + def _running_success_row() -> dict[str, object]: + wid = uuid4() + return { + "cancel_requested": 1, + "cancel_requested_ids": [uuid4()], + "cancel_requested_workers": [wid], + } + + async def test_deadlock_retry_succeeds_on_second_attempt(self) -> None: + """_cancel_where retries on DeadlockDetectedError and succeeds.""" + ps_row = self._ps_success_row() + running_row = self._running_empty_row() + pool, _ = self._mock_pool_and_conn( + fetch_rows=[ + asyncpg.DeadlockDetectedError(), + ps_row, + running_row, + ] + ) + + sql = MagicMock() + sql.insert_event = "INSERT INTO job_events VALUES ($1, $2, $3, $4)" + + result, _notify = await _cancel_where(pool, "taskq", sql, JobFilter(tags=("x",)), "test") + + assert result.cancelled_directly == 1 + assert result.cancel_requested == 0 + assert len(result.cancelled_ids) == 1 + + async def test_deadlock_retry_exhausted_raises(self) -> None: + """_cancel_where raises after 3 failed attempts.""" + pool, _ = self._mock_pool_and_conn( + fetch_rows=[asyncpg.DeadlockDetectedError()] * 3, + ) + + sql = MagicMock() + + with pytest.raises(asyncpg.DeadlockDetectedError): + await _cancel_where(pool, "taskq", sql, JobFilter(tags=("x",)), "test") + + async def test_no_deadlock_no_retry(self) -> None: + """_cancel_where succeeds immediately without retry.""" + ps_row = self._ps_success_row() + running_row = self._running_empty_row() + pool, conn = self._mock_pool_and_conn(fetch_rows=[ps_row, running_row]) + + sql = MagicMock() + sql.insert_event = "INSERT INTO job_events VALUES ($1, $2, $3, $4)" + + await _cancel_where(pool, "taskq", sql, JobFilter(tags=("x",)), "test") + + assert conn.fetchrow.call_count == 2 + + async def test_deadlock_during_executemany_retries_correctly(self) -> None: + """Deadlock on executemany (after fetchrow succeeds) retries the + whole transaction — no phantom IDs from the aborted attempt.""" + ps_row = self._ps_success_row() + running_row = self._running_empty_row() + # Attempt 1: fetchrow → ps_row, executemany → deadlock + # Attempt 2: fetchrow → ps_row, executemany → ok, fetchrow → running_row + pool, conn = self._mock_pool_and_conn(fetch_rows=[ps_row, ps_row, running_row]) + + call_count = [0] + + async def _flaky_executemany( + query: str, + args: list[tuple[object, ...]], + *a: object, + **kw: object, + ) -> None: + call_count[0] += 1 + if call_count[0] == 1: + raise asyncpg.DeadlockDetectedError() + + conn.executemany = _flaky_executemany + + sql = MagicMock() + sql.insert_event = "INSERT INTO job_events VALUES ($1, $2, $3, $4)" + + result, _notify = await _cancel_where(pool, "taskq", sql, JobFilter(tags=("x",)), "test") + + assert result.cancelled_directly == 1 + assert len(result.cancelled_ids) == 1 + + async def test_notify_target_filters_none_worker_id(self) -> None: + """A running job with NULL locked_by_worker is excluded from + notify_targets (no worker to NOTIFY).""" + jid = uuid4() + pool, _ = self._mock_pool_and_conn( + fetch_rows=[ + None, + { + "cancel_requested": 1, + "cancel_requested_ids": [jid], + "cancel_requested_workers": [None], + }, + ] + ) + + sql = MagicMock() + sql.insert_event = "INSERT INTO job_events VALUES ($1, $2, $3, $4)" + + result, notify = await _cancel_where(pool, "taskq", sql, JobFilter(tags=("x",)), "test") + + assert result.cancel_requested == 1 + assert len(notify) == 0 diff --git a/tests/test_filter_sql.py b/tests/test_filter_sql.py new file mode 100644 index 00000000..1c8a681b --- /dev/null +++ b/tests/test_filter_sql.py @@ -0,0 +1,164 @@ +"""Unit tests for the shared filter→SQL WHERE condition builder. + +Verifies that ``build_filter_conditions`` translates only the predicate +fields of ``JobFilter`` (queue, status, actor, identity_key, batch_id, +tags, active) into SQL fragments and parameters, and that cursor / +order_by are ignored (they are handled by the caller). +""" + +import re +from uuid import uuid4 + +from taskq.backend._filter_sql import build_filter_conditions +from taskq.backend._protocol import IdentityKey, JobFilter +from taskq.backend.statemachine import ACTIVE_STATUSES, TERMINAL_STATUSES + + +class TestBuildFilterConditions: + def test_empty_filter_produces_no_conditions(self) -> None: + result = build_filter_conditions(JobFilter()) + assert result.conditions == () + assert result.params == () + + def test_queue_filter(self) -> None: + result = build_filter_conditions(JobFilter(queue="default")) + assert len(result.conditions) == 1 + assert "queue" in result.conditions[0] + assert result.params == ("default",) + + def test_tags_filter(self) -> None: + result = build_filter_conditions(JobFilter(tags=("alpha", "beta"))) + assert len(result.conditions) == 1 + assert "tags" in result.conditions[0] + assert result.params == (["alpha", "beta"],) + + def test_batch_id_filter(self) -> None: + bid = uuid4() + result = build_filter_conditions(JobFilter(batch_id=bid)) + assert len(result.conditions) == 1 + assert "metadata" in result.conditions[0] + + def test_active_true_filter(self) -> None: + result = build_filter_conditions(JobFilter(active=True)) + assert len(result.conditions) == 1 + assert "status" in result.conditions[0] + assert set(result.params[0]) == set(ACTIVE_STATUSES) # type: ignore[arg-type] + + def test_active_false_filter(self) -> None: + result = build_filter_conditions(JobFilter(active=False)) + assert len(result.conditions) == 1 + assert "status" in result.conditions[0] + assert set(result.params[0]) == set(TERMINAL_STATUSES) # type: ignore[arg-type] + + def test_status_sequence_filter(self) -> None: + result = build_filter_conditions(JobFilter(status=["pending", "running"])) + assert len(result.conditions) == 1 + assert "ANY" in result.conditions[0] + assert result.params == (["pending", "running"],) + + def test_identity_key_filter(self) -> None: + key = IdentityKey("tenant-acme") + result = build_filter_conditions(JobFilter(identity_key=key)) + assert len(result.conditions) == 1 + assert "identity_key" in result.conditions[0] + assert result.params == (key,) + + def test_combined_filters(self) -> None: + result = build_filter_conditions( + JobFilter(queue="e2e", actor="my_actor", tags=("run-123",)), + ) + assert len(result.conditions) == 3 + + def test_cursor_and_order_by_ignored(self) -> None: + result = build_filter_conditions( + JobFilter(cursor="some-cursor", order_by=None), + ) + assert result.conditions == () + + def test_parameter_numbering_is_sequential(self) -> None: + """$N placeholders must be sequentially numbered and align + positionally with the params tuple.""" + result = build_filter_conditions(JobFilter(queue="q1", actor="a1", tags=("t1",))) + numbers: list[int] = [] + for c in result.conditions: + m = re.search(r"\$(\d+)", c) + assert m is not None, f"No $N placeholder in condition: {c!r}" + numbers.append(int(m.group(1))) + assert numbers == list(range(1, len(numbers) + 1)), numbers + assert result.params == ("q1", "a1", ["t1"]) + + +class TestSQLInjectionSafety: + """Verify user-supplied values are always parameterized, never + interpolated into the SQL condition string. + + Every condition must use ``$N`` positional binding — no raw user + value may appear in ``conditions``. Column names are hardcoded + literals, not derived from input. + """ + + def test_queue_with_sql_metacharacters_is_parameterized(self) -> None: + payload = "'; DROP TABLE jobs; --" + result = build_filter_conditions(JobFilter(queue=payload)) + assert result.params == (payload,) + for cond in result.conditions: + assert payload not in cond + assert "$" in cond + + def test_actor_with_sql_metacharacters_is_parameterized(self) -> None: + payload = "admin' OR '1'='1" + result = build_filter_conditions(JobFilter(actor=payload)) + assert result.params == (payload,) + for cond in result.conditions: + assert payload not in cond + + def test_tags_with_sql_metacharacters_are_parameterized(self) -> None: + payload = ("'; DELETE FROM jobs WHERE '1'='1",) + result = build_filter_conditions(JobFilter(tags=payload)) + assert result.params == (list(payload),) + for cond in result.conditions: + assert payload[0] not in cond + + def test_identity_key_with_sql_metacharacters_is_parameterized(self) -> None: + payload = IdentityKey("x'; DROP TABLE jobs; --") + result = build_filter_conditions(JobFilter(identity_key=payload)) + assert result.params == (payload,) + for cond in result.conditions: + assert str(payload) not in cond + + def test_conditions_only_contain_placeholders_and_column_names(self) -> None: + """Every condition string must be composed solely of known column + names, operators, and $N placeholders — never raw user input.""" + result = build_filter_conditions( + JobFilter( + queue="myqueue", + actor="myactor", + tags=("tag1", "tag2"), + batch_id=uuid4(), + ) + ) + allowed_substrings = { + "queue", + "status", + "actor", + "identity_key", + "metadata", + "tags", + " = $", + " = ANY($", + " @> $", + " && $", + "::jsonb", + "::text[]", + } + for cond in result.conditions: + stripped = cond + for s in allowed_substrings: + stripped = stripped.replace(s, "") + stripped = stripped.replace("$", "").replace("0", "").replace("1", "") + stripped = stripped.replace("2", "").replace("3", "").replace("4", "") + stripped = stripped.replace("5", "").replace("6", "").replace("7", "") + stripped = stripped.replace("8", "").replace("9", "") + assert stripped == "", ( + f"Unexpected content in condition: {cond!r} (residue: {stripped!r})" + ) diff --git a/tests/test_sub_job_tags.py b/tests/test_sub_job_tags.py new file mode 100644 index 00000000..d833c35a --- /dev/null +++ b/tests/test_sub_job_tags.py @@ -0,0 +1,298 @@ +"""Unit tests for SubJobEnqueuer.enqueue tags parameter and parent-tag inheritance.""" + +from datetime import UTC, datetime, timedelta +from uuid import UUID + +import pytest +from pydantic import BaseModel, TypeAdapter + +from taskq.actor import ActorRef +from taskq.backend._protocol import JobId +from taskq.client._enqueuer import SubJobEnqueuer, _parent_tags_var, set_parent_tags +from taskq.retry import RetryPolicy +from taskq.testing.clock import FakeClock +from taskq.testing.in_memory import InMemoryBackend + +_NOW = datetime(2025, 1, 1, tzinfo=UTC) + + +class _Payload(BaseModel): + value: str = "test" + + +class _Result(BaseModel): + ok: bool = True + + +def _make_actor_ref(name: str = "child") -> ActorRef[_Payload, _Result]: + async def _handler(payload: _Payload) -> _Result: + return _Result() + + return ActorRef( + name=name, + queue="default", + fn=_handler, + wants_ctx=False, + dependencies={}, + payload_type=_Payload, + result_adapter=TypeAdapter(_Result), + retry=RetryPolicy(), + result_ttl=None, + singleton=False, + unique_for=None, + max_pending=None, + ) + + +def _make_enqueuer(backend: InMemoryBackend | None = None) -> SubJobEnqueuer: + if backend is None: + backend = InMemoryBackend(clock=FakeClock(_NOW)) + return SubJobEnqueuer( + loop_scope_resolved=None, + worker_pool=object(), + backend=backend, + clock=FakeClock(_NOW), + ) + + +class TestSubJobExplicitTags: + async def test_explicit_tags_no_inheritance(self) -> None: + """tags= with inherit_tags=False sets only explicit tags.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["alpha", "beta"], + inherit_tags=False, + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("alpha", "beta") + + async def test_explicit_tags_no_parent_tags_inherit_true(self) -> None: + """inherit_tags=True (default) with no parent tags and explicit + tags returns only the explicit tags — no merge with empty parent.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["alpha"], + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("alpha",) + + async def test_tags_validated(self) -> None: + """Invalid tags raise ValueError.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + with pytest.raises(ValueError, match="invalid tag"): + await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["ab"], + ) + + async def test_tags_deduplicated(self) -> None: + """Duplicate tags are deduplicated.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["alpha", "alpha", "beta"], + inherit_tags=False, + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("alpha", "beta") + + +class TestSubJobTagInheritance: + async def test_inherit_parent_tags_default(self) -> None: + """With no explicit tags and default inherit_tags=True, sub-job inherits parent tags.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + token = set_parent_tags(("run-001", "tenant-acme")) + try: + handle = await enqueuer.enqueue(_make_actor_ref(), _Payload()) + finally: + _parent_tags_var.reset(token) + + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("run-001", "tenant-acme") + + async def test_inherit_and_merge_tags(self) -> None: + """Explicit tags merge with parent tags (parent first, deduped).""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + token = set_parent_tags(("run-001", "tenant-acme")) + try: + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["stage-2", "tenant-acme"], + ) + finally: + _parent_tags_var.reset(token) + + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("run-001", "tenant-acme", "stage-2") + + async def test_no_parent_tags_no_explicit_tags(self) -> None: + """With no parent tags and no explicit tags, sub-job has empty tags.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue(_make_actor_ref(), _Payload()) + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == () + + async def test_inherit_false_no_parent_tags(self) -> None: + """inherit_tags=False with no explicit tags -> empty tags.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + token = set_parent_tags(("run-001",)) + try: + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + inherit_tags=False, + ) + finally: + _parent_tags_var.reset(token) + + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == () + + async def test_inherit_false_with_explicit_tags(self) -> None: + """inherit_tags=False with explicit tags -> only explicit tags.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + token = set_parent_tags(("run-001",)) + try: + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=["custom-tag"], + inherit_tags=False, + ) + finally: + _parent_tags_var.reset(token) + + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("custom-tag",) + + async def test_empty_list_tags_with_inheritance(self) -> None: + """tags=[] with inherit_tags=True and parent tags -> inherits parent tags only. + + An empty list means "no additional tags to add" -- parent tags + are still inherited. This is the intuitive behavior: the caller + didn't add any new tags, so the sub-job carries what the parent + carried. + """ + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + token = set_parent_tags(("run-001", "tenant-acme")) + try: + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + tags=[], + ) + finally: + _parent_tags_var.reset(token) + + row = await backend.get(handle.job_id) + assert row is not None + assert row.tags == ("run-001", "tenant-acme") + + +class TestSubJobMissingFields: + async def test_schedule_to_close(self) -> None: + """schedule_to_close is accepted and stored on the row.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + deadline = _NOW + timedelta(hours=1) + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + schedule_to_close=deadline, + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.schedule_to_close == deadline + + async def test_start_to_close(self) -> None: + """start_to_close is accepted and stored on the row.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + start_to_close=timedelta(minutes=30), + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.start_to_close == timedelta(minutes=30) + + async def test_heartbeat_timeout(self) -> None: + """heartbeat_timeout is accepted and stored on the row.""" + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + handle = await enqueuer.enqueue( + _make_actor_ref(), + _Payload(), + heartbeat_timeout=timedelta(seconds=10), + ) + row = await backend.get(handle.job_id) + assert row is not None + assert row.heartbeat_timeout == timedelta(seconds=10) + + +class TestContextVarIsolation: + async def test_concurrent_jobs_separate_parent_tags(self) -> None: + """ContextVar ensures concurrent consumers don't share parent tags.""" + import asyncio + + backend = InMemoryBackend(clock=FakeClock(_NOW)) + enqueuer = _make_enqueuer(backend) + + async def enqueue_with_parent(parent_tags: tuple[str, ...]) -> UUID: + token = set_parent_tags(parent_tags) + try: + handle = await enqueuer.enqueue(_make_actor_ref(), _Payload()) + return handle.job_id + finally: + _parent_tags_var.reset(token) + + id1, id2 = await asyncio.gather( + enqueue_with_parent(("run-a",)), + enqueue_with_parent(("run-b",)), + ) + + row1 = await backend.get(JobId(id1)) + row2 = await backend.get(JobId(id2)) + assert row1 is not None + assert row2 is not None + assert row1.tags == ("run-a",) + assert row2.tags == ("run-b",)