oban 0.5.0__py3-none-any.whl

oban/_pruner.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from . import telemetry
+
+if TYPE_CHECKING:
+    from ._leader import Leader
+    from ._query import Query
+
+
+class Pruner:
+    """Manages periodic deletion of completed, cancelled, and discarded jobs.
+
+    This class is managed internally by Oban and shouldn't be constructed directly.
+    Instead, configure pruning via the Oban constructor:
+
+        >>> async with Oban(
+        ...     conn=conn,
+        ...     queues={"default": 10},
+        ...     pruner={"max_age": 86_400, "interval": 60.0, "limit": 20_000}
+        ... ) as oban:
+        ...     # Pruner runs automatically in the background
+    """
+
+    def __init__(
+        self,
+        *,
+        query: Query,
+        leader: Leader,
+        max_age: int = 86_400,
+        interval: float = 60.0,
+        limit: int = 20_000,
+    ) -> None:
+        self._leader = leader
+        self._max_age = max_age
+        self._interval = interval
+        self._limit = limit
+        self._query = query
+
+        self._loop_task = None
+
+        self._validate(max_age=max_age, interval=interval, limit=limit)
+
+    @staticmethod
+    def _validate(*, max_age: int, interval: float, limit: int) -> None:
+        if not isinstance(max_age, int):
+            raise TypeError(f"max_age must be an integer, got {max_age}")
+        if max_age <= 0:
+            raise ValueError(f"max_age must be positive, got {max_age}")
+
+        if not isinstance(interval, (int, float)):
+            raise TypeError(f"interval must be a number, got {interval}")
+        if interval <= 0:
+            raise ValueError(f"interval must be positive, got {interval}")
+
+        if not isinstance(limit, int):
+            raise TypeError(f"limit must be an integer, got {limit}")
+        if limit <= 0:
+            raise ValueError(f"limit must be positive, got {limit}")
+
+    async def start(self) -> None:
+        self._loop_task = asyncio.create_task(self._loop(), name="oban-pruner")
+
+    async def stop(self) -> None:
+        if not self._loop_task:
+            return
+
+        self._loop_task.cancel()
+
+        try:
+            await self._loop_task
+        except asyncio.CancelledError:
+            pass
+
+    async def _loop(self) -> None:
+        while True:
+            try:
+                await asyncio.sleep(self._interval)
+
+                if self._leader.is_leader:
+                    await self._prune()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                pass
+
+    async def _prune(self) -> None:
+        with telemetry.span("oban.pruner.prune", {}) as context:
+            pruned = await self._query.prune_jobs(self._max_age, self._limit)
+
+            context.add({"pruned_count": pruned})

oban/_query.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+import re
+from collections import defaultdict
+from contextlib import asynccontextmanager
+from dataclasses import replace
+from datetime import datetime, timezone
+from functools import cache
+from importlib.resources import files
+from typing import Any
+
+from psycopg.rows import class_row
+from psycopg.types.json import Jsonb
+from psycopg_pool import AsyncConnectionPool
+
+from ._executor import AckAction
+from ._extensions import use_ext
+from .job import Job, TIMESTAMP_FIELDS
+
+ACKABLE_FIELDS = [
+    "id",
+    "state",
+    "attempt_change",
+    "schedule_in",
+    "error",
+    "meta",
+]
+
+
+INSERTABLE_FIELDS = [
+    "args",
+    "inserted_at",
+    "max_attempts",
+    "meta",
+    "priority",
+    "queue",
+    "scheduled_at",
+    "state",
+    "tags",
+    "worker",
+]
+
+# The `Job` class has errors, but we only insert a single `error` at one time.
+JSON_FIELDS = ["args", "error", "errors", "meta", "tags"]
+
+
+UPDATABLE_FIELDS = [
+    "args",
+    "max_attempts",
+    "meta",
+    "priority",
+    "queue",
+    "scheduled_at",
+    "tags",
+    "worker",
+]
+
+
+async def _ack_jobs(query: Query, acks: list[AckAction]) -> list[int]:
+    async with query._pool.connection() as conn:
+        async with conn.transaction():
+            stmt = Query._load_file("ack_job.sql", query._prefix)
+            acked_ids = []
+
+            for ack in acks:
+                args = {
+                    field: Query._cast_type(field, getattr(ack, field))
+                    for field in ACKABLE_FIELDS
+                }
+
+                result = await conn.execute(stmt, args)
+                row = await result.fetchone()
+
+                if row:
+                    acked_ids.append(row[0])
+
+            return acked_ids
+
+
+async def _insert_jobs(query: Query, jobs: list[Job]) -> list[Job]:
+    async with query._pool.connection() as conn:
+        stmt = Query._load_file("insert_job.sql", query._prefix)
+        inserted = []
+
+        for job in jobs:
+            args = {
+                key: Query._cast_type(key, getattr(job, key))
+                for key in INSERTABLE_FIELDS
+            }
+
+            result = await conn.execute(stmt, args)
+            row = await result.fetchone()
+
+            inserted.append(
+                replace(
+                    job,
+                    id=row[0],
+                    inserted_at=row[1],
+                    queue=row[2],
+                    scheduled_at=row[3],
+                    state=row[4],
+                )
+            )
+
+        return inserted
+
+
+class Query:
+    @staticmethod
+    @cache
+    def _load_file(
+        path: str,
+        prefix: str = "public",
+        package: str = "oban.queries",
+        apply_prefix: bool = True,
+    ) -> str:
+        sql = files(package).joinpath(path).read_text(encoding="utf-8")
+
+        if apply_prefix:
+            return re.sub(
+                r"\b(oban_jobs|oban_leaders|oban_producers|oban_job_state|oban_state_to_bit)\b",
+                rf"{prefix}.\1",
+                sql,
+            )
+        else:
+            return sql
+
+    @staticmethod
+    def _cast_type(field: str, value: Any) -> Any:
+        if field in JSON_FIELDS and value is not None:
+            return Jsonb(value)
+
+        # Ensure timestamps are written as UTC rather than being implicitly cast to the current
+        # timezone. The database uses `TIMESTAMP WITHOUT TIME ZONE` and the value is automatically
+        # shifted when the zone is present.
+        if field in TIMESTAMP_FIELDS and value is not None:
+            return value.astimezone(timezone.utc).replace(tzinfo=None)
+
+        return value
+
+    def __init__(self, pool: AsyncConnectionPool, prefix: str = "public") -> None:
+        if not isinstance(pool, AsyncConnectionPool):
+            raise TypeError(f"Expected AsyncConnectionPool, got {type(pool).__name__}")
+
+        self._pool = pool
+        self._prefix = prefix
+
+    @property
+    def dsn(self) -> str:
+        return self._pool.conninfo
+
+    @asynccontextmanager
+    async def connection(self):
+        async with self._pool.connection() as conn:
+            yield conn
+
+    # Jobs
+
+    async def ack_jobs(self, acks: list[AckAction]) -> list[int]:
+        return await use_ext("query.ack_jobs", _ack_jobs, self, acks)
+
+    async def all_jobs(self, states: list[str]) -> list[Job]:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("all_jobs.sql", self._prefix)
+
+            async with conn.cursor(row_factory=class_row(Job)) as cur:
+                await cur.execute(stmt, {"states": states})
+                return await cur.fetchall()
+
+    async def cancel_many_jobs(self, ids: list[int]) -> tuple[int, list[int]]:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("cancel_many_jobs.sql", self._prefix)
+                args = {"ids": ids}
+
+                result = await conn.execute(stmt, args)
+                rows = await result.fetchall()
+
+                executing_ids = [row[0] for row in rows if row[1] == "executing"]
+
+                return len(rows), executing_ids
+
+    async def delete_many_jobs(self, ids: list[int]) -> int:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("delete_many_jobs.sql", self._prefix)
+                args = {"ids": ids}
+
+                result = await conn.execute(stmt, args)
+
+                return result.rowcount
+
+    async def get_job(self, job_id: int) -> Job:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("get_job.sql", self._prefix)
+
+            async with conn.cursor(row_factory=class_row(Job)) as cur:
+                await cur.execute(stmt, (job_id,))
+
+                return await cur.fetchone()
+
+    async def fetch_jobs(
+        self, demand: int, queue: str, node: str, uuid: str
+    ) -> list[Job]:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("fetch_jobs.sql", self._prefix)
+                args = {"queue": queue, "demand": demand, "attempted_by": [node, uuid]}
+
+                async with conn.cursor(row_factory=class_row(Job)) as cur:
+                    await cur.execute(stmt, args)
+
+                    return await cur.fetchall()
+
+    async def insert_jobs(self, jobs: list[Job]) -> list[Job]:
+        return await use_ext("query.insert_jobs", _insert_jobs, self, jobs)
+
+    async def prune_jobs(self, max_age: int, limit: int) -> int:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("prune_jobs.sql", self._prefix)
+                args = {"max_age": max_age, "limit": limit}
+
+                result = await conn.execute(stmt, args)
+
+                return result.rowcount
+
+    async def rescue_jobs(self, rescue_after: float) -> int:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("rescue_jobs.sql", self._prefix)
+                args = {"rescue_after": rescue_after}
+
+                result = await conn.execute(stmt, args)
+
+                return result.rowcount
+
+    async def retry_many_jobs(self, ids: list[int]) -> int:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("retry_many_jobs.sql", self._prefix)
+                args = {"ids": ids}
+
+                result = await conn.execute(stmt, args)
+
+                return result.rowcount
+
+    async def stage_jobs(
+        self, limit: int, queues: list[str], before: datetime | None = None
+    ) -> tuple[int, list[str]]:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("stage_jobs.sql", self._prefix)
+                args = {"limit": limit, "queues": queues, "before": before}
+
+                result = await conn.execute(stmt, args)
+                rows = await result.fetchall()
+                queues = [queue for (queue,) in rows]
+
+                return (len(rows), queues)
+
+    async def update_many_jobs(self, jobs: list[Job]) -> list[Job]:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                stmt = self._load_file("update_job.sql", self._prefix)
+                args = defaultdict(list)
+
+                for job in jobs:
+                    args["ids"].append(job.id)
+
+                    for key in UPDATABLE_FIELDS:
+                        args[key].append(self._cast_type(key, getattr(job, key)))
+
+                result = await conn.execute(stmt, dict(args))
+                rows = await result.fetchall()
+
+                return [
+                    replace(
+                        job,
+                        args=row[0],
+                        max_attempts=row[1],
+                        meta=row[2],
+                        priority=row[3],
+                        queue=row[4],
+                        scheduled_at=row[5],
+                        state=row[6],
+                        tags=row[7],
+                        worker=row[8],
+                    )
+                    for job, row in zip(jobs, rows)
+                ]
+
+    # Leadership
+
+    async def attempt_leadership(
+        self, name: str, node: str, ttl: int, is_leader: bool
+    ) -> bool:
+        async with self._pool.connection() as conn:
+            async with conn.transaction():
+                cleanup_stmt = self._load_file(
+                    "cleanup_expired_leaders.sql", self._prefix
+                )
+
+                await conn.execute(cleanup_stmt)
+
+                if is_leader:
+                    elect_stmt = self._load_file("reelect_leader.sql", self._prefix)
+                else:
+                    elect_stmt = self._load_file("elect_leader.sql", self._prefix)
+
+                args = {"name": name, "node": node, "ttl": ttl}
+                result = await conn.execute(elect_stmt, args)
+                rows = await result.fetchone()
+
+                return rows is not None and rows[0] == node
+
+    async def resign_leader(self, name: str, node: str) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("resign_leader.sql", self._prefix)
+            args = {"name": name, "node": node}
+
+            await conn.execute(stmt, args)
+
+    # Schema
+
+    async def install(self) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("install.sql", self._prefix)
+
+            await conn.execute(stmt)
+
+    async def reset(self) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("reset.sql", self._prefix)
+
+            await conn.execute(stmt)
+
+    async def uninstall(self) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("uninstall.sql", self._prefix)
+
+            await conn.execute(stmt)
+
+    async def verify_structure(self) -> list[str]:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("verify_structure.sql", apply_prefix=False)
+            args = {"prefix": self._prefix}
+            rows = await conn.execute(stmt, args)
+            results = await rows.fetchall()
+
+            return [table for (table,) in results]
+
+    # Producer
+
+    async def cleanup_expired_producers(self, max_age: float) -> int:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("cleanup_expired_producers.sql", self._prefix)
+            args = {"max_age": max_age}
+
+            result = await conn.execute(stmt, args)
+
+            return result.rowcount
+
+    async def delete_producer(self, uuid: str) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("delete_producer.sql", self._prefix)
+            args = {"uuid": uuid}
+
+            await conn.execute(stmt, args)
+
+    async def insert_producer(
+        self, uuid: str, name: str, node: str, queue: str, meta: dict[str, Any]
+    ) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("insert_producer.sql", self._prefix)
+            args = {
+                "uuid": uuid,
+                "name": name,
+                "node": node,
+                "queue": queue,
+                "meta": Jsonb(meta),
+            }
+
+            await conn.execute(stmt, args)
+
+    async def refresh_producers(self, uuids: list[str]) -> int:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("refresh_producers.sql", self._prefix)
+            args = {"uuids": uuids}
+
+            result = await conn.execute(stmt, args)
+
+            return result.rowcount
+
+    async def update_producer(self, uuid: str, meta: dict[str, Any]) -> None:
+        async with self._pool.connection() as conn:
+            stmt = self._load_file("update_producer.sql", self._prefix)
+            args = {"uuid": uuid, "meta": Jsonb(meta)}
+
+            await conn.execute(stmt, args)
+
+    # Notifier
+
+    async def notify(self, channel: str, payloads: list[str]) -> None:
+        async with self._pool.connection() as conn:
+            await conn.execute(
+                "SELECT pg_notify(%s, payload) FROM json_array_elements_text(%s::json) AS payload",
+                (channel, Jsonb(payloads)),
+            )

oban/_recorded.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import erlpack
+
+
+def _convert_bytes(value: Any) -> Any:
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    elif isinstance(value, dict):
+        return {_convert_bytes(key): _convert_bytes(val) for key, val in value.items()}
+    elif isinstance(value, list):
+        return [_convert_bytes(item) for item in value]
+    return value
+
+
+def encode_recorded(value: Any) -> str:
+    binary = erlpack.pack(value)
+    encoded = base64.b64encode(binary).decode("ascii")
+
+    return encoded.rstrip("=")
+
+
+def decode_recorded(encoded: str) -> Any:
+    padding = 4 - (len(encoded) % 4)
+
+    if padding != 4:
+        encoded = encoded + ("=" * padding)
+
+    binary = base64.b64decode(encoded)
+
+    return _convert_bytes(erlpack.unpack(binary))

oban/_refresher.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from . import telemetry
+
+if TYPE_CHECKING:
+    from ._producer import Producer
+    from ._leader import Leader
+    from ._query import Query
+
+
+class Refresher:
+    def __init__(
+        self,
+        *,
+        query: Query,
+        leader: Leader,
+        producers: dict[str, Producer],
+        interval: float = 15.0,
+        max_age: float = 60.0,
+    ) -> None:
+        self._leader = leader
+        self._interval = interval
+        self._max_age = max_age
+        self._producers = producers
+        self._query = query
+
+        self._loop_task = None
+
+        self._validate(interval=interval, max_age=max_age)
+
+    @staticmethod
+    def _validate(*, interval: float, max_age: float) -> None:
+        if not isinstance(interval, (int, float)):
+            raise TypeError(f"interval must be a number, got {interval}")
+        if interval <= 0:
+            raise ValueError(f"interval must be positive, got {interval}")
+
+        if not isinstance(max_age, (int, float)):
+            raise TypeError(f"max_age must be a number, got {max_age}")
+        if max_age <= 0:
+            raise ValueError(f"max_age must be positive, got {max_age}")
+
+    async def start(self) -> None:
+        self._loop_task = asyncio.create_task(self._loop(), name="oban-refresher")
+
+    async def stop(self) -> None:
+        if self._loop_task:
+            self._loop_task.cancel()
+
+            try:
+                await self._loop_task
+            except asyncio.CancelledError:
+                pass
+
+    async def _loop(self) -> None:
+        while True:
+            try:
+                await asyncio.sleep(self._interval)
+
+                await self._refresh()
+                await self._cleanup()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                pass
+
+    async def _refresh(self) -> None:
+        with telemetry.span("oban.refresher.refresh", {}) as context:
+            uuids = [producer._uuid for producer in self._producers.values()]
+
+            if uuids:
+                refreshed = await self._query.refresh_producers(uuids)
+            else:
+                refreshed = 0
+
+            context.add({"refreshed_count": refreshed})
+
+    async def _cleanup(self) -> None:
+        if not self._leader.is_leader:
+            return
+
+        with telemetry.span("oban.refresher.cleanup", {}) as context:
+            cleaned_up = await self._query.cleanup_expired_producers(self._max_age)
+
+            context.add({"cleanup_count": cleaned_up})