diff --git a/src/pipecat/bus/network/pgmq.py b/src/pipecat/bus/network/pgmq.py index 6a7dba637..8a4c76f3d 100644 --- a/src/pipecat/bus/network/pgmq.py +++ b/src/pipecat/bus/network/pgmq.py @@ -8,14 +8,15 @@ import asyncio import json -import re -import time -import uuid from loguru import logger from pipecat.bus.bus import TaskBus from pipecat.bus.messages import BusMessage +from pipecat.bus.network.pgmq_backends import ( + DirectPgmqBackend, + PgmqBackend, +) from pipecat.bus.serializers import JSONMessageSerializer from pipecat.bus.serializers.base import MessageSerializer @@ -27,81 +28,60 @@ except ModuleNotFoundError as e: # pragma: no cover - exercised only when extra raise Exception(f"Missing module: {e}") -_INVALID_CHANNEL_CHARS = re.compile(r"[^A-Za-z0-9_]") -_MAX_CHANNEL_LEN = 30 -_PEER_LIST_TTL_S = 1.0 - - -def _sanitize_channel(channel: str) -> str: - """Coerce an arbitrary channel string into a Postgres-safe identifier prefix. - - PGMQ queue names must match ``^[a-zA-Z_][a-zA-Z0-9_]*$`` and are bounded by - Postgres identifier limits (effective ~48 chars after PGMQ's internal - ``q_``/``a_`` table prefixes). This helper replaces invalid characters with - underscores, ensures the prefix doesn't start with a digit, and truncates - to leave headroom for an instance suffix. - - Args: - channel: User-supplied channel name (may contain colons, slashes, etc). - - Returns: - A sanitized prefix safe for use in PGMQ queue names. - """ - safe = _INVALID_CHANNEL_CHARS.sub("_", channel) - if safe and safe[0].isdigit(): - safe = f"q_{safe}" - return safe[:_MAX_CHANNEL_LEN] or "pipecat_bus" - - class PgmqBus(TaskBus): """Distributed task bus backed by PGMQ (PostgreSQL Message Queue). - Implements pub/sub fan-out on top of PGMQ's point-to-point queue semantics - by giving each ``PgmqBus`` instance its own queue and broadcasting on - publish. The reader long-polls its own queue and dispatches received - messages to local subscribers. + Pub/sub fan-out is implemented on top of PGMQ's point-to-point queue + semantics by giving each :class:`PgmqBus` instance its own queue and + broadcasting on publish. A reader long-polls the instance's queue and + dispatches received messages to local subscribers. - ``BusLocalMessage`` messages bypass PGMQ and are delivered directly to - local subscribers. + ``BusLocalMessage`` messages bypass the network entirely and are + delivered directly to local subscribers. - Requires the ``pgmq`` and ``asyncpg`` packages. Install with + Two backends are supported (see :mod:`pipecat.bus.network.pgmq_backends`): + + - :class:`DirectPgmqBackend` — calls ``pgmq.*`` directly and discovers + peers by queue-name prefix. Suitable when bus peers trust each other. + - :class:`IsolatedPgmqBackend` — calls SECURITY DEFINER Postgres + wrappers over an asyncpg pool. Suitable when peers should be isolated + and the channel name is the bus capability. + + Construct with either ``pgmq=PGMQueue`` (uses :class:`DirectPgmqBackend`) + or ``backend=PgmqBackend`` (any backend). The two are mutually exclusive. + + Requires the ``pgmq`` extra. Install with ``pip install pipecat-ai[pgmq]``. - The provided ``PGMQueue`` must already have its connection pool - initialized via ``await pgmq.init()`` before being passed. The adapter - does not own the client's lifetime and will not close it on stop. - - Notes: - Prefer the session-mode pooler (e.g. port 5432 in Supabase) when - one is available. Transaction-mode pooling (e.g. port 6543 in - Supabase) works in practice with this adapter because each PGMQ - call is a single SQL statement, but it logs benign "resetting - connection with an active transaction" warnings and is more - fragile around prepared statements. - The underlying pool must allow at least 2 concurrent connections - (one for the reader's long-poll, one for publishes); 4+ recommended - under load. - Example:: from pgmq.async_queue import PGMQueue pgmq = PGMQueue( - host="aws-0-us-east-1.pooler.supabase.com", + host="...", port="5432", database="postgres", - username="postgres.", + username="postgres", password="...", pool_size=4, ) await pgmq.init() bus = PgmqBus(pgmq=pgmq, channel="pipecat_acme") + + Notes: + Prefer a session-mode pooler when available. Transaction-mode + pooling works for direct ``pgmq.*`` calls but is fragile around the + long-poll inside the SECURITY DEFINER ``bus_subscribe`` wrapper used + by :class:`IsolatedPgmqBackend`. + The underlying connection pool must allow at least two concurrent + connections (one for the reader's long-poll, one for publishes). """ def __init__( self, *, - pgmq: PGMQueue, + pgmq: PGMQueue | None = None, + backend: PgmqBackend | None = None, serializer: MessageSerializer | None = None, channel: str = "pipecat_bus", visibility_timeout: int = 30, @@ -113,61 +93,74 @@ class PgmqBus(TaskBus): """Initialize the PgmqBus. Args: - pgmq: An initialized ``PGMQueue`` client. Call ``await pgmq.init()`` - before passing. - serializer: The `MessageSerializer` for encoding/decoding messages. - Defaults to `JSONMessageSerializer`. - channel: Channel prefix for queue names. Sanitized to alphanumeric - + underscore. Defaults to ``"pipecat_bus"``. - visibility_timeout: Seconds a read message stays invisible before - redelivery. Defaults to 30. + pgmq: An initialized ``PGMQueue`` client. Selects + :class:`DirectPgmqBackend`. Mutually exclusive with + ``backend``. + backend: A :class:`PgmqBackend` instance (e.g. + :class:`IsolatedPgmqBackend`, or a custom backend). + Mutually exclusive with ``pgmq``. + serializer: The :class:`MessageSerializer` for encoding/decoding + messages. Defaults to :class:`JSONMessageSerializer`. + channel: Channel name. With :class:`DirectPgmqBackend` this is + sanitized into a queue-name prefix. With + :class:`IsolatedPgmqBackend` it is the bus capability passed + to every wrapper call. + visibility_timeout: Seconds a read message stays invisible + before redelivery. Defaults to 30. batch_size: Maximum messages to fetch per read. Defaults to 10. poll_interval_ms: Long-poll check interval in milliseconds. - Defaults to 100. - max_poll_seconds: Maximum seconds the reader blocks per poll cycle. - Defaults to 5. - **kwargs: Additional arguments passed to `TaskBus`. + Defaults to 100. (Backend may ignore if it doesn't expose + this knob.) + max_poll_seconds: Maximum seconds the reader blocks per poll + cycle. Defaults to 5. + **kwargs: Additional arguments passed to :class:`TaskBus`. """ super().__init__(**kwargs) - self._pgmq = pgmq + if pgmq is not None and backend is not None: + raise ValueError("PgmqBus accepts pgmq= or backend=, not both") + if backend is not None: + self._backend: PgmqBackend = backend + elif pgmq is not None: + self._backend = DirectPgmqBackend(pgmq) + else: + raise ValueError("PgmqBus requires pgmq= or backend=") + self._serializer = serializer or JSONMessageSerializer() - self._safe_channel = _sanitize_channel(channel) + self._channel = channel self._visibility_timeout = visibility_timeout self._batch_size = batch_size self._poll_interval_ms = poll_interval_ms self._max_poll_seconds = max_poll_seconds self._queue_name: str | None = None self._reader_task: asyncio.Task | None = None - self._peer_cache: list[str] = [] - self._peer_cache_at: float = 0.0 async def start(self): - """Create this instance's queue and start the reader task.""" + """Join the channel via the backend and start the reader task.""" await super().start() - self._queue_name = f"{self._safe_channel}_{uuid.uuid4().hex[:12]}" - await self._pgmq.create_queue(self._queue_name) - logger.debug(f"{self}: created pgmq queue '{self._queue_name}'") + self._queue_name = await self._backend.join(self._channel) + logger.debug(f"{self}: joined channel via backend; queue='{self._queue_name}'") self._reader_task = self.create_task(self._reader_loop(), f"{self}::pgmq_reader") await asyncio.sleep(0) async def stop(self): - """Stop the reader task and drop this instance's queue.""" + """Stop the reader task and leave the channel.""" await super().stop() if self._reader_task: await self.cancel_task(self._reader_task) self._reader_task = None if self._queue_name: try: - await self._pgmq.drop_queue(self._queue_name) + await self._backend.leave(self._queue_name, channel=self._channel) except Exception: - logger.exception(f"{self}: failed to drop queue '{self._queue_name}'") + logger.exception(f"{self}: backend leave failed for queue '{self._queue_name}'") self._queue_name = None async def publish(self, message: BusMessage) -> None: - """Broadcast a message to every peer queue sharing the channel prefix. + """Broadcast a message to every peer on this channel. - Args: - message: The bus message to publish. + The backend handles peer discovery and fan-out; per-peer failures + are the backend's responsibility to absorb so the publish does not + raise. """ payload_bytes = self._serializer.serialize(message) try: @@ -176,39 +169,32 @@ class PgmqBus(TaskBus): logger.exception(f"{self}: failed to encode payload for pgmq") return - peers = await self._peer_queues() - logger.trace(f"{self}: publishing to {len(peers)} peer queue(s)") - for queue_name in peers: - try: - await self._pgmq.send(queue_name, payload) - except Exception: - logger.warning( - f"{self}: send to peer queue '{queue_name}' failed; invalidating cache" - ) - self._peer_cache_at = 0.0 + if self._queue_name is None: + logger.warning(f"{self}: publish called before start(); dropping message") + return - async def _peer_queues(self) -> list[str]: - now = time.monotonic() - if self._peer_cache and now - self._peer_cache_at < _PEER_LIST_TTL_S: - return self._peer_cache - names = await self._pgmq.list_queues() - prefix = f"{self._safe_channel}_" - self._peer_cache = [name for name in names if name.startswith(prefix)] - self._peer_cache_at = now - return self._peer_cache + try: + await self._backend.publish(self._channel, self._queue_name, payload) + except Exception: + logger.exception(f"{self}: backend publish failed") async def _reader_loop(self) -> None: while True: + if self._queue_name is None: + # Defensive: start() always sets this before launching the + # task, but cover the race where stop() clears it mid-loop. + return try: - messages = await self._pgmq.read_with_poll( - queue=self._queue_name, + messages = await self._backend.read( + self._queue_name, + channel=self._channel, vt=self._visibility_timeout, qty=self._batch_size, max_poll_seconds=self._max_poll_seconds, poll_interval_ms=self._poll_interval_ms, ) except Exception: - logger.exception(f"{self}: pgmq read failed; backing off") + logger.exception(f"{self}: backend read failed; backing off") await asyncio.sleep(1) continue @@ -222,6 +208,10 @@ class PgmqBus(TaskBus): logger.exception(f"{self}: failed to deserialize message {msg.msg_id}") finally: try: - await self._pgmq.delete(self._queue_name, msg.msg_id) + await self._backend.archive( + self._queue_name, + channel=self._channel, + msg_id=msg.msg_id, + ) except Exception: - logger.exception(f"{self}: failed to delete message {msg.msg_id}") + logger.exception(f"{self}: failed to archive message {msg.msg_id}") diff --git a/src/pipecat/bus/network/pgmq_backends.py b/src/pipecat/bus/network/pgmq_backends.py new file mode 100644 index 000000000..2a3eaad4b --- /dev/null +++ b/src/pipecat/bus/network/pgmq_backends.py @@ -0,0 +1,347 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Backends for :class:`pipecat.bus.network.pgmq.PgmqBus`. + +A backend owns the wire-side of bus operations: allocating a queue when a +bus instance joins a channel, fanning a published message out to channel +peers, long-polling for incoming messages, archiving them, and dropping +the queue when the bus stops. :class:`PgmqBus` is the orchestrator on top +of this surface. + +Two backends ship in this module: + +- :class:`DirectPgmqBackend` calls :class:`pgmq.async_queue.PGMQueue` + directly. Peer discovery is via ``list_queues()`` filtered by a channel + prefix. The channel name is encoded in the queue name and is therefore + enumerable by any role that can read ``pg_class``. Choose this backend + when bus peers trust each other (single-tenant deployments, internal + services). +- :class:`IsolatedPgmqBackend` calls a small set of SECURITY DEFINER + Postgres wrappers (``public.bus_join``, ``bus_publish``, + ``bus_subscribe``, ``bus_archive``, ``bus_leave``) over an + :mod:`asyncpg` pool. Queue names are server-allocated opaque UUIDs and + a server-side peer registry replaces ``list_queues`` discovery. Choose + this backend when bus peers should be isolated from each other and the + channel name itself is the bus capability. +""" + +import json +import re +import time +import uuid +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +from loguru import logger + +try: + from pgmq.async_queue import PGMQueue +except ModuleNotFoundError as e: # pragma: no cover - exercised only when extra is missing + logger.error(f"Exception: {e}") + logger.error("In order to use PgmqBus backends, you need to `pip install pipecat-ai[pgmq]`.") + raise Exception(f"Missing module: {e}") + + +_INVALID_CHANNEL_CHARS = re.compile(r"[^A-Za-z0-9_]") +_MAX_CHANNEL_LEN = 30 +_PEER_LIST_TTL_S = 1.0 + + +def _sanitize_channel(channel: str) -> str: + """Coerce an arbitrary channel string into a Postgres-safe identifier prefix. + + Only used by :class:`DirectPgmqBackend`, which builds queue names as + ``{channel}_{uuid}``. :class:`IsolatedPgmqBackend` does not embed the + channel in the queue name, so this helper does not apply there. + + Args: + channel: User-supplied channel name (may contain colons, slashes, etc). + + Returns: + A sanitized prefix safe for use as a PGMQ queue-name prefix. + """ + safe = _INVALID_CHANNEL_CHARS.sub("_", channel) + if safe and safe[0].isdigit(): + safe = f"q_{safe}" + return safe[:_MAX_CHANNEL_LEN] or "pipecat_bus" + + +@dataclass +class BackendMessage: + """A message returned by a backend's ``read`` call. + + Backends normalize whatever shape they get off the wire into this + minimal record so the bus orchestrator can stay backend-agnostic. + """ + + msg_id: int + message: dict + + +@runtime_checkable +class PgmqBackend(Protocol): + """Wire-side interface :class:`PgmqBus` delegates to. + + A backend instance is shared across the lifetime of a bus and may be + shared across multiple buses (e.g. one process running both an upstream + and downstream bus on different channels against the same database). + """ + + async def join(self, channel: str) -> str: + """Allocate a queue for this bus and register it on ``channel``. + + Returns: + The opaque queue name the bus should read from. + """ + ... + + async def publish( + self, + channel: str, + my_queue: str, + payload: dict, + ) -> None: + """Fan ``payload`` out to every peer queue on ``channel`` except ``my_queue``.""" + ... + + async def read( + self, + queue: str, + *, + channel: str, + vt: int, + qty: int, + max_poll_seconds: int, + poll_interval_ms: int, + ) -> list[BackendMessage]: + """Long-poll for messages on ``queue``. Returns an empty list on timeout.""" + ... + + async def archive( + self, + queue: str, + *, + channel: str, + msg_id: int, + ) -> bool: + """Acknowledge / archive a processed message.""" + ... + + async def leave(self, queue: str, *, channel: str) -> None: + """Drop the queue and unregister it from ``channel``.""" + ... + + +class DirectPgmqBackend: + """Backend that calls :class:`pgmq.async_queue.PGMQueue` directly. + + - Queue names are constructed client-side as ``{safe_channel}_{uuid12}``. + - Peer discovery uses ``pgmq.list_queues()`` filtered by channel prefix, + cached for :data:`_PEER_LIST_TTL_S` seconds. + - Per-peer ``send`` failures are caught; the cache is invalidated and the + fanout continues. The publish does not raise. + + The provided ``PGMQueue`` must already be initialized (``await pgmq.init()``). + The backend does not own the client's lifetime. + + Use this when bus peers trust each other. The channel name appears in + queue names visible to any role that can read ``pg_class``, so channels + are not secret. + """ + + def __init__(self, pgmq: PGMQueue): + """Initialize the backend with an already-initialized PGMQueue client.""" + self._pgmq = pgmq + # channel -> (cached_at, peer_queue_names) + self._peer_cache: dict[str, tuple[float, list[str]]] = {} + + async def join(self, channel: str) -> str: + """Create a queue named ``{safe_channel}_{uuid12}`` for this bus.""" + queue_name = f"{_sanitize_channel(channel)}_{uuid.uuid4().hex[:12]}" + await self._pgmq.create_queue(queue_name) + logger.debug(f"DirectPgmqBackend: created pgmq queue '{queue_name}'") + return queue_name + + async def publish(self, channel: str, my_queue: str, payload: dict) -> None: + """Send ``payload`` to every cached peer queue for ``channel``.""" + peers = await self._peers(channel) + logger.trace(f"DirectPgmqBackend: publishing to {len(peers)} peer queue(s)") + for queue_name in peers: + try: + await self._pgmq.send(queue_name, payload) + except Exception: + logger.warning( + f"DirectPgmqBackend: send to peer queue '{queue_name}' failed; " + "invalidating cache" + ) + self._peer_cache.pop(channel, None) + + async def read( + self, + queue: str, + *, + channel: str, + vt: int, + qty: int, + max_poll_seconds: int, + poll_interval_ms: int, + ) -> list[BackendMessage]: + """Long-poll ``queue`` via ``pgmq.read_with_poll``.""" + messages = await self._pgmq.read_with_poll( + queue=queue, + vt=vt, + qty=qty, + max_poll_seconds=max_poll_seconds, + poll_interval_ms=poll_interval_ms, + ) + return [BackendMessage(msg_id=m.msg_id, message=m.message) for m in messages or []] + + async def archive(self, queue: str, *, channel: str, msg_id: int) -> bool: + """Archive a processed message via ``pgmq.delete``.""" + return await self._pgmq.delete(queue, msg_id) + + async def leave(self, queue: str, *, channel: str) -> None: + """Drop the queue and invalidate the cached peer list for ``channel``.""" + await self._pgmq.drop_queue(queue) + self._peer_cache.pop(channel, None) + + async def _peers(self, channel: str) -> list[str]: + now = time.monotonic() + cached = self._peer_cache.get(channel) + if cached and now - cached[0] < _PEER_LIST_TTL_S: + return cached[1] + names = await self._pgmq.list_queues() + prefix = f"{_sanitize_channel(channel)}_" + peers = [n for n in names if n.startswith(prefix)] + self._peer_cache[channel] = (now, peers) + return peers + + +class IsolatedPgmqBackend: + """Backend that calls SECURITY DEFINER Postgres wrappers over asyncpg. + + Use this when bus peers should be isolated from each other and the + channel name is the bus capability. The backend never issues raw + ``pgmq.*`` calls; every operation goes through ``public.bus_*`` + wrappers, which enforce ``(queue_name, channel)`` membership against + a server-side peer registry table. + + Wire format (server-side SQL, defined out-of-band in the deployer's + migrations):: + + bus_join(p_channel text) RETURNS text + bus_publish(p_channel text, p_my_queue text, p_message jsonb) RETURNS bigint[] + bus_subscribe(p_my_queue text, p_channel text, p_vt int, + p_qty int, p_max_seconds int) + RETURNS TABLE(msg_id bigint, message jsonb) + bus_archive(p_my_queue text, p_channel text, p_msg_id bigint) RETURNS boolean + bus_leave(p_my_queue text, p_channel text) RETURNS void + + The ``asyncpg.Pool`` must allow at least two concurrent connections + (one held by the reader loop's long-poll, one for publishes). Pool + lifetime is owned by the caller — this backend does not close it. + + Notes: + - ``bus_subscribe`` long-polls inside a SECURITY DEFINER function. + Use a session-mode pooler; transaction-mode poolers may drop the + connection mid-poll. + - Payload serialization: the caller hands :class:`PgmqBus` a + :mod:`json`-encodable ``dict``; the backend forwards it as + ``jsonb`` via ``json.dumps`` because asyncpg does not auto-coerce + ``dict`` to ``jsonb``. + """ + + def __init__(self, pool: Any): + """Initialize the backend. + + Args: + pool: An :class:`asyncpg.pool.Pool` that the bus will use for all + wrapper calls. Typed as ``Any`` to keep :mod:`asyncpg` an + optional import. + """ + self._pool = pool + + async def join(self, channel: str) -> str: + """Call ``public.bus_join(channel)`` and return the server-allocated queue name.""" + async with self._pool.acquire() as conn: + queue_name = await conn.fetchval( + "SELECT public.bus_join($1)", + channel, + ) + if not queue_name: + raise RuntimeError("IsolatedPgmqBackend: bus_join returned no queue name") + logger.debug(f"IsolatedPgmqBackend: joined channel; queue='{queue_name}'") + return str(queue_name) + + async def publish(self, channel: str, my_queue: str, payload: dict) -> None: + """Call ``public.bus_publish`` to fan ``payload`` out to channel peers.""" + async with self._pool.acquire() as conn: + await conn.execute( + "SELECT public.bus_publish($1, $2, $3::jsonb)", + channel, + my_queue, + json.dumps(payload), + ) + + async def read( + self, + queue: str, + *, + channel: str, + vt: int, + qty: int, + max_poll_seconds: int, + poll_interval_ms: int, + ) -> list[BackendMessage]: + """Long-poll ``queue`` via ``public.bus_subscribe``. + + ``poll_interval_ms`` is honored server-side by ``pgmq.read_with_poll``'s + default; the wrapper does not expose it. + """ + async with self._pool.acquire() as conn: + rows = await conn.fetch( + "SELECT msg_id, message FROM public.bus_subscribe($1, $2, $3, $4, $5)", + queue, + channel, + vt, + qty, + max_poll_seconds, + ) + out: list[BackendMessage] = [] + for row in rows: + raw = row["message"] + message = json.loads(raw) if isinstance(raw, (bytes, str)) else raw + out.append(BackendMessage(msg_id=int(row["msg_id"]), message=message)) + return out + + async def archive(self, queue: str, *, channel: str, msg_id: int) -> bool: + """Call ``public.bus_archive`` to acknowledge a processed message.""" + async with self._pool.acquire() as conn: + result = await conn.fetchval( + "SELECT public.bus_archive($1, $2, $3)", + queue, + channel, + msg_id, + ) + return bool(result) + + async def leave(self, queue: str, *, channel: str) -> None: + """Call ``public.bus_leave`` to drop the queue and unregister it.""" + async with self._pool.acquire() as conn: + await conn.execute( + "SELECT public.bus_leave($1, $2)", + queue, + channel, + ) + + +__all__ = [ + "BackendMessage", + "DirectPgmqBackend", + "IsolatedPgmqBackend", + "PgmqBackend", +]