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 asyncio
|
||||||
import json
|
import json
|
||||||
import re
|
|
||||||
import time
|
|
||||||
import uuid
|
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from pipecat.bus.bus import TaskBus
|
from pipecat.bus.bus import TaskBus
|
||||||
from pipecat.bus.messages import BusMessage
|
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 import JSONMessageSerializer
|
||||||
from pipecat.bus.serializers.base import MessageSerializer
|
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}")
|
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):
|
class PgmqBus(TaskBus):
|
||||||
"""Distributed task bus backed by PGMQ (PostgreSQL Message Queue).
|
"""Distributed task bus backed by PGMQ (PostgreSQL Message Queue).
|
||||||
|
|
||||||
Implements pub/sub fan-out on top of PGMQ's point-to-point queue semantics
|
Pub/sub fan-out is implemented on top of PGMQ's point-to-point queue
|
||||||
by giving each ``PgmqBus`` instance its own queue and broadcasting on
|
semantics by giving each :class:`PgmqBus` instance its own queue and
|
||||||
publish. The reader long-polls its own queue and dispatches received
|
broadcasting on publish. A reader long-polls the instance's queue and
|
||||||
messages to local subscribers.
|
dispatches received messages to local subscribers.
|
||||||
|
|
||||||
``BusLocalMessage`` messages bypass PGMQ and are delivered directly to
|
``BusLocalMessage`` messages bypass the network entirely and are
|
||||||
local subscribers.
|
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]``.
|
``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::
|
Example::
|
||||||
|
|
||||||
from pgmq.async_queue import PGMQueue
|
from pgmq.async_queue import PGMQueue
|
||||||
|
|
||||||
pgmq = PGMQueue(
|
pgmq = PGMQueue(
|
||||||
host="aws-0-us-east-1.pooler.supabase.com",
|
host="...",
|
||||||
port="5432",
|
port="5432",
|
||||||
database="postgres",
|
database="postgres",
|
||||||
username="postgres.<project-ref>",
|
username="postgres",
|
||||||
password="...",
|
password="...",
|
||||||
pool_size=4,
|
pool_size=4,
|
||||||
)
|
)
|
||||||
await pgmq.init()
|
await pgmq.init()
|
||||||
bus = PgmqBus(pgmq=pgmq, channel="pipecat_acme")
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
pgmq: PGMQueue,
|
pgmq: PGMQueue | None = None,
|
||||||
|
backend: PgmqBackend | None = None,
|
||||||
serializer: MessageSerializer | None = None,
|
serializer: MessageSerializer | None = None,
|
||||||
channel: str = "pipecat_bus",
|
channel: str = "pipecat_bus",
|
||||||
visibility_timeout: int = 30,
|
visibility_timeout: int = 30,
|
||||||
@@ -113,61 +93,74 @@ class PgmqBus(TaskBus):
|
|||||||
"""Initialize the PgmqBus.
|
"""Initialize the PgmqBus.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
pgmq: An initialized ``PGMQueue`` client. Call ``await pgmq.init()``
|
pgmq: An initialized ``PGMQueue`` client. Selects
|
||||||
before passing.
|
:class:`DirectPgmqBackend`. Mutually exclusive with
|
||||||
serializer: The `MessageSerializer` for encoding/decoding messages.
|
``backend``.
|
||||||
Defaults to `JSONMessageSerializer`.
|
backend: A :class:`PgmqBackend` instance (e.g.
|
||||||
channel: Channel prefix for queue names. Sanitized to alphanumeric
|
:class:`IsolatedPgmqBackend`, or a custom backend).
|
||||||
+ underscore. Defaults to ``"pipecat_bus"``.
|
Mutually exclusive with ``pgmq``.
|
||||||
visibility_timeout: Seconds a read message stays invisible before
|
serializer: The :class:`MessageSerializer` for encoding/decoding
|
||||||
redelivery. Defaults to 30.
|
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.
|
batch_size: Maximum messages to fetch per read. Defaults to 10.
|
||||||
poll_interval_ms: Long-poll check interval in milliseconds.
|
poll_interval_ms: Long-poll check interval in milliseconds.
|
||||||
Defaults to 100.
|
Defaults to 100. (Backend may ignore if it doesn't expose
|
||||||
max_poll_seconds: Maximum seconds the reader blocks per poll cycle.
|
this knob.)
|
||||||
Defaults to 5.
|
max_poll_seconds: Maximum seconds the reader blocks per poll
|
||||||
**kwargs: Additional arguments passed to `TaskBus`.
|
cycle. Defaults to 5.
|
||||||
|
**kwargs: Additional arguments passed to :class:`TaskBus`.
|
||||||
"""
|
"""
|
||||||
super().__init__(**kwargs)
|
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._serializer = serializer or JSONMessageSerializer()
|
||||||
self._safe_channel = _sanitize_channel(channel)
|
self._channel = channel
|
||||||
self._visibility_timeout = visibility_timeout
|
self._visibility_timeout = visibility_timeout
|
||||||
self._batch_size = batch_size
|
self._batch_size = batch_size
|
||||||
self._poll_interval_ms = poll_interval_ms
|
self._poll_interval_ms = poll_interval_ms
|
||||||
self._max_poll_seconds = max_poll_seconds
|
self._max_poll_seconds = max_poll_seconds
|
||||||
self._queue_name: str | None = None
|
self._queue_name: str | None = None
|
||||||
self._reader_task: asyncio.Task | 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):
|
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()
|
await super().start()
|
||||||
self._queue_name = f"{self._safe_channel}_{uuid.uuid4().hex[:12]}"
|
self._queue_name = await self._backend.join(self._channel)
|
||||||
await self._pgmq.create_queue(self._queue_name)
|
logger.debug(f"{self}: joined channel via backend; queue='{self._queue_name}'")
|
||||||
logger.debug(f"{self}: created pgmq queue '{self._queue_name}'")
|
|
||||||
self._reader_task = self.create_task(self._reader_loop(), f"{self}::pgmq_reader")
|
self._reader_task = self.create_task(self._reader_loop(), f"{self}::pgmq_reader")
|
||||||
await asyncio.sleep(0)
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
async def stop(self):
|
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()
|
await super().stop()
|
||||||
if self._reader_task:
|
if self._reader_task:
|
||||||
await self.cancel_task(self._reader_task)
|
await self.cancel_task(self._reader_task)
|
||||||
self._reader_task = None
|
self._reader_task = None
|
||||||
if self._queue_name:
|
if self._queue_name:
|
||||||
try:
|
try:
|
||||||
await self._pgmq.drop_queue(self._queue_name)
|
await self._backend.leave(self._queue_name, channel=self._channel)
|
||||||
except Exception:
|
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
|
self._queue_name = None
|
||||||
|
|
||||||
async def publish(self, message: BusMessage) -> 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:
|
The backend handles peer discovery and fan-out; per-peer failures
|
||||||
message: The bus message to publish.
|
are the backend's responsibility to absorb so the publish does not
|
||||||
|
raise.
|
||||||
"""
|
"""
|
||||||
payload_bytes = self._serializer.serialize(message)
|
payload_bytes = self._serializer.serialize(message)
|
||||||
try:
|
try:
|
||||||
@@ -176,39 +169,32 @@ class PgmqBus(TaskBus):
|
|||||||
logger.exception(f"{self}: failed to encode payload for pgmq")
|
logger.exception(f"{self}: failed to encode payload for pgmq")
|
||||||
return
|
return
|
||||||
|
|
||||||
peers = await self._peer_queues()
|
if self._queue_name is None:
|
||||||
logger.trace(f"{self}: publishing to {len(peers)} peer queue(s)")
|
logger.warning(f"{self}: publish called before start(); dropping message")
|
||||||
for queue_name in peers:
|
return
|
||||||
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
|
|
||||||
|
|
||||||
async def _peer_queues(self) -> list[str]:
|
try:
|
||||||
now = time.monotonic()
|
await self._backend.publish(self._channel, self._queue_name, payload)
|
||||||
if self._peer_cache and now - self._peer_cache_at < _PEER_LIST_TTL_S:
|
except Exception:
|
||||||
return self._peer_cache
|
logger.exception(f"{self}: backend publish failed")
|
||||||
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
|
|
||||||
|
|
||||||
async def _reader_loop(self) -> None:
|
async def _reader_loop(self) -> None:
|
||||||
while True:
|
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:
|
try:
|
||||||
messages = await self._pgmq.read_with_poll(
|
messages = await self._backend.read(
|
||||||
queue=self._queue_name,
|
self._queue_name,
|
||||||
|
channel=self._channel,
|
||||||
vt=self._visibility_timeout,
|
vt=self._visibility_timeout,
|
||||||
qty=self._batch_size,
|
qty=self._batch_size,
|
||||||
max_poll_seconds=self._max_poll_seconds,
|
max_poll_seconds=self._max_poll_seconds,
|
||||||
poll_interval_ms=self._poll_interval_ms,
|
poll_interval_ms=self._poll_interval_ms,
|
||||||
)
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception(f"{self}: pgmq read failed; backing off")
|
logger.exception(f"{self}: backend read failed; backing off")
|
||||||
await asyncio.sleep(1)
|
await asyncio.sleep(1)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
@@ -222,6 +208,10 @@ class PgmqBus(TaskBus):
|
|||||||
logger.exception(f"{self}: failed to deserialize message {msg.msg_id}")
|
logger.exception(f"{self}: failed to deserialize message {msg.msg_id}")
|
||||||
finally:
|
finally:
|
||||||
try:
|
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:
|
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