From f0d2c94396eeea3717fe4e32b2dbcf063a80b0f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Artur=20Ku=C5=BAmi=C5=84ski?= Date: Thu, 16 Jul 2026 16:54:21 +0200 Subject: [PATCH 1/2] Feature: Add queue.pause() --- docs/howto/advanced.md | 1 + docs/howto/advanced/pause_queue.md | 91 ++++++++++ .../migrations/0042_pre_add_queue_pause.py | 17 ++ .../migrations/0043_post_add_queue_pause.py | 17 ++ procrastinate/jobs.py | 7 +- procrastinate/manager.py | 160 +++++++++++++++++- .../03.10.00_01_pre_add_queue_pause.sql | 86 ++++++++++ .../03.10.00_50_post_add_queue_pause.sql | 4 + procrastinate/sql/queries.sql | 21 ++- procrastinate/sql/schema.sql | 39 ++++- procrastinate/testing.py | 32 ++++ procrastinate/worker.py | 2 +- tests/acceptance/test_async.py | 130 +++++++++++++- tests/integration/test_manager.py | 100 +++++++++++ tests/unit/test_manager.py | 128 ++++++++++++++ tests/unit/test_worker.py | 32 ++++ 16 files changed, 858 insertions(+), 9 deletions(-) create mode 100644 docs/howto/advanced/pause_queue.md create mode 100644 procrastinate/contrib/django/migrations/0042_pre_add_queue_pause.py create mode 100644 procrastinate/contrib/django/migrations/0043_post_add_queue_pause.py create mode 100644 procrastinate/sql/migrations/03.10.00_01_pre_add_queue_pause.sql create mode 100644 procrastinate/sql/migrations/03.10.00_50_post_add_queue_pause.sql diff --git a/docs/howto/advanced.md b/docs/howto/advanced.md index 51884d9b5..2fd52b267 100644 --- a/docs/howto/advanced.md +++ b/docs/howto/advanced.md @@ -9,6 +9,7 @@ advanced/locks advanced/schedule advanced/priorities advanced/cancellation +advanced/pause_queue advanced/queueing_locks advanced/cron advanced/retry diff --git a/docs/howto/advanced/pause_queue.md b/docs/howto/advanced/pause_queue.md new file mode 100644 index 000000000..3a2b5b143 --- /dev/null +++ b/docs/howto/advanced/pause_queue.md @@ -0,0 +1,91 @@ +# Pause a queue + +We can pause a queue so that workers stop fetching its jobs, then resume it +later. This is useful to temporarily hold a queue's work, for example during a +maintenance window or while an external dependency is unavailable, without +stopping the worker process. + +Pausing only stops the *fetching* of new jobs: jobs already being processed run +to completion, and pending jobs stay in the `todo` state until the queue is +resumed. Workers keep fetching jobs from the other queues. + +The pause is stored in the database, so it is shared by every worker consuming +the queue and it survives a worker restart. + +## Pause a queue + +```python +# by using the sync method +app.job_manager.pause_queue("some_queue") +# or by using the async method +await app.job_manager.pause_queue_async("some_queue") +``` + +## Resume a queue + +```python +# by using the sync method +app.job_manager.resume_queue("some_queue") +# or by using the async method +await app.job_manager.resume_queue_async("some_queue") +``` + +Resuming a queue that is not paused does nothing. + +## Pause keys: several independent holders + +Every pause is held under a *pause key* (`"default"` when not specified). A +queue is paused as long as it holds at least one key, and each key must be +resumed for the queue to start working again. This lets independent processes +pause the same queue without stepping on each other: if a deploy script and a +maintenance task both pause a queue and the deploy finishes first, its resume +only releases its own key — the queue stays paused until the maintenance task +resumes too. + +```python +# deploy script +await app.job_manager.pause_queue_async("some_queue", pause_key="deploy") +... +await app.job_manager.resume_queue_async("some_queue", pause_key="deploy") + +# maintenance task, meanwhile +await app.job_manager.pause_queue_async("some_queue", pause_key="maintenance") +... +await app.job_manager.resume_queue_async("some_queue", pause_key="maintenance") +``` + +## Resume all keys at once + +A pause whose holder crashed before resuming (or was never going to resume) +stays in place until it is resumed under the same key. As an escape hatch, +`all_keys=True` removes every pause key from a queue, regardless of who holds +them: + +```python +app.job_manager.resume_queue("some_queue", all_keys=True) +``` + +## List paused queues + +```python +# by using the sync method +app.job_manager.list_paused_queues() +# or by using the async method +await app.job_manager.list_paused_queues_async() +``` + +This returns one `dict` per held pause key, with `queue_name`, `pause_key` and +`paused_at` keys — a queue paused by several holders appears once per key: + +```python +[ + {"queue_name": "some_queue", "pause_key": "deploy", "paused_at": ...}, + {"queue_name": "some_queue", "pause_key": "maintenance", "paused_at": ...}, +] +``` + +Both methods accept optional `queue` and `pause_key` filters: + +```python +await app.job_manager.list_paused_queues_async(queue="some_queue") +``` diff --git a/procrastinate/contrib/django/migrations/0042_pre_add_queue_pause.py b/procrastinate/contrib/django/migrations/0042_pre_add_queue_pause.py new file mode 100644 index 000000000..1c032ae6b --- /dev/null +++ b/procrastinate/contrib/django/migrations/0042_pre_add_queue_pause.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from django.db import migrations + +from .. import migrations_utils + + +class Migration(migrations.Migration): + operations = [ + migrations_utils.RunProcrastinateSQL( + name="03.10.00_01_pre_add_queue_pause.sql" + ), + ] + name = "0042_pre_add_queue_pause" + dependencies = [ + ("procrastinate", "0041_post_retry_failed_job"), + ] diff --git a/procrastinate/contrib/django/migrations/0043_post_add_queue_pause.py b/procrastinate/contrib/django/migrations/0043_post_add_queue_pause.py new file mode 100644 index 000000000..d191990c3 --- /dev/null +++ b/procrastinate/contrib/django/migrations/0043_post_add_queue_pause.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from django.db import migrations + +from .. import migrations_utils + + +class Migration(migrations.Migration): + operations = [ + migrations_utils.RunProcrastinateSQL( + name="03.10.00_50_post_add_queue_pause.sql" + ), + ] + name = "0043_post_add_queue_pause" + dependencies = [ + ("procrastinate", "0042_pre_add_queue_pause"), + ] diff --git a/procrastinate/jobs.py b/procrastinate/jobs.py index e124341b4..a567139bc 100644 --- a/procrastinate/jobs.py +++ b/procrastinate/jobs.py @@ -32,7 +32,12 @@ class AbortJobRequested(TypedDict): job_id: int -Notification: TypeAlias = JobInserted | AbortJobRequested +class QueueResumed(TypedDict): + type: Literal["queue_resumed"] + queue_name: str + + +Notification: TypeAlias = JobInserted | AbortJobRequested | QueueResumed def check_aware( diff --git a/procrastinate/manager.py b/procrastinate/manager.py index 862fc5113..db1e50a59 100644 --- a/procrastinate/manager.py +++ b/procrastinate/manager.py @@ -5,7 +5,7 @@ import logging import warnings from collections.abc import Awaitable, Iterable -from typing import Any, NoReturn, Protocol +from typing import Any, NoReturn, Protocol, TypedDict from procrastinate import connector, exceptions, sql, types, utils from procrastinate import jobs as jobs_module @@ -15,6 +15,12 @@ QUEUEING_LOCK_CONSTRAINT = "procrastinate_jobs_queueing_lock_idx_v1" +class PausedQueue(TypedDict): + queue_name: str + pause_key: str + paused_at: datetime.datetime + + class NotificationCallback(Protocol): def __call__( self, *, channel: str, notification: jobs_module.Notification @@ -454,6 +460,158 @@ async def cancel_job_by_id_async( assert result["id"] == job_id return True + def pause_queue(self, queue_name: str, *, pause_key: str = "default") -> None: + """ + Pause a queue under the given pause key, so that workers stop fetching + jobs from it. Jobs already being processed are not affected and run to + completion. Pending jobs stay in the ``todo`` state until the queue is + resumed. A queue is paused as long as it holds at least one pause key: + independent holders (e.g. a deploy script and a maintenance task) each + pause under their own key, and the queue only resumes once every key has + been resumed. Pausing a queue under a key it already holds does nothing. + + Parameters + ---------- + queue_name: + The name of the queue to pause + pause_key: + The key identifying the holder of the pause + """ + self.connector.get_sync_connector().execute_query( + query=sql.queries["pause_queue"], + queue_name=queue_name, + pause_key=pause_key, + ) + + async def pause_queue_async( + self, queue_name: str, *, pause_key: str = "default" + ) -> None: + """ + Pause a queue under the given pause key, so that workers stop fetching + jobs from it. Jobs already being processed are not affected and run to + completion. Pending jobs stay in the ``todo`` state until the queue is + resumed. A queue is paused as long as it holds at least one pause key: + independent holders (e.g. a deploy script and a maintenance task) each + pause under their own key, and the queue only resumes once every key has + been resumed. Pausing a queue under a key it already holds does nothing. + + Parameters + ---------- + queue_name: + The name of the queue to pause + pause_key: + The key identifying the holder of the pause + """ + await self.connector.execute_query_async( + query=sql.queries["pause_queue"], + queue_name=queue_name, + pause_key=pause_key, + ) + + def resume_queue( + self, queue_name: str, *, pause_key: str = "default", all_keys: bool = False + ) -> None: + """ + Remove the given pause key from a queue. Workers fetch its jobs again + once the queue holds no pause key at all. Resuming a key the queue does + not hold does nothing. With ``all_keys=True``, remove every pause key + from the queue at once, regardless of who holds them. + + Parameters + ---------- + queue_name: + The name of the queue to resume + pause_key: + The key identifying the holder of the pause + all_keys: + If ``True``, ignore ``pause_key`` and remove all pause keys from the + queue + """ + self.connector.get_sync_connector().execute_query( + query=sql.queries["resume_queue"], + queue_name=queue_name, + pause_key=pause_key, + all_keys=all_keys, + ) + + async def resume_queue_async( + self, queue_name: str, *, pause_key: str = "default", all_keys: bool = False + ) -> None: + """ + Remove the given pause key from a queue. Workers fetch its jobs again + once the queue holds no pause key at all. Resuming a key the queue does + not hold does nothing. With ``all_keys=True``, remove every pause key + from the queue at once, regardless of who holds them. + + Parameters + ---------- + queue_name: + The name of the queue to resume + pause_key: + The key identifying the holder of the pause + all_keys: + If ``True``, ignore ``pause_key`` and remove all pause keys from the + queue + """ + await self.connector.execute_query_async( + query=sql.queries["resume_queue"], + queue_name=queue_name, + pause_key=pause_key, + all_keys=all_keys, + ) + + def list_paused_queues( + self, queue: str | None = None, pause_key: str | None = None + ) -> Iterable[PausedQueue]: + """ + Sync version of `list_paused_queues_async` + """ + return [ + PausedQueue( + queue_name=row["queue_name"], + pause_key=row["pause_key"], + paused_at=row["paused_at"], + ) + for row in self.connector.get_sync_connector().execute_query_all( + query=sql.queries["list_paused_queues"], + queue_name=queue, + pause_key=pause_key, + ) + ] + + async def list_paused_queues_async( + self, queue: str | None = None, pause_key: str | None = None + ) -> Iterable[PausedQueue]: + """ + List the pauses currently held on queues. + + Parameters + ---------- + queue: + Filter by queue name + pause_key: + Filter by pause key + + Returns + ------- + : + One `PausedQueue` per held pause key, with ``queue_name``, + ``pause_key`` and ``paused_at`` keys, sorted by queue name then + pause key. A queue appears once per pause key it holds. + """ + return [ + PausedQueue( + queue_name=row["queue_name"], + pause_key=row["pause_key"], + paused_at=row["paused_at"], + ) + for row in await self.connector.execute_query_all_async( + query=sql.queries["list_paused_queues"], + queue_name=queue, + pause_key=pause_key, + ) + ] + def get_job_status( self, job_id: int, connection: Any | None = None ) -> jobs_module.Status: diff --git a/procrastinate/sql/migrations/03.10.00_01_pre_add_queue_pause.sql b/procrastinate/sql/migrations/03.10.00_01_pre_add_queue_pause.sql new file mode 100644 index 000000000..6ff15c320 --- /dev/null +++ b/procrastinate/sql/migrations/03.10.00_01_pre_add_queue_pause.sql @@ -0,0 +1,86 @@ +CREATE TABLE procrastinate_paused_queues ( + id bigint PRIMARY KEY GENERATED ALWAYS AS IDENTITY, + queue_name character varying(128) NOT NULL, + pause_key character varying(128) DEFAULT 'default' NOT NULL, + paused_at timestamp with time zone DEFAULT NOW() NOT NULL, + UNIQUE (queue_name, pause_key) +); + +CREATE FUNCTION procrastinate_fetch_job_v3( + target_queue_names character varying[], + p_worker_id bigint +) + RETURNS procrastinate_jobs + LANGUAGE plpgsql +AS $$ +DECLARE + found_jobs procrastinate_jobs; +BEGIN + WITH candidate AS ( + SELECT jobs.* + FROM procrastinate_jobs AS jobs + WHERE + -- reject the job if its lock has earlier or higher priority jobs + NOT EXISTS ( + SELECT 1 + FROM procrastinate_jobs AS other_jobs + WHERE + jobs.lock IS NOT NULL + AND other_jobs.lock = jobs.lock + AND ( + -- job with same lock is already running + other_jobs.status = 'doing' + OR + -- job with same lock is waiting and has higher priority (or same priority but was queued first) + ( + other_jobs.status = 'todo' + AND ( + other_jobs.priority > jobs.priority + OR ( + other_jobs.priority = jobs.priority + AND other_jobs.id < jobs.id + ) + ) + ) + ) + ) + AND jobs.status = 'todo' + AND (target_queue_names IS NULL OR jobs.queue_name = ANY( target_queue_names )) + AND (jobs.scheduled_at IS NULL OR jobs.scheduled_at <= now()) + -- reject the job if its queue is paused + AND NOT EXISTS ( + SELECT 1 + FROM procrastinate_paused_queues AS paused + WHERE paused.queue_name = jobs.queue_name + ) + ORDER BY jobs.priority DESC, jobs.id ASC LIMIT 1 + FOR UPDATE OF jobs SKIP LOCKED + ) + UPDATE procrastinate_jobs + SET status = 'doing', worker_id = p_worker_id + FROM candidate + WHERE procrastinate_jobs.id = candidate.id + RETURNING procrastinate_jobs.* INTO found_jobs; + + RETURN found_jobs; +END; +$$; + +CREATE FUNCTION procrastinate_notify_queue_resumed_v1() + RETURNS trigger + LANGUAGE plpgsql +AS $$ +DECLARE + payload TEXT; +BEGIN + SELECT json_build_object('type', 'queue_resumed', 'queue_name', OLD.queue_name)::text INTO payload; + PERFORM pg_notify('procrastinate_queue_v1#' || OLD.queue_name, payload); + PERFORM pg_notify('procrastinate_any_queue_v1', payload); + RETURN OLD; +END; +$$; + +CREATE TRIGGER procrastinate_paused_queues_notify_queue_resumed_v1 + AFTER DELETE ON procrastinate_paused_queues + FOR EACH ROW + EXECUTE PROCEDURE procrastinate_notify_queue_resumed_v1(); diff --git a/procrastinate/sql/migrations/03.10.00_50_post_add_queue_pause.sql b/procrastinate/sql/migrations/03.10.00_50_post_add_queue_pause.sql new file mode 100644 index 000000000..1ee4e086f --- /dev/null +++ b/procrastinate/sql/migrations/03.10.00_50_post_add_queue_pause.sql @@ -0,0 +1,4 @@ +-- Drop the previous fetch_job function now that the upgraded code calls +-- procrastinate_fetch_job_v3. + +DROP FUNCTION IF EXISTS procrastinate_fetch_job_v2(character varying[], bigint); diff --git a/procrastinate/sql/queries.sql b/procrastinate/sql/queries.sql index f54a315fe..75d8f6a1f 100644 --- a/procrastinate/sql/queries.sql +++ b/procrastinate/sql/queries.sql @@ -19,7 +19,7 @@ SELECT procrastinate_defer_periodic_job_v2(%(queue)s, %(lock)s, %(queueing_lock) -- fetch_job -- -- Get the first awaiting job SELECT id, status, task_name, priority, lock, queueing_lock, args, scheduled_at, queue_name, attempts, worker_id - FROM procrastinate_fetch_job_v2(%(queues)s::varchar[], %(worker_id)s); + FROM procrastinate_fetch_job_v3(%(queues)s::varchar[], %(worker_id)s); -- select_stalled_jobs_by_started -- -- Get running jobs that started more than a given time ago @@ -80,6 +80,25 @@ SELECT procrastinate_cancel_job_v1(%(job_id)s, %(abort)s, %(delete_job)s) AS id; -- Get the status of a job SELECT status FROM procrastinate_jobs WHERE id = %(job_id)s; +-- pause_queue -- +-- Pause a queue under the given pause key so that workers stop fetching its jobs +INSERT INTO procrastinate_paused_queues (queue_name, pause_key) + VALUES (%(queue_name)s, %(pause_key)s) + ON CONFLICT (queue_name, pause_key) DO NOTHING; + +-- resume_queue -- +-- Remove the given pause key (or all of them) from a queue so that workers fetch its jobs again +DELETE FROM procrastinate_paused_queues + WHERE queue_name = %(queue_name)s + AND (%(all_keys)s OR pause_key = %(pause_key)s); + +-- list_paused_queues -- +-- Get the pauses currently held on queues +SELECT queue_name, pause_key, paused_at FROM procrastinate_paused_queues + WHERE (%(queue_name)s::varchar IS NULL OR queue_name = %(queue_name)s::varchar) + AND (%(pause_key)s::varchar IS NULL OR pause_key = %(pause_key)s::varchar) + ORDER BY queue_name, pause_key; + -- retry_job -- -- Retry a job, changing it from "doing" to "todo" or from "failed" to "todo" SELECT procrastinate_retry_job_v2(%(job_id)s, %(retry_at)s, %(new_priority)s, %(new_queue_name)s, %(new_lock)s); diff --git a/procrastinate/sql/schema.sql b/procrastinate/sql/schema.sql index cd07fe49c..3b52655f5 100644 --- a/procrastinate/sql/schema.sql +++ b/procrastinate/sql/schema.sql @@ -94,6 +94,18 @@ CREATE TABLE procrastinate_events ( at timestamp with time zone DEFAULT NOW() NULL ); +-- A queue is paused as long as it has at least one row here: workers stop +-- fetching its jobs (already-running jobs keep going). Independent holders each +-- pause under their own pause_key and must all resume before the queue starts +-- working again. +CREATE TABLE procrastinate_paused_queues ( + id bigint PRIMARY KEY GENERATED ALWAYS AS IDENTITY, + queue_name character varying(128) NOT NULL, + pause_key character varying(128) DEFAULT 'default' NOT NULL, + paused_at timestamp with time zone DEFAULT NOW() NOT NULL, + UNIQUE (queue_name, pause_key) +); + -- Constraints & Indices -- this prevents from having several jobs with the same queueing lock in the "todo" state @@ -207,7 +219,7 @@ BEGIN END; $$; -CREATE FUNCTION procrastinate_fetch_job_v2( +CREATE FUNCTION procrastinate_fetch_job_v3( target_queue_names character varying[], p_worker_id bigint ) @@ -248,6 +260,12 @@ BEGIN AND jobs.status = 'todo' AND (target_queue_names IS NULL OR jobs.queue_name = ANY( target_queue_names )) AND (jobs.scheduled_at IS NULL OR jobs.scheduled_at <= now()) + -- reject the job if its queue is paused + AND NOT EXISTS ( + SELECT 1 + FROM procrastinate_paused_queues AS paused + WHERE paused.queue_name = jobs.queue_name + ) ORDER BY jobs.priority DESC, jobs.id ASC LIMIT 1 FOR UPDATE OF jobs SKIP LOCKED ) @@ -429,6 +447,20 @@ BEGIN END; $$; +CREATE FUNCTION procrastinate_notify_queue_resumed_v1() + RETURNS trigger + LANGUAGE plpgsql +AS $$ +DECLARE + payload TEXT; +BEGIN + SELECT json_build_object('type', 'queue_resumed', 'queue_name', OLD.queue_name)::text INTO payload; + PERFORM pg_notify('procrastinate_queue_v1#' || OLD.queue_name, payload); + PERFORM pg_notify('procrastinate_any_queue_v1', payload); + RETURN OLD; +END; +$$; + CREATE FUNCTION procrastinate_trigger_function_status_events_insert_v1() RETURNS trigger LANGUAGE plpgsql @@ -574,6 +606,11 @@ CREATE TRIGGER procrastinate_jobs_notify_queue_job_aborted_v1 FOR EACH ROW WHEN ((old.abort_requested = false AND new.abort_requested = true AND new.status = 'doing'::procrastinate_job_status)) EXECUTE PROCEDURE procrastinate_notify_queue_abort_job_v1(); +CREATE TRIGGER procrastinate_paused_queues_notify_queue_resumed_v1 + AFTER DELETE ON procrastinate_paused_queues + FOR EACH ROW + EXECUTE PROCEDURE procrastinate_notify_queue_resumed_v1(); + CREATE TRIGGER procrastinate_trigger_status_events_update_v1 AFTER UPDATE OF status ON procrastinate_jobs FOR EACH ROW diff --git a/procrastinate/testing.py b/procrastinate/testing.py index 310988636..a540b9f88 100644 --- a/procrastinate/testing.py +++ b/procrastinate/testing.py @@ -57,6 +57,7 @@ def reset(self) -> None: self.periodic_defers: dict[tuple[str, str], int] = {} self.table_exists = True self.states: list[str] = [] + self.queue_pauses: dict[tuple[str, str], datetime.datetime] = {} def get_sync_connector(self) -> connector.BaseConnector: return self @@ -288,12 +289,14 @@ async def fetch_job_one( ) -> dict[str, Any]: assert worker_id in self.workers, f"Worker {worker_id} not found" + paused = self.paused_queues() filtered_jobs = [ job for job in self.jobs.values() if ( job["status"] == "todo" and (queues is None or job["queue_name"] in queues) + and job["queue_name"] not in paused and (not job["scheduled_at"] or job["scheduled_at"] <= utils.utcnow()) and job["lock"] not in self.current_locks ) @@ -351,6 +354,35 @@ async def cancel_job_one( async def get_job_status_one(self, job_id: int) -> dict[str, Any]: return {"status": self.jobs[job_id]["status"]} + def paused_queues(self) -> set[str]: + return {queue_name for queue_name, _ in self.queue_pauses} + + async def pause_queue_run(self, queue_name: str, pause_key: str) -> None: + self.queue_pauses.setdefault((queue_name, pause_key), utils.utcnow()) + + async def resume_queue_run( + self, queue_name: str, pause_key: str, all_keys: bool + ) -> None: + deleted = False + for key in list(self.queue_pauses): + if key[0] == queue_name and (all_keys or key[1] == pause_key): + del self.queue_pauses[key] + deleted = True + if deleted: + await self._notify( + queue_name, {"type": "queue_resumed", "queue_name": queue_name} + ) + + async def list_paused_queues_all( + self, queue_name: str | None, pause_key: str | None + ) -> list[dict[str, Any]]: + return [ + {"queue_name": name, "pause_key": key, "paused_at": paused_at} + for (name, key), paused_at in sorted(self.queue_pauses.items()) + if (queue_name is None or name == queue_name) + and (pause_key is None or key == pause_key) + ] + async def retry_job_run( self, job_id: int, diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 7dbfa93d3..7dd29422b 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -475,7 +475,7 @@ async def run(self): async def _handle_notification( self, *, channel: str, notification: jobs.Notification ): - if notification["type"] == "job_inserted": + if notification["type"] in ("job_inserted", "queue_resumed"): self._new_job_event.set() elif notification["type"] == "abort_job_requested": self._handle_abort_jobs_requested([notification["job_id"]]) diff --git a/tests/acceptance/test_async.py b/tests/acceptance/test_async.py index 38b56fbb6..3f091960f 100644 --- a/tests/acceptance/test_async.py +++ b/tests/acceptance/test_async.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import time import pytest @@ -24,16 +25,23 @@ async def async_app(request, psycopg_connector, connection_params): yield app -async def wait_for_job_status( - app: app_module.App, job_id: int, status: Status, timeout: float = 5 -): +async def wait_until(condition, timeout: float = 5): async def poll(): - while await app.job_manager.get_job_status_async(job_id) != status: + while not await condition(): await asyncio.sleep(0.02) await asyncio.wait_for(poll(), timeout) +async def wait_for_job_status( + app: app_module.App, job_id: int, status: Status, timeout: float = 5 +): + async def has_status(): + return await app.job_manager.get_job_status_async(job_id) == status + + await wait_until(has_status, timeout=timeout) + + async def test_defer(async_app: app_module.App): sum_results = [] product_results = [] @@ -128,6 +136,120 @@ async def sum_task(a, b): assert sum_results == [7] +@pytest.mark.skip_before_version("3.10.0") +async def test_pause_queue(async_app: app_module.App): + sum_results = [] + + @async_app.task(queue="default", name="sum_task") + async def sum_task(a, b): + sum_results.append(a + b) + + await async_app.job_manager.pause_queue_async("default") + job_id = await sum_task.defer_async(a=1, b=2) + + # The queue is paused: the worker fetches nothing and the job stays todo. + await async_app.run_worker_async(queues=["default"], wait=False) + assert sum_results == [] + assert await async_app.job_manager.get_job_status_async(job_id) == Status.TODO + rows = await async_app.job_manager.list_paused_queues_async() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("default", "default") + ] + + # After resuming, the worker processes the job. + await async_app.job_manager.resume_queue_async("default") + await async_app.run_worker_async(queues=["default"], wait=False) + assert sum_results == [3] + assert await async_app.job_manager.get_job_status_async(job_id) == Status.SUCCEEDED + assert await async_app.job_manager.list_paused_queues_async() == [] + + +@pytest.mark.skip_before_version("3.10.0") +async def test_pause_queue_does_not_interrupt_running_job(async_app: app_module.App): + may_finish = asyncio.Event() + + @async_app.task(queue="default", name="blocking_task") + async def blocking_task(): + await may_finish.wait() + + job_id = await blocking_task.defer_async() + + worker_task = asyncio.create_task( + async_app.run_worker_async(queues=["default"], wait=False) + ) + await wait_for_job_status(async_app, job_id, Status.DOING) + + # Pause while the job is running: it must run to completion, not be aborted. + await async_app.job_manager.pause_queue_async("default") + may_finish.set() + + await asyncio.wait_for(worker_task, timeout=2) + assert await async_app.job_manager.get_job_status_async(job_id) == Status.SUCCEEDED + + +@pytest.mark.skip_before_version("3.10.0") +async def test_resume_queue_wakes_waiting_worker(async_app: app_module.App): + results = [] + + @async_app.task(queue="default", name="sum_task") + async def sum_task(a, b): + results.append(a + b) + + ping_ids = [] + + @async_app.task(queue="ping", name="ping_task") + async def ping_task(): + ping_ids.append(1) + + await async_app.job_manager.pause_queue_async("default", pause_key="deploy") + await async_app.job_manager.pause_queue_async("default", pause_key="maintenance") + await sum_task.defer_async(a=1, b=2) + + # A long polling interval so the worker can only pick the job back up through + # the resume notification, not by polling. + worker_task = asyncio.create_task( + async_app.run_worker_async( + queues=["default", "ping"], + wait=True, + fetch_job_polling_interval=30, + listen_notify=True, + ) + ) + + async def has_ping(): + return bool(ping_ids) + + async def wait_until_listening(timeout: float = 10): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + await ping_task.defer_async() + with contextlib.suppress(asyncio.TimeoutError): + await wait_until(has_ping, timeout=0.5) + return + pytest.fail("worker never started listening for notifications") + + await wait_until_listening() + assert results == [] + + # One pause key is still held, so the queue stays paused: the worker may be + # woken by the resume notification, but it still fetches nothing. + await async_app.job_manager.resume_queue_async("default", pause_key="deploy") + await asyncio.sleep(0.5) + assert results == [] + + # Resuming the last key notifies the idle worker, which picks the job up. + await async_app.job_manager.resume_queue_async("default", pause_key="maintenance") + + async def job_processed(): + return results == [3] + + await wait_until(job_processed, timeout=5) + + worker_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await worker_task + + async def test_no_job_to_cancel_found(async_app: app_module.App): @async_app.task(queue="default", name="example_task") def example_task(): diff --git a/tests/integration/test_manager.py b/tests/integration/test_manager.py index 61e913343..a7a621e65 100644 --- a/tests/integration/test_manager.py +++ b/tests/integration/test_manager.py @@ -126,6 +126,106 @@ async def test_fetch_job_no_result( ) +async def test_pause_and_resume_queue(pg_job_manager, deferred_job_factory, worker_id): + job = await deferred_job_factory(queue="queue_a") + + await pg_job_manager.pause_queue_async("queue_a") + rows = await pg_job_manager.list_paused_queues_async() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("queue_a", "default") + ] + assert rows[0]["paused_at"] is not None + assert await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) is None + + await pg_job_manager.resume_queue_async("queue_a") + assert await pg_job_manager.list_paused_queues_async() == [] + fetched = await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched.id == job.id + + +async def test_pause_queue_is_idempotent(pg_job_manager): + await pg_job_manager.pause_queue_async("queue_a") + rows = await pg_job_manager.list_paused_queues_async() + await pg_job_manager.pause_queue_async("queue_a") + assert await pg_job_manager.list_paused_queues_async() == rows + + await pg_job_manager.resume_queue_async("queue_a") + await pg_job_manager.resume_queue_async("queue_a") + assert await pg_job_manager.list_paused_queues_async() == [] + + +async def test_pause_queue_with_multiple_keys( + pg_job_manager, deferred_job_factory, worker_id +): + job = await deferred_job_factory(queue="queue_a") + + await pg_job_manager.pause_queue_async("queue_a", pause_key="deploy") + await pg_job_manager.pause_queue_async("queue_a", pause_key="maintenance") + + await pg_job_manager.resume_queue_async("queue_a", pause_key="deploy") + rows = await pg_job_manager.list_paused_queues_async() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("queue_a", "maintenance") + ] + assert await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) is None + + await pg_job_manager.resume_queue_async("queue_a", pause_key="maintenance") + assert await pg_job_manager.list_paused_queues_async() == [] + fetched = await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched.id == job.id + + +async def test_resume_queue_all_keys(pg_job_manager, deferred_job_factory, worker_id): + job = await deferred_job_factory(queue="queue_a") + + await pg_job_manager.pause_queue_async("queue_a", pause_key="deploy") + await pg_job_manager.pause_queue_async("queue_a", pause_key="maintenance") + await pg_job_manager.pause_queue_async("queue_b") + + await pg_job_manager.resume_queue_async("queue_a", all_keys=True) + rows = await pg_job_manager.list_paused_queues_async() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("queue_b", "default") + ] + fetched = await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched.id == job.id + + +async def test_pause_queue_only_affects_paused_queue( + pg_job_manager, deferred_job_factory, worker_id +): + await deferred_job_factory(queue="queue_a") + job_b = await deferred_job_factory(queue="queue_b") + + await pg_job_manager.pause_queue_async("queue_a") + + fetched = await pg_job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched.id == job_b.id + + +async def test_fetch_job_scoped_to_paused_queue_returns_none( + pg_job_manager, deferred_job_factory, worker_id +): + await deferred_job_factory(queue="queue_a") + await pg_job_manager.pause_queue_async("queue_a") + + # A worker restricted to the paused queue fetches nothing. + assert ( + await pg_job_manager.fetch_job(queues=["queue_a"], worker_id=worker_id) is None + ) + + +def test_pause_and_resume_queue_sync(pg_job_manager): + pg_job_manager.pause_queue("queue_a") + rows = pg_job_manager.list_paused_queues() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("queue_a", "default") + ] + + pg_job_manager.resume_queue("queue_a") + assert pg_job_manager.list_paused_queues() == [] + + @pytest.mark.parametrize( "filter_args", [ diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py index a0fd78ef9..7ff9a0567 100644 --- a/tests/unit/test_manager.py +++ b/tests/unit/test_manager.py @@ -804,3 +804,131 @@ async def test_unsupported_connector_raises_query_one_async(): c = BaseConnector() with pytest.raises(exceptions.ConnectorException, match="does not support"): await c.execute_query_one_async_with_connection(object(), "SELECT 1") + + +def test_pause_queue(job_manager, connector): + job_manager.pause_queue(queue_name="foo") + assert connector.queries[-1] == ( + "pause_queue", + {"queue_name": "foo", "pause_key": "default"}, + ) + assert connector.paused_queues() == {"foo"} + + +async def test_pause_queue_async(job_manager, connector): + await job_manager.pause_queue_async(queue_name="foo", pause_key="deploy") + assert connector.queries[-1] == ( + "pause_queue", + {"queue_name": "foo", "pause_key": "deploy"}, + ) + assert connector.paused_queues() == {"foo"} + + +def test_resume_queue(job_manager, connector): + job_manager.pause_queue(queue_name="foo") + job_manager.resume_queue(queue_name="foo") + assert connector.queries[-1] == ( + "resume_queue", + {"queue_name": "foo", "pause_key": "default", "all_keys": False}, + ) + assert connector.paused_queues() == set() + + +async def test_resume_queue_async(job_manager, connector): + await job_manager.pause_queue_async(queue_name="foo") + await job_manager.resume_queue_async(queue_name="foo") + assert connector.queries[-1] == ( + "resume_queue", + {"queue_name": "foo", "pause_key": "default", "all_keys": False}, + ) + assert connector.paused_queues() == set() + + +def test_resume_queue_only_releases_own_key(job_manager, connector): + job_manager.pause_queue(queue_name="foo", pause_key="deploy") + job_manager.pause_queue(queue_name="foo", pause_key="maintenance") + + job_manager.resume_queue(queue_name="foo", pause_key="deploy") + assert connector.paused_queues() == {"foo"} + + job_manager.resume_queue(queue_name="foo", pause_key="maintenance") + assert connector.paused_queues() == set() + + +def test_resume_queue_all_keys(job_manager, connector): + job_manager.pause_queue(queue_name="foo", pause_key="deploy") + job_manager.pause_queue(queue_name="foo", pause_key="maintenance") + job_manager.pause_queue(queue_name="bar") + + job_manager.resume_queue(queue_name="foo", all_keys=True) + assert connector.paused_queues() == {"bar"} + + +def test_list_paused_queues(job_manager, connector): + job_manager.pause_queue(queue_name="foo") + job_manager.pause_queue(queue_name="bar", pause_key="deploy") + rows = job_manager.list_paused_queues() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("bar", "deploy"), + ("foo", "default"), + ] + + +async def test_list_paused_queues_async(job_manager, connector): + await job_manager.pause_queue_async(queue_name="foo") + rows = await job_manager.list_paused_queues_async() + assert [(row["queue_name"], row["pause_key"]) for row in rows] == [ + ("foo", "default") + ] + + +def test_list_paused_queues_filters(job_manager, connector): + job_manager.pause_queue(queue_name="foo", pause_key="deploy") + job_manager.pause_queue(queue_name="foo", pause_key="maintenance") + job_manager.pause_queue(queue_name="bar", pause_key="deploy") + + assert [ + (row["queue_name"], row["pause_key"]) + for row in job_manager.list_paused_queues(queue="foo") + ] == [("foo", "deploy"), ("foo", "maintenance")] + + assert [ + (row["queue_name"], row["pause_key"]) + for row in job_manager.list_paused_queues(pause_key="deploy") + ] == [("bar", "deploy"), ("foo", "deploy")] + + assert [ + (row["queue_name"], row["pause_key"]) + for row in job_manager.list_paused_queues(queue="bar", pause_key="maintenance") + ] == [] + + +async def test_fetch_job_skips_paused_queue(job_manager, job_factory, worker_id): + job = job_factory(id=None, queue="paused") + await job_manager.defer_job_async(job=job) + + await job_manager.pause_queue_async(queue_name="paused") + assert await job_manager.fetch_job(queues=None, worker_id=worker_id) is None + + await job_manager.resume_queue_async(queue_name="paused") + fetched = await job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched is not None + assert fetched.queue == "paused" + + +async def test_fetch_job_skips_queue_until_all_keys_resumed( + job_manager, job_factory, worker_id +): + job = job_factory(id=None, queue="paused") + await job_manager.defer_job_async(job=job) + + await job_manager.pause_queue_async(queue_name="paused", pause_key="deploy") + await job_manager.pause_queue_async(queue_name="paused", pause_key="maintenance") + + await job_manager.resume_queue_async(queue_name="paused", pause_key="deploy") + assert await job_manager.fetch_job(queues=None, worker_id=worker_id) is None + + await job_manager.resume_queue_async(queue_name="paused", pause_key="maintenance") + fetched = await job_manager.fetch_job(queues=None, worker_id=worker_id) + assert fetched is not None + assert fetched.queue == "paused" diff --git a/tests/unit/test_worker.py b/tests/unit/test_worker.py index bb949b35d..9ba5a1c54 100644 --- a/tests/unit/test_worker.py +++ b/tests/unit/test_worker.py @@ -234,6 +234,38 @@ async def perform_job(): complete_tasks.set() +async def test_worker_run_fetches_job_on_queue_resumed_notification(worker, app: App): + complete_tasks = asyncio.Event() + + @app.task(queue="paused") + async def perform_job(): + await complete_tasks.wait() + + await app.job_manager.pause_queue_async("paused") + await perform_job.defer_async() + + await start_worker(worker) + + connector = cast(InMemoryConnector, app.connector) + fetch_count = len([query for query in connector.queries if query[0] == "fetch_job"]) + + await asyncio.sleep(0.01) + assert ( + len([query for query in connector.queries if query[0] == "fetch_job"]) + == fetch_count + ) + + await app.job_manager.resume_queue_async("paused") + await asyncio.sleep(0.01) + + assert ( + len([query for query in connector.queries if query[0] == "fetch_job"]) + == fetch_count + 1 + ) + + complete_tasks.set() + + @pytest.mark.parametrize( "worker", [({"fetch_job_polling_interval": 0.05})], From d1c94e278673e16077d39765364dc28b5ef2608b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Artur=20Ku=C5=BAmi=C5=84ski?= Date: Thu, 16 Jul 2026 16:54:21 +0200 Subject: [PATCH 2/2] Add queue.pause django changes --- docs/howto/django/models.md | 5 +-- .../migrations/0044_add_paused_queue_model.py | 31 +++++++++++++++++++ procrastinate/contrib/django/models.py | 20 ++++++++++++ .../integration/contrib/django/test_models.py | 22 +++++++++++++ 4 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 procrastinate/contrib/django/migrations/0044_add_paused_queue_model.py diff --git a/docs/howto/django/models.md b/docs/howto/django/models.md index cbcfd2972..544869247 100644 --- a/docs/howto/django/models.md +++ b/docs/howto/django/models.md @@ -1,6 +1,6 @@ # Interact with Procrastinate tables as Django models -Procrastinate exposes 3 of its internal tables as Django models. You can use +Procrastinate exposes some of its internal tables as Django models. You can use them to query the state of your jobs. They're also exposed in the Django admin. :::{note} @@ -15,6 +15,7 @@ from procrastinate.contrib.django.models import ( ProcrastinateEvent, ProcrastinatePeriodicDefer, ProcrastinateWorker, + ProcrastinatePausedQueue, ) ProcrastinateJob.objects.filter(task_name="mytask").count() @@ -30,7 +31,7 @@ or events through the ORM. ```{eval-rst} .. automodule:: procrastinate.contrib.django.models - :members: ProcrastinateJob, ProcrastinateEvent, ProcrastinatePeriodicDefer, ProcrastinateWorker + :members: ProcrastinateJob, ProcrastinateEvent, ProcrastinatePeriodicDefer, ProcrastinateWorker, ProcrastinatePausedQueue ``` diff --git a/procrastinate/contrib/django/migrations/0044_add_paused_queue_model.py b/procrastinate/contrib/django/migrations/0044_add_paused_queue_model.py new file mode 100644 index 000000000..eea77cccb --- /dev/null +++ b/procrastinate/contrib/django/migrations/0044_add_paused_queue_model.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from django.db import migrations, models + +import procrastinate.contrib.django.models + + +class Migration(migrations.Migration): + operations = [ + migrations.CreateModel( + name="ProcrastinatePausedQueue", + fields=[ + ("id", models.BigAutoField(primary_key=True, serialize=False)), + ("queue_name", models.CharField(max_length=128)), + ("pause_key", models.CharField(max_length=128)), + ("paused_at", models.DateTimeField()), + ], + options={ + "db_table": "procrastinate_paused_queues", + "managed": False, + }, + bases=( + procrastinate.contrib.django.models.ProcrastinateReadOnlyModelMixin, + models.Model, + ), + ), + ] + name = "0044_add_paused_queue_model" + dependencies = [ + ("procrastinate", "0043_post_add_queue_pause"), + ] diff --git a/procrastinate/contrib/django/models.py b/procrastinate/contrib/django/models.py index d1739aa8b..19aaefa8b 100644 --- a/procrastinate/contrib/django/models.py +++ b/procrastinate/contrib/django/models.py @@ -165,3 +165,23 @@ class Meta: # type: ignore managed = False db_table = "procrastinate_periodic_defers" unique_together = [("task_name", "periodic_id", "defer_timestamp")] + + +class ProcrastinatePausedQueue(ProcrastinateReadOnlyModelMixin, models.Model): + id = models.BigAutoField(primary_key=True) + queue_name = models.CharField(max_length=128) + pause_key = models.CharField(max_length=128) + paused_at = models.DateTimeField() + + objects = ProcrastinateReadOnlyManager() + + class Meta: # type: ignore + managed = False + db_table = "procrastinate_paused_queues" + unique_together = [("queue_name", "pause_key")] + + def __str__(self) -> str: + return ( + f"Queue {self.queue_name} - " + f"Paused with key {self.pause_key} at {self.paused_at}" + ) diff --git a/tests/integration/contrib/django/test_models.py b/tests/integration/contrib/django/test_models.py index d3b302f67..3acad26bc 100644 --- a/tests/integration/contrib/django/test_models.py +++ b/tests/integration/contrib/django/test_models.py @@ -93,6 +93,28 @@ def test_procrastinate_job__no_delete(db): models.ProcrastinateJob().delete() +def test_procrastinate_paused_queue(db): + procrastinate.contrib.django.app.job_manager.pause_queue("foo", pause_key="deploy") + now = datetime.datetime.now(datetime.timezone.utc) + one_sec = datetime.timedelta(seconds=1) + paused = models.ProcrastinatePausedQueue.objects.values().get(queue_name="foo") + paused_at = paused.pop("paused_at") + paused.pop("id") + assert paused == {"queue_name": "foo", "pause_key": "deploy"} + assert now - one_sec < paused_at < now + one_sec + + +def test_procrastinate_paused_queue__resumed(db): + procrastinate.contrib.django.app.job_manager.pause_queue("foo") + procrastinate.contrib.django.app.job_manager.resume_queue("foo") + assert not models.ProcrastinatePausedQueue.objects.exists() + + +def test_procrastinate_paused_queue__no_create(db): + with pytest.raises(procrastinate.contrib.django.exceptions.ReadOnlyModel): + models.ProcrastinatePausedQueue.objects.create(queue_name="foo") + + def test_procrastinate_event(db): job_id = procrastinate.contrib.django.app.configure_task("test_task").defer( a=1, b=2