python-pq 0.8.0__tar.gz → 0.9.0__tar.gz

{python_pq-0.8.0 → python_pq-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-pq
-Version: 0.8.0
+Version: 0.9.0
 Summary: Postgres-backed job queue for Python
 Keywords: postgres,job-queue,task-queue,background-jobs,worker
 Author: ricwo
@@ -143,6 +143,19 @@ pq.upsert(send_email, to="a@b.com", client_id="welcome-email")
 pq.upsert(send_email, to="new@b.com", client_id="welcome-email")
 ```
 
+If the task is running, the new version is queued and runs after the current run ends, so one `client_id` does not run twice at the same time. This lets a task schedule its own continuation:
+
+```python
+def backfill(offset: int) -> None:
+    more = process_rows(offset, limit=1000)
+    if more:
+        with PQ(DATABASE_URL) as client:
+            # Runs after this run ends, never next to it
+            client.upsert(backfill, offset=offset + 1000, client_id="backfill")
+```
+
+If the worker running the task dies, the queued version runs after the stale-task reaper picks up the row (`stale_task_timeout`, default 1 hour). Keep `stale_task_timeout` longer than your longest run: if the reaper picks up a run that is still alive, the queued version can start next to it.
+
 ## Periodic Tasks
 
 ### Intervals

{python_pq-0.8.0 → python_pq-0.9.0}/README.md

@@ -107,6 +107,19 @@ pq.upsert(send_email, to="a@b.com", client_id="welcome-email")
 pq.upsert(send_email, to="new@b.com", client_id="welcome-email")
 ```
 
+If the task is running, the new version is queued and runs after the current run ends, so one `client_id` does not run twice at the same time. This lets a task schedule its own continuation:
+
+```python
+def backfill(offset: int) -> None:
+    more = process_rows(offset, limit=1000)
+    if more:
+        with PQ(DATABASE_URL) as client:
+            # Runs after this run ends, never next to it
+            client.upsert(backfill, offset=offset + 1000, client_id="backfill")
+```
+
+If the worker running the task dies, the queued version runs after the stale-task reaper picks up the row (`stale_task_timeout`, default 1 hour). Keep `stale_task_timeout` longer than your longest run: if the reaper picks up a run that is still alive, the queued version can start next to it.
+
 ## Periodic Tasks
 
 ### Intervals

{python_pq-0.8.0 → python_pq-0.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "python-pq"
-version = "0.8.0"
+version = "0.9.0"
 description = "Postgres-backed job queue for Python"
 readme = "README.md"
 authors = [

{python_pq-0.8.0 → python_pq-0.9.0}/src/pq/__init__.py

@@ -5,7 +5,7 @@ from pq.models import Periodic, Task, TaskStatus
 from pq.priority import Priority
 from pq.worker import PostExecuteHook, PreExecuteHook, TaskTimeoutError
 
-__version__ = "0.8.0"
+__version__ = "0.9.0"
 
 __all__ = [
     "PQ",

{python_pq-0.8.0 → python_pq-0.9.0}/src/pq/client.py

@@ -13,9 +13,9 @@ from croniter.croniter import CroniterBadCronError
 from loguru import logger
 import sqlalchemy as sa
 from sqlalchemy import create_engine, delete, func, select, update
-from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.dialects.postgresql import JSONB, insert
 from sqlalchemy.engine import Engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import InstrumentedAttribute, sessionmaker
 
 from pq.models import Base, Periodic, Task, TaskStatus
 from pq.priority import Priority
@@ -230,6 +230,16 @@ class PQ:
         Behaves like enqueue(), but on conflict for client_id, updates all fields.
         Status resets to PENDING, attempts to 0, and timestamps are cleared.
 
+        If the existing task is RUNNING, it is not changed. The new version
+        is stored on the row and runs after the current run ends (the worker
+        or the stale reaper re-queues the row with it). A second upsert while
+        the row is still RUNNING replaces the stored version, so only the
+        latest one runs. One ``client_id`` does not run twice at the same
+        time, as long as the stale reaper's timeout is longer than the
+        longest run. If the worker running the task died, the stored version
+        runs only after the stale reaper re-queues the row: with the reaper
+        disabled (``stale_task_timeout=None``) it never runs.
+
         Args:
             task: Callable function to execute.
             *args: Positional arguments to pass to the handler.
@@ -255,6 +265,41 @@ class PQ:
         if run_at is None:
             run_at = datetime.now(UTC)
 
+        # In ``ON CONFLICT DO UPDATE``, ``Task.<column>`` refers to the
+        # existing row. A RUNNING row keeps every column as it is (the run
+        # in progress keeps its payload, ``started_at``, ``attempts`` and
+        # ``max_runtime_seconds``); the new version is parked in
+        # ``requeue`` and applied when the run ends. Any other status gets
+        # the plain overwrite. All in one statement, so the decision is
+        # made under the row lock that ``ON CONFLICT`` takes.
+        running = Task.status == TaskStatus.RUNNING
+        # ``run_at`` is added in SQL as ``timestamptz``, so Postgres
+        # serializes it as ISO 8601 with an offset. A naive ``run_at`` is
+        # then read in this session's time zone, the same as on the direct
+        # path, and not later in the session of whoever applies the version.
+        requeue = sa.literal(
+            {
+                "name": name,
+                "payload": payload,
+                "priority": int(priority),
+                "max_runtime_seconds": max_runtime_seconds,
+            },
+            JSONB,
+        ).op("||", return_type=JSONB)(
+            sa.func.jsonb_build_object(
+                "run_at", sa.cast(sa.literal(run_at), sa.DateTime(timezone=True))
+            )
+        )
+
+        def unless_running[T](
+            value: T, column: InstrumentedAttribute[T]
+        ) -> sa.ColumnElement[T]:
+            # ``sa.literal(..., column.type)`` binds the value through the
+            # column's type. Required for the ``task_status`` enum, which
+            # stores member names (``'PENDING'``): a bare enum member in a
+            # ``CASE`` would bind ``'pending'`` and fail.
+            return sa.case((running, column), else_=sa.literal(value, column.type))
+
         stmt = (
             insert(Task)
             .values(
@@ -269,16 +314,19 @@ class PQ:
             .on_conflict_do_update(
                 index_elements=["client_id"],
                 set_={
-                    "name": name,
-                    "payload": payload,
-                    "priority": priority,
-                    "status": TaskStatus.PENDING,
-                    "run_at": run_at,
-                    "max_runtime_seconds": max_runtime_seconds,
-                    "attempts": 0,
-                    "started_at": None,
-                    "completed_at": None,
-                    "error": None,
+                    "name": unless_running(name, Task.name),
+                    "payload": unless_running(payload, Task.payload),
+                    "priority": unless_running(int(priority), Task.priority),
+                    "status": unless_running(TaskStatus.PENDING, Task.status),
+                    "run_at": unless_running(run_at, Task.run_at),
+                    "max_runtime_seconds": unless_running(
+                        max_runtime_seconds, Task.max_runtime_seconds
+                    ),
+                    "attempts": unless_running(0, Task.attempts),
+                    "started_at": unless_running(None, Task.started_at),
+                    "completed_at": unless_running(None, Task.completed_at),
+                    "error": unless_running(None, Task.error),
+                    "requeue": sa.case((running, requeue), else_=sa.null()),
                 },
             )
             .returning(Task.id)
@@ -582,13 +630,17 @@ class PQ:
             return result.rowcount
 
     def reap_stale_tasks(self, threshold: timedelta) -> int:
-        """Mark stale RUNNING tasks as FAILED.
+        """Mark stale RUNNING tasks as FAILED, or re-queue them.
 
         When a worker dies mid-execution (e.g. pod restart, OOM on the worker
         process), in-flight tasks stay RUNNING forever because no parent
         process remains to update their status. This method detects those
         orphaned rows and transitions them to FAILED.
 
+        A stale row that holds a version parked by ``upsert()`` (see
+        ``Task.requeue``) is re-queued as PENDING with that version instead
+        of being failed, so the parked version is not lost.
+
         Tasks enqueued with a per-task ``max_runtime_seconds`` override get a
         proportionally larger reaper window: the effective threshold for
         a row is ``max(threshold, max_runtime_seconds * 2)``. This avoids
@@ -605,7 +657,7 @@ class PQ:
                 ``timedelta(seconds=max_runtime_seconds * 2)``.
 
         Returns:
-            Number of tasks reaped.
+            Number of tasks reaped (failed plus re-queued).
         """
         now = datetime.now(UTC)
         threshold_seconds = threshold.total_seconds()
@@ -649,6 +701,7 @@ class PQ:
                 .where(
                     Task.status == TaskStatus.RUNNING,
                     Task.started_at < stale_cutoff,
+                    Task.requeue.is_(None),
                 )
                 .values(
                     status=TaskStatus.FAILED,
@@ -669,7 +722,31 @@ class PQ:
                     f" started_at={started_at},"
                     f" max_runtime_seconds={max_runtime_seconds})"
                 )
-            return len(reaped)
+
+            # Stale rows with a parked upsert: re-queue them with the
+            # parked version. ``FOR UPDATE`` serializes with a concurrent
+            # ``upsert()`` (which may replace ``requeue``) and with the
+            # worker's own end-of-run update.
+            requeued = list(
+                session.execute(
+                    select(Task)
+                    .where(
+                        Task.status == TaskStatus.RUNNING,
+                        Task.started_at < stale_cutoff,
+                        Task.requeue.is_not(None),
+                    )
+                    .with_for_update()
+                ).scalars()
+            )
+            for task in requeued:
+                logger.warning(
+                    f"Reaped stale task '{task.name}' (id={task.id},"
+                    f" started_at={task.started_at},"
+                    f" max_runtime_seconds={task.max_runtime_seconds});"
+                    " re-queued by a parked upsert"
+                )
+                task.apply_requeue()
+            return len(reaped) + len(requeued)
 
     def run_worker(
         self,
@@ -687,7 +764,9 @@ class PQ:
         On SIGTERM or SIGINT the worker stops claiming new tasks, waits up
         to ``drain_timeout`` seconds for in-flight tasks to finish, then
         kills any still-running task and marks its row FAILED with an
-        explicit shutdown error. Interrupted tasks are not re-queued.
+        explicit shutdown error. Interrupted tasks are not re-queued,
+        unless an ``upsert()`` stored a new version while they ran: then the
+        row is re-queued with that version.
 
         Args:
             concurrency: Maximum number of tasks to process simultaneously.

python_pq-0.9.0/src/pq/migrations/versions/20261007T120000Z_d4e5f6a7b8c9_add_task_requeue.py

@@ -0,0 +1,43 @@
+"""add requeue to pq_tasks
+
+Revision ID: d4e5f6a7b8c9
+Revises: aee3e8e7e647
+Create Date: 2026-10-07 12:00:00 Z
+
+Holds the next version of a one-off task when ``upsert()`` hits a row
+that is RUNNING. The current run keeps its row unchanged; whoever ends
+the run (worker or stale reaper) re-queues the row from this column
+instead of closing it. This makes sure one ``client_id`` never runs
+twice at the same time (ricwo/pq#27).
+
+Nullable without a default, so on PostgreSQL 11+ ``ADD COLUMN`` is a
+catalog-only change (no table rewrite). Existing rows get NULL, which
+means "no parked version" — the previous behaviour.
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+
+# revision identifiers, used by Alembic.
+revision: str = "d4e5f6a7b8c9"
+down_revision: Union[str, Sequence[str], None] = "aee3e8e7e647"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Add the nullable ``requeue`` column to ``pq_tasks``."""
+    op.add_column(
+        "pq_tasks",
+        sa.Column("requeue", postgresql.JSONB(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    """Remove the ``requeue`` column from ``pq_tasks``."""
+    op.drop_column("pq_tasks", "requeue")

{python_pq-0.8.0 → python_pq-0.9.0}/src/pq/models.py

@@ -79,6 +79,37 @@ class Task(Base):
     # ``max_runtime_seconds * 2``). Useful for occasionally-long tasks
     # in a fleet whose default is sized for typical short work.
     max_runtime_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)
+    # Next version of this task, parked by an ``upsert()`` that arrived while
+    # the row was RUNNING: name, payload, priority, run_at (ISO 8601) and
+    # max_runtime_seconds. Whoever ends the current run (worker or reaper)
+    # re-queues the row with it instead of closing the row. NULL when no
+    # upsert is parked (the common case). ``none_as_null`` makes the ORM
+    # write SQL NULL for ``None`` (plain ``JSONB`` writes JSON ``null``,
+    # which ``IS NULL`` does not match).
+    requeue: Mapped[dict[str, Any] | None] = mapped_column(
+        JSONB(none_as_null=True), nullable=True
+    )
+
+    def apply_requeue(self) -> None:
+        """Turn this RUNNING row into a fresh PENDING task from ``requeue``.
+
+        The caller must hold a row lock (``FOR UPDATE``) and guarantee
+        that ``requeue`` is not None.
+        """
+        spec = self.requeue
+        if spec is None:
+            raise ValueError(f"Task id={self.id} has no parked requeue version")
+        self.name = spec["name"]
+        self.payload = spec["payload"]
+        self.priority = spec["priority"]
+        self.run_at = datetime.fromisoformat(spec["run_at"])
+        self.max_runtime_seconds = spec["max_runtime_seconds"]
+        self.status = TaskStatus.PENDING
+        self.attempts = 0
+        self.started_at = None
+        self.completed_at = None
+        self.error = None
+        self.requeue = None
 
 
 class Periodic(Base):

{python_pq-0.8.0 → python_pq-0.9.0}/src/pq/worker.py

@@ -123,6 +123,9 @@ class _ChildSlot:
     name: str
     start_time: float
     is_periodic: bool
+    # ``started_at`` written by the claim of a one-off run. Identifies this
+    # run on the row (see ``_finish_one_off``). None for periodic slots.
+    started_at: datetime | None = None
     periodic_max_concurrent: int | None = None
 
 
@@ -569,7 +572,8 @@ def _maybe_reap_stale(
 
     Args:
         pq: PQ client instance.
-        stale_task_timeout: RUNNING tasks older than this are marked FAILED.
+        stale_task_timeout: RUNNING tasks older than this are marked FAILED
+            (or re-queued, if an upsert parked a new version on them).
             ``None`` disables reaping.
         reaper_interval: Seconds between reaper checks.
         last_reap: Mutable list containing last reap timestamp.
@@ -581,9 +585,15 @@ def _maybe_reap_stale(
     if now - last_reap[0] < reaper_interval:
         return
 
-    reaped = pq.reap_stale_tasks(stale_task_timeout)
-    if reaped:
-        logger.info(f"Reaped {reaped} stale RUNNING task(s)")
+    # A reaper error must not stop the worker loop (in concurrent mode that
+    # would also skip the drain of in-flight children). Retry next interval.
+    try:
+        reaped = pq.reap_stale_tasks(stale_task_timeout)
+    except Exception as e:
+        logger.error(f"Error reaping stale tasks: {e}")
+    else:
+        if reaped:
+            logger.info(f"Reaped {reaped} stale RUNNING task(s)")
 
     last_reap[0] = now
 
@@ -612,7 +622,9 @@ def run_worker(
     still-running task children and marks their rows FAILED with an explicit
     shutdown error. Interrupted tasks are NOT re-queued — pq's at-most-once
     semantics are preserved; applications that need redelivery must
-    re-enqueue such tasks themselves.
+    re-enqueue such tasks themselves. The one exception: if an ``upsert()``
+    stored a new version on the row while it ran, the row is re-queued with
+    that new version (the interrupted run itself is not retried).
 
     Args:
         pq: PQ client instance.
@@ -625,9 +637,12 @@ def run_worker(
         retention_days: Days to keep completed/failed tasks. Default: 7.
             Set to 0 to disable automatic cleanup.
         cleanup_interval: Seconds between cleanup runs. Default: 3600 (1 hour).
-        stale_task_timeout: Mark RUNNING tasks older than this as FAILED.
+        stale_task_timeout: Mark RUNNING tasks older than this as FAILED
+            (or re-queue them, if an upsert parked a new version on them).
             Catches orphaned tasks whose worker died mid-execution.
-            Default: 1 hour. Set to ``None`` to disable.
+            Default: 1 hour. Set to ``None`` to disable. Must be longer
+            than the longest run: a live run that is reaped can overlap
+            with the version re-queued after it.
         drain_timeout: Seconds to wait for in-flight tasks when shutting
             down on SIGTERM/SIGINT. Default: 20. Must be set below the
             orchestrator's termination grace period, with headroom for the
@@ -756,6 +771,66 @@ def run_worker_once(
     return False
 
 
+def _finish_one_off(
+    pq: PQ,
+    task_id: int,
+    started_at: datetime | None,
+    name: str,
+    status: TaskStatus,
+    error_msg: str | None,
+) -> None:
+    """Record the end of a one-off run.
+
+    Only touches the row if it is still RUNNING *with this run's*
+    ``started_at`` (the value written by its claim). So a reaper's FAILED
+    verdict (or a cancel) is not overwritten, and neither is a later run
+    of the same row: if the reaper re-queued the row while this run was
+    still alive and another worker claimed it, ``started_at`` no longer
+    matches. ``None`` never matches (a RUNNING row always has
+    ``started_at``), so the row is left alone. If an ``upsert()`` parked a
+    new version on the row while it ran (``Task.requeue``), the row is
+    re-queued as PENDING with that version instead of being closed. The
+    outcome of this run is then only in the worker log, the same as when
+    an upsert overwrites a COMPLETED or FAILED row.
+
+    ``FOR UPDATE`` serializes this read-then-write with a concurrent
+    ``upsert()``, whose ``ON CONFLICT DO UPDATE`` locks the same row: a
+    parked version is either seen here, or the upsert runs after this
+    commit and hits a non-RUNNING row, which it overwrites directly.
+    """
+    try:
+        with pq.session() as session:
+            task = session.execute(
+                select(Task)
+                .where(
+                    Task.id == task_id,
+                    Task.status == TaskStatus.RUNNING,
+                    Task.started_at == started_at,
+                )
+                .with_for_update()
+            ).scalar_one_or_none()
+            if task is None:
+                logger.warning(
+                    f"Task '{name}' (id={task_id}) was no longer RUNNING"
+                    " for this run when the worker tried to record its result"
+                    " — likely reaped or canceled"
+                )
+                return
+            if task.requeue is not None:
+                logger.info(
+                    f"Task '{name}' (id={task_id}) ended as {status.value};"
+                    " re-queued with the version upserted while it ran"
+                )
+                task.apply_requeue()
+                return
+            task.status = status
+            task.completed_at = datetime.now(UTC)
+            if error_msg:
+                task.error = error_msg
+    except Exception as e:
+        logger.error(f"Error updating task status: {e}")
+
+
 def _process_one_off_task(
     pq: PQ,
     *,
@@ -798,9 +873,11 @@ def _process_one_off_task(
             if task is None:
                 return False
 
-            # Mark as running
+            # Mark as running. ``started_at`` also identifies this run when
+            # Phase 3 records its result (see ``_finish_one_off``).
+            started_at = datetime.now(UTC)
             task.status = TaskStatus.RUNNING
-            task.started_at = datetime.now(UTC)
+            task.started_at = started_at
             task.attempts += 1
 
             # Get task data for execution (before session closes)
@@ -869,27 +946,9 @@ def _process_one_off_task(
     elapsed = time.perf_counter() - start
 
     # Phase 3: Update task status (guarded — only if still RUNNING, so we
-    # don't overwrite a reaper's FAILED verdict for an orphaned task)
-    try:
-        with pq.session() as session:
-            values: dict[str, object] = {
-                "status": status,
-                "completed_at": datetime.now(UTC),
-            }
-            if error_msg:
-                values["error"] = error_msg
-            result = session.execute(
-                update(Task)
-                .where(Task.id == task_id, Task.status == TaskStatus.RUNNING)
-                .values(**values)
-            )
-            if result.rowcount == 0:
-                logger.warning(
-                    f"Task '{name}' (id={task_id}) was no longer RUNNING"
-                    " when Phase 3 tried to update — likely reaped"
-                )
-    except Exception as e:
-        logger.error(f"Error updating task status: {e}")
+    # don't overwrite a reaper's FAILED verdict for an orphaned task; and
+    # re-queued instead of closed if an upsert was parked on the row)
+    _finish_one_off(pq, task_id, started_at, name, status, error_msg)
 
     # Log result
     if status == TaskStatus.COMPLETED:
@@ -1092,8 +1151,9 @@ def _claim_and_fork_one_off(
             if task is None:
                 return None
 
+            started_at = datetime.now(UTC)
             task.status = TaskStatus.RUNNING
-            task.started_at = datetime.now(UTC)
+            task.started_at = started_at
             task.attempts += 1
 
             name = task.name
@@ -1128,15 +1188,7 @@ def _claim_and_fork_one_off(
         )
     except Exception as e:
         logger.error(f"Error starting task '{name}': {e}")
-        try:
-            with pq.session() as session:
-                t = session.get(Task, task_id)
-                if t:
-                    t.status = TaskStatus.FAILED
-                    t.completed_at = datetime.now(UTC)
-                    t.error = str(e)
-        except Exception as update_err:
-            logger.error(f"Error updating task status: {update_err}")
+        _finish_one_off(pq, task_id, started_at, name, TaskStatus.FAILED, str(e))
         return None
 
     return _ChildSlot(
@@ -1146,6 +1198,7 @@ def _claim_and_fork_one_off(
         name=name,
         start_time=time.perf_counter(),
         is_periodic=False,
+        started_at=started_at,
     )
 
 
@@ -1342,29 +1395,14 @@ def _reap_and_update(pq: PQ, slot: _ChildSlot) -> None:
         if result.exit_kind == "timeout":
             error_msg = f"Timed out after {elapsed:.3f} s"
 
-        try:
-            with pq.session() as session:
-                values: dict[str, object] = {
-                    "status": result.task_status,
-                    "completed_at": datetime.now(UTC),
-                }
-                if error_msg:
-                    values["error"] = error_msg
-                row_result = session.execute(
-                    update(Task)
-                    .where(
-                        Task.id == slot.task_id,
-                        Task.status == TaskStatus.RUNNING,
-                    )
-                    .values(**values)
-                )
-                if row_result.rowcount == 0:
-                    logger.warning(
-                        f"Task '{slot.name}' (id={slot.task_id}) was no longer"
-                        " RUNNING when reap tried to update — likely reaped"
-                    )
-        except Exception as e:
-            logger.error(f"Error updating task status: {e}")
+        _finish_one_off(
+            pq,
+            slot.task_id,
+            slot.started_at,
+            slot.name,
+            result.task_status,
+            error_msg,
+        )
 
         if result.task_status == TaskStatus.COMPLETED:
             logger.debug(f"Task '{slot.name}' completed in {elapsed:.3f} s")