diff --git a/docs/guides/cron.md b/docs/guides/cron.md index 03323bb..5b1530d 100644 --- a/docs/guides/cron.md +++ b/docs/guides/cron.md @@ -40,6 +40,29 @@ Expressions use standard five-field cron syntax (`minute hour day-of-month month - `queue.delete_cron(task, "*/1 * * * *")` removes a cron; pass the same task and expression you registered. - `sheppy cron list` shows registered crons from the CLI. +## Declarative cron jobs (pyproject.toml) + +Instead of calling `add_cron()`, you can declare cron jobs in your project's `pyproject.toml`: + +```toml +[[tool.sheppy.cron]] +task = "myapp.tasks:cleanup" +expression = "0 3 * * *" +args = [30] + +[[tool.sheppy.cron]] +task = "myapp.tasks:backup" +expression = "0 4 * * sun" +``` + +Every worker started with `sheppy work` from that directory reads the file and reconciles continuously: + +- crons declared in the file are created and marked as file-managed +- remove an entry from the file and the worker deletes that cron within seconds +- crons added through `queue.add_cron()` are left alone, even when identical to a declaration + +Entries take optional `args`, `kwargs`, and `queue` (which defaults to the worker's first queue). Invalid entries are logged and skipped, and if the file is missing or broken the worker keeps the current state instead of deleting anything. Use `Worker(cron_config_file=...)` for the programmatic path. + ## When to use cron vs `schedule()` Use `queue.schedule()` for one-off future tasks ("send this email in 30 minutes"). Use `add_cron()` for recurring work ("nightly cleanup at 03:00"). A cron registration persists until you delete it; a scheduled task fires once. diff --git a/pyproject.toml b/pyproject.toml index d64ae1a..99d2d32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "croniter>=6.0.0", "pydantic>=2.1.1", "redis>=6.0.0,<9.0.0", + "tomli>=2.0.0; python_version < '3.11'", "typer>=0.19.0", ] diff --git a/src/sheppy/_sync_queue.py b/src/sheppy/_sync_queue.py index 0032c53..ecf54ec 100644 --- a/src/sheppy/_sync_queue.py +++ b/src/sheppy/_sync_queue.py @@ -180,6 +180,38 @@ def retry(self, task: Task | UUID | str, at: datetime | timedelta | None = None, """ return self._run_coro(self._queue.retry(task, at, force)) + def cancel(self, task: Task | UUID | str) -> Task: + """Cancel a pending or scheduled task. + + Args: + task: Instance of a Task or its ID. + + Returns: + The updated Task instance with status 'cancelled'. + + Raises: + TaskCancellationError: If the task cannot be cancelled - either it + was already claimed by a worker, it already finished, or it + does not exist. + """ + return self._run_coro(self._queue.cancel(task)) + + def delete(self, task: Task | UUID | str) -> bool: + """Hard-delete a finished task's stored metadata. + + Only tasks in a terminal state (completed, failed, crashed, cancelled) can be deleted. + + Args: + task: Instance of a Task or its ID. + + Returns: + True if the task existed and was deleted, False if it was not found. + + Raises: + ValueError: If the task has not finished yet. + """ + return self._run_coro(self._queue.delete(task)) + def size(self) -> int: """Get number of pending tasks in the queue. diff --git a/src/sheppy/_utils/cron_config.py b/src/sheppy/_utils/cron_config.py new file mode 100644 index 0000000..4e9423a --- /dev/null +++ b/src/sheppy/_utils/cron_config.py @@ -0,0 +1,86 @@ +""" +This file contains utility functions meant for internal use only. Expect breaking changes if you use them directly. +""" + +import logging +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +if sys.version_info >= (3, 11): + import tomllib +else: # Python 3.10 + import tomli as tomllib + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class CronDeclaration: + task: str + expression: str + args: tuple[Any, ...] = () + kwargs: dict[str, Any] = field(default_factory=dict) + queue: str | None = None + + +def _parse_entry(index: int, entry: Any) -> CronDeclaration | None: + if not isinstance(entry, dict): + logger.warning(f"[tool.sheppy.cron] entry #{index} is not a table, skipping") + return None + + task = entry.get("task") + expression = entry.get("expression") + + if not isinstance(task, str) or ":" not in task: + logger.warning(f"[tool.sheppy.cron] entry #{index} has a missing or invalid 'task' (expected 'module:function'), skipping") + return None + + if not isinstance(expression, str) or not expression.strip(): + logger.warning(f"[tool.sheppy.cron] entry #{index} ({task}) has a missing or invalid 'expression', skipping") + return None + + args = entry.get("args", []) + if not isinstance(args, list): + logger.warning(f"[tool.sheppy.cron] entry #{index} ({task}) has non-list 'args', skipping") + return None + + kwargs = entry.get("kwargs", {}) + if not isinstance(kwargs, dict): + logger.warning(f"[tool.sheppy.cron] entry #{index} ({task}) has non-table 'kwargs', skipping") + return None + + queue = entry.get("queue") + if queue is not None and not isinstance(queue, str): + logger.warning(f"[tool.sheppy.cron] entry #{index} ({task}) has a non-string 'queue', skipping") + return None + + return CronDeclaration(task=task, expression=expression, args=tuple(args), kwargs=dict(kwargs), queue=queue) + + +def load_cron_declarations(path: str | Path) -> list[CronDeclaration] | None: + path = Path(path) + + try: + data = tomllib.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + logger.debug(f"cron config file not found: {path}") + return None + except (OSError, tomllib.TOMLDecodeError) as e: + logger.error(f"cannot parse cron config file {path}: {e}") + return None + + entries = data.get("tool", {}).get("sheppy", {}).get("cron", []) + + if not isinstance(entries, list): + logger.error(f"[tool.sheppy.cron] in {path} must be an array of tables") + return None + + declarations = [] + for index, entry in enumerate(entries, 1): + declaration = _parse_entry(index, entry) + if declaration is not None: + declarations.append(declaration) + + return declarations diff --git a/src/sheppy/backend/base.py b/src/sheppy/backend/base.py index 0b4acf7..c49b042 100644 --- a/src/sheppy/backend/base.py +++ b/src/sheppy/backend/base.py @@ -17,7 +17,7 @@ def resolve_metadata_ttl(task_data: dict[str, Any], *, ttl: int | None, error_tt task_ttl: TTLValue = config.get("ttl", "inherit") task_error_ttl: TTLValue = config.get("error_ttl", "inherit") - if task_data.get("status") in ("failed", "crashed"): + if task_data.get("status") in ("failed", "crashed", "cancelled"): for candidate in (task_error_ttl, task_ttl, error_ttl, ttl): if not isinstance(candidate, str): # skip "inherit" return candidate @@ -77,6 +77,14 @@ async def pop_scheduled(self, queue_name: str, now: datetime | None = None) -> l async def store_result(self, queue_name: str, task_data: dict[str, Any]) -> bool: pass + @abstractmethod + async def cancel(self, queue_name: str, task_id: str) -> dict[str, Any] | None: + pass + + @abstractmethod + async def delete_task(self, queue_name: str, task_id: str) -> bool: + pass + @abstractmethod async def get_results(self, queue_name: str, task_ids: list[str], timeout: float | None = None) -> dict[str, dict[str, Any]]: pass diff --git a/src/sheppy/backend/memory.py b/src/sheppy/backend/memory.py index 7a4bba7..4ae0fa3 100644 --- a/src/sheppy/backend/memory.py +++ b/src/sheppy/backend/memory.py @@ -253,6 +253,51 @@ async def store_result(self, queue_name: str, task_data: dict[str, Any]) -> bool return True + async def cancel(self, queue_name: str, task_id: str) -> dict[str, Any] | None: + self._check_connected() + + async with self._locks[queue_name]: + self._purge_expired(queue_name) + task_data = self._task_metadata[queue_name].get(task_id) + + if not task_data or task_data.get("finished_at") is not None: + return None + + # task was already claimed by a worker + removed = False + if task_id in self._pending[queue_name]: + self._pending[queue_name].remove(task_id) + removed = True + else: + scheduled = self._scheduled[queue_name] + for i, scheduled_task in enumerate(scheduled): + if scheduled_task.task_id == task_id: + scheduled.pop(i) + heapq.heapify(scheduled) + removed = True + break + + if not removed: + return None + + task_data["status"] = "cancelled" + task_data["finished_at"] = datetime.now(timezone.utc).isoformat() + self._set_expiry(queue_name, task_data) + + return deepcopy(task_data) + + async def delete_task(self, queue_name: str, task_id: str) -> bool: + self._check_connected() + + async with self._locks[queue_name]: + if task_id not in self._task_metadata[queue_name]: + return False + + del self._task_metadata[queue_name][task_id] + self._task_expiry[queue_name].pop(task_id, None) + + return True + async def get_results(self, queue_name: str, task_ids: list[str], timeout: float | None = None) -> dict[str,dict[str, Any]]: self._check_connected() diff --git a/src/sheppy/backend/redis.py b/src/sheppy/backend/redis.py index 10736b6..797ef47 100644 --- a/src/sheppy/backend/redis.py +++ b/src/sheppy/backend/redis.py @@ -1,7 +1,7 @@ import asyncio import contextlib import json -from datetime import datetime +from datetime import datetime, timezone from time import time from typing import Any, cast @@ -108,6 +108,10 @@ def _pending_tasks_key(self, queue_name: str) -> str: """Queued tasks to be processed (stream)""" return f"sheppy:pending:{queue_name}" + def _pending_index_key(self, queue_name: str) -> str: + """Index of unclaimed pending tasks: task_id -> stream message_id (hash)""" + return f"sheppy:pending_ids:{queue_name}" + def _finished_tasks_key(self, queue_name: str) -> str: """Notifications about finished tasks (stream)""" return f"sheppy:finished:{queue_name}" @@ -165,6 +169,7 @@ async def append(self, queue_name: str, tasks: list[dict[str, Any]], unique: boo try: async with self.client.pipeline(transaction=False) as pipe: pipe.hsetnx(self._queues_registry_key(), queue_name, "{}") + xadd_positions = [] for t in to_queue: _task_data = json.dumps(t) @@ -173,12 +178,21 @@ async def append(self, queue_name: str, tasks: list[dict[str, Any]], unique: boo pipe.set(f"{tasks_metadata_key}:{t['id']}", _task_data) # add to pending stream + xadd_positions.append(len(pipe.command_stack)) pipe.xadd(pending_tasks_key, {"data": _task_data}) - await pipe.execute() + res = await pipe.execute() except Exception as e: raise BackendError(f"Failed to enqueue task: {e}") from e + # record stream message ids in the pending index (used by cancel()) + if to_queue: + mapping = {} + for t, pos in zip(to_queue, xadd_positions, strict=True): + message_id = res[pos] + mapping[t["id"]] = message_id.decode() if isinstance(message_id, bytes) else message_id + await self.client.hset(self._pending_index_key(queue_name), mapping=mapping) + return success async def pop(self, queue_name: str, limit: int = 1, timeout: float | None = None) -> list[dict[str, Any]]: @@ -212,6 +226,8 @@ async def pop(self, queue_name: str, limit: int = 1, timeout: float | None = Non self._pending_messages[task_data["id"]] = (queue_name, message_id.decode()) tasks.append(task_data) + await self.client.hdel(self._pending_index_key(queue_name), *[t["id"] for t in tasks]) + return tasks except Exception as e: @@ -246,6 +262,7 @@ async def clear(self, queue_name: str) -> int: count += 1 await self.client.xtrim(pending_tasks_key, maxlen=0) + await self.client.delete(self._pending_index_key(queue_name)) await self.client.delete(scheduled_key) await self.client.hdel(self._queues_registry_key(), queue_name) await self.client.delete(self._rate_limit_key(queue_name)) @@ -354,6 +371,67 @@ async def store_result(self, queue_name: str, task_data: dict[str, Any]) -> bool except Exception as e: raise BackendError(f"Failed to store task result: {e}") from e + async def cancel(self, queue_name: str, task_id: str) -> dict[str, Any] | None: + tasks_metadata_key = self._tasks_metadata_key(queue_name) + scheduled_key = self._scheduled_tasks_key(queue_name) + pending_tasks_key = self._pending_tasks_key(queue_name) + pending_index_key = self._pending_index_key(queue_name) + + await self._ensure_consumer_group(pending_tasks_key) + + raw = await self.client.get(f"{tasks_metadata_key}:{task_id}") + if not raw: + return None + + task_data: dict[str, Any] = json.loads(raw) + + if task_data.get("finished_at") is not None: + return None + + if task_data.get("status") == "scheduled": + removed = await self.client.zrem(scheduled_key, task_id) + if removed > 0: + return await self._finalize_cancellation(queue_name, task_data) + # else: scheduled task just got into pending queue + + raw_message_id = await self.client.hget(pending_index_key, task_id) + if raw_message_id is None: + # already claimed by a worker (or never queued) -> cancellation fails + return None + + message_id = raw_message_id.decode() if isinstance(raw_message_id, bytes) else raw_message_id + + # XDEL is atomic, so only one of two concurrent cancellations wins + removed = await self.client.xdel(pending_tasks_key, message_id) + await self.client.hdel(pending_index_key, task_id) + + if removed <= 0: + # stale index entry -> the message was already removed + return None + + return await self._finalize_cancellation(queue_name, task_data) + + async def _finalize_cancellation(self, queue_name: str, task_data: dict[str, Any]) -> dict[str, Any]: + task_data["status"] = "cancelled" + task_data["finished_at"] = datetime.now(timezone.utc).isoformat() + + async with self.client.pipeline(transaction=True) as pipe: + pipe.set( + f"{self._tasks_metadata_key(queue_name)}:{task_data['id']}", + json.dumps(task_data), + ex=resolve_metadata_ttl(task_data, ttl=self.ttl, error_ttl=self.error_ttl), + ) + # notify waiters (wait_for) that the task reached a terminal state + min_id = f"{int((time() - self._results_stream_ttl) * 1000)}-0" + pipe.xadd(self._finished_tasks_key(queue_name), {"task_id": task_data["id"]}, minid=min_id) + await pipe.execute() + + return task_data + + async def delete_task(self, queue_name: str, task_id: str) -> bool: + tasks_metadata_key = self._tasks_metadata_key(queue_name) + return bool(await self.client.delete(f"{tasks_metadata_key}:{task_id}")) + async def get_stats(self, queue_name: str) -> dict[str, int]: scheduled_tasks_key = self._scheduled_tasks_key(queue_name) pending_tasks_key = self._pending_tasks_key(queue_name) diff --git a/src/sheppy/cli/cli.py b/src/sheppy/cli/cli.py index 6f350ec..45d9006 100644 --- a/src/sheppy/cli/cli.py +++ b/src/sheppy/cli/cli.py @@ -8,6 +8,8 @@ from .commands.cron.list import list_crons from .commands.queue.list import list_queues from .commands.task.add import add +from .commands.task.cancel import cancel +from .commands.task.delete import delete from .commands.task.info import info from .commands.task.list import list_tasks from .commands.task.retry import retry @@ -41,6 +43,8 @@ def callback( task_app.command()(test) task_app.command()(add) task_app.command()(schedule) +task_app.command()(cancel) +task_app.command()(delete) app.add_typer(task_app, name="task") diff --git a/src/sheppy/cli/commands/task/cancel.py b/src/sheppy/cli/commands/task/cancel.py new file mode 100644 index 0000000..e47af35 --- /dev/null +++ b/src/sheppy/cli/commands/task/cancel.py @@ -0,0 +1,59 @@ +import asyncio +import os +import sys +from typing import Annotated +from uuid import UUID + +import typer + +from sheppy import Queue +from sheppy._config import config +from sheppy.exceptions import TaskCancellationError +from sheppy.queue import _create_backend_from_url + +from ...utils import OutputFormat, console, print_json, task_status_label + + +def cancel( + task_id: Annotated[str, typer.Argument(help="Task ID to cancel")], + queue: Annotated[str, typer.Option("--queue", "-q", help="Queue name. Env: SHEPPY_QUEUE")] = config.queue_list[0], + backend_url: Annotated[str | None, typer.Option("--backend-url", "-u", help="Backend URL. Env: SHEPPY_BACKEND_URL")] = config.backend_url, + format_output: Annotated[OutputFormat, typer.Option("--format", help="Output format")] = OutputFormat.table, +) -> None: + """Cancel a pending or scheduled task so it is never executed.""" + + cwd = os.getcwd() + if cwd not in sys.path: + sys.path.insert(0, cwd) + + async def _cancel(backend_url: str | None) -> None: + if backend_url is None: + backend_url = "redis://127.0.0.1:6379" + backend_instance = _create_backend_from_url(backend_url) + q = Queue(backend_instance, queue) + + try: + uuid_obj = UUID(task_id) + except ValueError: + console.print("[red]Error: Task ID must be UUID format[/red]") + raise typer.Exit(1) from None + + try: + task = await q.cancel(uuid_obj) + except TaskCancellationError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from None + + if format_output == OutputFormat.json: + task_dict = task.model_dump(mode='json') + task_dict["queue"] = queue + task_dict["queue_status"] = task_status_label(task) + task_dict["cancelled"] = True + print_json(task_dict) + return + + console.print(f"[green]✓ Task {task_id} has been cancelled[/green]") + console.print(f" Function: [blue]{task.spec.func}[/blue]") + console.print(f" Status: [magenta]{task.status}[/magenta]") + + asyncio.run(_cancel(backend_url)) diff --git a/src/sheppy/cli/commands/task/delete.py b/src/sheppy/cli/commands/task/delete.py new file mode 100644 index 0000000..174f0d2 --- /dev/null +++ b/src/sheppy/cli/commands/task/delete.py @@ -0,0 +1,56 @@ +import asyncio +import os +import sys +from typing import Annotated +from uuid import UUID + +import typer + +from sheppy import Queue +from sheppy._config import config +from sheppy.queue import _create_backend_from_url + +from ...utils import OutputFormat, console, print_json + + +def delete( + task_id: Annotated[str, typer.Argument(help="Task ID to delete")], + queue: Annotated[str, typer.Option("--queue", "-q", help="Queue name. Env: SHEPPY_QUEUE")] = config.queue_list[0], + backend_url: Annotated[str | None, typer.Option("--backend-url", "-u", help="Backend URL. Env: SHEPPY_BACKEND_URL")] = config.backend_url, + format_output: Annotated[OutputFormat, typer.Option("--format", help="Output format")] = OutputFormat.table, +) -> None: + """Hard-delete a finished task's metadata (cancel pending/scheduled tasks first).""" + + cwd = os.getcwd() + if cwd not in sys.path: + sys.path.insert(0, cwd) + + async def _delete(backend_url: str | None) -> None: + if backend_url is None: + backend_url = "redis://127.0.0.1:6379" + backend_instance = _create_backend_from_url(backend_url) + q = Queue(backend_instance, queue) + + try: + uuid_obj = UUID(task_id) + except ValueError: + console.print("[red]Error: Task ID must be UUID format[/red]") + raise typer.Exit(1) from None + + try: + deleted = await q.delete(uuid_obj) + except ValueError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from None + + if not deleted: + console.print(f"[red]Error: Task {task_id} not found in queue '{queue}'[/red]") + raise typer.Exit(1) + + if format_output == OutputFormat.json: + print_json({"task_id": task_id, "queue": queue, "deleted": True}) + return + + console.print(f"[green]✓ Task {task_id} has been deleted[/green]") + + asyncio.run(_delete(backend_url)) diff --git a/src/sheppy/cli/commands/work.py b/src/sheppy/cli/commands/work.py index 078dc78..f97cf2c 100644 --- a/src/sheppy/cli/commands/work.py +++ b/src/sheppy/cli/commands/work.py @@ -67,6 +67,7 @@ def work( console.print(f" Max concurrent tasks: [yellow]{max_concurrent}[/yellow]") if prestart: console.print(f" Prestart hook: [yellow]{prestart}[/yellow]") + console.print() _start_worker(queues, backend_instance, max_concurrent, max_prefetch, _log_level, diff --git a/src/sheppy/cli/utils.py b/src/sheppy/cli/utils.py index cf5d4f1..9493868 100644 --- a/src/sheppy/cli/utils.py +++ b/src/sheppy/cli/utils.py @@ -1,9 +1,13 @@ +import json import logging from datetime import datetime, timezone from enum import Enum +from typing import Any from rich.console import Console +from sheppy.models import Task + console = Console() @@ -54,3 +58,21 @@ def humanize_datetime(dt: datetime | None, now: datetime | None = None) -> str: time_string = f"{int(abs_delta)} second" + ("s" if abs_delta >= 2 else "") return f"{time_string} ago" if is_past else f"in {time_string}" + + +def print_json(data: Any) -> None: + console.print(json.dumps(data, indent=2, default=str), markup=False, highlight=False, soft_wrap=True, crop=False) + + +def task_status_label(task: Task) -> str: + if task.status == 'completed': + return "completed" + if task.status == 'crashed': + return "crashed" + if task.status == 'retrying': # FIXME: doesn't work + return "retrying" + if task.exception: + return "failed" + if task.status == 'scheduled': + return "scheduled" + return task.status diff --git a/src/sheppy/exceptions.py b/src/sheppy/exceptions.py index e7ea93f..247f348 100644 --- a/src/sheppy/exceptions.py +++ b/src/sheppy/exceptions.py @@ -14,6 +14,10 @@ class WorkerCrashedError(Exception): pass +class TaskCancellationError(Exception): + """Raised when a task cannot be cancelled (already claimed by a worker, already finished, or not found).""" + + class TaskFailedError(Exception): """Fallback exception for task failures whose original exception class cannot be reconstructed.""" diff --git a/src/sheppy/models.py b/src/sheppy/models.py index 7af355a..7cc4722 100644 --- a/src/sheppy/models.py +++ b/src/sheppy/models.py @@ -328,6 +328,9 @@ def add(x: int, y: int) -> int: workflow_id: UUID | None = None """UUID|None: ID of the workflow this task belongs to (if created within a workflow).""" + cron_id: UUID | None = None + """UUID|None: ID of the CronTask that created this job (temporary)""" + # caller: str | None = None # worker: str | None = None @@ -400,7 +403,9 @@ class TaskCron(BaseModel): expression: Cron expression defining the schedule, e.g. "*/5 * * * *" for every 5 minutes. spec: Task specification config: Task configuration - + managed_by: Origin of the cron definition. None means programmatic + (added via `Queue.add_cron()`); "pyproject" means declared in a + pyproject.toml file and reconciled by workers. Note: - You should not create TaskCron instances directly. Instead, use the `add_cron` method of the Queue class to create a cron definition. - `args` and `kwargs` in `spec` must be JSON serializable. @@ -442,6 +447,9 @@ def say_hello(to: str) -> str: config: TaskConfig """Task configuration""" + managed_by: str | None = None + """str|None: Origin of the cron definition; None for programmatic, "pyproject" for declarative.""" + # enabled: bool = True # last_run: AwareDatetime | None = None # next_run: AwareDatetime | None = None @@ -509,5 +517,6 @@ def create_task(self, start: datetime) -> Task: return Task( id=uuid5(TASK_CRON_NS, str(self.deterministic_id) + str(start.timestamp())), spec=self.spec.model_copy(deep=True), - config=self.config.model_copy(deep=True) + config=self.config.model_copy(deep=True), + cron_id=self.deterministic_id, ) diff --git a/src/sheppy/queue.py b/src/sheppy/queue.py index d203ee7..d9b1879 100644 --- a/src/sheppy/queue.py +++ b/src/sheppy/queue.py @@ -6,6 +6,7 @@ from ._config import config from ._workflow import Workflow, WorkflowResult, WorkflowRunner from .backend.base import Backend +from .exceptions import TaskCancellationError from .models import Task, TaskCron from .task_factory import TaskFactory @@ -335,6 +336,72 @@ async def retry(self, task: Task | UUID | str, at: datetime | timedelta | None = success = await self.backend.append(self.name, [_task.model_dump(mode="json")], unique=False) return success[0] + async def cancel(self, task: Task | UUID | str) -> Task: + """Cancel a pending or scheduled task. + + Args: + task: Instance of a Task or its ID. + + Returns: + The updated Task instance with status 'cancelled'. + + Raises: + TaskCancellationError: If the task cannot be cancelled - either it + was already claimed by a worker, it already finished, or it + does not exist. + + Example: + ```python + q = Queue(...) + + await q.schedule(task, timedelta(minutes=10)) + + cancelled = await q.cancel(task) + assert cancelled.status == 'cancelled' + ``` + """ + await self.__ensure_backend_is_connected() + + task_id = str(task.id if isinstance(task, Task) else task) + task_data = await self.backend.cancel(self.name, task_id) + + if task_data is None: + raise TaskCancellationError(f"Task {task_id} could not be cancelled.") + + return Task.model_validate(task_data) + + async def delete(self, task: Task | UUID | str) -> bool: + """Hard-delete a finished task's stored metadata. + + Only tasks in a terminal state (completed, failed, crashed, cancelled) can be deleted. + + Args: + task: Instance of a Task or its ID. + + Returns: + True if the task existed and was deleted, False if it was not found. + + Raises: + ValueError: If the task has not finished yet. + + Example: + ```python + q = Queue(...) + + await q.cancel(task) + deleted = await q.delete(task) + assert deleted is True + ``` + """ + _task = await self.get_task(task) + if not _task: + return False + + if _task.finished_at is None: + raise ValueError("Only finished tasks can be deleted, cancel the task first") + + return await self.backend.delete_task(self.name, str(_task.id)) + async def size(self) -> int: """Get number of pending tasks in the queue. @@ -428,6 +495,14 @@ async def get_crons(self) -> list[TaskCron]: await self.__ensure_backend_is_connected() return [TaskCron.model_validate(tc) for tc in await self.backend.get_crons(self.name)] + async def _store_cron(self, cron: TaskCron) -> bool: + await self.__ensure_backend_is_connected() + return await self.backend.add_cron(self.name, str(cron.deterministic_id), cron.model_dump(mode="json")) + + async def _delete_task_cron(self, cron: TaskCron) -> bool: + await self.__ensure_backend_is_connected() + return await self.backend.delete_cron(self.name, str(cron.deterministic_id)) + async def _pop_pending(self, limit: int = 1, timeout: float | None = None) -> list[Task]: """Get next task to process. Internal method used by workers. diff --git a/src/sheppy/task_factory.py b/src/sheppy/task_factory.py index 59be40b..1f10114 100644 --- a/src/sheppy/task_factory.py +++ b/src/sheppy/task_factory.py @@ -92,11 +92,12 @@ def create_task(func: Callable[..., Any], return _task @staticmethod - def create_cron_from_task(task: Task, cron_expression: str) -> TaskCron: + def create_cron_from_task(task: Task, cron_expression: str, managed_by: str | None = None) -> TaskCron: return TaskCron( expression=cron_expression, spec=task.spec.model_copy(deep=True), config=task.config.model_copy(deep=True), + managed_by=managed_by, ) diff --git a/src/sheppy/worker.py b/src/sheppy/worker.py index e67e698..10dbd36 100644 --- a/src/sheppy/worker.py +++ b/src/sheppy/worker.py @@ -1,13 +1,17 @@ import asyncio import contextlib import logging +import os import signal from collections.abc import Callable from functools import partial from typing import Any, cast +from uuid import UUID from pydantic import BaseModel +from ._utils.cron_config import load_cron_declarations +from ._utils.functions import resolve_function from ._utils.task_execution import ( LoggingMiddleware, TaskChainingMiddleware, @@ -23,6 +27,7 @@ TaskProcessorProtocol, ) from .queue import Queue +from .task_factory import TaskFactory logger = logging.getLogger(__name__) @@ -56,7 +61,6 @@ class Worker: enable_job_processing: If True, enables job processing. Default is True. enable_scheduler: If True, enables the scheduler to enqueue scheduled tasks. Default is True. enable_cron_manager: If True, enables the cron manager to handle cron jobs. Default is True. - Note: The SHEPPY_* environment variables are read by the `sheppy work` CLI command, not by this class. Configure the class explicitly. @@ -137,6 +141,7 @@ def __init__( self.enable_job_processing = enable_job_processing self.enable_scheduler = enable_scheduler self.enable_cron_manager = enable_cron_manager + self._cron_config_file = "pyproject.toml" if os.path.exists("pyproject.toml") else None self._work_queue_tasks: list[asyncio.Task[None]] = [] self._scheduler_task: asyncio.Task[None] | None = None @@ -277,9 +282,64 @@ async def _run_scheduler(self, poll_interval: float) -> None: logger.info(SCHEDULER_PREFIX + "stopped") + async def _reconcile_declared_crons(self) -> None: + if self._cron_config_file is None: + return + + declarations = load_cron_declarations(self._cron_config_file) + if declarations is None: + return + + declared_by_queue: dict[str, dict[UUID, TaskCron]] = {q.name: {} for q in self.queues} + + for decl in declarations: + queue_name = decl.queue or self.queues[0].name + if queue_name not in declared_by_queue: + logger.warning(CRON_MANAGER_PREFIX + f"Cron declaration for {decl.task!r} targets queue {queue_name!r} which this worker does not serve, skipping") + continue + + try: + func = resolve_function(decl.task) + task = TaskFactory.create_task(func, decl.args, decl.kwargs, retry=0, retry_delay=None, + middleware=None, timeout=None, retry_on_timeout=None, retry_on_crash=None) + cron = TaskFactory.create_cron_from_task(task, decl.expression, managed_by="pyproject") + except Exception as e: + logger.error(CRON_MANAGER_PREFIX + f"Invalid cron declaration for {decl.task!r}: {e}") + continue + + declared_by_queue[queue_name][cron.deterministic_id] = cron + + for queue in self.queues: + declared = declared_by_queue[queue.name] + + scheduled = await queue.get_scheduled() + + for existing in await queue.get_crons(): + det_id = existing.deterministic_id + if det_id in declared: + continue # already registered (either declared or programmatic) + + if existing.managed_by == "pyproject": + logger.info(CRON_MANAGER_PREFIX + f"Removing undeclared cron {existing.id} ({existing.spec.func})") + await queue._delete_task_cron(existing) + + cron_id = existing.deterministic_id + + for t in scheduled: + if t.cron_id == cron_id: + # todo: do for loop and batch cancel? + logger.info(CRON_MANAGER_PREFIX + f"Cancelling residual task {t.id}") + await queue.cancel(t) + + + for cron in declared.values(): + await queue._store_cron(cron) + async def _run_cron_manager(self, poll_interval: float) -> None: logger.info(CRON_MANAGER_PREFIX + "started") + await self._reconcile_declared_crons() + while not self._shutdown_event.is_set(): try: for queue in self.queues: @@ -453,6 +513,12 @@ async def _run_worker_loop(self, queue: Queue, oneshot: bool = False) -> None: async def process_task(self, task: Task, queue: Queue) -> Task: async with self._task_semaphore: + fresh = await queue.get_task(task.id) + if fresh is None or fresh.status == 'cancelled': + await queue.backend.acknowledge(queue.name, [str(task.id)]) + return fresh if fresh is not None else task + task = fresh + try: _, task = await self._task_processor.process_task(task, queue, self.worker_id) except MiddlewareError: diff --git a/tests/contract/test_cancel.py b/tests/contract/test_cancel.py new file mode 100644 index 0000000..55ea3ee --- /dev/null +++ b/tests/contract/test_cancel.py @@ -0,0 +1,254 @@ +from datetime import timedelta +from uuid import uuid4 + +import pytest + +from sheppy import Backend, Queue, RedisBackend, Worker +from sheppy.exceptions import TaskCancellationError +from tests.dependencies import ( + assert_is_completed, + simple_async_task, + simple_sync_task, +) + + +@pytest.fixture(params=["async_task", "sync_task"]) +def task_fn(request): + if request.param == "async_task": + return simple_async_task + + if request.param == "sync_task": + return simple_sync_task + + raise NotImplementedError + + +async def test_cancel_pending_task(task_fn, queue: Queue, worker: Worker): + worker.enable_scheduler = False + worker.enable_cron_manager = False + + t = task_fn(1, 2) + await queue.add(t) + assert await queue.size() == 1 + + cancelled = await queue.cancel(t) + + assert cancelled.status == 'cancelled' + assert cancelled.finished_at is not None + assert await queue.size() == 0 + + # metadata is kept + stored = await queue.get_task(t) + assert stored is not None + assert stored.status == 'cancelled' + assert stored.finished_at is not None + assert stored.result is None + + # the cancelled task must never be executed + await worker.work(oneshot=True) + + stored = await queue.get_task(t) + assert stored is not None + assert stored.status == 'cancelled' + assert stored.result is None + + +async def test_cancel_pending_task_by_id(task_fn, queue: Queue): + t1 = task_fn(1, 2) + t2 = task_fn(3, 4) + await queue.add([t1, t2]) + + cancelled = await queue.cancel(t1.id) + assert cancelled.status == 'cancelled' + + cancelled = await queue.cancel(str(t2.id)) + assert cancelled.status == 'cancelled' + + +async def test_cancel_scheduled_task(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.schedule(t, timedelta(minutes=10)) + assert len(await queue.get_scheduled()) == 1 + + cancelled = await queue.cancel(t) + + assert cancelled.status == 'cancelled' + assert cancelled.finished_at is not None + assert await queue.get_scheduled() == [] + assert await queue.size() == 0 + + stored = await queue.get_task(t) + assert stored is not None + assert stored.status == 'cancelled' + + +async def test_cancel_claimed_task_fails(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.add(t) + + # simulate a worker claiming the task + claimed = await queue._pop_pending() + assert len(claimed) == 1 + + with pytest.raises(TaskCancellationError): + await queue.cancel(t) + + +async def test_cancel_completed_task_fails(task_fn, queue: Queue, worker: Worker): + worker.enable_scheduler = False + worker.enable_cron_manager = False + + t = task_fn(1, 2) + await queue.add(t) + await worker.work(1) + + processed = await queue.get_task(t) + assert_is_completed(processed) + + with pytest.raises(TaskCancellationError): + await queue.cancel(t) + + +async def test_cancel_twice_fails(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.add(t) + + await queue.cancel(t) + + with pytest.raises(TaskCancellationError): + await queue.cancel(t) + + +async def test_cancel_nonexistent_task_fails(queue: Queue): + with pytest.raises(TaskCancellationError): + await queue.cancel(uuid4()) + + +async def test_wait_for_returns_cancelled_task(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.add(t) + + await queue.cancel(t) + + finished = await queue.wait_for(t) + assert finished is not None + assert finished.status == 'cancelled' + + +async def test_delete_cancelled_task(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.add(t) + await queue.cancel(t) + + assert await queue.delete(t) is True + assert await queue.get_task(t) is None + + # already deleted + assert await queue.delete(t) is False + + +async def test_delete_completed_task(task_fn, queue: Queue, worker: Worker): + worker.enable_scheduler = False + worker.enable_cron_manager = False + + t = task_fn(1, 2) + await queue.add(t) + await worker.work(1) + + assert await queue.delete(t) is True + assert await queue.get_task(t) is None + + +async def test_delete_unfinished_task_fails(task_fn, queue: Queue): + t = task_fn(1, 2) + await queue.add(t) + + with pytest.raises(ValueError, match="cancel the task first"): + await queue.delete(t) + + # scheduled tasks cannot be deleted either + t2 = task_fn(1, 2) + await queue.schedule(t2, timedelta(minutes=10)) + + with pytest.raises(ValueError, match="cancel the task first"): + await queue.delete(t2) + + +async def test_delete_nonexistent_task(queue: Queue): + assert await queue.delete(uuid4()) is False + + +#Redis-specific tests for the task_id -> message_id index that cancel() relies on +class TestPendingIndexRedis: + + async def test_index_lifecycle(self, queue: Queue, backend: Backend): + if not isinstance(backend, RedisBackend): + pytest.skip("pending index is Redis-specific") + + index_key = f"sheppy:pending_ids:{queue.name}" + + t = simple_async_task(1, 2) + await queue.add(t) + + # enqueued tasks are indexed + assert await backend.client.hget(index_key, str(t.id)) is not None + + # claiming the task removes it from the index (making it non-cancellable) + await queue._pop_pending() + assert await backend.client.hget(index_key, str(t.id)) is None + + async def test_index_removed_on_cancel(self, queue: Queue, backend: Backend): + if not isinstance(backend, RedisBackend): + pytest.skip("pending index is Redis-specific") + + index_key = f"sheppy:pending_ids:{queue.name}" + + t = simple_async_task(1, 2) + await queue.add(t) + await queue.cancel(t) + + assert await backend.client.hget(index_key, str(t.id)) is None + + async def test_cancelled_task_never_executed_even_if_claimed(self, queue: Queue, worker: Worker, backend: Backend): + if not isinstance(backend, RedisBackend): + pytest.skip("this race is Redis-specific") + + index_key = f"sheppy:pending_ids:{queue.name}" + + t = simple_async_task(1, 2) + await queue.add(t) + + # grab the message id, then claim the task (pop removes the index entry) + message_id = await backend.client.hget(index_key, str(t.id)) + assert message_id is not None + claimed = await queue._pop_pending() + assert len(claimed) == 1 + + # simulate the cancel landing in the window between the worker's + # XREADGROUP and the index cleanup (index entry still present) + await backend.client.hset(index_key, str(t.id), message_id) + + cancelled = await queue.cancel(t) + assert cancelled.status == 'cancelled' + + # the worker must skip execution and leave the cancelled state intact + result = await worker.process_task(claimed[0], queue) + assert result.status == 'cancelled' + assert result.result is None + + stored = await queue.get_task(t) + assert stored is not None + assert stored.status == 'cancelled' + assert stored.result is None + + +async def test_cancelled_task_uses_error_ttl(queue: Queue, backend: Backend): + if not isinstance(backend, RedisBackend): + pytest.skip("TTL inspection is Redis-specific") + + t = simple_async_task(1, 2) + await queue.add(t) + await queue.cancel(t) + + ttl = await backend.client.ttl(f"sheppy:tasks:{queue.name}:{t.id}") + assert ttl > 0 diff --git a/tests/contract/test_declared_cron.py b/tests/contract/test_declared_cron.py new file mode 100644 index 0000000..dd68105 --- /dev/null +++ b/tests/contract/test_declared_cron.py @@ -0,0 +1,107 @@ +import asyncio +import contextlib +from pathlib import Path + +from sheppy import Queue, Worker +from sheppy.backend.base import Backend +from tests.conftest import TEST_QUEUE_NAME +from tests.dependencies import simple_sync_task, simple_sync_task_no_param + +DECLARED_TOML = """ +[[tool.sheppy.cron]] +task = "tests.dependencies:simple_sync_task" +expression = "0 3 * * *" +args = [1, 2] + +[[tool.sheppy.cron]] +task = "tests.dependencies:simple_sync_task_no_param" +expression = "*/5 * * * *" +""" + +DECLARED_TOML_REDUCED = """ +[[tool.sheppy.cron]] +task = "tests.dependencies:simple_sync_task_no_param" +expression = "*/5 * * * *" +""" + + +def _make_worker(backend: Backend, cron_config_file: str) -> Worker: + worker = Worker(TEST_QUEUE_NAME, backend, enable_job_processing=False, enable_scheduler=False) + worker._cron_config_file = cron_config_file + worker._cron_polling_interval = 0.001 + return worker + + +async def _run_worker_briefly(worker: Worker, seconds: float = 0.1) -> None: + worker_task = asyncio.create_task(worker.work()) + await asyncio.sleep(seconds) + worker_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await worker_task + + +async def test_declared_crons_are_created_and_reconciled(queue: Queue, worker_backend: Backend, tmp_path: Path) -> None: + config_file = tmp_path / "pyproject.toml" + config_file.write_text(DECLARED_TOML) + + worker = _make_worker(worker_backend, str(config_file)) + await _run_worker_briefly(worker) + + crons = await queue.get_crons() + assert len(crons) == 2 + assert all(cron.managed_by == "pyproject" for cron in crons) + functions = {cron.spec.func for cron in crons} + assert functions == {"tests.dependencies:simple_sync_task", "tests.dependencies:simple_sync_task_no_param"} + + # drop one entry from the file. The next cycle must remove only that cron + config_file.write_text(DECLARED_TOML_REDUCED) + + worker = _make_worker(worker_backend, str(config_file)) + await _run_worker_briefly(worker) + + crons = await queue.get_crons() + assert [cron.spec.func for cron in crons] == ["tests.dependencies:simple_sync_task_no_param"] + + +async def test_programmatic_crons_survive_reconciliation(queue: Queue, worker_backend: Backend, tmp_path: Path) -> None: + config_file = tmp_path / "pyproject.toml" + config_file.write_text(DECLARED_TOML) + + # programmatic cron for a task that is not declared in the file + await queue.add_cron(simple_sync_task(5, 5), "0 12 * * *") + + # programmatic cron identical to a declaration: it must not be hijacked or deleted + await queue.add_cron(simple_sync_task(1, 2), "0 3 * * *") + + worker = _make_worker(worker_backend, str(config_file)) + await _run_worker_briefly(worker) + + crons = {(cron.spec.func, cron.spec.args): cron for cron in await queue.get_crons()} + assert len(crons) == 3 + + # the identical one stays programmatic + identical = crons[("tests.dependencies:simple_sync_task", (1, 2))] + assert identical.managed_by is None + + # now remove the declaration from the file. The programmatic twin must survive + config_file.write_text(""" +[[tool.sheppy.cron]] +task = "tests.dependencies:simple_sync_task_no_param" +expression = "*/5 * * * *" +""") + + worker = _make_worker(worker_backend, str(config_file)) + await _run_worker_briefly(worker) + + crons = {(cron.spec.func, cron.spec.args): cron for cron in await queue.get_crons()} + assert ("tests.dependencies:simple_sync_task", (1, 2)) in crons + assert ("tests.dependencies:simple_sync_task", (5, 5)) in crons + + +async def test_missing_config_file_keeps_state(queue: Queue, worker_backend: Backend, tmp_path: Path) -> None: + await queue.add_cron(simple_sync_task_no_param(), "*/5 * * * *") + + worker = _make_worker(worker_backend, str(tmp_path / "does-not-exist.toml")) + await _run_worker_briefly(worker) + + assert len(await queue.get_crons()) == 1 diff --git a/tests/unit/test_cron_config.py b/tests/unit/test_cron_config.py new file mode 100644 index 0000000..af2cd01 --- /dev/null +++ b/tests/unit/test_cron_config.py @@ -0,0 +1,97 @@ +from sheppy._utils.cron_config import load_cron_declarations + +VALID_TOML = """ +[[tool.sheppy.cron]] +task = "myapp.tasks:cleanup" +expression = "0 3 * * *" +args = [30] +kwargs = {dry_run = true} +queue = "maintenance" + +[[tool.sheppy.cron]] +task = "myapp.tasks:ping" +expression = "*/5 * * * *" +""" + + +def test_load_valid_declarations(tmp_path): + path = tmp_path / "pyproject.toml" + path.write_text(VALID_TOML) + + declarations = load_cron_declarations(path) + + assert declarations is not None + assert len(declarations) == 2 + + first, second = declarations + assert first.task == "myapp.tasks:cleanup" + assert first.expression == "0 3 * * *" + assert first.args == (30,) + assert first.kwargs == {"dry_run": True} + assert first.queue == "maintenance" + + assert second.task == "myapp.tasks:ping" + assert second.args == () + assert second.kwargs == {} + assert second.queue is None + + +def test_missing_file_returns_none(tmp_path): + assert load_cron_declarations(tmp_path / "nope.toml") is None + + +def test_invalid_toml_returns_none(tmp_path): + path = tmp_path / "pyproject.toml" + path.write_text("[[[not toml") + assert load_cron_declarations(path) is None + + +def test_cron_section_not_a_list_returns_none(tmp_path): + path = tmp_path / "pyproject.toml" + path.write_text('[tool.sheppy]\ncron = "nope"\n') + assert load_cron_declarations(path) is None + + +def test_no_cron_section_returns_empty_list(tmp_path): + path = tmp_path / "pyproject.toml" + path.write_text('[project]\nname = "something"\n') + assert load_cron_declarations(path) == [] + + +def test_invalid_entries_are_skipped(tmp_path): + path = tmp_path / "pyproject.toml" + path.write_text(""" +[[tool.sheppy.cron]] +task = "myapp.tasks:ok" +expression = "* * * * *" + +[[tool.sheppy.cron]] +expression = "* * * * *" + +[[tool.sheppy.cron]] +task = "not-a-module-path" +expression = "* * * * *" + +[[tool.sheppy.cron]] +task = "myapp.tasks:no_expression" + +[[tool.sheppy.cron]] +task = "myapp.tasks:bad_args" +expression = "* * * * *" +args = "not-a-list" + +[[tool.sheppy.cron]] +task = "myapp.tasks:bad_kwargs" +expression = "* * * * *" +kwargs = [1, 2] + +[[tool.sheppy.cron]] +task = "myapp.tasks:bad_queue" +expression = "* * * * *" +queue = 5 +""") + + declarations = load_cron_declarations(path) + + assert declarations is not None + assert [d.task for d in declarations] == ["myapp.tasks:ok"] diff --git a/tests/unit/test_ttl_resolution.py b/tests/unit/test_ttl_resolution.py index 9792a19..7a44455 100644 --- a/tests/unit/test_ttl_resolution.py +++ b/tests/unit/test_ttl_resolution.py @@ -28,19 +28,19 @@ def test_no_expiry_when_nothing_configured(self): assert resolve_metadata_ttl(task_data("failed"), ttl=None, error_ttl="inherit") is None def test_backend_ttl_applies_to_regular_statuses(self): - for status in ("completed", "cancelled", "unknown"): + for status in ("completed", "unknown"): assert resolve_metadata_ttl(task_data(status), ttl=100, error_ttl=10) == 100 def test_backend_error_ttl_applies_to_error_statuses(self): - for status in ("failed", "crashed"): + for status in ("failed", "crashed", "cancelled"): assert resolve_metadata_ttl(task_data(status), ttl=100, error_ttl=10) == 10 def test_backend_error_ttl_inherit_falls_back_to_backend_ttl(self): - for status in ("failed", "crashed"): + for status in ("failed", "crashed", "cancelled"): assert resolve_metadata_ttl(task_data(status), ttl=100, error_ttl="inherit") == 100 def test_backend_error_ttl_none_disables_expiry(self): - for status in ("failed", "crashed"): + for status in ("failed", "crashed", "cancelled"): assert resolve_metadata_ttl(task_data(status), ttl=100, error_ttl=None) is None def test_backend_ttl_none_disables_expiry(self):