oban 0.5.0__py3-none-any.whl

oban/_notifier.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import gzip
+import inspect
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Callable, Protocol, runtime_checkable
+from uuid import uuid4
+
+import orjson
+from psycopg import AsyncConnection, OperationalError
+
+from ._backoff import jittery_exponential
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ._query import Query
+
+
+def encode_payload(payload: dict) -> str:
+    """Encode a dict payload to an efficient format for publishing.
+
+    Args:
+        payload: Dict to encode
+
+    Returns:
+        Base64-encoded gzipped JSON string
+    """
+    dumped = orjson.dumps(payload)
+    zipped = gzip.compress(dumped)
+
+    return base64.b64encode(zipped).decode("ascii")
+
+
+def decode_payload(payload: str) -> dict:
+    """Decode a payload string to a dict.
+
+    Handles both compressed (base64+gzip) and plain JSON payloads.
+
+    Args:
+        payload: Payload string to decode
+
+    Returns:
+        Decoded dict
+    """
+    # Payloads created by SQL queries won't be encoded or compressed.
+    if payload.startswith("{"):
+        return orjson.loads(payload)
+    else:
+        decoded = base64.b64decode(payload)
+        unzipped = gzip.decompress(decoded)
+
+        return orjson.loads(unzipped.decode("utf-8"))
+
+
+@runtime_checkable
+class Notifier(Protocol):
+    """Protocol for pub/sub notification systems.
+
+    Notifiers enable real-time communication between Oban components using
+    channels. The default implementation uses PostgreSQL LISTEN/NOTIFY.
+    """
+
+    async def start(self) -> None:
+        """Start the notifier and establish necessary connections."""
+        ...
+
+    async def stop(self) -> None:
+        """Stop the notifier and clean up resources."""
+        ...
+
+    async def listen(
+        self, channel: str, callback: Callable[[str, dict], Any], wait: bool = True
+    ) -> str:
+        """Register a callback for a channel.
+
+        Args:
+            channels: Channel name to listen on
+            callback: Sync or async function called when notification received.
+                      Receives (channel, payload) as arguments where payload is a dict.
+                      Async callbacks are executed in the background without blocking.
+            wait: If True, blocks until the subscription is fully established and ready to
+                  receive notifications. If False, returns immediately after registering.
+                  Defaults to True for test reliability.
+
+        Returns:
+            Token (UUID string) used to unregister this subscription
+
+        Example:
+            >>> async def handler(channel: str, payload: dict):
+            ...     print(f"Received on {channel}: {payload}")
+            ...
+            >>> token = await notifier.listen("insert", handler)
+        """
+        ...
+
+    async def unlisten(self, token: str) -> None:
+        """Unregister a subscription by token.
+
+        Args:
+            token: The token returned from listen()
+
+        Example:
+            >>> token = await notifier.listen("insert", handler)
+            >>> await notifier.unlisten(token)
+        """
+        ...
+
+    async def notify(self, channel: str, payloads: dict | list[dict]) -> None:
+        """Send one or more notifications to a channel.
+
+        Args:
+            channel: Channel name to send notification on
+            payloads: Payload dict or list of payload dicts to send
+
+        Example:
+            >>> await notifier.notify("insert", {"queue": "default"})
+            >>> await notifier.notify("insert", [{"queue": "default"}, {"queue": "mailers"}])
+        """
+        ...
+
+
+class PostgresNotifier:
+    """PostgreSQL-based notifier using LISTEN/NOTIFY.
+
+    Maintains a dedicated connection for receiving notifications and dispatches
+    them to registered callbacks. Automatically reconnects on connection loss.
+    """
+
+    def __init__(
+        self,
+        *,
+        query: Query,
+        prefix: str = "public",
+        beat_interval: float = 30.0,
+        connect_timeout: float = 5.0,
+        notify_timeout: float = 0.1,
+    ) -> None:
+        self._query = query
+        self._prefix = prefix
+        self._beat_interval = beat_interval
+        self._connect_timeout = connect_timeout
+        self._notify_timeout = notify_timeout
+
+        self._pending_listen = set()
+        self._pending_unlisten = set()
+        self._listen_events = {}
+        self._subscriptions = defaultdict(dict)
+        self._tokens = {}
+
+        self._conn = None
+        self._loop_task = None
+        self._beat_task = None
+        self._reconnect_attempts = 0
+
+    def _to_full_channel(self, channel: str) -> str:
+        return f"{self._prefix}.oban_{channel}"
+
+    def _from_full_channel(self, full_channel: str) -> str:
+        (_prefix, channel) = full_channel.split(".", 1)
+
+        return channel[5:]
+
+    async def start(self) -> None:
+        await self._connect()
+
+    async def stop(self) -> None:
+        if not self._loop_task or not self._beat_task or not self._conn:
+            return
+
+        self._loop_task.cancel()
+        self._beat_task.cancel()
+
+        await asyncio.gather(self._loop_task, self._beat_task, return_exceptions=True)
+
+        try:
+            await self._conn.close()
+        except Exception:
+            pass
+
+    async def listen(
+        self, channel: str, callback: Callable[[str, dict], Any], wait: bool = True
+    ) -> str:
+        token = str(uuid4())
+
+        if channel not in self._subscriptions:
+            self._pending_listen.add(channel)
+            self._listen_events[channel] = asyncio.Event()
+
+        self._tokens[token] = channel
+        self._subscriptions[channel][token] = callback
+
+        if wait and channel in self._listen_events:
+            await self._listen_events[channel].wait()
+
+        return token
+
+    async def unlisten(self, token: str) -> None:
+        if token not in self._tokens:
+            return
+
+        channel = self._tokens.pop(token)
+
+        if channel in self._subscriptions:
+            self._subscriptions[channel].pop(token, None)
+
+            if len(self._subscriptions[channel]) == 0:
+                del self._subscriptions[channel]
+                self._pending_unlisten.add(channel)
+
+    async def notify(self, channel: str, payloads: dict | list[dict]) -> None:
+        if isinstance(payloads, dict):
+            payloads = [payloads]
+
+        channel = self._to_full_channel(channel)
+        encoded = [encode_payload(payload) for payload in payloads]
+
+        await self._query.notify(channel, encoded)
+
+    async def _connect(self) -> None:
+        dsn = self._query.dsn
+
+        self._conn = await asyncio.wait_for(
+            AsyncConnection.connect(dsn, autocommit=True),
+            timeout=self._connect_timeout,
+        )
+
+        for channel in list(self._subscriptions.keys()):
+            self._pending_listen.add(channel)
+
+        self._loop_task = asyncio.create_task(self._loop(), name="oban-notifier-loop")
+        self._beat_task = asyncio.create_task(self._beat(), name="oban-notifier-beat")
+
+        self._reconnect_attempts = 0
+
+    async def _reconnect(self) -> None:
+        self._reconnect_attempts += 1
+
+        delay = jittery_exponential(self._reconnect_attempts, max_pow=4)
+
+        await asyncio.sleep(delay)
+
+        try:
+            await self._connect()
+        except (OSError, OperationalError, asyncio.TimeoutError):
+            asyncio.create_task(self._reconnect())
+
+    async def _loop(self) -> None:
+        if not self._conn:
+            return
+
+        while True:
+            try:
+                await self._process_pending()
+                await self._process_notifications()
+            except asyncio.CancelledError:
+                raise
+            except (OSError, OperationalError):
+                self._conn = None
+                asyncio.create_task(self._reconnect())
+                break
+
+    async def _process_pending(self) -> None:
+        if not self._conn:
+            return
+
+        for channel in list(self._pending_listen):
+            full_channel = self._to_full_channel(channel)
+            await self._conn.execute(f'LISTEN "{full_channel}"')
+
+            self._pending_listen.discard(channel)
+
+            if channel in self._listen_events:
+                self._listen_events[channel].set()
+                del self._listen_events[channel]
+
+        for channel in list(self._pending_unlisten):
+            full_channel = self._to_full_channel(channel)
+            await self._conn.execute(f'UNLISTEN "{full_channel}"')
+
+            self._pending_unlisten.discard(channel)
+
+    async def _process_notifications(self) -> None:
+        if not self._conn:
+            return
+
+        gen = self._conn.notifies(timeout=self._notify_timeout)
+
+        async for notify in gen:
+            await self._dispatch(notify)
+
+    async def _dispatch(self, notify) -> None:
+        channel = self._from_full_channel(notify.channel)
+        payload = decode_payload(notify.payload)
+
+        if channel in self._subscriptions:
+            for callback in self._subscriptions[channel].values():
+                try:
+                    if inspect.iscoroutinefunction(callback):
+                        asyncio.create_task(callback(channel, payload))
+                    else:
+                        callback(channel, payload)
+                except Exception:
+                    logger.exception(
+                        "Error in notifier callback for channel %s with payload %s",
+                        channel,
+                        payload,
+                    )
+
+    async def _beat(self) -> None:
+        while True:
+            try:
+                await asyncio.sleep(self._beat_interval)
+
+                if self._conn and not self._conn.closed:
+                    await self._conn.execute("SELECT 1")
+
+            except asyncio.CancelledError:
+                raise
+            except (OSError, OperationalError):
+                pass

oban/_producer.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from . import telemetry
+from ._executor import Executor
+from ._extensions import use_ext
+from .job import Job
+
+if TYPE_CHECKING:
+    from ._notifier import Notifier
+    from ._query import Query
+
+logger = logging.getLogger(__name__)
+
+
+def _init(producer: Producer) -> dict:
+    return {"local_limit": producer._limit, "paused": producer._paused}
+
+
+def _pause(producer: Producer) -> dict:
+    producer._paused = True
+
+    return {"paused": True}
+
+
+def _resume(producer: Producer) -> dict:
+    producer._paused = False
+
+    return {"paused": False}
+
+
+def _scale(producer: Producer, **kwargs: Any) -> dict:
+    if "limit" in kwargs:
+        producer._limit = kwargs["limit"]
+
+    return {"local_limit": producer._limit}
+
+
+def _validate(*, queue: str, limit: int, **extra) -> None:
+    if not isinstance(queue, str):
+        raise TypeError(f"queue must be a string, got {queue}")
+    if not queue.strip():
+        raise ValueError("queue must not be blank")
+
+    if limit < 1:
+        raise ValueError(f"Queue '{queue}' limit must be positive")
+
+
+async def _get_jobs(producer: Producer) -> list[Job]:
+    demand = producer._limit - len(producer._running_jobs)
+
+    if demand > 0:
+        return await producer._query.fetch_jobs(
+            demand=demand,
+            queue=producer._queue,
+            node=producer._node,
+            uuid=producer._uuid,
+        )
+    else:
+        return []
+
+
+class LocalDispatcher:
+    def dispatch(self, producer: Producer, job: Job) -> asyncio.Task:
+        return asyncio.create_task(producer._execute(job))
+
+
+@dataclass(frozen=True, slots=True)
+class QueueInfo:
+    """Information about a queue's runtime state."""
+
+    limit: int
+    """The concurrency limit for this queue."""
+
+    meta: dict[str, Any]
+    """Queue metadata including extension options (global_limit, partition, etc.)."""
+
+    node: str
+    """The node name where this queue is running."""
+
+    paused: bool
+    """Whether the queue is currently paused."""
+
+    queue: str
+    """The queue name."""
+
+    running: list[int]
+    """List of currently executing job IDs."""
+
+    started_at: datetime | None
+    """When the queue was started."""
+
+
+class Producer:
+    def __init__(
+        self,
+        *,
+        debounce_interval: float = 0.005,
+        dispatcher: Any = None,
+        limit: int = 10,
+        paused: bool = False,
+        queue: str = "default",
+        name: str,
+        node: str,
+        notifier: Notifier,
+        query: Query,
+        **extra,
+    ) -> None:
+        self._debounce_interval = debounce_interval
+        self._dispatcher = dispatcher or LocalDispatcher()
+        self._extra = extra
+        self._limit = limit
+        self._name = name
+        self._node = node
+        self._notifier = notifier
+        self._paused = paused
+        self._query = query
+        self._queue = queue
+
+        self._validate()
+
+        self._init_lock = asyncio.Lock()
+        self._last_fetch_time = 0.0
+        self._listen_token = None
+        self._loop_task = None
+        self._notified = asyncio.Event()
+        self._pending_acks = []
+        self._running_jobs = {}
+        self._started_at = None
+        self._uuid = str(uuid4())
+
+    def _validate(self, **opts) -> None:
+        params = {"queue": self._queue, "limit": self._limit, **self._extra}
+        merged = {**params, **opts}
+
+        use_ext("producer.validate", _validate, **merged)
+
+    async def start(self) -> None:
+        async with self._init_lock:
+            self._started_at = datetime.now(timezone.utc)
+
+            await self._query.insert_producer(
+                uuid=self._uuid,
+                name=self._name,
+                node=self._node,
+                queue=self._queue,
+                meta=use_ext("producer.init", _init, self),
+            )
+
+            self._listen_token = await self._notifier.listen(
+                "signal", self._on_signal, wait=False
+            )
+
+            self._loop_task = asyncio.create_task(
+                self._loop(), name=f"oban-producer-{self._queue}"
+            )
+
+    async def stop(self) -> None:
+        async with self._init_lock:
+            if not self._listen_token or not self._loop_task:
+                return
+
+            self._loop_task.cancel()
+
+            await self._notifier.unlisten(self._listen_token)
+
+            running_tasks = [task for (_job, task) in self._running_jobs.values()]
+
+            if running_tasks:
+                try:
+                    now = datetime.now(timezone.utc).isoformat()
+                    await self._query.update_producer(
+                        uuid=self._uuid,
+                        meta={"paused": True, "shutdown_started_at": now},
+                    )
+                except Exception:
+                    logger.debug(
+                        "Failed to update producer %s during shutdown", self._uuid
+                    )
+
+            await asyncio.gather(
+                self._loop_task,
+                *running_tasks,
+                return_exceptions=True,
+            )
+
+            try:
+                await self._query.delete_producer(self._uuid)
+            except Exception:
+                logger.debug(
+                    "Failed to delete producer %s during shutdown, will be cleaned up asynchronously",
+                    self._uuid,
+                )
+
+    def notify(self) -> None:
+        self._notified.set()
+
+    async def pause(self) -> None:
+        meta = use_ext("producer.pause", _pause, self)
+
+        await self._query.update_producer(uuid=self._uuid, meta=meta)
+
+    async def resume(self) -> None:
+        meta = use_ext("producer.resume", _resume, self)
+
+        await self._query.update_producer(uuid=self._uuid, meta=meta)
+
+        self.notify()
+
+    async def scale(self, **kwargs: Any) -> None:
+        self._validate(**kwargs)
+
+        meta = use_ext("producer.scale", _scale, self, **kwargs)
+
+        await self._query.update_producer(uuid=self._uuid, meta=meta)
+
+        self.notify()
+
+    def check(self) -> QueueInfo:
+        return QueueInfo(
+            limit=self._limit,
+            meta=self._extra,
+            node=self._node,
+            paused=self._paused,
+            queue=self._queue,
+            running=list(self._running_jobs.keys()),
+            started_at=self._started_at,
+        )
+
+    async def _loop(self) -> None:
+        while True:
+            try:
+                await asyncio.wait_for(self._notified.wait(), timeout=1.0)
+            except asyncio.TimeoutError:
+                continue
+            except asyncio.CancelledError:
+                break
+
+            self._notified.clear()
+
+            try:
+                await self._debounce()
+                await self._produce()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                logger.exception("Error in producer for queue %s", self._queue)
+
+    async def _debounce(self) -> None:
+        now = asyncio.get_event_loop().time()
+        elapsed = now - self._last_fetch_time
+
+        if elapsed < self._debounce_interval:
+            await asyncio.sleep(self._debounce_interval - elapsed)
+
+        self._last_fetch_time = asyncio.get_event_loop().time()
+
+    async def _produce(self) -> None:
+        if self._paused or (self._limit - len(self._running_jobs)) <= 0:
+            return
+
+        _ack = await self._ack_jobs()
+        jobs = await self._get_jobs()
+
+        for job in jobs:
+            task = self._dispatcher.dispatch(self, job)
+            task.add_done_callback(
+                lambda _, job_id=job.id: self._on_job_complete(job_id)
+            )
+
+            self._running_jobs[job.id] = (job, task)
+
+    async def _ack_jobs(self):
+        with telemetry.span("oban.producer.ack", {"queue": self._queue}) as context:
+            if self._pending_acks:
+                acked_ids = await self._query.ack_jobs(self._pending_acks)
+                acked_set = set(acked_ids)
+
+                self._pending_acks = [
+                    ack for ack in self._pending_acks if ack.id not in acked_set
+                ]
+            else:
+                acked_ids = []
+
+            context.add({"count": len(acked_ids)})
+
+    async def _get_jobs(self):
+        with telemetry.span("oban.producer.get", {"queue": self._queue}) as context:
+            jobs = await use_ext("producer.get_jobs", _get_jobs, self)
+
+            context.add({"count": len(jobs)})
+
+            return jobs
+
+    async def _execute(self, job: Job) -> None:
+        job._cancellation = asyncio.Event()
+
+        executor = await Executor(job=job, safe=True).execute()
+
+        self._pending_acks.append(executor.action)
+
+    def _on_job_complete(self, job_id: int) -> None:
+        self._running_jobs.pop(job_id, None)
+
+        self.notify()
+
+    async def _on_signal(self, _channel: str, payload: dict) -> None:
+        ident = payload.get("ident", "any")
+        queue = payload.get("queue", "*")
+
+        if queue != "*" and queue != self._queue:
+            return
+
+        if ident != "any" and ident != f"{self._name}.{self._node}":
+            return
+
+        match payload.get("action"):
+            case "pause":
+                await self.pause()
+            case "resume":
+                await self.resume()
+            case "pkill":
+                job_id = payload["job_id"]
+
+                if job_id in self._running_jobs:
+                    (job, _task) = self._running_jobs[job_id]
+
+                    job._cancellation.set()