Split PgmqBus into orchestrator + pluggable backends

Move the wire-side of PGMQ operations into a new
``pipecat.bus.network.pgmq_backends`` module with a ``PgmqBackend``
Protocol, a ``DirectPgmqBackend`` (peers discovered by queue prefix),
and an ``IsolatedPgmqBackend`` (SECURITY DEFINER ``public.bus_*``
wrappers over an asyncpg pool). ``PgmqBus`` now delegates join,
publish, read, archive, and leave to the configured backend.

Construct ``PgmqBus`` with either ``pgmq=PGMQueue`` (uses
``DirectPgmqBackend``) or ``backend=PgmqBackend`` (any backend); the
two are mutually exclusive.
This commit is contained in:
Jon Taylor
2026-05-13 23:52:40 -07:00
committed by Aleix Conchillo Flaqué
parent a63abc41b6
commit d757d8d06d
2 changed files with 441 additions and 104 deletions

View File

@@ -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.<project-ref>",
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}")

View File

@@ -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",
]