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:
committed by
Aleix Conchillo Flaqué
parent
a63abc41b6
commit
d757d8d06d
@@ -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}")
|
||||
|
||||
347
src/pipecat/bus/network/pgmq_backends.py
Normal file
347
src/pipecat/bus/network/pgmq_backends.py
Normal 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",
|
||||
]
|
||||
Reference in New Issue
Block a user