percival 0.1.0__py3-none-any.whl

percival/__init__.py

@@ -0,0 +1,67 @@
+"""Client side of Percival. Installed in the Celery project, not in the panel.
+
+import percival
+percival.init(app, redis_url="redis://...", code_version=os.environ["GIT_SHA"])
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+__version__ = "0.1.0"
+
+from percival.config import UNKNOWN_VERSION, Settings
+from percival.dedupe import once_per
+from percival.runtime import ClientFactory, Runtime
+
+log = logging.getLogger(__name__)
+
+__all__ = ["Settings", "__version__", "init", "once_per", "shutdown"]
+
+
+def init(
+    app: Any,
+    redis_url: str | None = None,
+    code_version: str | None = None,
+    *,
+    client_factory: ClientFactory | None = None,
+    **options: Any,
+) -> Runtime:
+    """Wire percival into a celery app. Safe to call more than once.
+
+    Any field of :class:`percival.Settings` can be passed as a keyword, e.g.
+    ``redact_kwargs=[...]``, ``stale_action="drop"``, ``capture_args=False``.
+    ``client_factory`` exists for tests, you shouldn't need it.
+    """
+    # imported here so `import percival` stays cheap for things like the CLI
+    from percival import gate, heartbeat, metadata, runtime, signals
+
+    settings = Settings.build(redis_url=redis_url, code_version=code_version, **options)
+    if settings.code_version == UNKNOWN_VERSION:
+        log.warning(
+            "percival: no code_version set. pass code_version= or set PERCIVAL_CODE_VERSION, "
+            "otherwise schedules that require a version will be blocked on this worker"
+        )
+
+    rt = runtime.configure(settings, client_factory=client_factory)
+    signals.connect()
+    heartbeat.connect()
+    metadata.connect()
+    gate.connect()
+
+    # if the worker is already up (init called late), gate the tasks now
+    if getattr(app, "finalized", False):
+        gate.install(app)
+    return rt
+
+
+def shutdown() -> None:
+    """Undo init(). Mostly useful in tests."""
+    from percival import gate, heartbeat, metadata, runtime, signals
+
+    signals.disconnect()
+    heartbeat.disconnect()
+    metadata.disconnect()
+    gate.disconnect()
+    runtime.reset()

percival/beat.py

@@ -0,0 +1,91 @@
+"""RedBeat scheduler with per-schedule timezones and a heartbeat.
+
+    celery -A proj beat -S percival.beat.Scheduler
+
+It's stock RedBeat underneath. Entries stay in RedBeat's format, so plain
+``redbeat.RedBeatScheduler`` still reads them (crontabs just fall back to the
+app timezone). What this adds:
+
+- crontabs with a ``percival_timezone`` header are evaluated in that zone
+- a heartbeat key so the panel knows beat is alive, when it last ticked and
+  who holds the lock
+- a shorter default loop interval, so a newly created schedule is picked up
+  in seconds instead of up to 5 minutes
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import socket
+import time
+from typing import Any
+
+from celery.schedules import crontab
+from redbeat.schedulers import RedBeatScheduler, RedBeatSchedulerEntry
+
+from percival import __version__, runtime
+from percival.naming import Headers
+from percival.schedules import zoned_crontab
+
+log = logging.getLogger(__name__)
+
+DEFAULT_MAX_INTERVAL = 15
+
+
+def apply_zone(sched: Any, options: dict[str, Any] | None, app: Any = None) -> Any:
+    headers = (options or {}).get("headers") or {}
+    zone = headers.get(Headers.TIMEZONE)
+    if not zone or not isinstance(sched, crontab) or isinstance(sched, zoned_crontab):
+        return sched
+    try:
+        return zoned_crontab.from_crontab(sched, zone, app=app)
+    except Exception:
+        log.warning("percival: unknown timezone %r, using the app timezone", zone)
+        return sched
+
+
+class Entry(RedBeatSchedulerEntry):  # type: ignore[misc]
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.schedule = apply_zone(self.schedule, self.options, self.app)  # type: ignore[has-type]
+
+
+class Scheduler(RedBeatScheduler):  # type: ignore[misc]
+    Entry = Entry
+
+    def __init__(self, app: Any, *args: Any, **kwargs: Any) -> None:
+        if kwargs.get("max_interval") is None and not app.conf.beat_max_loop_interval:
+            kwargs["max_interval"] = DEFAULT_MAX_INTERVAL
+        super().__init__(app, *args, **kwargs)
+        self._started_at = time.time()
+
+    def tick(self, *args: Any, **kwargs: Any) -> Any:
+        result = super().tick(*args, **kwargs)
+        self._heartbeat()
+        return result
+
+    def _heartbeat(self) -> None:
+        rt = runtime.get_optional()
+        if rt is None:
+            return
+        payload = {
+            "hostname": socket.gethostname(),
+            "pid": os.getpid(),
+            "scheduler": f"{type(self).__module__}.{type(self).__name__}",
+            "started_at": self._started_at,
+            "last_tick": time.time(),
+            "max_interval": self.max_interval,
+            "lock_key": self.lock_key,
+            "code_version": rt.settings.code_version,
+            "client_version": __version__,
+        }
+        try:
+            rt.redis.get().set(
+                rt.keys.beat,
+                json.dumps(payload),
+                ex=int(max(self.max_interval * 3, 60)),
+            )
+        except Exception:
+            log.debug("percival: beat heartbeat failed", exc_info=True)

percival/cli.py

@@ -0,0 +1,78 @@
+"""``percival`` command line. Only ``import`` for now."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import os
+import sys
+from typing import Any
+
+from percival.naming import NAME
+
+
+def load_app(spec: str) -> Any:
+    """``proj.celery:app``, or just ``proj.celery`` if the app is called app/celery."""
+    # same as celery -A: make the current directory importable
+    if os.getcwd() not in sys.path:
+        sys.path.insert(0, os.getcwd())
+
+    module_name, _, attr = spec.partition(":")
+    module = importlib.import_module(module_name)
+    if attr:
+        return getattr(module, attr)
+    for candidate in ("app", "celery"):
+        if hasattr(module, candidate):
+            return getattr(module, candidate)
+    raise SystemExit(f"couldn't find a celery app in {module_name}, try {module_name}:app")
+
+
+def cmd_import(args: argparse.Namespace) -> int:
+    from percival import importer
+
+    app = load_app(args.app)
+    importer.configure_redbeat(app, args.redis_url, args.key_prefix)
+    planned = importer.plan(app, require_version=args.require_version, only=args.only or None)
+    print(importer.render(planned))
+
+    if not args.apply:
+        if any(p.action in (importer.CREATE, importer.UPDATE) for p in planned):
+            print("\ndry run, nothing written. run again with --apply to write these.")
+        return 0
+
+    written = importer.apply(app, planned)
+    print(f"\nwrote {written} entries to redbeat.")
+    if written:
+        print(
+            "note: beat re-syncs beat_schedule from config on every start, so remove the\n"
+            "imported entries from your config or they will overwrite changes made in the panel."
+        )
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog=NAME)
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    imp = sub.add_parser("import", help="copy beat_schedule from your app into redbeat")
+    imp.add_argument("--app", required=True, help="celery app, e.g. proj.celery:app")
+    imp.add_argument("--redis-url", help="redis used by redbeat (default: redbeat_redis_url)")
+    imp.add_argument("--key-prefix", help="redbeat_key_prefix if you changed it")
+    imp.add_argument(
+        "--require-version",
+        help="set the required code version on every imported entry",
+    )
+    imp.add_argument("--only", nargs="+", metavar="NAME", help="import just these entries")
+    imp.add_argument("--apply", action="store_true", help="actually write (default is dry run)")
+    imp.set_defaults(func=cmd_import)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    result: int = args.func(args)
+    return result
+
+
+if __name__ == "__main__":
+    sys.exit(main())

percival/config.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass, replace
+from typing import Any, Literal
+
+from percival.naming import ENV_PREFIX, NAME
+
+DEFAULT_REDACT = (
+    "password",
+    "passwd",
+    "secret",
+    "token",
+    "api_key",
+    "apikey",
+    "authorization",
+    "cookie",
+)
+
+UNKNOWN_VERSION = "unknown"
+
+StaleAction = Literal["requeue", "drop"]
+
+
+def _env(name: str) -> str | None:
+    value = os.environ.get(ENV_PREFIX + name)
+    return value if value else None
+
+
+def _env_bool(name: str) -> bool | None:
+    value = _env(name)
+    if value is None:
+        return None
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+@dataclass(frozen=True)
+class Settings:
+    redis_url: str = "redis://localhost:6379/0"
+    code_version: str = UNKNOWN_VERSION
+    key_prefix: str = NAME
+
+    # what goes into run events
+    capture_args: bool = True
+    redact_kwargs: tuple[str, ...] = DEFAULT_REDACT
+    max_value_chars: int = 2_000
+    max_traceback_chars: int = 20_000
+
+    # what the stale-worker gate does with a blocked message
+    stale_action: StaleAction = "requeue"
+    requeue_delay: float = 5.0
+    max_requeues: int = 50
+
+    heartbeat_interval: float = 30.0
+    heartbeat_ttl: int = 90
+
+    events_maxlen: int = 100_000
+    publish_source: bool = True
+    source_max_lines: int = 300
+
+    # once_per behaviour when redis is down; False means skip the run
+    dedupe_fail_open: bool = False
+
+    @classmethod
+    def build(cls, **overrides: Any) -> Settings:
+        """Explicit arguments win, then PERCIVAL_* env vars, then defaults."""
+        values: dict[str, Any] = {}
+
+        env_map: dict[str, Any] = {
+            "redis_url": _env("REDIS_URL"),
+            "code_version": _env("CODE_VERSION"),
+            "key_prefix": _env("KEY_PREFIX"),
+            "capture_args": _env_bool("CAPTURE_ARGS"),
+            "dedupe_fail_open": _env_bool("DEDUPE_FAIL_OPEN"),
+            "stale_action": _env("STALE_ACTION"),
+        }
+        values.update({k: v for k, v in env_map.items() if v is not None})
+        values.update({k: v for k, v in overrides.items() if v is not None})
+
+        redact = values.get("redact_kwargs")
+        if redact is not None and not isinstance(redact, tuple):
+            values["redact_kwargs"] = _as_tuple(redact)
+
+        settings = cls(**values)
+        settings.validate()
+        return settings
+
+    def validate(self) -> None:
+        if self.stale_action not in ("requeue", "drop"):
+            raise ValueError(f"stale_action must be 'requeue' or 'drop', got {self.stale_action!r}")
+        if self.heartbeat_ttl <= self.heartbeat_interval:
+            raise ValueError("heartbeat_ttl has to be larger than heartbeat_interval")
+
+    def with_(self, **changes: Any) -> Settings:
+        return replace(self, **changes)
+
+
+def _as_tuple(items: Iterable[str] | str) -> tuple[str, ...]:
+    if isinstance(items, str):
+        return tuple(s.strip() for s in items.split(",") if s.strip())
+    return tuple(items)

percival/dedupe.py

@@ -0,0 +1,183 @@
+"""``@once_per``: run a task at most once per key per time window.
+
+    @app.task
+    @once_per(key=lambda user_id, **kw: user_id, window="1d")
+    def send_digest(user_id): ...
+
+Backed by ``SET NX EX`` in redis. Fixed windows ("5m", "1h", "1d") are
+aligned to the unix epoch, so "1d" means one UTC calendar day. For anything
+else, like the recipient's local date, pass a callable as ``window`` and it
+becomes the bucket.
+
+If the task raises (including celery's Retry) the key is released again,
+otherwise a failed attempt would block its own retry for the whole window.
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+import re
+import time
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any, TypeVar
+
+from celery import Task, current_task
+from celery.exceptions import Ignore
+
+from percival import runtime
+from percival.events import Status, build_event, emit
+
+log = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}
+_WINDOW_RE = re.compile(r"^\s*(\d+)\s*([smhdw])\s*$")
+
+# ttl for callable windows when none is given, generous on purpose
+DEFAULT_CALLABLE_TTL = 2 * 86400
+
+
+def parse_window(window: str) -> int:
+    match = _WINDOW_RE.match(window)
+    if not match:
+        raise ValueError(f"can't parse window {window!r}, use something like '5m', '1h', '1d'")
+    seconds = int(match.group(1)) * _UNITS[match.group(2)]
+    if seconds <= 0:
+        raise ValueError("window has to be longer than zero")
+    return seconds
+
+
+def once_per(
+    key: Callable[..., Any] | None = None,
+    window: str | Callable[..., Any] = "1d",
+    *,
+    name: str | None = None,
+    ttl: int | None = None,
+    fail_open: bool | None = None,
+    release_on_error: bool = True,
+) -> Callable[[F], F]:
+    # validate up front so a typo fails at import, not at 3am
+    window_seconds = parse_window(window) if isinstance(window, str) else None
+    if window_seconds is None and not callable(window):
+        raise TypeError("window must be a string like '1d' or a callable")
+
+    def decorator(fn: F) -> F:
+        fallback_name = f"{fn.__module__}.{fn.__qualname__}"
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            rt = runtime.current()
+            task = _running_task()
+            task_name = name or (task.name if task is not None else fallback_name)
+
+            # bind=True tasks get self first, the key function shouldn't see it
+            user_args = args[1:] if args and isinstance(args[0], Task) else args
+            key_value = str(key(*user_args, **kwargs)) if key is not None else "_"
+
+            if window_seconds is not None:
+                bucket = str(int(time.time() // window_seconds))
+                expire = ttl or window_seconds
+            else:
+                bucket = str(window(*user_args, **kwargs))  # type: ignore[operator]
+                expire = ttl or DEFAULT_CALLABLE_TTL
+
+            redis_key = rt.keys.once(task_name, key_value, bucket)
+            holder = getattr(getattr(task, "request", None), "id", None) or "direct"
+            should_fail_open = rt.settings.dedupe_fail_open if fail_open is None else fail_open
+
+            try:
+                acquired = bool(rt.redis.get().set(redis_key, holder, nx=True, ex=expire))
+            except Exception as exc:
+                if should_fail_open:
+                    log.warning(
+                        "percival: once_per can't reach redis, running %s anyway",
+                        task_name,
+                        exc_info=True,
+                    )
+                    return fn(*args, **kwargs)
+                log.error(
+                    "percival: once_per can't reach redis, skipping %s", task_name, exc_info=True
+                )
+                _report(
+                    rt,
+                    Status.DEDUPE_UNAVAILABLE,
+                    task,
+                    task_name,
+                    fn,
+                    user_args,
+                    kwargs,
+                    redis_key,
+                    exc_type=type(exc).__name__,
+                    exc_message=str(exc),
+                )
+                return _skip(task)
+
+            if not acquired:
+                _report(
+                    rt, Status.SKIPPED_DUPLICATE, task, task_name, fn, user_args, kwargs, redis_key
+                )
+                return _skip(task)
+
+            try:
+                return fn(*args, **kwargs)
+            except BaseException:
+                if release_on_error:
+                    _release(rt, redis_key, holder)
+                raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def _running_task() -> Any:
+    task = current_task
+    if not task:
+        return None
+    request = getattr(task, "request", None)
+    if request is None or getattr(request, "called_directly", True):
+        return None
+    return task
+
+
+def _skip(task: Any) -> None:
+    # inside a worker, Ignore keeps celery from also recording a "success"
+    if task is not None:
+        raise Ignore()
+    return None
+
+
+def _release(rt: runtime.Runtime, redis_key: str, holder: str) -> None:
+    try:
+        client = rt.redis.get()
+        # only delete it if it's still ours
+        current = client.get(redis_key)
+        if isinstance(current, bytes):
+            current = current.decode()
+        if current == holder:
+            client.delete(redis_key)
+    except Exception:
+        log.warning("percival: could not release once_per key %s", redis_key, exc_info=True)
+
+
+def _report(
+    rt: runtime.Runtime,
+    status: str,
+    task: Any,
+    task_name: str,
+    fn: Callable[..., Any],
+    args: Any,
+    kwargs: Any,
+    redis_key: str,
+    **extra: Any,
+) -> None:
+    subject = task if task is not None else SimpleNamespace(name=task_name, request=None, run=fn)
+    emit(
+        rt,
+        build_event(
+            rt, status, task=subject, args=args, kwargs=kwargs, dedupe_key=redis_key, **extra
+        ),
+    )

percival/events.py

@@ -0,0 +1,117 @@
+"""Run events, pushed to a Redis Stream that the panel consumes."""
+
+from __future__ import annotations
+
+import json
+import logging
+import socket
+import time
+from typing import Any
+
+from percival.naming import Headers
+from percival.redaction import capture_args
+from percival.runtime import Runtime
+
+log = logging.getLogger(__name__)
+
+
+class Status:
+    STARTED = "started"
+    SUCCESS = "success"
+    FAILURE = "failure"
+    RETRY = "retry"
+    REVOKED = "revoked"
+    SKIPPED_DUPLICATE = "skipped_duplicate"
+    BLOCKED_STALE_WORKER = "blocked_stale_worker"
+    # once_per could not reach redis and fail_open is off, so the run was skipped
+    DEDUPE_UNAVAILABLE = "dedupe_unavailable"
+
+
+def emit(rt: Runtime, event: dict[str, Any]) -> None:
+    """Best effort. A broken redis must never take a task down with it."""
+    try:
+        payload = json.dumps(event, default=repr, separators=(",", ":"))
+        rt.redis.get().xadd(
+            rt.keys.events,
+            {"data": payload},
+            maxlen=rt.settings.events_maxlen,
+            approximate=True,
+        )
+    except Exception:
+        log.warning(
+            "percival: could not record %s event for %s",
+            event.get("status"),
+            event.get("task_id"),
+            exc_info=True,
+        )
+
+
+def get_header(request: Any, name: str) -> Any:
+    """Custom headers end up in different places depending on the message
+    protocol and celery version, so look in both."""
+    if request is None:
+        return None
+    headers = getattr(request, "headers", None)
+    if isinstance(headers, dict) and name in headers:
+        return headers[name]
+    return getattr(request, name, None)
+
+
+def build_event(
+    rt: Runtime,
+    status: str,
+    *,
+    task: Any,
+    request: Any = None,
+    args: Any = None,
+    kwargs: Any = None,
+    started_at: float | None = None,
+    **extra: Any,
+) -> dict[str, Any]:
+    settings = rt.settings
+    now = time.time()
+    if request is None:
+        request = getattr(task, "request", None)
+
+    event: dict[str, Any] = {
+        "status": status,
+        "task_id": getattr(request, "id", None),
+        "task": getattr(task, "name", None),
+        "schedule": get_header(request, Headers.SCHEDULE),
+        "required_version": get_header(request, Headers.REQUIRED_VERSION),
+        "worker": getattr(request, "hostname", None) or rt.hostname or socket.gethostname(),
+        "code_version": settings.code_version,
+        "retries": getattr(request, "retries", None),
+        "queue": _queue_of(request),
+        "ts": now,
+    }
+    if started_at is not None:
+        event["started_at"] = started_at
+        if status != Status.STARTED:
+            event["finished_at"] = now
+            event["duration"] = round(now - started_at, 6)
+
+    if settings.capture_args:
+        if args is None and kwargs is None and request is not None:
+            args = getattr(request, "args", None)
+            kwargs = getattr(request, "kwargs", None)
+        event.update(
+            capture_args(
+                getattr(task, "run", None),
+                args,
+                kwargs,
+                settings.redact_kwargs,
+                settings.max_value_chars,
+            )
+        )
+
+    event.update({k: v for k, v in extra.items() if v is not None})
+    return event
+
+
+def _queue_of(request: Any) -> str | None:
+    info = getattr(request, "delivery_info", None)
+    if isinstance(info, dict):
+        queue = info.get("routing_key") or info.get("queue")
+        return str(queue) if queue else None
+    return None