From e4bfcaf10c63c5e0bc07e96f79da7a5195d7e087 Mon Sep 17 00:00:00 2001 From: mducros Date: Mon, 18 May 2026 16:48:31 +0200 Subject: [PATCH 01/44] feat(Broker): add PgmqBroker --- pyproject.toml | 10 +- remoulade/__init__.py | 2 + remoulade/brokers/__init__.py | 7 + remoulade/brokers/pgmq.py | 191 ++++++++++++++++++ remoulade/errors.py | 4 + tests/conftest.py | 35 ++++ tests/docker-compose.yml | 6 +- .../postgres/01-create-pgmq-extension.sql | 1 + tests/test_pgmq.py | 184 +++++++++++++++++ 9 files changed, 433 insertions(+), 7 deletions(-) create mode 100644 remoulade/brokers/pgmq.py create mode 100644 tests/docker/postgres/01-create-pgmq-extension.sql create mode 100644 tests/test_pgmq.py diff --git a/pyproject.toml b/pyproject.toml index 8a173cfe9..0aa26fa0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,8 @@ classifiers = [ rabbitmq = ["amqpstorm>=2.6,<3"] redis = ["redis>=7.0.0"] server = ["flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec"] -postgres = ["sqlalchemy>=1.4.29,<2", "psycopg2>=2.9.11"] +postgres = ["sqlalchemy>=2.0,<3", "psycopg2>=2.9.11"] +pgmq = ["pgmq[sqlalchemy]>=1.1.1"] pydantic = ["pydantic>=2.12", "simplejson"] limits = ["limits~=5.3.0"] tracing = ["opentelemetry-api>=1.20"] @@ -34,8 +35,9 @@ dev = [ "flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec", - "sqlalchemy>=1.4.29,<2", + "sqlalchemy>=2.0,<3", "psycopg2>=2.9.11", + "pgmq[sqlalchemy]>=1.1.1", "pydantic>=2.12", "simplejson", "limits~=5.3.0", @@ -48,7 +50,7 @@ dev = [ # Linting "ruff", "mypy~=1.18.2", - "sqlalchemy[mypy]>=1.4.29,<2", + "sqlalchemy[mypy]>=2.0,<3", "types-redis", "types-python-dateutil", "types-simplejson", @@ -98,8 +100,6 @@ testpaths = ["tests"] asyncio_mode = "auto" markers = ["confirm_delivery", "group_transaction"] filterwarnings = [ - "error::sqlalchemy.exc.RemovedIn20Warning", - "error::sqlalchemy.exc.MovedIn20Warning", ] [tool.ruff] diff --git a/remoulade/__init__.py b/remoulade/__init__.py index 4d193704c..78da3d860 100644 --- a/remoulade/__init__.py +++ b/remoulade/__init__.py @@ -33,6 +33,7 @@ QueueJoinTimeout, QueueNotFound, RemouladeError, + UnsupportedMessageEncoding, ) from .generic import GenericActor from .logging import get_logger @@ -71,6 +72,7 @@ # Errors "RemouladeError", "Result", + "UnsupportedMessageEncoding", # Workers "Worker", "actor", diff --git a/remoulade/brokers/__init__.py b/remoulade/brokers/__init__.py index 1d583ac71..61763f308 100644 --- a/remoulade/brokers/__init__.py +++ b/remoulade/brokers/__init__.py @@ -14,3 +14,10 @@ # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . + +from .local import LocalBroker +from .pgmq import PgmqBroker +from .rabbitmq import RabbitmqBroker +from .stub import StubBroker + +__all__ = ["LocalBroker", "PgmqBroker", "RabbitmqBroker", "StubBroker"] diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py new file mode 100644 index 000000000..1633d4de7 --- /dev/null +++ b/remoulade/brokers/pgmq.py @@ -0,0 +1,191 @@ +# This file is a part of Remoulade. +# +# Copyright (C) 2017,2018 CLEARTYPE SRL +# +# Remoulade is free software; you can redistribute it and/or modify it +# under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or (at +# your option) any later version. +# +# Remoulade is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public +# License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . +import json +import os +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from threading import local +from typing import TYPE_CHECKING + +from pgmq import SQLAlchemyPGMQueue +from sqlalchemy import bindparam, create_engine, text +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.engine import Engine +from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker +from sqlalchemy.orm.session import sessionmaker + +from ..broker import Broker +from ..errors import QueueNotFound, UnsupportedMessageEncoding +from ..state.backends.postgres import DEFAULT_POSTGRES_URI + +if TYPE_CHECKING: + from ..message import Message + from ..middleware import Middleware + + +DELAYED_SEND_SQL = text( + """ + SELECT * + FROM pgmq.send( + queue_name => :queue_name, + msg => :message, + delay => :delay + ) + """ +).bindparams(bindparam("message", type_=JSONB)) + + +class PgmqBroker(Broker): + """A broker backed by PostgreSQL via the PGMQ extension.""" + + def __init__( + self, + *, + url: str | None = None, + middleware: list["Middleware"] | None = None, + sessionmaker: sessionmaker | None = None, + engine: Engine | None = None, + group_transaction: bool = False, + ): + super().__init__(middleware=middleware) + + if engine is not None and sessionmaker is not None: + raise ValueError("Only one of engine or sessionmaker can be provided.") + + self.url = ( + url + or os.getenv("REMOULADE_PGMQ_URL") + or os.getenv("REMOULADE_POSTGRESQL_URL") + or DEFAULT_POSTGRES_URI + ) + self.state = local() + self.group_transaction = group_transaction + self._owns_engine = engine is None and sessionmaker is None + + if sessionmaker is not None: + bound_engine = sessionmaker.kw.get("bind") + if bound_engine is None or not isinstance(bound_engine, Engine): + raise ValueError("sessionmaker must be bound to a SQLAlchemy engine.") + self.engine = bound_engine + self.sessionmaker = sessionmaker + else: + self.engine = engine or create_engine(self.url, pool_pre_ping=True) + self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) + + self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False) + + @contextmanager + def tx(self): + if self._has_transaction: + self.state.transaction_depth += 1 + try: + yield + finally: + self.state.transaction_depth -= 1 + return + + with self.sessionmaker.begin() as session: + self.state.transaction_depth = 1 + self.state.transaction_session = session + self.state.transaction_connection = session.connection() + try: + yield + finally: + self.state.transaction_depth = 0 + self.state.transaction_session = None + self.state.transaction_connection = None + + @property + def _has_transaction(self) -> bool: + return getattr(self.state, "transaction_connection", None) is not None + + @property + def _current_connection(self): + return getattr(self.state, "transaction_connection", None) + + def close(self): + if self._owns_engine: + self.engine.dispose() + + def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000): + raise NotImplementedError("PgmqBroker.consume() will be implemented in jalon 2.") + + def declare_queue(self, queue_name: str) -> None: + if queue_name in self.queues: + return + + self.client.validate_queue_name(queue_name, conn=self._current_connection) + self.emit_before("declare_queue", queue_name) + self.client.create_queue(queue_name, conn=self._current_connection) + self.queues[queue_name] = None + self.emit_after("declare_queue", queue_name) + + def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": + return message + + def _encode_message(self, message: "Message") -> dict: + try: + payload = json.loads(message.encode().decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise UnsupportedMessageEncoding( + "PgmqBroker only supports encoders that serialize messages as JSON objects." + ) from exc + + if not isinstance(payload, dict): + raise UnsupportedMessageEncoding("PgmqBroker requires message encoders to produce JSON objects.") + + return payload + + def _send_with_delay(self, queue_name: str, payload: dict, delay: int) -> None: + visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) + parameters = {"queue_name": queue_name, "message": payload, "delay": visible_at} + + if self._current_connection is not None: + self._current_connection.execute(DELAYED_SEND_SQL, parameters) + return + + with self.sessionmaker.begin() as session: + session.connection().execute(DELAYED_SEND_SQL, parameters) + + def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": + if message.queue_name not in self.queues: + raise QueueNotFound(message.queue_name) + + payload = self._encode_message(message) + + if delay is None: + self.client.send(message.queue_name, payload, conn=self._current_connection) + else: + self._send_with_delay(message.queue_name, payload, delay) + + return message + + def flush(self, queue_name: str) -> None: + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + + self.client.purge(queue_name, conn=self._current_connection) + + def purge_queue(self, queue_name: str) -> None: + self.flush(queue_name) + + def flush_all(self) -> None: + for queue_name in self.queues: + self.flush(queue_name) + + def join(self, queue_name: str, *, timeout: int | None = None) -> None: + raise NotImplementedError("PgmqBroker.join() will be implemented in jalon 3.") diff --git a/remoulade/errors.py b/remoulade/errors.py index 271d06d24..10e4dda46 100644 --- a/remoulade/errors.py +++ b/remoulade/errors.py @@ -64,6 +64,10 @@ class MessageNotDelivered(ConnectionError): """Raised when a message has not been delivered.""" +class UnsupportedMessageEncoding(BrokerError): + """Raised when the current message encoder does not produce JSON compatible payloads.""" + + class NoResultBackend(BrokerError): """Raised when trying to access the result backend on a broker without it""" diff --git a/tests/conftest.py b/tests/conftest.py index 1cbd25c11..75bc3a014 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,11 +13,13 @@ from sqlalchemy.inspection import inspect from sqlalchemy.orm.session import sessionmaker from sqlalchemy.pool import NullPool +from sqlalchemy.sql import text import remoulade from remoulade import Worker from remoulade.api import app from remoulade.brokers.local import LocalBroker +from remoulade.brokers.pgmq import PgmqBroker from remoulade.brokers.rabbitmq import RabbitmqBroker from remoulade.brokers.stub import StubBroker from remoulade.cancel import backends as cl_backends @@ -73,6 +75,22 @@ def check_postgres(client): return version_exists and states_exists and len(versions) == 1 and versions[0].version == DB_VERSION +def check_pgmq(client): + try: + with client.begin() as session: + session.execute(text("CREATE EXTENSION IF NOT EXISTS pgmq")) + session.execute(text("SELECT pgmq.validate_queue_name('remoulade_test_queue')")) + except Exception as e: + raise e from e if CI else pytest.skip("No connection to PostgreSQL/PGMQ server.") + + +def cleanup_pgmq_queues(client): + with client.begin() as session: + queue_names = [row[0] for row in session.execute(text("SELECT queue_name FROM pgmq.list_queues()")).all()] + for queue_name in queue_names: + session.execute(text("SELECT pgmq.drop_queue(:queue_name)"), {"queue_name": queue_name}) + + @pytest.fixture def check_postgres_begin(): client = remoulade.get_broker().get_state_backend().client @@ -113,6 +131,23 @@ def rabbitmq_broker(request): broker.close() +@pytest.fixture() +def pgmq_broker(): + db_string = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" + engine = create_engine(db_string, poolclass=NullPool) + client = sessionmaker(bind=engine) + check_pgmq(client) + cleanup_pgmq_queues(client) + + broker = PgmqBroker(sessionmaker=client) + broker.emit_after("process_boot") + remoulade.set_broker(broker) + yield broker + cleanup_pgmq_queues(client) + broker.emit_before("process_stop") + broker.close() + + @pytest.fixture() def local_broker(): broker = LocalBroker() diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index aaf02c9c8..2800823c4 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -8,10 +8,12 @@ services: ports: - "5784:5672" postgres: - image: postgres + image: ghcr.io/pgmq/pg18-pgmq:latest ports: - "5544:5432" environment: POSTGRES_USER: remoulade POSTGRES_HOST_AUTH_METHOD: trust - POSTGRES_DB: test \ No newline at end of file + POSTGRES_DB: test + volumes: + - ./docker/postgres/01-create-pgmq-extension.sql:/docker-entrypoint-initdb.d/01-create-pgmq-extension.sql:ro diff --git a/tests/docker/postgres/01-create-pgmq-extension.sql b/tests/docker/postgres/01-create-pgmq-extension.sql new file mode 100644 index 000000000..b484e86e5 --- /dev/null +++ b/tests/docker/postgres/01-create-pgmq-extension.sql @@ -0,0 +1 @@ +CREATE EXTENSION IF NOT EXISTS pgmq; diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py new file mode 100644 index 000000000..2cfb63918 --- /dev/null +++ b/tests/test_pgmq.py @@ -0,0 +1,184 @@ +import json +import time +from unittest.mock import Mock + +import pytest +from sqlalchemy import JSON, Column, Integer, MetaData, Table, create_engine, func +from sqlalchemy.sql import select, text + +import remoulade +from remoulade import Message, Middleware, UnsupportedMessageEncoding +from remoulade.brokers.pgmq import PgmqBroker + + +class RecorderMiddleware(Middleware): + def __init__(self): + self.before_messages = [] + self.after_messages = [] + + def before_enqueue(self, broker, message, delay): + self.before_messages.append((message, delay)) + + def after_enqueue(self, broker, message, delay, exception=None): + self.after_messages.append((message, delay, exception)) + + +@pytest.fixture +def sqlite_pgmq_broker(): + broker = PgmqBroker(engine=create_engine("sqlite://"), middleware=[]) + remoulade.set_broker(broker) + yield broker + broker.close() + + +def _count_messages(broker, queue_name="default"): + queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") + with broker.sessionmaker.begin() as session: + return session.execute(select(func.count()).select_from(queue_table)).scalar_one() + + +def _first_payload(broker, queue_name="default"): + queue_table = Table( + f"q_{queue_name}", + MetaData(), + Column("msg_id", Integer), + Column("message", JSON), + schema="pgmq", + ) + with broker.sessionmaker.begin() as session: + row = session.execute(select(queue_table.c.message).order_by(queue_table.c.msg_id).limit(1)).one() + return row[0] + + +def _queue_exists(broker, queue_name): + with broker.sessionmaker.begin() as session: + query = text("SELECT EXISTS(SELECT 1 FROM pgmq.list_queues() WHERE queue_name = :queue_name)") + return session.execute(query, {"queue_name": queue_name}).scalar_one() + + +def _expected_payload(message): + return json.loads(message.encode().decode("utf-8")) + + +def test_pgmq_broker_prefers_pgmq_url_from_environment(monkeypatch): + monkeypatch.setenv("REMOULADE_PGMQ_URL", "postgresql://pgmq-user@localhost/pgmq") + monkeypatch.setenv("REMOULADE_POSTGRESQL_URL", "postgresql://postgres-user@localhost/postgres") + + broker = PgmqBroker(engine=create_engine("sqlite://")) + + assert broker.url == "postgresql://pgmq-user@localhost/pgmq" + + +def test_pgmq_broker_declares_queue_idempotently(sqlite_pgmq_broker): + sqlite_pgmq_broker.client.validate_queue_name = Mock() + sqlite_pgmq_broker.client.create_queue = Mock() + sqlite_pgmq_broker.declare_queue("default") + sqlite_pgmq_broker.declare_queue("default") + + assert sqlite_pgmq_broker.get_declared_queues() == {"default"} + assert sqlite_pgmq_broker.get_declared_delay_queues() == set() + assert sqlite_pgmq_broker.client.validate_queue_name.call_count == 1 + assert sqlite_pgmq_broker.client.create_queue.call_count == 1 + + +def test_pgmq_broker_rejects_too_long_queue_names(sqlite_pgmq_broker): + sqlite_pgmq_broker.client.validate_queue_name = Mock(side_effect=ValueError("queue name too long")) + + with pytest.raises(ValueError): + sqlite_pgmq_broker.declare_queue("q" * 49) + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) + pgmq_broker.declare_queue(message.queue_name) + + pgmq_broker.enqueue(message) + + assert _count_messages(pgmq_broker) == 1 + assert _first_payload(pgmq_broker) == _expected_payload(message) + + +def test_pgmq_broker_rejects_non_json_encoders(sqlite_pgmq_broker, pickle_encoder): + message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + sqlite_pgmq_broker.client.validate_queue_name = Mock() + sqlite_pgmq_broker.client.create_queue = Mock() + sqlite_pgmq_broker.declare_queue(message.queue_name) + sqlite_pgmq_broker.client.send = Mock() + + with pytest.raises(UnsupportedMessageEncoding): + sqlite_pgmq_broker.enqueue(message) + + sqlite_pgmq_broker.client.send.assert_not_called() + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_broker_uses_native_visibility_delay_without_delay_queue(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) + pgmq_broker.declare_queue(message.queue_name) + + pgmq_broker.enqueue(message, delay=250) + + assert pgmq_broker.client.read("default", vt=1) is None + assert not _queue_exists(pgmq_broker, "default.DQ") + + time.sleep(0.35) + delayed = pgmq_broker.client.read("default", vt=1) + + assert delayed is not None + assert delayed.message == _expected_payload(message) + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_broker_transactions_commit_and_rollback_messages(pgmq_broker): + @remoulade.actor + def do_work(): + return 1 + + pgmq_broker.declare_actor(do_work) + + with pgmq_broker.tx(): + do_work.send() + + with pytest.raises(ValueError), pgmq_broker.tx(): + do_work.send() + raise ValueError("rollback") + + assert _count_messages(pgmq_broker) == 1 + + +def test_pgmq_broker_flush_purges_messages(sqlite_pgmq_broker): + sqlite_pgmq_broker.client.validate_queue_name = Mock() + sqlite_pgmq_broker.client.create_queue = Mock() + sqlite_pgmq_broker.client.purge = Mock() + + sqlite_pgmq_broker.declare_queue("default") + sqlite_pgmq_broker.flush("default") + + sqlite_pgmq_broker.client.purge.assert_called_once_with("default", conn=None) + + +def test_pgmq_broker_middleware_receives_standard_messages(sqlite_pgmq_broker): + middleware = RecorderMiddleware() + sqlite_pgmq_broker.client.validate_queue_name = Mock() + sqlite_pgmq_broker.client.create_queue = Mock() + sqlite_pgmq_broker.client.send = Mock() + sqlite_pgmq_broker.add_middleware(middleware) + + @remoulade.actor + def do_work(): + return 1 + + sqlite_pgmq_broker.declare_actor(do_work) + message = do_work.send() + + assert middleware.before_messages == [(message, None)] + assert middleware.after_messages == [(message, None, None)] + + +def test_pgmq_broker_worker_entrypoints_fail_explicitly_until_jalon_two(sqlite_pgmq_broker): + with pytest.raises(NotImplementedError, match="jalon 2"): + sqlite_pgmq_broker.consume("default") + + with pytest.raises(NotImplementedError, match="jalon 3"): + sqlite_pgmq_broker.join("default") From 783802d9b82cc5fc5c377f81283355cf0e8ccf6c Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 20 May 2026 17:06:43 +0200 Subject: [PATCH 02/44] feat(Broker): partionned queue --- local_pgmq_broker.py | 28 +++++++++++++++++ remoulade/brokers/pgmq.py | 65 +++++++++++++-------------------------- tests/conftest.py | 2 +- tests/test_pgmq.py | 61 +++++++++++++++++++++++++++++++----- 4 files changed, 104 insertions(+), 52 deletions(-) create mode 100644 local_pgmq_broker.py diff --git a/local_pgmq_broker.py b/local_pgmq_broker.py new file mode 100644 index 000000000..1ce33d2dd --- /dev/null +++ b/local_pgmq_broker.py @@ -0,0 +1,28 @@ +from sqlalchemy import text + +import remoulade +from remoulade import Message +from remoulade.brokers.pgmq import PgmqBroker + +# Just for local test +url = "postgresql://remoulade@localhost:5544/test" +broker = PgmqBroker(url=url, middleware=[]) + +remoulade.set_broker(broker) +broker.declare_queue("default") + +msg = Message( + queue_name="default", + actor_name="demo.add", + args=(1, 2), + kwargs={}, + options={}, +) +broker.enqueue(msg) + +with broker.sessionmaker.begin() as session: + row = session.execute(text("SELECT msg_id, message FROM pgmq.q_default ORDER BY msg_id DESC LIMIT 1")).one() + print("msg_id:", row.msg_id) + print("payload:", row.message) + +broker.close() diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 1633d4de7..2695365ee 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -15,7 +15,6 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . import json -import os from contextlib import contextmanager from datetime import UTC, datetime, timedelta from threading import local @@ -24,13 +23,10 @@ from pgmq import SQLAlchemyPGMQueue from sqlalchemy import bindparam, create_engine, text from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker -from sqlalchemy.orm.session import sessionmaker from ..broker import Broker from ..errors import QueueNotFound, UnsupportedMessageEncoding -from ..state.backends.postgres import DEFAULT_POSTGRES_URI if TYPE_CHECKING: from ..message import Message @@ -55,71 +51,44 @@ class PgmqBroker(Broker): def __init__( self, *, - url: str | None = None, + url: str, middleware: list["Middleware"] | None = None, - sessionmaker: sessionmaker | None = None, - engine: Engine | None = None, group_transaction: bool = False, + partition_archive_on_queue_init: bool = False, + archive_partition_interval: int | str = "1 day", + archive_retention_interval: int | str = "7 days", ): super().__init__(middleware=middleware) - if engine is not None and sessionmaker is not None: - raise ValueError("Only one of engine or sessionmaker can be provided.") - - self.url = ( - url - or os.getenv("REMOULADE_PGMQ_URL") - or os.getenv("REMOULADE_POSTGRESQL_URL") - or DEFAULT_POSTGRES_URI - ) + self.url = url self.state = local() self.group_transaction = group_transaction - self._owns_engine = engine is None and sessionmaker is None - - if sessionmaker is not None: - bound_engine = sessionmaker.kw.get("bind") - if bound_engine is None or not isinstance(bound_engine, Engine): - raise ValueError("sessionmaker must be bound to a SQLAlchemy engine.") - self.engine = bound_engine - self.sessionmaker = sessionmaker - else: - self.engine = engine or create_engine(self.url, pool_pre_ping=True) - self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) + self.partition_archive_on_queue_init = partition_archive_on_queue_init + self.archive_partition_interval = archive_partition_interval + self.archive_retention_interval = archive_retention_interval + + self.engine = create_engine(self.url, pool_pre_ping=True) + self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False) @contextmanager def tx(self): - if self._has_transaction: - self.state.transaction_depth += 1 - try: - yield - finally: - self.state.transaction_depth -= 1 - return - with self.sessionmaker.begin() as session: - self.state.transaction_depth = 1 self.state.transaction_session = session self.state.transaction_connection = session.connection() try: yield finally: - self.state.transaction_depth = 0 self.state.transaction_session = None self.state.transaction_connection = None - @property - def _has_transaction(self) -> bool: - return getattr(self.state, "transaction_connection", None) is not None - @property def _current_connection(self): return getattr(self.state, "transaction_connection", None) def close(self): - if self._owns_engine: - self.engine.dispose() + self.engine.dispose() def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000): raise NotImplementedError("PgmqBroker.consume() will be implemented in jalon 2.") @@ -130,7 +99,15 @@ def declare_queue(self, queue_name: str) -> None: self.client.validate_queue_name(queue_name, conn=self._current_connection) self.emit_before("declare_queue", queue_name) - self.client.create_queue(queue_name, conn=self._current_connection) + if self.partition_archive_on_queue_init: + self.client.create_partitioned_queue( + queue_name, + partition_interval=self.archive_partition_interval, + retention_interval=self.archive_retention_interval, + conn=self._current_connection, + ) + else: + self.client.create_queue(queue_name, conn=self._current_connection) self.queues[queue_name] = None self.emit_after("declare_queue", queue_name) diff --git a/tests/conftest.py b/tests/conftest.py index 75bc3a014..95ebdc2b6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -139,7 +139,7 @@ def pgmq_broker(): check_pgmq(client) cleanup_pgmq_queues(client) - broker = PgmqBroker(sessionmaker=client) + broker = PgmqBroker(url=db_string) broker.emit_after("process_boot") remoulade.set_broker(broker) yield broker diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 2cfb63918..9f70aced6 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -3,7 +3,7 @@ from unittest.mock import Mock import pytest -from sqlalchemy import JSON, Column, Integer, MetaData, Table, create_engine, func +from sqlalchemy import JSON, Column, Integer, MetaData, Table, func from sqlalchemy.sql import select, text import remoulade @@ -25,7 +25,7 @@ def after_enqueue(self, broker, message, delay, exception=None): @pytest.fixture def sqlite_pgmq_broker(): - broker = PgmqBroker(engine=create_engine("sqlite://"), middleware=[]) + broker = PgmqBroker(url="sqlite://", middleware=[]) remoulade.set_broker(broker) yield broker broker.close() @@ -60,13 +60,60 @@ def _expected_payload(message): return json.loads(message.encode().decode("utf-8")) -def test_pgmq_broker_prefers_pgmq_url_from_environment(monkeypatch): - monkeypatch.setenv("REMOULADE_PGMQ_URL", "postgresql://pgmq-user@localhost/pgmq") - monkeypatch.setenv("REMOULADE_POSTGRESQL_URL", "postgresql://postgres-user@localhost/postgres") +def test_pgmq_broker_uses_provided_url(): + broker_url = "postgresql://pgmq-user@localhost/pgmq" + broker = PgmqBroker(url=broker_url) - broker = PgmqBroker(engine=create_engine("sqlite://")) + assert broker.url == broker_url - assert broker.url == "postgresql://pgmq-user@localhost/pgmq" + +def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): + broker = PgmqBroker( + url="postgresql://pgmq-user@localhost/pgmq", + middleware=[], + partition_archive_on_queue_init=True, + ) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.create_queue = Mock() + + broker.declare_queue("default") + + broker.client.create_partitioned_queue.assert_called_once_with( + "default", + partition_interval="1 day", + retention_interval="7 days", + conn=None, + ) + broker.client.create_queue.assert_not_called() + + +def test_pgmq_broker_uses_non_partitioned_queue_creation_when_disabled(): + broker = PgmqBroker( + url="postgresql://pgmq-user@localhost/pgmq", + middleware=[], + partition_archive_on_queue_init=False, + ) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.create_queue = Mock() + + broker.declare_queue("default") + + broker.client.create_queue.assert_called_once_with("default", conn=None) + broker.client.create_partitioned_queue.assert_not_called() + + +def test_pgmq_broker_uses_non_partitioned_queue_creation_on_non_postgresql(): + broker = PgmqBroker(url="sqlite://", middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.create_queue = Mock() + + broker.declare_queue("default") + + broker.client.create_queue.assert_called_once_with("default", conn=None) + broker.client.create_partitioned_queue.assert_not_called() def test_pgmq_broker_declares_queue_idempotently(sqlite_pgmq_broker): From d18e47f0b2aa9693daa647e57787b3ffde17abe7 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 21 May 2026 10:36:49 +0200 Subject: [PATCH 03/44] feat(Pgmq): add consumer --- local_pgmq_consumer.py | 35 +++++++++++++++ remoulade/broker.py | 1 + remoulade/brokers/pgmq.py | 79 ++++++++++++++++++++++++++++++-- remoulade/worker.py | 14 +++--- tests/test_pgmq.py | 94 +++++++++++++++++++++++++++++++++++++-- 5 files changed, 211 insertions(+), 12 deletions(-) create mode 100644 local_pgmq_consumer.py diff --git a/local_pgmq_consumer.py b/local_pgmq_consumer.py new file mode 100644 index 000000000..9c362f714 --- /dev/null +++ b/local_pgmq_consumer.py @@ -0,0 +1,35 @@ +import time + +import remoulade +from remoulade import Worker +from remoulade.brokers.pgmq import PgmqBroker + +# Same DSN as local_pgmq_broker.py +URL = "postgresql://remoulade@localhost:5544/test" + +broker = PgmqBroker(url=URL, middleware=[]) +remoulade.set_broker(broker) + + +@remoulade.actor(actor_name="demo.add", queue_name="default") +def add(x: int, y: int) -> int: + result = x + y + print(f"[demo.add] {x} + {y} = {result}") + return result + + +remoulade.declare_actors([add]) + + +if __name__ == "__main__": + worker = Worker(broker, queues={"default"}, worker_threads=1, worker_timeout=500) + worker.start() + print("PGMQ local consumer started on queue 'default'. Press Ctrl+C to stop.") + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + pass + finally: + worker.stop() + broker.close() diff --git a/remoulade/broker.py b/remoulade/broker.py index 52bfa97f3..57b46ec1e 100644 --- a/remoulade/broker.py +++ b/remoulade/broker.py @@ -199,6 +199,7 @@ def __init__(self, middleware: "Iterable[Middleware] | None" = None): self.actors: dict[str, Actor] = {} self.queues: dict[str, Queue | None] = {} self.delay_queues: set[str] = set() + self.supports_native_delay = False self.actor_options: set[str] = set() self.middleware: list[Middleware] = [] diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 2695365ee..2afef924a 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -15,6 +15,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . import json +from collections import deque from contextlib import contextmanager from datetime import UTC, datetime, timedelta from threading import local @@ -25,11 +26,11 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker -from ..broker import Broker +from ..broker import Broker, Consumer, MessageProxy from ..errors import QueueNotFound, UnsupportedMessageEncoding +from ..message import Message if TYPE_CHECKING: - from ..message import Message from ..middleware import Middleware @@ -69,6 +70,7 @@ def __init__( self.engine = create_engine(self.url, pool_pre_ping=True) self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) + self.supports_native_delay = True self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False) @@ -90,8 +92,10 @@ def _current_connection(self): def close(self): self.engine.dispose() - def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000): - raise NotImplementedError("PgmqBroker.consume() will be implemented in jalon 2.") + def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PgmqConsumer": + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + return _PgmqConsumer(self, queue_name=queue_name, prefetch=prefetch, timeout=timeout) def declare_queue(self, queue_name: str) -> None: if queue_name in self.queues: @@ -166,3 +170,70 @@ def flush_all(self) -> None: def join(self, queue_name: str, *, timeout: int | None = None) -> None: raise NotImplementedError("PgmqBroker.join() will be implemented in jalon 3.") + + +class _PgmqConsumer(Consumer): + def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int): + self.broker = broker + self.client = broker.client + self.queue_name = queue_name + self.prefetch = max(prefetch, 1) + self.timeout = max(timeout, 0) + self.messages = deque() + + @property + def max_poll_seconds(self) -> int: + return max((self.timeout + 999) // 1000, 1) + + @property + def poll_interval_ms(self) -> int: + return max(min(self.timeout, 1000), 1) + + def ack(self, message): + if not isinstance(message, _PgmqMessage): + raise ValueError("It must be a PgmqMessage") + self.client.delete(self.queue_name, message._pgmq_message.msg_id) + + def nack(self, message): + if not isinstance(message, _PgmqMessage): + raise ValueError("It must be a PgmqMessage") + self.client.archive(self.queue_name, message._pgmq_message.msg_id) + + def requeue(self, messages): + msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] + if msg_ids: + self.client.set_vt(self.queue_name, msg_ids, 0) + + def __next__(self): + if self.messages: + return _PgmqMessage(self.messages.popleft()) + + messages = self.client.read_with_poll( + self.queue_name, + qty=self.prefetch, + max_poll_seconds=self.max_poll_seconds, + poll_interval_ms=self.poll_interval_ms, + ) + if not messages: + return None + + self.messages.extend(messages) + return _PgmqMessage(self.messages.popleft()) + + def close(self): + return + + +class _PgmqMessage(MessageProxy): + def __init__(self, pgmq_message): + payload = pgmq_message.message + if not isinstance(payload, dict): + raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") + + try: + message = Message(**payload) + except TypeError as exc: + raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc + + super().__init__(message) + self._pgmq_message = pgmq_message diff --git a/remoulade/worker.py b/remoulade/worker.py index a9cc4b748..9dff76a6c 100644 --- a/remoulade/worker.py +++ b/remoulade/worker.py @@ -171,7 +171,8 @@ def join(self): """ while True: for consumer in self.consumers.values(): - consumer.delay_queue.join() + if consumer.uses_in_memory_delay: + consumer.delay_queue.join() self.work_queue.join() @@ -179,7 +180,7 @@ def join(self): # joining on the work queue then it should be safe to exit. # This could still miss stuff but the chances are slim. for consumer in self.consumers.values(): - if consumer.delay_queue.unfinished_tasks: + if consumer.uses_in_memory_delay and consumer.delay_queue.unfinished_tasks: break else: if self.work_queue.unfinished_tasks: @@ -241,6 +242,7 @@ def __init__(self, *, broker, queue_name, prefetch, work_queue, worker_timeout): self.queue_name = queue_name self.work_queue = work_queue self.worker_timeout = worker_timeout + self.uses_in_memory_delay = not broker.supports_native_delay self.delay_queue = PriorityQueue() # type: PriorityQueue def run(self): @@ -270,7 +272,8 @@ def run(self): elif self.paused: break - self.handle_delayed_messages() + if self.uses_in_memory_delay: + self.handle_delayed_messages() if not self.running: break @@ -323,7 +326,7 @@ def handle_message(self, message): work queue. """ try: - if "eta" in message.options: + if self.uses_in_memory_delay and "eta" in message.options: self.logger.debug("Pushing message %r onto delay queue.", message.message_id) self.broker.emit_before("delay_message", message) self.delay_queue.put((message.options.get("eta", 0), message)) @@ -388,7 +391,8 @@ def close(self): """Close this consumer thread and its underlying connection.""" try: if self.consumer: - self.requeue_messages(m for _, m in iter_queue(self.delay_queue)) + if self.uses_in_memory_delay: + self.requeue_messages(m for _, m in iter_queue(self.delay_queue)) self.consumer.close() except ConnectionError: pass diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 9f70aced6..a05f98434 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -7,7 +7,7 @@ from sqlalchemy.sql import select, text import remoulade -from remoulade import Message, Middleware, UnsupportedMessageEncoding +from remoulade import Message, Middleware, QueueNotFound, UnsupportedMessageEncoding, Worker from remoulade.brokers.pgmq import PgmqBroker @@ -50,6 +50,12 @@ def _first_payload(broker, queue_name="default"): return row[0] +def _count_archived_messages(broker, queue_name="default"): + archive_table = Table(f"a_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") + with broker.sessionmaker.begin() as session: + return session.execute(select(func.count()).select_from(archive_table)).scalar_one() + + def _queue_exists(broker, queue_name): with broker.sessionmaker.begin() as session: query = text("SELECT EXISTS(SELECT 1 FROM pgmq.list_queues() WHERE queue_name = :queue_name)") @@ -223,9 +229,91 @@ def do_work(): assert middleware.after_messages == [(message, None, None)] -def test_pgmq_broker_worker_entrypoints_fail_explicitly_until_jalon_two(sqlite_pgmq_broker): - with pytest.raises(NotImplementedError, match="jalon 2"): +def test_pgmq_broker_consume_fails_when_queue_was_not_declared(sqlite_pgmq_broker): + with pytest.raises(QueueNotFound): sqlite_pgmq_broker.consume("default") + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_consumer_reads_messages_and_acks_with_delete(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(42,), kwargs={}, options={}) + pgmq_broker.declare_queue(message.queue_name) + pgmq_broker.enqueue(message) + + consumer = pgmq_broker.consume("default", prefetch=2, timeout=200) + consumed_message = next(consumer) + assert consumed_message is not None + assert consumed_message.message_id == message.message_id + + consumer.ack(consumed_message) + consumer.close() + + assert _count_messages(pgmq_broker) == 0 + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_consumer_nack_archives_messages(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) + pgmq_broker.declare_queue(message.queue_name) + pgmq_broker.enqueue(message) + + consumer = pgmq_broker.consume("default", prefetch=1, timeout=200) + consumed_message = next(consumer) + + assert consumed_message is not None + consumer.nack(consumed_message) + consumer.close() + + assert _count_messages(pgmq_broker) == 0 + assert _count_archived_messages(pgmq_broker) == 1 + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_consumer_requeue_restores_visibility_with_set_vt(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) + pgmq_broker.declare_queue(message.queue_name) + pgmq_broker.enqueue(message) + + consumer = pgmq_broker.consume("default", prefetch=1, timeout=200) + consumed_message = next(consumer) + + assert consumed_message is not None + consumer.requeue([consumed_message]) + + replayed_message = next(consumer) + assert replayed_message is not None + assert replayed_message.message_id == message.message_id + + consumer.ack(replayed_message) + consumer.close() + + assert _count_messages(pgmq_broker) == 0 + + +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_worker_processes_native_delayed_messages_without_delay_queue(pgmq_broker): + seen = [] + + @remoulade.actor + def do_work(value): + seen.append(value) + + pgmq_broker.declare_actor(do_work) + worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=2) + worker.start() + try: + do_work.send_with_options(args=(3,), delay=150) + deadline = time.monotonic() + 10 + while not seen and time.monotonic() < deadline: + time.sleep(0.05) + assert seen == [3] + worker.join() + finally: + worker.stop() + + assert not _queue_exists(pgmq_broker, "default.DQ") + + +def test_pgmq_broker_join_is_not_implemented_until_jalon_three(sqlite_pgmq_broker): with pytest.raises(NotImplementedError, match="jalon 3"): sqlite_pgmq_broker.join("default") From 86a9c864caa85f692b76a847997f66f7b52e5054 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 21 May 2026 15:40:00 +0200 Subject: [PATCH 04/44] feat(Pgmq): add listen notify --- local_pgmq_broker.py | 25 ++++-- local_pgmq_consumer.py | 21 +++-- pyproject.toml | 5 +- remoulade/brokers/pgmq.py | 184 +++++++++++++++++++++++++++++++++----- tests/test_pgmq.py | 147 ++++++++++++++++++++++++++++++ 5 files changed, 345 insertions(+), 37 deletions(-) diff --git a/local_pgmq_broker.py b/local_pgmq_broker.py index 1ce33d2dd..d5468a33e 100644 --- a/local_pgmq_broker.py +++ b/local_pgmq_broker.py @@ -1,3 +1,6 @@ +from collections.abc import Iterable +from typing import Any, cast + from sqlalchemy import text import remoulade @@ -5,16 +8,23 @@ from remoulade.brokers.pgmq import PgmqBroker # Just for local test -url = "postgresql://remoulade@localhost:5544/test" -broker = PgmqBroker(url=url, middleware=[]) +URL = "postgresql://remoulade@localhost:5544/test" +QUEUE_NAME = "default" +ACTOR_NAME = "demo.add" +broker = PgmqBroker( + url=URL, + middleware=[], + listen_notify_enabled=True, +) remoulade.set_broker(broker) -broker.declare_queue("default") +broker.declare_queue(QUEUE_NAME) -msg = Message( - queue_name="default", - actor_name="demo.add", - args=(1, 2), +msg_args: Iterable[Any] = cast(Iterable[Any], (1, 2)) +msg: Message = Message( + queue_name=QUEUE_NAME, + actor_name=ACTOR_NAME, + args=msg_args, kwargs={}, options={}, ) @@ -24,5 +34,6 @@ row = session.execute(text("SELECT msg_id, message FROM pgmq.q_default ORDER BY msg_id DESC LIMIT 1")).one() print("msg_id:", row.msg_id) print("payload:", row.message) + print("listen_notify_channel:", f"pgmq.q_{QUEUE_NAME}.INSERT") broker.close() diff --git a/local_pgmq_consumer.py b/local_pgmq_consumer.py index 9c362f714..e63f2a550 100644 --- a/local_pgmq_consumer.py +++ b/local_pgmq_consumer.py @@ -6,12 +6,19 @@ # Same DSN as local_pgmq_broker.py URL = "postgresql://remoulade@localhost:5544/test" - -broker = PgmqBroker(url=URL, middleware=[]) +QUEUE_NAME = "default" +ACTOR_NAME = "demo.add" + +broker = PgmqBroker( + url=URL, + middleware=[], + listen_notify_enabled=True, +) remoulade.set_broker(broker) +broker.declare_queue(QUEUE_NAME) -@remoulade.actor(actor_name="demo.add", queue_name="default") +@remoulade.actor(actor_name=ACTOR_NAME, queue_name=QUEUE_NAME) def add(x: int, y: int) -> int: result = x + y print(f"[demo.add] {x} + {y} = {result}") @@ -22,9 +29,13 @@ def add(x: int, y: int) -> int: if __name__ == "__main__": - worker = Worker(broker, queues={"default"}, worker_threads=1, worker_timeout=500) + probe = broker.consume(QUEUE_NAME, prefetch=1, timeout=30_000) + listener_mode = "LISTEN/NOTIFY" if probe._listener_available else "polling fallback" + probe.close() + + worker = Worker(broker, queues={QUEUE_NAME}, worker_threads=1, worker_timeout=30_000) worker.start() - print("PGMQ local consumer started on queue 'default'. Press Ctrl+C to stop.") + print(f"PGMQ local consumer started on queue '{QUEUE_NAME}' ({listener_mode}). Press Ctrl+C to stop.") try: while True: time.sleep(1) diff --git a/pyproject.toml b/pyproject.toml index 0aa26fa0e..1b6da4a48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,8 +23,8 @@ classifiers = [ rabbitmq = ["amqpstorm>=2.6,<3"] redis = ["redis>=7.0.0"] server = ["flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec"] -postgres = ["sqlalchemy>=2.0,<3", "psycopg2>=2.9.11"] -pgmq = ["pgmq[sqlalchemy]>=1.1.1"] +postgres = ["sqlalchemy>=1.4.29", "psycopg2>=2.9.11"] +pgmq = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1"] pydantic = ["pydantic>=2.12", "simplejson"] limits = ["limits~=5.3.0"] tracing = ["opentelemetry-api>=1.20"] @@ -38,6 +38,7 @@ dev = [ "sqlalchemy>=2.0,<3", "psycopg2>=2.9.11", "pgmq[sqlalchemy]>=1.1.1", + "psycopg>=3.2", "pydantic>=2.12", "simplejson", "limits~=5.3.0", diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 2afef924a..65dfd5657 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -16,14 +16,19 @@ # along with this program. If not, see . import json from collections import deque +from collections.abc import Iterable, Iterator from contextlib import contextmanager from datetime import UTC, datetime, timedelta -from threading import local -from typing import TYPE_CHECKING +from threading import Event, Thread, local +from typing import TYPE_CHECKING, Any +import psycopg from pgmq import SQLAlchemyPGMQueue +from psycopg import sql as psycopg_sql +from pydantic import BaseModel, ConfigDict, ValidationError from sqlalchemy import bindparam, create_engine, text from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.engine import make_url from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker from ..broker import Broker, Consumer, MessageProxy @@ -34,6 +39,16 @@ from ..middleware import Middleware +PgmqPayload = dict[str, Any] + + +class _PgmqQueueMessage(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + msg_id: int + message: PgmqPayload + + DELAYED_SEND_SQL = text( """ SELECT * @@ -44,6 +59,7 @@ ) """ ).bindparams(bindparam("message", type_=JSONB)) +LISTEN_NOTIFY_THROTTLE_MS = 250 class PgmqBroker(Broker): @@ -58,7 +74,8 @@ def __init__( partition_archive_on_queue_init: bool = False, archive_partition_interval: int | str = "1 day", archive_retention_interval: int | str = "7 days", - ): + listen_notify_enabled: bool = True, + ) -> None: super().__init__(middleware=middleware) self.url = url @@ -67,6 +84,7 @@ def __init__( self.partition_archive_on_queue_init = partition_archive_on_queue_init self.archive_partition_interval = archive_partition_interval self.archive_retention_interval = archive_retention_interval + self.listen_notify_enabled = listen_notify_enabled self.engine = create_engine(self.url, pool_pre_ping=True) self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) @@ -75,7 +93,7 @@ def __init__( self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False) @contextmanager - def tx(self): + def tx(self) -> Iterator[None]: with self.sessionmaker.begin() as session: self.state.transaction_session = session self.state.transaction_connection = session.connection() @@ -86,10 +104,31 @@ def tx(self): self.state.transaction_connection = None @property - def _current_connection(self): + def _current_connection(self) -> Any | None: return getattr(self.state, "transaction_connection", None) - def close(self): + def _build_listener_conninfo(self) -> str: + # psycopg expects a plain PostgreSQL URL and rejects SQLAlchemy-style driver suffixes. + return make_url(self.url).set(drivername="postgresql").render_as_string(hide_password=False) + + def _try_enable_notify(self, queue_name: str) -> None: + if not self.listen_notify_enabled: + return + + try: + self.client.enable_notify( + queue_name, + throttle_interval_ms=LISTEN_NOTIFY_THROTTLE_MS, + conn=self._current_connection, + ) + except Exception: + self.logger.warning( + "Failed to enable LISTEN/NOTIFY for queue %s; consumer will fall back to polling.", + queue_name, + exc_info=True, + ) + + def close(self) -> None: self.engine.dispose() def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PgmqConsumer": @@ -112,13 +151,14 @@ def declare_queue(self, queue_name: str) -> None: ) else: self.client.create_queue(queue_name, conn=self._current_connection) + self._try_enable_notify(queue_name) self.queues[queue_name] = None self.emit_after("declare_queue", queue_name) def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": return message - def _encode_message(self, message: "Message") -> dict: + def _encode_message(self, message: "Message") -> PgmqPayload: try: payload = json.loads(message.encode().decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: @@ -131,7 +171,7 @@ def _encode_message(self, message: "Message") -> dict: return payload - def _send_with_delay(self, queue_name: str, payload: dict, delay: int) -> None: + def _send_with_delay(self, queue_name: str, payload: PgmqPayload, delay: int) -> None: visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) parameters = {"queue_name": queue_name, "message": payload, "delay": visible_at} @@ -173,13 +213,21 @@ def join(self, queue_name: str, *, timeout: int | None = None) -> None: class _PgmqConsumer(Consumer): - def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int): + def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int) -> None: self.broker = broker self.client = broker.client self.queue_name = queue_name self.prefetch = max(prefetch, 1) self.timeout = max(timeout, 0) - self.messages = deque() + self.messages: deque[_PgmqQueueMessage] = deque() + self._notify_event = Event() + self._listener_stop = Event() + self._listener_thread: Thread | None = None + self._listener_connection: psycopg.Connection[Any] | None = None + self._listener_available = False + + if self.broker.listen_notify_enabled: + self._start_listener() @property def max_poll_seconds(self) -> int: @@ -189,49 +237,139 @@ def max_poll_seconds(self) -> int: def poll_interval_ms(self) -> int: return max(min(self.timeout, 1000), 1) - def ack(self, message): + @property + def wait_timeout_seconds(self) -> float: + return self.timeout / 1000 if self.timeout > 0 else 0 + + def _normalize_messages(self, messages: Any | list[Any] | None) -> list[_PgmqQueueMessage]: + if messages is None: + return [] + raw_messages = messages if isinstance(messages, list) else [messages] + normalized_messages: list[_PgmqQueueMessage] = [] + for message in raw_messages: + try: + normalized_messages.append(_PgmqQueueMessage.model_validate(message, from_attributes=True)) + except ValidationError as exc: + raise UnsupportedMessageEncoding( + "PGMQ messages must expose 'msg_id' and a JSONB-serializable object in 'message'." + ) from exc + return normalized_messages + + def _read_immediate(self) -> list[_PgmqQueueMessage]: + return self._normalize_messages(self.client.read(self.queue_name, qty=self.prefetch)) + + def _read_with_poll(self) -> list[_PgmqQueueMessage]: + return self._normalize_messages( + self.client.read_with_poll( + self.queue_name, + qty=self.prefetch, + max_poll_seconds=self.max_poll_seconds, + poll_interval_ms=self.poll_interval_ms, + ) + ) + + def _start_listener(self) -> None: + channel = f"pgmq.q_{self.queue_name}.INSERT" + try: + conninfo = self.broker._build_listener_conninfo() + self._listener_connection = psycopg.connect(conninfo, autocommit=True) + self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) + self._listener_available = True + except Exception: + self._listener_available = False + self._listener_connection = None + self.broker.logger.warning( + "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling.", + self.queue_name, + exc_info=True, + ) + return + + self._listener_thread = Thread( + target=self._run_listener, + name=f"pgmq-listener-{self.queue_name}", + daemon=True, + ) + self._listener_thread.start() + + def _run_listener(self) -> None: + while not self._listener_stop.is_set(): + try: + if self._listener_connection is None: + break + for _ in self._listener_connection.notifies(timeout=0.5, stop_after=1): + self._notify_event.set() + except Exception: + if self._listener_stop.is_set(): + break + self._listener_available = False + self._notify_event.set() + self.broker.logger.warning( + "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling.", + self.queue_name, + exc_info=True, + ) + break + + self._listener_available = False + + def ack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self.client.delete(self.queue_name, message._pgmq_message.msg_id) - def nack(self, message): + def nack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self.client.archive(self.queue_name, message._pgmq_message.msg_id) - def requeue(self, messages): + def requeue(self, messages: Iterable["MessageProxy"]) -> None: msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] if msg_ids: self.client.set_vt(self.queue_name, msg_ids, 0) - def __next__(self): + def __next__(self) -> "_PgmqMessage | None": if self.messages: return _PgmqMessage(self.messages.popleft()) - messages = self.client.read_with_poll( - self.queue_name, - qty=self.prefetch, - max_poll_seconds=self.max_poll_seconds, - poll_interval_ms=self.poll_interval_ms, - ) + if self._listener_available: + self._notify_event.clear() + messages = self._read_immediate() + if not messages: + self._notify_event.wait(self.wait_timeout_seconds) + messages = self._read_with_poll() if not self._listener_available else self._read_immediate() + else: + messages = self._read_with_poll() + if not messages: return None self.messages.extend(messages) return _PgmqMessage(self.messages.popleft()) - def close(self): + def close(self) -> None: + self._listener_stop.set() + self._notify_event.set() + if self._listener_connection is not None: + try: + self._listener_connection.close() + except Exception as e: + self.broker.logger.error("Listener not joinable: %s", str(e)) + finally: + self._listener_connection = None + if self._listener_thread is not None: + self._listener_thread.join(timeout=1) return class _PgmqMessage(MessageProxy): - def __init__(self, pgmq_message): + def __init__(self, pgmq_message: _PgmqQueueMessage) -> None: payload = pgmq_message.message if not isinstance(payload, dict): raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") try: - message = Message(**payload) + message: Message = Message(**payload) except TypeError as exc: raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index a05f98434..458ffa130 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -1,4 +1,5 @@ import json +import threading import time from unittest.mock import Mock @@ -110,16 +111,41 @@ def test_pgmq_broker_uses_non_partitioned_queue_creation_when_disabled(): broker.client.create_partitioned_queue.assert_not_called() +def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_queue = Mock() + broker.client.enable_notify = Mock() + + broker.declare_queue("default") + + broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) + + +def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_queue = Mock() + broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) + + broker.declare_queue("default") + + assert "default" in broker.queues + assert "Failed to enable LISTEN/NOTIFY" in caplog.text + + def test_pgmq_broker_uses_non_partitioned_queue_creation_on_non_postgresql(): broker = PgmqBroker(url="sqlite://", middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.create_queue = Mock() + broker.client.enable_notify = Mock() broker.declare_queue("default") broker.client.create_queue.assert_called_once_with("default", conn=None) broker.client.create_partitioned_queue.assert_not_called() + broker.client.enable_notify.assert_not_called() def test_pgmq_broker_declares_queue_idempotently(sqlite_pgmq_broker): @@ -234,6 +260,97 @@ def test_pgmq_broker_consume_fails_when_queue_was_not_declared(sqlite_pgmq_broke sqlite_pgmq_broker.consume("default") +def test_pgmq_consumer_uses_notification_path_when_listener_is_available(monkeypatch): + def _fake_start_listener(self): + self._listener_available = True + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + payload = _expected_payload(message) + broker.client.read = Mock(side_effect=[None, Mock(msg_id=1, message=payload)]) + broker.client.read_with_poll = Mock(return_value=[]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumer._notify_event.set() + + consumed = next(consumer) + + assert consumed is not None + assert consumed.message_id == message.message_id + assert broker.client.read.call_count == 2 + broker.client.read_with_poll.assert_not_called() + consumer.close() + + +def test_pgmq_consumer_falls_back_to_polling_when_listener_is_unavailable(monkeypatch): + def _fake_start_listener(self): + self._listener_available = False + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) + payload = _expected_payload(message) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=2, message=payload)]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumed = next(consumer) + + assert consumed is not None + assert consumed.message_id == message.message_id + broker.client.read.assert_not_called() + broker.client.read_with_poll.assert_called_once_with( + "default", + qty=1, + max_poll_seconds=1, + poll_interval_ms=200, + ) + consumer.close() + + +def test_pgmq_consumer_falls_back_to_polling_when_listener_stops_during_wait(monkeypatch): + def _fake_start_listener(self): + self._listener_available = True + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(3,), kwargs={}, options={}) + payload = _expected_payload(message) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=3, message=payload)]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + + def _wait_and_stop(_timeout): + consumer._listener_available = False + return False + + consumer._notify_event.wait = Mock(side_effect=_wait_and_stop) + + consumed = next(consumer) + + assert consumed is not None + assert consumed.message_id == message.message_id + broker.client.read.assert_called_once_with("default", qty=1) + broker.client.read_with_poll.assert_called_once_with( + "default", + qty=1, + max_poll_seconds=1, + poll_interval_ms=200, + ) + consumer.close() + + @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_consumer_reads_messages_and_acks_with_delete(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(42,), kwargs={}, options={}) @@ -314,6 +431,36 @@ def do_work(value): assert not _queue_exists(pgmq_broker, "default.DQ") +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_consumer_listener_wakes_on_enqueue_with_listen_notify(pgmq_broker): + message = Message(queue_name="default", actor_name="do_work", args=(99,), kwargs={}, options={}) + pgmq_broker.declare_queue(message.queue_name) + consumer = pgmq_broker.consume(message.queue_name, prefetch=1, timeout=1500) + + if not consumer._listener_available: + pytest.skip("LISTEN/NOTIFY listener unavailable in this environment.") + + consumed_messages = [] + + def _consume_once(): + consumed_messages.append(next(consumer)) + + thread = threading.Thread(target=_consume_once) + thread.start() + try: + time.sleep(0.15) + pgmq_broker.enqueue(message) + thread.join(timeout=3) + assert not thread.is_alive() + assert consumed_messages + consumed = consumed_messages[0] + assert consumed is not None + assert consumed.message_id == message.message_id + consumer.ack(consumed) + finally: + consumer.close() + + def test_pgmq_broker_join_is_not_implemented_until_jalon_three(sqlite_pgmq_broker): with pytest.raises(NotImplementedError, match="jalon 3"): sqlite_pgmq_broker.join("default") From 0ad1c3aae49e05ab41a65f2b71a62e6816370662 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 21 May 2026 16:27:09 +0200 Subject: [PATCH 05/44] fix(Workflow): tests --- .github/workflows/tests.yml | 2 +- tests/test_pgmq.py | 107 +----------------------------------- 2 files changed, 2 insertions(+), 107 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3391ead5c..f7a4efb1b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -28,7 +28,7 @@ jobs: ports: - 5784:5672 postgres: - image: postgres:16 + image: ghcr.io/pgmq/pg18-pgmq:latest ports: - 5544:5432 env: diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 458ffa130..a06de7b92 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -8,30 +8,10 @@ from sqlalchemy.sql import select, text import remoulade -from remoulade import Message, Middleware, QueueNotFound, UnsupportedMessageEncoding, Worker +from remoulade import Message, Worker from remoulade.brokers.pgmq import PgmqBroker -class RecorderMiddleware(Middleware): - def __init__(self): - self.before_messages = [] - self.after_messages = [] - - def before_enqueue(self, broker, message, delay): - self.before_messages.append((message, delay)) - - def after_enqueue(self, broker, message, delay, exception=None): - self.after_messages.append((message, delay, exception)) - - -@pytest.fixture -def sqlite_pgmq_broker(): - broker = PgmqBroker(url="sqlite://", middleware=[]) - remoulade.set_broker(broker) - yield broker - broker.close() - - def _count_messages(broker, queue_name="default"): queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") with broker.sessionmaker.begin() as session: @@ -134,39 +114,6 @@ def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): assert "Failed to enable LISTEN/NOTIFY" in caplog.text -def test_pgmq_broker_uses_non_partitioned_queue_creation_on_non_postgresql(): - broker = PgmqBroker(url="sqlite://", middleware=[]) - broker.client.validate_queue_name = Mock() - broker.client.create_partitioned_queue = Mock() - broker.client.create_queue = Mock() - broker.client.enable_notify = Mock() - - broker.declare_queue("default") - - broker.client.create_queue.assert_called_once_with("default", conn=None) - broker.client.create_partitioned_queue.assert_not_called() - broker.client.enable_notify.assert_not_called() - - -def test_pgmq_broker_declares_queue_idempotently(sqlite_pgmq_broker): - sqlite_pgmq_broker.client.validate_queue_name = Mock() - sqlite_pgmq_broker.client.create_queue = Mock() - sqlite_pgmq_broker.declare_queue("default") - sqlite_pgmq_broker.declare_queue("default") - - assert sqlite_pgmq_broker.get_declared_queues() == {"default"} - assert sqlite_pgmq_broker.get_declared_delay_queues() == set() - assert sqlite_pgmq_broker.client.validate_queue_name.call_count == 1 - assert sqlite_pgmq_broker.client.create_queue.call_count == 1 - - -def test_pgmq_broker_rejects_too_long_queue_names(sqlite_pgmq_broker): - sqlite_pgmq_broker.client.validate_queue_name = Mock(side_effect=ValueError("queue name too long")) - - with pytest.raises(ValueError): - sqlite_pgmq_broker.declare_queue("q" * 49) - - @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) @@ -178,19 +125,6 @@ def test_pgmq_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(pgmq_b assert _first_payload(pgmq_broker) == _expected_payload(message) -def test_pgmq_broker_rejects_non_json_encoders(sqlite_pgmq_broker, pickle_encoder): - message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) - sqlite_pgmq_broker.client.validate_queue_name = Mock() - sqlite_pgmq_broker.client.create_queue = Mock() - sqlite_pgmq_broker.declare_queue(message.queue_name) - sqlite_pgmq_broker.client.send = Mock() - - with pytest.raises(UnsupportedMessageEncoding): - sqlite_pgmq_broker.enqueue(message) - - sqlite_pgmq_broker.client.send.assert_not_called() - - @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_broker_uses_native_visibility_delay_without_delay_queue(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) @@ -226,40 +160,6 @@ def do_work(): assert _count_messages(pgmq_broker) == 1 -def test_pgmq_broker_flush_purges_messages(sqlite_pgmq_broker): - sqlite_pgmq_broker.client.validate_queue_name = Mock() - sqlite_pgmq_broker.client.create_queue = Mock() - sqlite_pgmq_broker.client.purge = Mock() - - sqlite_pgmq_broker.declare_queue("default") - sqlite_pgmq_broker.flush("default") - - sqlite_pgmq_broker.client.purge.assert_called_once_with("default", conn=None) - - -def test_pgmq_broker_middleware_receives_standard_messages(sqlite_pgmq_broker): - middleware = RecorderMiddleware() - sqlite_pgmq_broker.client.validate_queue_name = Mock() - sqlite_pgmq_broker.client.create_queue = Mock() - sqlite_pgmq_broker.client.send = Mock() - sqlite_pgmq_broker.add_middleware(middleware) - - @remoulade.actor - def do_work(): - return 1 - - sqlite_pgmq_broker.declare_actor(do_work) - message = do_work.send() - - assert middleware.before_messages == [(message, None)] - assert middleware.after_messages == [(message, None, None)] - - -def test_pgmq_broker_consume_fails_when_queue_was_not_declared(sqlite_pgmq_broker): - with pytest.raises(QueueNotFound): - sqlite_pgmq_broker.consume("default") - - def test_pgmq_consumer_uses_notification_path_when_listener_is_available(monkeypatch): def _fake_start_listener(self): self._listener_available = True @@ -459,8 +359,3 @@ def _consume_once(): consumer.ack(consumed) finally: consumer.close() - - -def test_pgmq_broker_join_is_not_implemented_until_jalon_three(sqlite_pgmq_broker): - with pytest.raises(NotImplementedError, match="jalon 3"): - sqlite_pgmq_broker.join("default") From 3999582e9b0b72cdc23698486601ebec224b55dc Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 28 May 2026 11:05:12 +0200 Subject: [PATCH 06/44] fix(pgmq): support encoder --- remoulade/brokers/pgmq.py | 40 +++++++++++-------------------------- tests/test_pgmq.py | 42 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 28 deletions(-) diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 65dfd5657..ba7c354d4 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -24,8 +24,8 @@ import psycopg from pgmq import SQLAlchemyPGMQueue +from pgmq.messages import Message as PgmqQueueMessage from psycopg import sql as psycopg_sql -from pydantic import BaseModel, ConfigDict, ValidationError from sqlalchemy import bindparam, create_engine, text from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.engine import make_url @@ -38,17 +38,8 @@ if TYPE_CHECKING: from ..middleware import Middleware - PgmqPayload = dict[str, Any] - -class _PgmqQueueMessage(BaseModel): - model_config = ConfigDict(extra="ignore", from_attributes=True) - - msg_id: int - message: PgmqPayload - - DELAYED_SEND_SQL = text( """ SELECT * @@ -219,7 +210,7 @@ def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeou self.queue_name = queue_name self.prefetch = max(prefetch, 1) self.timeout = max(timeout, 0) - self.messages: deque[_PgmqQueueMessage] = deque() + self.messages: deque[PgmqQueueMessage] = deque() self._notify_event = Event() self._listener_stop = Event() self._listener_thread: Thread | None = None @@ -241,24 +232,15 @@ def poll_interval_ms(self) -> int: def wait_timeout_seconds(self) -> float: return self.timeout / 1000 if self.timeout > 0 else 0 - def _normalize_messages(self, messages: Any | list[Any] | None) -> list[_PgmqQueueMessage]: + def _normalize_messages(self, messages: PgmqQueueMessage | list[PgmqQueueMessage] | None) -> list[PgmqQueueMessage]: if messages is None: return [] - raw_messages = messages if isinstance(messages, list) else [messages] - normalized_messages: list[_PgmqQueueMessage] = [] - for message in raw_messages: - try: - normalized_messages.append(_PgmqQueueMessage.model_validate(message, from_attributes=True)) - except ValidationError as exc: - raise UnsupportedMessageEncoding( - "PGMQ messages must expose 'msg_id' and a JSONB-serializable object in 'message'." - ) from exc - return normalized_messages - - def _read_immediate(self) -> list[_PgmqQueueMessage]: + return [messages] if not isinstance(messages, list) else messages + + def _read_immediate(self) -> list[PgmqQueueMessage]: return self._normalize_messages(self.client.read(self.queue_name, qty=self.prefetch)) - def _read_with_poll(self) -> list[_PgmqQueueMessage]: + def _read_with_poll(self) -> list[PgmqQueueMessage]: return self._normalize_messages( self.client.read_with_poll( self.queue_name, @@ -316,7 +298,7 @@ def _run_listener(self) -> None: def ack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") - self.client.delete(self.queue_name, message._pgmq_message.msg_id) + self.client.archive(self.queue_name, message._pgmq_message.msg_id) def nack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): @@ -363,13 +345,15 @@ def close(self) -> None: class _PgmqMessage(MessageProxy): - def __init__(self, pgmq_message: _PgmqQueueMessage) -> None: + def __init__(self, pgmq_message: PgmqQueueMessage) -> None: payload = pgmq_message.message if not isinstance(payload, dict): raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") try: - message: Message = Message(**payload) + # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder) + # can rehydrate actor args/kwargs to their typed schemas. + message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8")) except TypeError as exc: raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index a06de7b92..046b2981b 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -4,6 +4,7 @@ from unittest.mock import Mock import pytest +from pydantic import BaseModel from sqlalchemy import JSON, Column, Integer, MetaData, Table, func from sqlalchemy.sql import select, text @@ -251,6 +252,47 @@ def _wait_and_stop(_timeout): consumer.close() +def test_pgmq_consumer_decodes_payload_with_global_encoder(monkeypatch, pydantic_encoder): + class InputSchema(BaseModel): + value: int + + def _fake_start_listener(self): + self._listener_available = False + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + remoulade.set_broker(broker) + broker.client.validate_queue_name = Mock() + broker.client.create_queue = Mock() + broker.client.enable_notify = Mock() + + @remoulade.actor(actor_name="typed.actor", queue_name="default") + def typed_actor(payload: InputSchema): + return payload.value + + broker.declare_actor(typed_actor) + + message = Message( + queue_name="default", + actor_name="typed.actor", + args=(InputSchema(value=42),), + kwargs={}, + options={}, + ) + payload = _expected_payload(message) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=1, message=payload)]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumed = next(consumer) + + assert consumed is not None + assert isinstance(consumed.args[0], InputSchema) + assert consumed.args[0].value == 42 + consumer.close() + + @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_consumer_reads_messages_and_acks_with_delete(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(42,), kwargs={}, options={}) From 2221af28b3f11b3e5546b3950c3af0853376ea54 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 28 May 2026 14:32:51 +0200 Subject: [PATCH 07/44] feat(pgmq-broker): add join method --- remoulade/brokers/pgmq.py | 59 ++++++++++++++++++++++++++++++++++++--- tests/test_pgmq.py | 35 ++++++++++++++++++++--- 2 files changed, 86 insertions(+), 8 deletions(-) diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index ba7c354d4..1e5fa8ff6 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -15,6 +15,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . import json +import time from collections import deque from collections.abc import Iterable, Iterator from contextlib import contextmanager @@ -26,13 +27,13 @@ from pgmq import SQLAlchemyPGMQueue from pgmq.messages import Message as PgmqQueueMessage from psycopg import sql as psycopg_sql -from sqlalchemy import bindparam, create_engine, text +from sqlalchemy import Column, Integer, MetaData, Table, bindparam, create_engine, func, select, text from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.engine import make_url from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker from ..broker import Broker, Consumer, MessageProxy -from ..errors import QueueNotFound, UnsupportedMessageEncoding +from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding from ..message import Message if TYPE_CHECKING: @@ -199,8 +200,58 @@ def flush_all(self) -> None: for queue_name in self.queues: self.flush(queue_name) - def join(self, queue_name: str, *, timeout: int | None = None) -> None: - raise NotImplementedError("PgmqBroker.join() will be implemented in jalon 3.") + def _count_enqueued_messages(self, queue_name: str) -> int: + queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") + count_query = select(func.count()).select_from(queue_table) + + with self.sessionmaker.begin() as session: + return session.execute(count_query).scalar_one() + + def join( + self, + queue_name: str, + min_successes: int = 10, + idle_time: int = 100, + *, + timeout: int | None = None, + ) -> None: + """Wait for all the messages on the given queue to be processed. + + This method checks the full PGMQ queue table and therefore waits for + all states: visible messages, invisible in-flight messages and native + delayed messages. + + Parameters: + queue_name(str): The queue to wait on. + min_successes(int): The minimum number of times the queue should be + observed as empty. + idle_time(int): The number of milliseconds to wait between checks. + timeout(Optional[int]): The max amount of time, in milliseconds, to + wait on this queue. + + Raises: + QueueNotFound: If the given queue was never declared. + QueueJoinTimeout: When the timeout elapses. + """ + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + + deadline = time.monotonic() + timeout / 1000 if timeout is not None else None + successes = 0 + + while successes < min_successes: + if deadline and time.monotonic() >= deadline: + raise QueueJoinTimeout(queue_name) + + total_messages = self._count_enqueued_messages(queue_name) + if total_messages == 0: + successes += 1 + if successes >= min_successes: + return + else: + successes = 0 + + time.sleep(idle_time / 1000) class _PgmqConsumer(Consumer): diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 046b2981b..8ee219329 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -9,7 +9,7 @@ from sqlalchemy.sql import select, text import remoulade -from remoulade import Message, Worker +from remoulade import Message, QueueJoinTimeout, Worker from remoulade.brokers.pgmq import PgmqBroker @@ -362,9 +362,7 @@ def do_work(value): worker.start() try: do_work.send_with_options(args=(3,), delay=150) - deadline = time.monotonic() + 10 - while not seen and time.monotonic() < deadline: - time.sleep(0.05) + pgmq_broker.join(do_work.queue_name, timeout=10_000) assert seen == [3] worker.join() finally: @@ -373,6 +371,35 @@ def do_work(value): assert not _queue_exists(pgmq_broker, "default.DQ") +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_broker_join_times_out_while_processing_invisible_message(pgmq_broker): + started = threading.Event() + release = threading.Event() + + @remoulade.actor + def do_work(): + started.set() + release.wait(timeout=5) + + pgmq_broker.declare_actor(do_work) + + worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=1) + worker.start() + try: + do_work.send() + assert started.wait(timeout=2) + + with pytest.raises(QueueJoinTimeout): + pgmq_broker.join(do_work.queue_name, timeout=100) + + release.set() + pgmq_broker.join(do_work.queue_name, timeout=5_000) + worker.join() + finally: + release.set() + worker.stop() + + @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_consumer_listener_wakes_on_enqueue_with_listen_notify(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(99,), kwargs={}, options={}) From 4c7df94347dd6cd0e740860f4d09c6a72a2e6a19 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 28 May 2026 14:50:14 +0200 Subject: [PATCH 08/44] feat(api): remove superbowl route --- remoulade/api/main.py | 154 +--------------- tests/state/test_state_api.py | 335 +++------------------------------- tests/test_cancel.py | 36 ---- tests/test_scheduler.py | 138 -------------- 4 files changed, 34 insertions(+), 629 deletions(-) diff --git a/remoulade/api/main.py b/remoulade/api/main.py index b0a9831d2..3efd151f3 100644 --- a/remoulade/api/main.py +++ b/remoulade/api/main.py @@ -1,158 +1,14 @@ """This file describe the API to get the state of messages""" -import sys -from typing import Any, TypedDict - from flask import Flask -from flask_apispec import marshal_with -from marshmallow import Schema, ValidationError, fields, validate, validates_schema +from marshmallow import ValidationError from werkzeug.exceptions import HTTPException -import remoulade -from remoulade import get_broker -from remoulade.errors import NoResultBackend, NoStateBackend, RemouladeError -from remoulade.result import Result -from remoulade.results import ResultMissing +from remoulade.errors import RemouladeError -from .apispec import add_swagger, validate_schema -from .scheduler import scheduler_bp, scheduler_routes -from .state import messages_bp, messages_routes +from .apispec import add_swagger app = Flask(__name__) -app.register_blueprint(scheduler_bp) -app.register_blueprint(messages_bp) - - -class MessageSchema(Schema): - """ - Class to validate post data in /messages - """ - - actor_name = fields.Str(validate=validate.Length(min=1), required=True) - args = fields.List(fields.Raw(), allow_none=True) - kwargs = fields.Dict(allow_none=True) - options = fields.Dict(allow_none=True) - delay = fields.Float(validate=validate.Range(min=1), allow_none=True) - - @validates_schema - def validate_actor_name(self, data, **kwargs): - actor_name = data.get("actor_name") - if actor_name not in remoulade.get_broker().actors: - raise ValidationError(f"No actor named {actor_name} exists") - - -class ResponseSchema(Schema): - result = fields.Raw(allow_none=True) - error = fields.Str(allow_none=True) - - -class ArgsSchema(Schema): - name = fields.Str() - type = fields.Str() - default = fields.Str(allow_none=True) - - -class ActorSchema(Schema): - name = fields.Str() - queue_name = fields.Str() - alternative_queues = fields.List(fields.Str()) - priority = fields.Int() - args = fields.List(fields.Nested(ArgsSchema)) - - -class ActorResponseSchema(Schema): - result = fields.List(fields.Nested(ActorSchema)) - - -class OptionsResponseSchema(Schema): - options = fields.List(fields.Str()) - - -@app.route("/messages/cancel/", methods=["POST"]) -@marshal_with(ResponseSchema) -def cancel_message(message_id): - broker = get_broker() - cancel_backend = broker.get_cancel_backend() - try: - # If a single message in a composition is canceled, we cancel the whole composition - state_backend = broker.get_state_backend() - states_count = state_backend.get_states_count(selected_composition_ids=[message_id]) - state = state_backend.get_state(message_id) - if states_count == 0 and state is None: - raise ValidationError(f"This message id or composition id {message_id} does not exist.") - if state and state.composition_id: - message_id = state.composition_id - except NoStateBackend: - pass - cancel_backend.cancel([message_id]) - return {"result": "ok"} - - -@app.route("/messages/requeue/", methods=["POST"]) -@marshal_with(ResponseSchema) -def requeue_message(message_id): - broker = get_broker() - backend = broker.get_state_backend() - state = backend.get_state(message_id) - if state is None: - raise ValidationError(f"No message with id {message_id} exists") - actor = broker.get_actor(state.actor_name) - payload = {"args": state.args, "kwargs": state.kwargs} - pipe_target = state.options.get("pipe_target") - if pipe_target is None: - actor.send_with_options(**payload, **state.options) - return {"result": "ok"} - else: - return {"error": "requeue message in a pipeline not supported"}, 400 - - -@app.route("/messages/result/") -@marshal_with(ResponseSchema) -def get_results(message_id): - from ..message import get_encoder - - max_size = 1e4 - try: - result = Result[Any](message_id=message_id).get() - encoded_result = get_encoder().encode(result).decode("utf-8") - size_result = sys.getsizeof(encoded_result) - if size_result >= max_size: - encoded_result = f"The result is too big {size_result / 1e6}M" - return {"result": encoded_result} - except ResultMissing: - return {"result": "result is missing"} - except NoResultBackend: - return {"result": "no result backend"} - except (UnicodeDecodeError, TypeError): - return {"result": "non serializable result"} - - -@app.route("/messages", methods=["POST"]) -@marshal_with(ResponseSchema) -@validate_schema(MessageSchema) -def enqueue_message(**kwargs): - actor = get_broker().get_actor(kwargs.pop("actor_name")) - options = kwargs.pop("options") or {} - actor.send_with_options(**kwargs, **options) - return {"result": "ok"} - - -class GroupMessagesT(TypedDict): - group_id: str - messages: list[dict] - - -@app.route("/actors") -@marshal_with(ActorResponseSchema) -def get_actors(): - return {"result": [actor.as_dict() for actor in get_broker().actors.values()]} - - -@app.route("/options") -@marshal_with(OptionsResponseSchema) -def get_options(): - broker = get_broker() - return {"options": list(broker.actor_options)} @app.errorhandler(RemouladeError) @@ -170,6 +26,4 @@ def validation_error(e): return {"error": e.normalized_messages()}, 400 -routes = [cancel_message, requeue_message, get_results, enqueue_message, get_actors, get_options] - -add_swagger(app, {"main": routes, "scheduler": scheduler_routes, "messages": messages_routes}) +add_swagger(app, {"main": []}) diff --git a/tests/state/test_state_api.py b/tests/state/test_state_api.py index 0911a4bdf..34d8a165e 100644 --- a/tests/state/test_state_api.py +++ b/tests/state/test_state_api.py @@ -1,294 +1,42 @@ import datetime -import json -from datetime import date -from operator import itemgetter -from random import choice -from unittest import mock -from unittest.mock import MagicMock -from zoneinfo import ZoneInfo import pytest from dateutil.parser import parse -import remoulade -from remoulade import set_scheduler from remoulade.api.main import app -from remoulade.cancel import Cancel -from remoulade.message import Message -from remoulade.scheduler import ScheduledJob from remoulade.state import State, StateStatusesEnum -class TestMessageStateAPI: - """Class Responsible to do the test of the API""" - - def test_no_messages(self, stub_broker, state_middleware, api_client): - res = api_client.post("/messages/states") - assert res.status_code == 200 - assert len(res.json["data"]) == 0 - - def test_invalid_url(self, stub_broker, state_middleware, api_client): - res = api_client.get("/invalid_url/") - assert res.status_code == 404 - - def test_invalid_state_message(self, stub_broker, api_client): - res = api_client.post( - "/messages/states", data=json.dumps({"status": "invalid_state"}), content_type="application/json" - ) - assert res.status_code == 400 - - @pytest.mark.parametrize("status", [StateStatusesEnum.Skipped, StateStatusesEnum.Success]) - def test_get_state_by_name(self, status, stub_broker, state_middleware, api_client): - state = State("1", status) - state_middleware.backend.set_state(state, ttl=1000) - args = {"selected_statuses": [status.value]} - res = api_client.post("/messages/states", data=json.dumps(args), content_type="application/json") - assert res.json == {"count": 1, "data": [state.as_dict()]} - - def test_get_by_message_id(self, stub_broker, state_middleware, api_client): - message_id = "1" - state = State(message_id, StateStatusesEnum.Pending) - state_middleware.backend.set_state(state, ttl=1000) - res = api_client.get(f"/messages/states/{message_id}") - assert res.json == state.as_dict() - - @pytest.mark.parametrize("n", [0, 10, 50, 100]) - def test_get_states_by_name(self, n, stub_broker, state_middleware, api_client): - # generate random test - random_states = [] - for i in range(n): - random_state = State(f"id{i}", choice(list(StateStatusesEnum))) # random state - random_states.append(random_state.as_dict()) - state_middleware.backend.set_state(random_state, ttl=1000) - - # check storage - res = api_client.post("/messages/states", data=json.dumps({"size": 100}), content_type="application/json") - states = res.json["data"] - assert len(states) == n - # check integrity - for state in states: - assert state in random_states - - def test_request_cancel(self, stub_broker, cancel_backend): - with app.test_client() as client: - message_id = "1" - stub_broker.add_middleware(Cancel(backend=cancel_backend)) - client.post(f"/messages/cancel/{message_id}") - assert cancel_backend.is_canceled(message_id, None) - - def test_request_cancel_without_backend(self, stub_broker, api_client): - res = api_client.post("/messages/cancel/{}".format("id")) - # if there is not cancel backend should raise an error - assert res.json["error"] == "The default broker doesn't have a cancel backend." - assert res.status_code == 500 - - def test_no_scheduler(self, stub_broker, api_client): - set_scheduler(None) - res = api_client.get("/scheduled/jobs") - assert res.status_code == 400 - - def test_scheduled_jobs(self, scheduler, api_client, do_work, frozen_datetime): - timezone = ZoneInfo("Europe/Paris") - scheduler.schedule = [ - ScheduledJob( - actor_name=do_work.actor_name, daily_time=(datetime.datetime.now(timezone)).time(), tz="Europe/Paris" - ) - ] - scheduler.sync_config() - - res = api_client.get("/scheduled/jobs") - jobs = res.json["result"] - assert jobs == [ - { - "hash": jobs[0]["hash"], - "actor_name": "do_work", - "args": [], - "daily_time": "01:00:00", - "enabled": True, - "interval": 86400, - "iso_weekday": None, - "kwargs": {}, - "last_queued": None, - "tz": "Europe/Paris", - } - ] - - def test_enqueue_message(self, stub_broker, do_work, api_client): - data = { - "actor_name": "do_work", - "args": ["1", "2"], - "kwargs": {}, - "options": {"time_limit": 1000}, - "delay": None, - } - stub_broker.enqueue = MagicMock() - stub_broker.join(do_work.queue_name) - res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") - message = Message( - queue_name="default", - actor_name="do_work", - args=("1", "2"), - kwargs={}, - options={"time_limit": 1000}, - message_id=mock.ANY, - message_timestamp=mock.ANY, - ) - assert stub_broker.enqueue.call_count == 1 - assert stub_broker.enqueue.call_args == mock.call(message, delay=None) - assert res.status_code == 200 - - @pytest.mark.parametrize( - "actor_name,error", - [ - (None, "Field may not be null."), - ("", "Shorter than minimum length 1."), - (111, "Not a valid string."), - ([], "Not a valid string."), - ], - ) - def test_invalid_actor_name_to_enqueue(self, actor_name, error, stub_broker, do_work, api_client): - data = { - "actor_name": actor_name, - } - res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") - validation_error = res.json["error"] - assert validation_error["actor_name"] == [error] - assert res.status_code == 400 - - @pytest.mark.parametrize( - "delay,error", [(-1, "Must be greater than or equal to 1."), ("str", "Not a valid number.")] - ) - def test_invalid_delay_to_enqueue(self, delay, error, stub_broker, do_work, api_client): - data = { - "actor_name": "do_work", - "delay": delay, - } - res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") - validation_error = res.json["error"] - assert validation_error["delay"] == [error] - assert res.status_code == 400 - - def test_enqueue_invalid_actor_name(self, stub_broker, api_client): - data = {"actor_name": "invalid_actor"} - res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") - - assert res.status_code == 400 - - def test_get_declared_actors(self, stub_broker, do_work, api_client): - @remoulade.actor(queue_name="foo", priority=10) - def do_job(): - pass - - stub_broker.declare_actor(do_job) - res = api_client.get("/actors") - expected = [ - {"name": "do_work", "priority": 0, "queue_name": "default"}, - {"name": "do_job", "priority": 10, "queue_name": "foo"}, - ].sort(key=itemgetter("name")) - assert res.json["result"].sort(key=itemgetter("name")) == expected - - def test_filter_messages(self, stub_broker, api_client): - state = State("some_message_id") - state_backend = stub_broker.get_state_backend() - state_backend.set_state(state, ttl=1000) - data = {"sort_column": "message_id", "selected_message_ids": ["some_message_id"]} - res = api_client.post("/messages/states", data=json.dumps(data), content_type="application/json") - assert res.json == {"count": 1, "data": [state.as_dict()]} - - @pytest.mark.parametrize("offset", [0, 1, 5, 100]) - def test_get_states_offset(self, offset, stub_broker, api_client, state_middleware): - for i in range(0, 10): - state_middleware.backend.set_state(State(f"id{i}"), ttl=1000) - res = api_client.post( - "/messages/states", data=json.dumps({"offset": offset, "size": 50}), content_type="application/json" - ) - if offset >= 10: - assert res.json["data"] == [] - else: - assert len(res.json["data"]) + offset == 10 - - @pytest.mark.parametrize("size", [1, 5, 100]) - def test_get_states_page_size(self, size, stub_broker, api_client, state_middleware): - for i in range(0, 10): - state_middleware.backend.set_state(State(f"id{i}"), ttl=1000) - res = api_client.post("/messages/states", data=json.dumps({"size": size}), content_type="application/json") - if size >= 10: - assert res.json["count"] == 10 - else: - assert len(res.json["data"]) == size - - def test_not_raise_error_with_pickle_and_non_serializable( - self, pickle_encoder, stub_broker, api_client, state_middleware - ): - state = State("id1", args=["some_status", date(2020, 12, 12)]) - state_middleware.backend.set_state(state, ttl=1000) - res = api_client.post("messages/states") - assert res.json == { - "count": 1, - "data": [{"args": ["some_status", "Sat, 12 Dec 2020 00:00:00 GMT"], "message_id": "id1"}], - } - - def test_requeue_message(self, stub_broker, do_work, api_client, state_middleware): - stub_broker.enqueue = MagicMock() - state = State("id1", StateStatusesEnum.Success, actor_name="do_work", options={"time_limit": 1000}) - state_middleware.backend.set_state(state, ttl=1000) - res = api_client.post("messages/requeue/id1") - message = Message( - queue_name="default", - actor_name="do_work", - args=(), - kwargs={}, - options={"time_limit": 1000}, - message_id=mock.ANY, - message_timestamp=mock.ANY, - ) - assert stub_broker.enqueue.call_count == 1 - assert stub_broker.enqueue.call_args == mock.call(message, delay=None) - assert res.status_code == 200 - - def test_requeue_invalid_id(self, stub_broker, api_client, state_middleware): - res = api_client.post("messages/requeue/invalid_id") - assert res.status_code == 400 - - def test_requeue_message_with_pipeline(self, stub_broker, do_work, api_client, state_middleware): - stub_broker.enqueue = MagicMock() - state = State("id1", StateStatusesEnum.Success, actor_name="do_work", options={"pipe_target": "some_pipe"}) - state_middleware.backend.set_state(state, ttl=1000) - res = api_client.post("messages/requeue/id1") - assert res.json == {"error": "requeue message in a pipeline not supported"} - assert res.status_code == 400 - - def test_get_result_message(self, stub_broker, stub_worker, api_client, state_middleware, result_middleware): - @remoulade.actor(store_results=True) - def do_work(): - return 42 - - stub_broker.declare_actor(do_work) - message = do_work.send() - stub_broker.join(do_work.queue_name) - stub_worker.join() - res = api_client.get(f"/messages/result/{message.message_id}") - assert res.json["result"] == "42" - - def test_no_result_backend(self, stub_broker, stub_worker, api_client, do_work): - message = do_work.send() - res = api_client.get(f"/messages/result/{message.message_id}") - assert res.json["result"] == "no result backend" - - def test_result_not_serializable( - self, pickle_encoder, stub_broker, stub_worker, api_client, state_middleware, result_middleware - ): - @remoulade.actor(store_results=True) - def do_work(): - return {"date": date(2020, 10, 10)} - - stub_broker.declare_actor(do_work) - message = do_work.send() - stub_broker.join(do_work.queue_name) - stub_worker.join() - res = api_client.get(f"/messages/result/{message.message_id}") - assert res.json["result"] == "non serializable result" - +@pytest.fixture +def local_api_client(): + with app.test_client() as client: + yield client + + +@pytest.mark.parametrize( + "method,path", + [ + ("post", "/messages/states"), + ("get", "/messages/states/some-id"), + ("delete", "/messages/states"), + ("post", "/messages/cancel/some-id"), + ("post", "/messages/requeue/some-id"), + ("get", "/messages/result/some-id"), + ("post", "/messages"), + ("get", "/actors"), + ("get", "/options"), + ("get", "/scheduled/jobs"), + ("post", "/scheduled/jobs"), + ("put", "/scheduled/jobs/some-id"), + ("delete", "/scheduled/jobs/some-id"), + ], +) +def test_superbowl_routes_are_removed(local_api_client, method, path): + response = getattr(local_api_client, method)(path) + assert response.status_code == 404 + + +class TestMessageStateBackend: def test_select_actors(self, stub_broker, postgres_state_middleware): backend = postgres_state_middleware.backend for i in range(2): @@ -357,26 +105,3 @@ def test_clean_not_started(self, stub_broker, postgres_state_middleware): res = backend.get_states() assert len(res) == 1 assert res[0].message_id == "id0" - - def test_clean_route(self, stub_broker, postgres_state_middleware): - client = app.test_client() - backend = postgres_state_middleware.backend - backend.set_state(State("id0")) - - assert len(backend.get_states()) == 1 - client.delete("/messages/states") - assert len(backend.get_states()) == 0 - - def test_cant_sort_by_args_kwargs_options(self, stub_broker, state_middleware, api_client): - res = api_client.post( - "/messages/states", data=json.dumps({"sort_column": "args"}), content_type="application/json" - ) - assert res.status_code == 400 - res = api_client.post( - "/messages/states", data=json.dumps({"sort_column": "kwargs"}), content_type="application/json" - ) - assert res.status_code == 400 - res = api_client.post( - "/messages/states", data=json.dumps({"sort_column": "options"}), content_type="application/json" - ) - assert res.status_code == 400 diff --git a/tests/test_cancel.py b/tests/test_cancel.py index af20f745e..fec9c9b95 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -89,32 +89,6 @@ def do_fail(arg): assert cancel_backend.is_canceled("", "mocked_composition_id") == cancel -def test_compositions_are_canceled_on_message_cancel(stub_broker, cancel_backend, state_middleware, api_client): - # Given a cancel backend - # And a broker with the cancel middleware - stub_broker.add_middleware(Cancel(backend=cancel_backend)) - - # And an actor - @remoulade.actor - def do_work(arg=None): - return 1 - - # And those actors are declared - stub_broker.declare_actor(do_work) - - message_to_cancel = do_work.message() - - # And a group that I enqueue - g = group([message_to_cancel | do_work.message(), do_work.message()], cancel_on_error=True) - g.run() - - # When I cancel a message of this group - api_client.post("messages/cancel/" + message_to_cancel.message_id) - - # The whole composition should be canceled - assert cancel_backend.is_canceled("", g.group_id) - - def test_cannot_cancel_on_error_if_no_cancel(stub_broker): # Given an actor @remoulade.actor() @@ -188,13 +162,3 @@ def do_work(): # It messages should not have runMessageSchema assert calls_count == 0 - - -def test_raise_error_if_unknown_id(stub_broker, cancel_backend, api_client): - # Given a cancel middleware - stub_broker.add_middleware(Cancel(backend=cancel_backend)) - - # If I try to cancel a id that is not a id of a message or composition - res = api_client.post("messages/cancel/invalid_id") - - assert res.status_code == 400 diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 2ebea8e75..e0d4abe45 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1,5 +1,4 @@ import datetime -import json import threading import time from zoneinfo import ZoneInfo @@ -260,143 +259,6 @@ def test_delete_job(scheduler): assert scheduler.get_redis_schedule() == {} -def test_get_scheduled_jobs(scheduler, api_client): - scheduler.schedule = [ScheduledJob(actor_name="do_work")] - scheduler.sync_config() - res = api_client.get("/scheduled/jobs") - assert res.status_code == 200 - assert res.json == { - "result": [ - { - "hash": scheduler.schedule[0].get_hash(), - "actor_name": "do_work", - "args": [], - "daily_time": None, - "enabled": True, - "interval": 86400, - "iso_weekday": None, - "kwargs": {}, - "last_queued": None, - "tz": "UTC", - } - ] - } - - -def test_api_add_job(scheduler, api_client, do_work): - res = api_client.post( - "/scheduled/jobs", data=json.dumps({"actor_name": "do_work"}), content_type="application/json" - ) - assert res.status_code == 200 - assert res.json == { - "result": [ - { - "hash": res.json["result"][0]["hash"], - "actor_name": "do_work", - "args": [], - "daily_time": None, - "enabled": True, - "interval": 86400, - "iso_weekday": None, - "kwargs": {}, - "last_queued": None, - "tz": "UTC", - } - ] - } - - -def test_api_delete_job(scheduler, api_client): - scheduler.schedule = [ScheduledJob(actor_name="do_work")] - scheduler.sync_config() - res = api_client.delete(f"/scheduled/jobs/{scheduler.schedule[0].get_hash()}") - assert res.status_code == 200 - assert res.json == {"result": []} - - -def test_update_job(scheduler, api_client, do_work): - scheduler.schedule = [ScheduledJob(actor_name="do_work")] - scheduler.sync_config() - res = api_client.put( - f"/scheduled/jobs/{scheduler.schedule[0].get_hash()}", - data=json.dumps({"actor_name": "do_work", "enabled": False}), - content_type="application/json", - ) - assert res.status_code == 200 - assert res.json == { - "result": [ - { - "hash": res.json["result"][0]["hash"], - "actor_name": "do_work", - "args": [], - "daily_time": None, - "enabled": False, - "interval": 86400, - "iso_weekday": None, - "kwargs": {}, - "last_queued": None, - "tz": "UTC", - } - ] - } - - -def test_api_update_jobs(scheduler, api_client, do_work): - scheduler.schedule = [ScheduledJob(actor_name="do_work"), ScheduledJob(actor_name="do_other_work")] - scheduler.sync_config() - res = api_client.put( - "/scheduled/jobs", - data=json.dumps( - { - "jobs": { - scheduler.schedule[0].get_hash(): {"actor_name": "do_work", "enabled": False}, - scheduler.schedule[1].get_hash(): {"actor_name": "do_work", "enabled": False, "interval": "55"}, - } - } - ), - content_type="application/json", - ) - assert res.status_code == 200 - assert len(res.json["result"]) == 2 - for i in range(2): - assert not res.json["result"][i]["enabled"] - - -def test_daily_time_wrong_interval(scheduler, api_client): - res = api_client.post( - "/scheduled/jobs", - data=json.dumps({"actor_name": "do_work", "daily_time": "00:00", "interval": 1000}), - content_type="application/json", - ) - assert res.status_code == 400 - - -def test_invalid_timezone(scheduler, api_client): - res = api_client.post( - "/scheduled/jobs", - data=json.dumps({"actor_name": "do_work", "tz": "invalid_tz"}), - content_type="application/json", - ) - assert res.status_code == 400 - - -def test_invalid_actor_name(scheduler, api_client): - res = api_client.post( - "/scheduled/jobs", data=json.dumps({"actor_name": "invalid_actor"}), content_type="application/json" - ) - assert res.status_code == 400 - - -def test_tz_aware_last_queued(scheduler, api_client, do_work): - res = api_client.post( - "scheduled/jobs", - data=json.dumps({"actor_name": "do_work", "last_queued": "2020-10-10 10:00:00Z"}), - content_type="application/json", - ) - - assert res.status_code == 400 - - class InputArg(BaseModel): data: str From 39443d24acad11a7b821c658c27830f4d8c2836e Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 28 May 2026 15:46:37 +0200 Subject: [PATCH 09/44] fix(pgmq-broker): url format --- remoulade/brokers/pgmq.py | 111 +++++++++++++++++++++++++++++++++----- tests/test_pgmq.py | 92 +++++++++++++++++++++++++++---- 2 files changed, 180 insertions(+), 23 deletions(-) diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 1e5fa8ff6..15511ec3c 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -1,6 +1,6 @@ # This file is a part of Remoulade. # -# Copyright (C) 2017,2018 CLEARTYPE SRL +# Copyright (C) 2017,2018 WIREMIND SAS # # Remoulade is free software; you can redistribute it and/or modify it # under the terms of the GNU Lesser General Public License as published by @@ -20,7 +20,7 @@ from collections.abc import Iterable, Iterator from contextlib import contextmanager from datetime import UTC, datetime, timedelta -from threading import Event, Thread, local +from threading import Event, Lock, Thread, local from typing import TYPE_CHECKING, Any import psycopg @@ -29,7 +29,6 @@ from psycopg import sql as psycopg_sql from sqlalchemy import Column, Integer, MetaData, Table, bindparam, create_engine, func, select, text from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.engine import make_url from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker from ..broker import Broker, Consumer, MessageProxy @@ -67,8 +66,31 @@ def __init__( archive_partition_interval: int | str = "1 day", archive_retention_interval: int | str = "7 days", listen_notify_enabled: bool = True, + visibility_timeout: int = 30, + heartbeat_interval: float = 10.0, ) -> None: + """Initialize a PostgreSQL-backed broker using the PGMQ extension. + + Parameters: + url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg. + middleware(list[Middleware] | None): Middleware stack applied to this broker. + group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. + partition_archive_on_queue_init(bool): If True, creates partitioned queues when declaring queues. + archive_partition_interval(int | str): Partition interval passed to PGMQ when creating partitioned queues. + archive_retention_interval(int | str): Retention interval passed to PGMQ for archive partitions. + listen_notify_enabled(bool): If True, enables LISTEN/NOTIFY to wake consumers when new messages arrive. + visibility_timeout(int): Message visibility timeout in seconds after read; must be greater than 0. + heartbeat_interval(float): Heartbeat interval in seconds used to extend in-flight message visibility; + must be greater than 0 and lower than visibility_timeout. + + Raises: + ValueError: If visibility_timeout or heartbeat_interval values are invalid. + """ super().__init__(middleware=middleware) + if visibility_timeout <= 0: + raise ValueError("visibility_timeout must be greater than 0") + if heartbeat_interval <= 0 or heartbeat_interval >= visibility_timeout: + raise ValueError("heartbeat_interval must be greater than 0 and lower than visibility_timeout") self.url = url self.state = local() @@ -77,12 +99,14 @@ def __init__( self.archive_partition_interval = archive_partition_interval self.archive_retention_interval = archive_retention_interval self.listen_notify_enabled = listen_notify_enabled + self.visibility_timeout = visibility_timeout + self.heartbeat_interval = heartbeat_interval self.engine = create_engine(self.url, pool_pre_ping=True) self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) self.supports_native_delay = True - self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False) + self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=True, vt=self.visibility_timeout) @contextmanager def tx(self) -> Iterator[None]: @@ -99,10 +123,6 @@ def tx(self) -> Iterator[None]: def _current_connection(self) -> Any | None: return getattr(self.state, "transaction_connection", None) - def _build_listener_conninfo(self) -> str: - # psycopg expects a plain PostgreSQL URL and rejects SQLAlchemy-style driver suffixes. - return make_url(self.url).set(drivername="postgresql").render_as_string(hide_password=False) - def _try_enable_notify(self, queue_name: str) -> None: if not self.listen_notify_enabled: return @@ -256,20 +276,36 @@ def join( class _PgmqConsumer(Consumer): def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int) -> None: + """Initialize a consumer for a PGMQ queue. + + Parameters: + broker(PgmqBroker): Broker instance that owns the queue and database client. + queue_name(str): Name of the declared queue to consume from. + prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1. + timeout(int): Idle wait timeout in milliseconds when polling for messages; values lower than 0 are coerced to + 0. A value of 0 performs non-blocking reads. + """ self.broker = broker self.client = broker.client self.queue_name = queue_name self.prefetch = max(prefetch, 1) self.timeout = max(timeout, 0) + self.visibility_timeout = broker.visibility_timeout + self.heartbeat_interval = broker.heartbeat_interval self.messages: deque[PgmqQueueMessage] = deque() self._notify_event = Event() self._listener_stop = Event() self._listener_thread: Thread | None = None self._listener_connection: psycopg.Connection[Any] | None = None self._listener_available = False + self._heartbeat_stop = Event() + self._heartbeat_thread: Thread | None = None + self._heartbeat_msg_ids_lock = Lock() + self._heartbeat_msg_ids: set[int] = set() if self.broker.listen_notify_enabled: self._start_listener() + self._start_heartbeat() @property def max_poll_seconds(self) -> int: @@ -289,23 +325,60 @@ def _normalize_messages(self, messages: PgmqQueueMessage | list[PgmqQueueMessage return [messages] if not isinstance(messages, list) else messages def _read_immediate(self) -> list[PgmqQueueMessage]: - return self._normalize_messages(self.client.read(self.queue_name, qty=self.prefetch)) + return self._normalize_messages( + self.client.read( + self.queue_name, + vt=self.visibility_timeout, + qty=self.prefetch, + ) + ) def _read_with_poll(self) -> list[PgmqQueueMessage]: return self._normalize_messages( self.client.read_with_poll( self.queue_name, + vt=self.visibility_timeout, qty=self.prefetch, max_poll_seconds=self.max_poll_seconds, poll_interval_ms=self.poll_interval_ms, ) ) + def _start_heartbeat(self) -> None: + self._heartbeat_thread = Thread( + target=self._run_heartbeat, + name=f"pgmq-heartbeat-{self.queue_name}", + daemon=True, + ) + self._heartbeat_thread.start() + + def _run_heartbeat(self) -> None: + while not self._heartbeat_stop.wait(self.heartbeat_interval): + with self._heartbeat_msg_ids_lock: + msg_ids = list(self._heartbeat_msg_ids) + if not msg_ids: + continue + try: + self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout) + except Exception: + self.broker.logger.warning( + "Failed to extend visibility timeout heartbeat for queue %s.", + self.queue_name, + exc_info=True, + ) + + def _register_heartbeat_msg(self, message: "_PgmqMessage") -> None: + with self._heartbeat_msg_ids_lock: + self._heartbeat_msg_ids.add(message._pgmq_message.msg_id) + + def _unregister_heartbeat_msg_id(self, msg_id: int) -> None: + with self._heartbeat_msg_ids_lock: + self._heartbeat_msg_ids.discard(msg_id) + def _start_listener(self) -> None: channel = f"pgmq.q_{self.queue_name}.INSERT" try: - conninfo = self.broker._build_listener_conninfo() - self._listener_connection = psycopg.connect(conninfo, autocommit=True) + self._listener_connection = psycopg.connect(self.broker.url, autocommit=True) self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) self._listener_available = True except Exception: @@ -349,21 +422,30 @@ def _run_listener(self) -> None: def ack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") + self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) self.client.archive(self.queue_name, message._pgmq_message.msg_id) def nack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") + self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) self.client.archive(self.queue_name, message._pgmq_message.msg_id) def requeue(self, messages: Iterable["MessageProxy"]) -> None: msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] if msg_ids: + for msg_id in msg_ids: + self._unregister_heartbeat_msg_id(msg_id) self.client.set_vt(self.queue_name, msg_ids, 0) + def _build_message(self, pgmq_message: PgmqQueueMessage) -> "_PgmqMessage": + message = _PgmqMessage(pgmq_message) + self._register_heartbeat_msg(message) + return message + def __next__(self) -> "_PgmqMessage | None": if self.messages: - return _PgmqMessage(self.messages.popleft()) + return self._build_message(self.messages.popleft()) if self._listener_available: self._notify_event.clear() @@ -378,10 +460,11 @@ def __next__(self) -> "_PgmqMessage | None": return None self.messages.extend(messages) - return _PgmqMessage(self.messages.popleft()) + return self._build_message(self.messages.popleft()) def close(self) -> None: self._listener_stop.set() + self._heartbeat_stop.set() self._notify_event.set() if self._listener_connection is not None: try: @@ -392,6 +475,8 @@ def close(self) -> None: self._listener_connection = None if self._listener_thread is not None: self._listener_thread.join(timeout=1) + if self._heartbeat_thread is not None: + self._heartbeat_thread.join(timeout=1) return diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 8ee219329..19f6c031f 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -1,4 +1,5 @@ import json +import os import threading import time from unittest.mock import Mock @@ -12,6 +13,8 @@ from remoulade import Message, QueueJoinTimeout, Worker from remoulade.brokers.pgmq import PgmqBroker +TEST_PGMQ_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" + def _count_messages(broker, queue_name="default"): queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") @@ -49,7 +52,7 @@ def _expected_payload(message): def test_pgmq_broker_uses_provided_url(): - broker_url = "postgresql://pgmq-user@localhost/pgmq" + broker_url = TEST_PGMQ_URL broker = PgmqBroker(url=broker_url) assert broker.url == broker_url @@ -57,7 +60,7 @@ def test_pgmq_broker_uses_provided_url(): def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): broker = PgmqBroker( - url="postgresql://pgmq-user@localhost/pgmq", + url=TEST_PGMQ_URL, middleware=[], partition_archive_on_queue_init=True, ) @@ -78,7 +81,7 @@ def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): def test_pgmq_broker_uses_non_partitioned_queue_creation_when_disabled(): broker = PgmqBroker( - url="postgresql://pgmq-user@localhost/pgmq", + url=TEST_PGMQ_URL, middleware=[], partition_archive_on_queue_init=False, ) @@ -93,7 +96,7 @@ def test_pgmq_broker_uses_non_partitioned_queue_creation_when_disabled(): def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_queue = Mock() broker.client.enable_notify = Mock() @@ -104,7 +107,7 @@ def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_queue = Mock() broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) @@ -167,7 +170,7 @@ def _fake_start_listener(self): monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) @@ -193,7 +196,7 @@ def _fake_start_listener(self): monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) @@ -209,6 +212,35 @@ def _fake_start_listener(self): broker.client.read.assert_not_called() broker.client.read_with_poll.assert_called_once_with( "default", + vt=30, + qty=1, + max_poll_seconds=1, + poll_interval_ms=200, + ) + consumer.close() + + +def test_pgmq_consumer_uses_broker_visibility_timeout_for_reads(monkeypatch): + def _fake_start_listener(self): + self._listener_available = False + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[], visibility_timeout=17) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(5,), kwargs={}, options={}) + payload = _expected_payload(message) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=5, message=payload)]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumed = next(consumer) + + assert consumed is not None + broker.client.read_with_poll.assert_called_once_with( + "default", + vt=17, qty=1, max_poll_seconds=1, poll_interval_ms=200, @@ -222,7 +254,7 @@ def _fake_start_listener(self): monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(3,), kwargs={}, options={}) @@ -242,9 +274,10 @@ def _wait_and_stop(_timeout): assert consumed is not None assert consumed.message_id == message.message_id - broker.client.read.assert_called_once_with("default", qty=1) + broker.client.read.assert_called_once_with("default", vt=30, qty=1) broker.client.read_with_poll.assert_called_once_with( "default", + vt=30, qty=1, max_poll_seconds=1, poll_interval_ms=200, @@ -252,6 +285,45 @@ def _wait_and_stop(_timeout): consumer.close() +def test_pgmq_consumer_heartbeat_extends_inflight_message_visibility(monkeypatch): + def _fake_start_listener(self): + self._listener_available = False + + monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + + broker = PgmqBroker( + url=TEST_PGMQ_URL, + middleware=[], + visibility_timeout=2, + heartbeat_interval=0.05, + ) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(7,), kwargs={}, options={}) + payload = _expected_payload(message) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=7, message=payload)]) + broker.client.set_vt = Mock() + broker.client.archive = Mock() + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumed = next(consumer) + assert consumed is not None + + deadline = time.monotonic() + 1.0 + while broker.client.set_vt.call_count == 0 and time.monotonic() < deadline: + time.sleep(0.01) + + assert broker.client.set_vt.call_count >= 1 + queue_name, msg_ids, vt = broker.client.set_vt.call_args.args + assert queue_name == "default" + assert msg_ids == [7] + assert vt == 2 + + consumer.ack(consumed) + consumer.close() + + def test_pgmq_consumer_decodes_payload_with_global_encoder(monkeypatch, pydantic_encoder): class InputSchema(BaseModel): value: int @@ -261,7 +333,7 @@ def _fake_start_listener(self): monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url="postgresql://pgmq-user@localhost/pgmq", middleware=[]) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) remoulade.set_broker(broker) broker.client.validate_queue_name = Mock() broker.client.create_queue = Mock() From 23fb4e165ad6258cb040d0e9614e17c2a25758cd Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 29 May 2026 10:57:09 +0200 Subject: [PATCH 10/44] fix(app): works without pgmq extra --- remoulade/brokers/__init__.py | 7 ------- remoulade/brokers/pgmq.py | 2 +- tests/docker-compose.yml | 2 +- tests/state/test_state_api.py | 23 ----------------------- 4 files changed, 2 insertions(+), 32 deletions(-) diff --git a/remoulade/brokers/__init__.py b/remoulade/brokers/__init__.py index 61763f308..1d583ac71 100644 --- a/remoulade/brokers/__init__.py +++ b/remoulade/brokers/__init__.py @@ -14,10 +14,3 @@ # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . - -from .local import LocalBroker -from .pgmq import PgmqBroker -from .rabbitmq import RabbitmqBroker -from .stub import StubBroker - -__all__ = ["LocalBroker", "PgmqBroker", "RabbitmqBroker", "StubBroker"] diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 15511ec3c..927762bb2 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -106,7 +106,7 @@ def __init__( self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) self.supports_native_delay = True - self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=True, vt=self.visibility_timeout) + self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False, vt=self.visibility_timeout) @contextmanager def tx(self) -> Iterator[None]: diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index 2800823c4..12173d68e 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -8,7 +8,7 @@ services: ports: - "5784:5672" postgres: - image: ghcr.io/pgmq/pg18-pgmq:latest + image: ghcr.io/pgmq/pg18-pgmq:v1.10.0 ports: - "5544:5432" environment: diff --git a/tests/state/test_state_api.py b/tests/state/test_state_api.py index 34d8a165e..69b2a13dd 100644 --- a/tests/state/test_state_api.py +++ b/tests/state/test_state_api.py @@ -13,29 +13,6 @@ def local_api_client(): yield client -@pytest.mark.parametrize( - "method,path", - [ - ("post", "/messages/states"), - ("get", "/messages/states/some-id"), - ("delete", "/messages/states"), - ("post", "/messages/cancel/some-id"), - ("post", "/messages/requeue/some-id"), - ("get", "/messages/result/some-id"), - ("post", "/messages"), - ("get", "/actors"), - ("get", "/options"), - ("get", "/scheduled/jobs"), - ("post", "/scheduled/jobs"), - ("put", "/scheduled/jobs/some-id"), - ("delete", "/scheduled/jobs/some-id"), - ], -) -def test_superbowl_routes_are_removed(local_api_client, method, path): - response = getattr(local_api_client, method)(path) - assert response.status_code == 404 - - class TestMessageStateBackend: def test_select_actors(self, stub_broker, postgres_state_middleware): backend = postgres_state_middleware.backend From 2cf1621fdafddb1e2a203c763f94f4c99a8229c9 Mon Sep 17 00:00:00 2001 From: mducros Date: Tue, 2 Jun 2026 15:01:46 +0200 Subject: [PATCH 11/44] fix(pgmq): simplify code --- .github/workflows/tests.yml | 2 +- examples/composition/composition/actors.py | 4 +- local_pgmq_broker.py | 4 +- local_pgmq_consumer.py | 1 - pyproject.toml | 3 +- remoulade/api/state.py | 14 +- remoulade/brokers/pgmq.py | 126 ++++----- remoulade/brokers/rabbitmq.py | 16 +- remoulade/state/backends/__init__.py | 13 +- remoulade/state/backends/postgres.py | 250 ------------------ tests/conftest.py | 31 +-- .../postgres/01-create-pgmq-extension.sql | 2 + tests/middleware/test_message_state.py | 4 - tests/state/test_backend.py | 31 --- tests/state/test_postgres.py | 65 ----- tests/test_pgmq.py | 63 +++-- 16 files changed, 109 insertions(+), 520 deletions(-) delete mode 100644 remoulade/state/backends/postgres.py delete mode 100644 tests/state/test_postgres.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f7a4efb1b..7f73b5cf5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -28,7 +28,7 @@ jobs: ports: - 5784:5672 postgres: - image: ghcr.io/pgmq/pg18-pgmq:latest + image: ghcr.io/pgmq/pg18-pgmq:v1.10.0 ports: - 5544:5432 env: diff --git a/examples/composition/composition/actors.py b/examples/composition/composition/actors.py index 58ad693ef..a3ab5c607 100644 --- a/examples/composition/composition/actors.py +++ b/examples/composition/composition/actors.py @@ -8,7 +8,7 @@ from remoulade.results import Results from remoulade.results.backends import RedisBackend from remoulade.state import MessageState -from remoulade.state.backends import PostgresBackend +from remoulade.state.backends import RedisBackend as RedisStateBackend encoder = PickleEncoder() backend = RedisBackend(encoder=encoder) @@ -16,7 +16,7 @@ broker.add_middleware(Results(backend=backend)) remoulade.set_broker(broker) remoulade.set_encoder(encoder) -remoulade.get_broker().add_middleware(MessageState(backend=PostgresBackend())) +remoulade.get_broker().add_middleware(MessageState(backend=RedisStateBackend())) @remoulade.actor(store_results=True) diff --git a/local_pgmq_broker.py b/local_pgmq_broker.py index d5468a33e..857a96a7e 100644 --- a/local_pgmq_broker.py +++ b/local_pgmq_broker.py @@ -12,10 +12,10 @@ QUEUE_NAME = "default" ACTOR_NAME = "demo.add" + broker = PgmqBroker( url=URL, middleware=[], - listen_notify_enabled=True, ) remoulade.set_broker(broker) broker.declare_queue(QUEUE_NAME) @@ -30,7 +30,7 @@ ) broker.enqueue(msg) -with broker.sessionmaker.begin() as session: +with broker.client.session() as session: row = session.execute(text("SELECT msg_id, message FROM pgmq.q_default ORDER BY msg_id DESC LIMIT 1")).one() print("msg_id:", row.msg_id) print("payload:", row.message) diff --git a/local_pgmq_consumer.py b/local_pgmq_consumer.py index e63f2a550..8ff921b94 100644 --- a/local_pgmq_consumer.py +++ b/local_pgmq_consumer.py @@ -12,7 +12,6 @@ broker = PgmqBroker( url=URL, middleware=[], - listen_notify_enabled=True, ) remoulade.set_broker(broker) broker.declare_queue(QUEUE_NAME) diff --git a/pyproject.toml b/pyproject.toml index 1b6da4a48..b78986d59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,8 +23,7 @@ classifiers = [ rabbitmq = ["amqpstorm>=2.6,<3"] redis = ["redis>=7.0.0"] server = ["flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec"] -postgres = ["sqlalchemy>=1.4.29", "psycopg2>=2.9.11"] -pgmq = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1"] +postgres = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1"] pydantic = ["pydantic>=2.12", "simplejson"] limits = ["limits~=5.3.0"] tracing = ["opentelemetry-api>=1.20"] diff --git a/remoulade/api/state.py b/remoulade/api/state.py index 4bb82f086..91bb868a0 100644 --- a/remoulade/api/state.py +++ b/remoulade/api/state.py @@ -5,7 +5,6 @@ from remoulade import get_broker from remoulade.state import State, StateStatusesEnum -from remoulade.state.backends import PostgresBackend from .apispec import validate_schema @@ -81,17 +80,6 @@ def get_states(**kwargs): return {"data": data, "count": backend.get_states_count(**kwargs)} -@messages_bp.route("/states", methods=["DELETE"]) -@doc(tags=["state"]) -@validate_schema(DeleteSchema) -def clean_states(**kwargs): - backend = get_broker().get_state_backend() - if not isinstance(backend, PostgresBackend): - return {"error": "deleting states is only supported by the PostgresBackend"}, 400 - get_broker().get_state_backend().clean(**kwargs) - return {"result": "ok"} - - @messages_bp.route("/states/") @doc(tags=["state"]) @marshal_with(StateSchema) @@ -103,4 +91,4 @@ def get_state(message_id): return data.as_dict() -messages_routes = [get_states, clean_states, get_state] +messages_routes = [get_states, get_state] diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 927762bb2..61351cb86 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -21,15 +21,13 @@ from contextlib import contextmanager from datetime import UTC, datetime, timedelta from threading import Event, Lock, Thread, local -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import psycopg from pgmq import SQLAlchemyPGMQueue from pgmq.messages import Message as PgmqQueueMessage from psycopg import sql as psycopg_sql -from sqlalchemy import Column, Integer, MetaData, Table, bindparam, create_engine, func, select, text -from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker +from sqlalchemy import Column, Integer, MetaData, Table, func, select from ..broker import Broker, Consumer, MessageProxy from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding @@ -39,17 +37,6 @@ from ..middleware import Middleware PgmqPayload = dict[str, Any] - -DELAYED_SEND_SQL = text( - """ - SELECT * - FROM pgmq.send( - queue_name => :queue_name, - msg => :message, - delay => :delay - ) - """ -).bindparams(bindparam("message", type_=JSONB)) LISTEN_NOTIFY_THROTTLE_MS = 250 @@ -62,10 +49,8 @@ def __init__( url: str, middleware: list["Middleware"] | None = None, group_transaction: bool = False, - partition_archive_on_queue_init: bool = False, - archive_partition_interval: int | str = "1 day", - archive_retention_interval: int | str = "7 days", - listen_notify_enabled: bool = True, + archive_partition_interval: str = "1 day", + archive_retention_interval: str = "7 days", visibility_timeout: int = 30, heartbeat_interval: float = 10.0, ) -> None: @@ -75,10 +60,8 @@ def __init__( url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg. middleware(list[Middleware] | None): Middleware stack applied to this broker. group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. - partition_archive_on_queue_init(bool): If True, creates partitioned queues when declaring queues. archive_partition_interval(int | str): Partition interval passed to PGMQ when creating partitioned queues. archive_retention_interval(int | str): Retention interval passed to PGMQ for archive partitions. - listen_notify_enabled(bool): If True, enables LISTEN/NOTIFY to wake consumers when new messages arrive. visibility_timeout(int): Message visibility timeout in seconds after read; must be greater than 0. heartbeat_interval(float): Heartbeat interval in seconds used to extend in-flight message visibility; must be greater than 0 and lower than visibility_timeout. @@ -95,28 +78,22 @@ def __init__( self.url = url self.state = local() self.group_transaction = group_transaction - self.partition_archive_on_queue_init = partition_archive_on_queue_init self.archive_partition_interval = archive_partition_interval self.archive_retention_interval = archive_retention_interval - self.listen_notify_enabled = listen_notify_enabled self.visibility_timeout = visibility_timeout self.heartbeat_interval = heartbeat_interval - - self.engine = create_engine(self.url, pool_pre_ping=True) - self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine) self.supports_native_delay = True - self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False, vt=self.visibility_timeout) + self.client = SQLAlchemyPGMQueue(conn_string=url, init_extension=False, vt=self.visibility_timeout) + @override @contextmanager def tx(self) -> Iterator[None]: - with self.sessionmaker.begin() as session: - self.state.transaction_session = session - self.state.transaction_connection = session.connection() + with self.client.engine.begin() as connection: + self.state.transaction_connection = connection try: yield finally: - self.state.transaction_session = None self.state.transaction_connection = None @property @@ -124,9 +101,6 @@ def _current_connection(self) -> Any | None: return getattr(self.state, "transaction_connection", None) def _try_enable_notify(self, queue_name: str) -> None: - if not self.listen_notify_enabled: - return - try: self.client.enable_notify( queue_name, @@ -140,33 +114,34 @@ def _try_enable_notify(self, queue_name: str) -> None: exc_info=True, ) + @override def close(self) -> None: - self.engine.dispose() + self.client.dispose() + @override def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PgmqConsumer": if queue_name not in self.queues: raise QueueNotFound(queue_name) return _PgmqConsumer(self, queue_name=queue_name, prefetch=prefetch, timeout=timeout) + @override def declare_queue(self, queue_name: str) -> None: if queue_name in self.queues: return self.client.validate_queue_name(queue_name, conn=self._current_connection) self.emit_before("declare_queue", queue_name) - if self.partition_archive_on_queue_init: - self.client.create_partitioned_queue( - queue_name, - partition_interval=self.archive_partition_interval, - retention_interval=self.archive_retention_interval, - conn=self._current_connection, - ) - else: - self.client.create_queue(queue_name, conn=self._current_connection) + self.client.create_partitioned_queue( + queue_name, + partition_interval=self.archive_partition_interval, + retention_interval=self.archive_retention_interval, + conn=self._current_connection, + ) self._try_enable_notify(queue_name) self.queues[queue_name] = None self.emit_after("declare_queue", queue_name) + @override def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": return message @@ -183,39 +158,30 @@ def _encode_message(self, message: "Message") -> PgmqPayload: return payload - def _send_with_delay(self, queue_name: str, payload: PgmqPayload, delay: int) -> None: - visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) - parameters = {"queue_name": queue_name, "message": payload, "delay": visible_at} - - if self._current_connection is not None: - self._current_connection.execute(DELAYED_SEND_SQL, parameters) - return - - with self.sessionmaker.begin() as session: - session.connection().execute(DELAYED_SEND_SQL, parameters) - + @override def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": if message.queue_name not in self.queues: raise QueueNotFound(message.queue_name) payload = self._encode_message(message) - - if delay is None: - self.client.send(message.queue_name, payload, conn=self._current_connection) - else: - self._send_with_delay(message.queue_name, payload, delay) + visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) if delay is not None else None + self.client.send( + message.queue_name, + payload, + conn=self._current_connection, + delay=visible_at, + ) return message + @override def flush(self, queue_name: str) -> None: if queue_name not in self.queues: raise QueueNotFound(queue_name) self.client.purge(queue_name, conn=self._current_connection) - def purge_queue(self, queue_name: str) -> None: - self.flush(queue_name) - + @override def flush_all(self) -> None: for queue_name in self.queues: self.flush(queue_name) @@ -224,9 +190,10 @@ def _count_enqueued_messages(self, queue_name: str) -> int: queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") count_query = select(func.count()).select_from(queue_table) - with self.sessionmaker.begin() as session: + with self.client.session() as session: return session.execute(count_query).scalar_one() + @override def join( self, queue_name: str, @@ -288,8 +255,8 @@ def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeou self.broker = broker self.client = broker.client self.queue_name = queue_name - self.prefetch = max(prefetch, 1) - self.timeout = max(timeout, 0) + self.prefetch = prefetch + self.timeout = timeout self.visibility_timeout = broker.visibility_timeout self.heartbeat_interval = broker.heartbeat_interval self.messages: deque[PgmqQueueMessage] = deque() @@ -303,18 +270,9 @@ def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeou self._heartbeat_msg_ids_lock = Lock() self._heartbeat_msg_ids: set[int] = set() - if self.broker.listen_notify_enabled: - self._start_listener() + self._start_listener() self._start_heartbeat() - @property - def max_poll_seconds(self) -> int: - return max((self.timeout + 999) // 1000, 1) - - @property - def poll_interval_ms(self) -> int: - return max(min(self.timeout, 1000), 1) - @property def wait_timeout_seconds(self) -> float: return self.timeout / 1000 if self.timeout > 0 else 0 @@ -339,8 +297,8 @@ def _read_with_poll(self) -> list[PgmqQueueMessage]: self.queue_name, vt=self.visibility_timeout, qty=self.prefetch, - max_poll_seconds=self.max_poll_seconds, - poll_interval_ms=self.poll_interval_ms, + max_poll_seconds=max((self.timeout + 999) // 1000, 1), + poll_interval_ms=max(min(self.timeout, 1000), 1), ) ) @@ -419,18 +377,21 @@ def _run_listener(self) -> None: self._listener_available = False + @override def ack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) self.client.archive(self.queue_name, message._pgmq_message.msg_id) + @override def nack(self, message: "MessageProxy") -> None: if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) self.client.archive(self.queue_name, message._pgmq_message.msg_id) + @override def requeue(self, messages: Iterable["MessageProxy"]) -> None: msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] if msg_ids: @@ -443,16 +404,18 @@ def _build_message(self, pgmq_message: PgmqQueueMessage) -> "_PgmqMessage": self._register_heartbeat_msg(message) return message + @override def __next__(self) -> "_PgmqMessage | None": if self.messages: return self._build_message(self.messages.popleft()) if self._listener_available: - self._notify_event.clear() + self._notify_event.wait(self.wait_timeout_seconds) messages = self._read_immediate() if not messages: - self._notify_event.wait(self.wait_timeout_seconds) - messages = self._read_with_poll() if not self._listener_available else self._read_immediate() + messages = self._read_with_poll() + self._listener_available = False + self._notify_event.clear() else: messages = self._read_with_poll() @@ -462,6 +425,7 @@ def __next__(self) -> "_PgmqMessage | None": self.messages.extend(messages) return self._build_message(self.messages.popleft()) + @override def close(self) -> None: self._listener_stop.set() self._heartbeat_stop.set() diff --git a/remoulade/brokers/rabbitmq.py b/remoulade/brokers/rabbitmq.py index b60b1e299..7b5003aa9 100644 --- a/remoulade/brokers/rabbitmq.py +++ b/remoulade/brokers/rabbitmq.py @@ -21,7 +21,7 @@ from functools import partial from queue import Empty, Full, LifoQueue from threading import Lock, local -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from amqpstorm import AMQPChannelError, AMQPConnectionError, AMQPError, Channel, UriConnection from amqpstorm.compatibility import urlparse @@ -181,6 +181,7 @@ def clear_channel_pools(self): self.channel_pools["confirm_delivery"].clear() self.channel_pools["no_confirm_delivery"].clear() + @override def close(self) -> None: """Close all open RabbitMQ connections.""" @@ -193,6 +194,7 @@ def close(self) -> None: self.clear_channel_pools() self.logger.debug("Channels and connections closed.") + @override def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 5000) -> "_RabbitmqConsumer": """Create a new consumer for a queue. @@ -227,6 +229,7 @@ def _declare_rabbitmq_queues(self): self._declare_dq_queue(channel, queue_name) self._declare_xq_queue(channel, queue_name) + @override def declare_queue(self, queue_name: str) -> None: """Declare a queue. Has no effect if a queue with the given name already exists. @@ -283,6 +286,7 @@ def _declare_xq_queue(self, channel, queue_name): arguments["x-queue-type"] = "quorum" return channel.queue.declare(queue=xq_name(queue_name), durable=True, arguments=arguments) + @override def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": if delay is not None: message_eta = current_millis() + delay @@ -291,6 +295,7 @@ def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message return message + @override @contextmanager def tx(self): with self.get_channel_pool(confirm_delivery=False).acquire() as channel, channel.tx: @@ -312,6 +317,7 @@ def _get_channel(self, confirm_delivery: bool): with self.get_channel_pool(confirm_delivery).acquire() as channel: yield channel + @override def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": """Enqueue a message. @@ -393,6 +399,7 @@ def get_queue_message_counts(self, queue_name: str): xq_queue_response["message_count"], ) + @override def flush(self, queue_name: str) -> None: """Drop all the messages from a queue. @@ -406,11 +413,13 @@ def flush(self, queue_name: str) -> None: with self.default_channel_pool.acquire() as channel: channel.queue.purge(name) + @override def flush_all(self) -> None: """Drop all messages from all declared queues.""" for queue_name in self.queues: self.flush(queue_name) + @override def join( self, queue_name: str, min_successes: int = 10, idle_time: int = 100, *, timeout: int | None = None ) -> None: @@ -458,6 +467,7 @@ def __init__(self, connection, queue_name, prefetch, timeout): except (AMQPConnectionError, AMQPChannelError) as e: raise ConnectionClosed(e) from None + @override def ack(self, message): try: message.ack() @@ -466,6 +476,7 @@ def ack(self, message): except Exception: # pragma: no cover self.logger.error("Failed to ack message.", exc_info=True) + @override def nack(self, message): try: message.nack(requeue=False) @@ -474,11 +485,13 @@ def nack(self, message): except Exception: # pragma: no cover self.logger.error("Failed to nack message.", exc_info=True) + @override def requeue(self, messages): """RabbitMQ automatically re-enqueues unacked messages when consumers disconnect so this is a no-op. """ + @override def __next__(self): """Return None if no value after timeout seconds""" try: @@ -494,6 +507,7 @@ def __next__(self): except (AMQPConnectionError, AMQPChannelError) as e: raise ConnectionClosed(e) from None + @override def close(self): with suppress(AMQPConnectionError, AMQPChannelError): self.channel.close() diff --git a/remoulade/state/backends/__init__.py b/remoulade/state/backends/__init__.py index 384ce0295..567912bf6 100644 --- a/remoulade/state/backends/__init__.py +++ b/remoulade/state/backends/__init__.py @@ -1,14 +1,3 @@ -try: - from .postgres import PostgresBackend -except ImportError: # pragma: no cover - import warnings - - warnings.warn( - "PostgresBackend is not available. Run `pip install remoulade[postgres]` to add support for that backend.", - ImportWarning, - stacklevel=2, - ) - try: from .redis import RedisBackend from .stub import StubBackend @@ -21,4 +10,4 @@ stacklevel=2, ) -__all__ = ["PostgresBackend", "RedisBackend", "StubBackend"] +__all__ = ["RedisBackend", "StubBackend"] diff --git a/remoulade/state/backends/postgres.py b/remoulade/state/backends/postgres.py deleted file mode 100644 index eb8a85278..000000000 --- a/remoulade/state/backends/postgres.py +++ /dev/null @@ -1,250 +0,0 @@ -import datetime -import os -import sys -import threading -from typing import TypeVar - -from sqlalchemy import ( - Column, - DateTime, - Float, - LargeBinary, - SmallInteger, - String, - create_engine, - distinct, - inspect, - or_, - text, -) -from sqlalchemy.orm import declarative_base -from sqlalchemy.orm.session import sessionmaker -from sqlalchemy.sql import func -from sqlalchemy.sql.functions import coalesce, count, max, min # noqa: A004 - -from remoulade import Encoder -from remoulade.state import State, StateBackend - -Base = declarative_base() - -DEFAULT_POSTGRES_URI = "postgresql://remoulade@localhost:5432/remoulade" -DB_VERSION = 3 -T = TypeVar("T", bound="StoredState") - - -class StoredState(Base): - __tablename__ = "states" - - message_id = Column(String(length=36), primary_key=True, index=True) - status = Column(String(length=10), index=True) - actor_name = Column(String(length=79), index=True) - args = Column(LargeBinary) - kwargs = Column(LargeBinary) - options = Column(LargeBinary) - priority = Column(SmallInteger) - progress = Column(Float) - enqueued_datetime = Column(DateTime(timezone=True), index=True) - started_datetime = Column(DateTime(timezone=True), index=True) - end_datetime = Column(DateTime(timezone=True), index=True) - queue_name = Column(String(length=60)) - composition_id = Column(String) - - def as_state(self, encoder: Encoder) -> State: - state_dict = {} - mapper = inspect(StoredState) - for column in mapper.attrs: - column_value = getattr(self, column.key) - if column_value is None: - continue - if column.key in ["args", "kwargs", "options"]: - column_value = encoder.decode(column_value) - state_dict[column.key] = column_value - return State.from_dict(state_dict) - - @classmethod - def from_state(cls: type[T], state: State, max_size: int, encoder: Encoder) -> T: - state_dict = state.as_dict() - for key in ["args", "kwargs", "options"]: - if key in state_dict: - encoded_value = encoder.encode(state_dict[key]) - state_dict[key] = encoded_value if sys.getsizeof(encoded_value) <= max_size else None - return cls(**state_dict) - - -class StateVersion(Base): - __tablename__ = "version" - - version = Column(SmallInteger, primary_key=True) - - -def filter_query( - *, - query, - selected_actors: list[str] | None, - selected_statuses: list[str] | None, - selected_message_ids: list[str] | None, - selected_composition_ids: list[str] | None, - start_datetime: datetime.datetime | None, - end_datetime: datetime.datetime | None, -): - if selected_actors is not None: - query = query.filter(StoredState.actor_name.in_(selected_actors)) - if selected_statuses is not None: - query = query.filter(StoredState.status.in_(selected_statuses)) - if selected_message_ids is not None: - query = query.filter(StoredState.message_id.in_(selected_message_ids)) - if selected_composition_ids is not None: - query = query.filter(StoredState.composition_id.in_(selected_composition_ids)) - if start_datetime is not None: - query = query.filter(StoredState.enqueued_datetime >= start_datetime) - if end_datetime is not None: - query = query.filter(StoredState.enqueued_datetime <= end_datetime) - - return query - - -class PostgresBackend(StateBackend): - def __init__( - self, - *, - namespace: str = "remoulade-state", - encoder: Encoder | None = None, - client: sessionmaker | None = None, - max_size: int = 2000000, - url: str | None = None, - future: bool = False, - ): - self.url = url or os.getenv("REMOULADE_POSTGRESQL_URL") or DEFAULT_POSTGRES_URI - super().__init__(namespace=namespace, encoder=encoder) - self.client = client or sessionmaker(create_engine(self.url, pool_pre_ping=True, future=future)) - self.init_db() - self.max_size = max_size - self.lock = threading.Lock() - - def init_db(self): - with self.client.begin() as session: - bind = session.get_bind() - insp = inspect(bind) - - if not insp.has_table("version"): - Base.metadata.create_all(bind=bind, tables=[StateVersion.__table__]) - - state_version = session.query(StateVersion).first() - if state_version is None: - session.add(StateVersion(version=DB_VERSION)) - - if not insp.has_table("states"): - Base.metadata.create_all(bind=bind, tables=[StoredState.__table__]) - elif state_version is None or state_version.version != DB_VERSION: - StoredState.__table__.drop(bind=bind) - Base.metadata.create_all(bind=bind, tables=[StoredState.__table__]) - if state_version is not None: - state_version.version = DB_VERSION - - def get_state(self, message_id: str): - with self.client.begin() as session: - state = session.query(StoredState).filter_by(message_id=message_id).first() - if state is None: - return None - return state.as_state(self.encoder) - - def set_state(self, state: State, ttl=3600): - with self.lock, self.client.begin() as session: - session.merge(StoredState.from_state(state, self.max_size, self.encoder)) - - def get_states( - self, - *, - size: int | None = None, - offset: int = 0, - selected_actors: list[str] | None = None, - selected_statuses: list[str] | None = None, - selected_message_ids: list[str] | None = None, - selected_composition_ids: list[str] | None = None, - start_datetime: datetime.datetime | None = None, - end_datetime: datetime.datetime | None = None, - sort_column: str | None = None, - sort_direction: str | None = None, - ): - sort_column = sort_column or "enqueued_datetime" - sort_direction = sort_direction or "desc" - - with self.client.begin() as session: - query = session.query(StoredState) - query = filter_query( - query=query, - selected_actors=selected_actors, - selected_statuses=selected_statuses, - selected_message_ids=selected_message_ids, - selected_composition_ids=selected_composition_ids, - start_datetime=start_datetime, - end_datetime=end_datetime, - ) - if size is not None: - query = query.subquery() - query_group = ( - session.query( - max(query.c.composition_id).label("grouped_composition_id"), - max(query.c.message_id).label("grouped_message_id"), - max(query.c.status).label("grouped_status"), - max(query.c.actor_name).label("grouped_actor_name"), - max(query.c.priority).label("grouped_priority"), - func.avg(query.c.progress).label("grouped_progress"), - min(query.c.enqueued_datetime).label("grouped_enqueued_datetime"), - min(query.c.started_datetime).label("grouped_started_datetime"), - max(query.c.end_datetime).label("grouped_end_datetime"), - max(query.c.queue_name).label("grouped_queue_name"), - ) - .group_by(coalesce(query.c.composition_id, query.c.message_id)) - .order_by(text(f"grouped_{sort_column} {sort_direction}")) - ) - query_group = query_group.offset(offset).limit(size).subquery() - query = ( - session.query(StoredState) - .select_from(StoredState) - .join( - query_group, - or_( - StoredState.message_id == query_group.c.grouped_message_id, - StoredState.composition_id == query_group.c.grouped_composition_id, - ), - ) - ) - query = query.order_by(text(f"{sort_column} {sort_direction}")) - return [state_model.as_state(self.encoder) for state_model in query] - - def get_states_count( - self, - *, - selected_actors: list[str] | None = None, - selected_statuses: list[str] | None = None, - selected_messages_ids: list[str] | None = None, - selected_composition_ids: list[str] | None = None, - start_datetime: datetime.datetime | None = None, - end_datetime: datetime.datetime | None = None, - **kwargs, - ): - with self.client.begin() as session: - query = session.query(count(distinct(coalesce(StoredState.composition_id, StoredState.message_id)))) - - query = filter_query( - query=query, - selected_actors=selected_actors, - selected_statuses=selected_statuses, - selected_message_ids=selected_messages_ids, - selected_composition_ids=selected_composition_ids, - start_datetime=start_datetime, - end_datetime=end_datetime, - ) - return query.first()[0] - - def clean(self, max_age: int | None = None, not_started: bool = False): - with self.client.begin() as session: - query = session.query(StoredState) - if max_age: - now = datetime.datetime.now(datetime.UTC) - min_datetime = now - datetime.timedelta(minutes=max_age) - query = session.query(StoredState).filter(StoredState.end_datetime < min_datetime) - if not_started: - query = session.query(StoredState).filter(StoredState.started_datetime.is_(None)) - query.delete() diff --git a/tests/conftest.py b/tests/conftest.py index 95ebdc2b6..11400d613 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,6 @@ import redis from freezegun import freeze_time from sqlalchemy.engine import create_engine -from sqlalchemy.inspection import inspect from sqlalchemy.orm.session import sessionmaker from sqlalchemy.pool import NullPool from sqlalchemy.sql import text @@ -34,7 +33,6 @@ MessageState, backends as st_backends, ) -from remoulade.state.backends.postgres import DB_VERSION, StateVersion logfmt = "[%(asctime)s] [%(threadName)s] [%(name)s] [%(levelname)s] %(message)s" logging.basicConfig(level=logging.INFO, format=logfmt) @@ -65,16 +63,6 @@ def check_redis(client): client.flushall() -def check_postgres(client): - with client.begin() as session: - insp = inspect(session.get_bind()) - version_exists = insp.has_table("version") - states_exists = insp.has_table("states") - if version_exists: - versions = session.query(StateVersion).all() - return version_exists and states_exists and len(versions) == 1 and versions[0].version == DB_VERSION - - def check_pgmq(client): try: with client.begin() as session: @@ -91,13 +79,6 @@ def cleanup_pgmq_queues(client): session.execute(text("SELECT pgmq.drop_queue(:queue_name)"), {"queue_name": queue_name}) -@pytest.fixture -def check_postgres_begin(): - client = remoulade.get_broker().get_state_backend().client - if not check_postgres(client): - pytest.skip("Postgres Database is not in the proper state. Database initialisation is probably incorrect.") - - @pytest.fixture() def stub_broker(): broker = StubBroker() @@ -235,16 +216,8 @@ def stub_state_backend(): @pytest.fixture -def postgres_state_backend(): - db_string = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" - backend = st_backends.PostgresBackend(client=sessionmaker(create_engine(db_string, poolclass=NullPool))) - backend.clean() - return backend - - -@pytest.fixture -def state_backends(postgres_state_backend, redis_state_backend, stub_state_backend): - return {"postgres": postgres_state_backend, "redis": redis_state_backend, "stub": stub_state_backend} +def state_backends(redis_state_backend, stub_state_backend): + return {"redis": redis_state_backend, "stub": stub_state_backend} @pytest.fixture(params=["postgres", "redis", "stub"]) diff --git a/tests/docker/postgres/01-create-pgmq-extension.sql b/tests/docker/postgres/01-create-pgmq-extension.sql index b484e86e5..3afb0e3ff 100644 --- a/tests/docker/postgres/01-create-pgmq-extension.sql +++ b/tests/docker/postgres/01-create-pgmq-extension.sql @@ -1 +1,3 @@ +CREATE SCHEMA IF NOT EXISTS partman; +CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman; CREATE EXTENSION IF NOT EXISTS pgmq; diff --git a/tests/middleware/test_message_state.py b/tests/middleware/test_message_state.py index ee66000f9..fad062c3e 100644 --- a/tests/middleware/test_message_state.py +++ b/tests/middleware/test_message_state.py @@ -8,7 +8,6 @@ from remoulade.cancel import Cancel from remoulade.middleware import Middleware, SkipMessage from remoulade.state.backend import State, StateStatusesEnum -from remoulade.state.backends import PostgresBackend from remoulade.state.middleware import MessageState from tests.conftest import mock_func @@ -104,9 +103,6 @@ def before_process_message(self, broker, message): @pytest.mark.parametrize("ttl, result_type", [pytest.param(1000, State), pytest.param(1, type(None))]) def test_expiration_data_backend(self, ttl, result_type, stub_broker, state_backend): - if isinstance(state_backend, PostgresBackend): - pytest.skip("Skipping this test as there is no expiration on PostgresBackend") - @remoulade.actor def wait(): pass diff --git a/tests/state/test_backend.py b/tests/state/test_backend.py index 59c8b6972..5ea04c87e 100644 --- a/tests/state/test_backend.py +++ b/tests/state/test_backend.py @@ -1,7 +1,4 @@ -import pytest - from remoulade.state import State, StateStatusesEnum -from remoulade.state.backends import PostgresBackend class TestStateBackend: @@ -31,31 +28,3 @@ def test_count_messages(self, stub_broker, state_middleware): backend.set_state(State(f"id{i}")) assert backend.get_states_count() == 3 - - def test_count_compositions(self, stub_broker, state_middleware): - if not isinstance(state_middleware.backend, PostgresBackend): - pytest.skip() - - backend = state_middleware.backend - - for i in range(3): - for j in range(2): - backend.set_state(State(f"id{i * j}", composition_id=f"id{j}")) - - assert backend.get_states_count() == 2 - - def test_sort_with_offset(self, stub_broker, state_middleware): - if not isinstance(state_middleware.backend, PostgresBackend): - pytest.skip() - backend = state_middleware.backend - for i in range(8): - backend.set_state(State(f"id{i}", actor_name=f"{3 + 4 * (i // 4) - i % 4}")) - - res = backend.get_states(size=3, sort_column="actor_name", sort_direction="desc") - assert res[0].actor_name == "7" - assert res[1].actor_name == "6" - assert res[2].actor_name == "5" - res = backend.get_states(size=3, sort_column="actor_name", sort_direction="asc") - assert res[0].actor_name == "0" - assert res[1].actor_name == "1" - assert res[2].actor_name == "2" diff --git a/tests/state/test_postgres.py b/tests/state/test_postgres.py deleted file mode 100644 index 4f2acc8d2..000000000 --- a/tests/state/test_postgres.py +++ /dev/null @@ -1,65 +0,0 @@ -from remoulade.state.backends.postgres import DB_VERSION, StateVersion, StoredState -from tests.conftest import check_postgres - - -def test_no_changes(stub_broker, postgres_state_middleware, check_postgres_begin): - backend = postgres_state_middleware.backend - client = backend.client - with client.begin() as session: - session.add(StoredState(message_id="id")) - - backend.init_db() - assert check_postgres(client) - with client.begin() as session: - assert len(session.query(StoredState).all()) == 1 - - -def test_create_tables(stub_broker, postgres_state_middleware, check_postgres_begin): - backend = postgres_state_middleware.backend - client = backend.client - with client.begin() as session: - engine = session.get_bind() - StoredState.__table__.drop(bind=engine) - StateVersion.__table__.drop(bind=engine) - - backend.init_db() - assert check_postgres(client) - - -def test_change_version(stub_broker, postgres_state_middleware, check_postgres_begin): - backend = postgres_state_middleware.backend - client = backend.client - with client.begin() as session: - version = session.query(StateVersion).first() - version.version = DB_VERSION + 1 - session.add(StoredState(message_id="id")) - - backend.init_db() - assert check_postgres(client) - with client.begin() as session: - assert len(session.query(StoredState).all()) == 0 - - -def test_no_version(stub_broker, postgres_state_middleware, check_postgres_begin): - backend = postgres_state_middleware.backend - client = backend.client - with client.begin() as session: - engine = session.get_bind() - StateVersion.__table__.drop(bind=engine) - session.add(StoredState(message_id="id")) - - backend.init_db() - assert check_postgres(client) - with client.begin() as session: - assert len(session.query(StoredState).all()) == 0 - - -def test_no_states(stub_broker, postgres_state_middleware, check_postgres_begin): - backend = postgres_state_middleware.backend - client = backend.client - with client.begin() as session: - engine = session.get_bind() - StoredState.__table__.drop(bind=engine) - - backend.init_db() - assert check_postgres(client) diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 19f6c031f..98a3639f5 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -18,7 +18,7 @@ def _count_messages(broker, queue_name="default"): queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") - with broker.sessionmaker.begin() as session: + with broker.client.session() as session: return session.execute(select(func.count()).select_from(queue_table)).scalar_one() @@ -30,19 +30,19 @@ def _first_payload(broker, queue_name="default"): Column("message", JSON), schema="pgmq", ) - with broker.sessionmaker.begin() as session: + with broker.client.session() as session: row = session.execute(select(queue_table.c.message).order_by(queue_table.c.msg_id).limit(1)).one() return row[0] def _count_archived_messages(broker, queue_name="default"): archive_table = Table(f"a_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") - with broker.sessionmaker.begin() as session: + with broker.client.session() as session: return session.execute(select(func.count()).select_from(archive_table)).scalar_one() def _queue_exists(broker, queue_name): - with broker.sessionmaker.begin() as session: + with broker.client.session() as session: query = text("SELECT EXISTS(SELECT 1 FROM pgmq.list_queues() WHERE queue_name = :queue_name)") return session.execute(query, {"queue_name": queue_name}).scalar_one() @@ -62,7 +62,6 @@ def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): broker = PgmqBroker( url=TEST_PGMQ_URL, middleware=[], - partition_archive_on_queue_init=True, ) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() @@ -79,37 +78,27 @@ def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): broker.client.create_queue.assert_not_called() -def test_pgmq_broker_uses_non_partitioned_queue_creation_when_disabled(): - broker = PgmqBroker( - url=TEST_PGMQ_URL, - middleware=[], - partition_archive_on_queue_init=False, - ) - broker.client.validate_queue_name = Mock() - broker.client.create_partitioned_queue = Mock() - broker.client.create_queue = Mock() - - broker.declare_queue("default") - - broker.client.create_queue.assert_called_once_with("default", conn=None) - broker.client.create_partitioned_queue.assert_not_called() - - def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.client.validate_queue_name = Mock() - broker.client.create_queue = Mock() + broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() broker.declare_queue("default") + broker.client.create_partitioned_queue.assert_called_once_with( + "default", + partition_interval="1 day", + retention_interval="7 days", + conn=None, + ) broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) broker.client.validate_queue_name = Mock() - broker.client.create_queue = Mock() + broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) broker.declare_queue("default") @@ -118,6 +107,27 @@ def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): assert "Failed to enable LISTEN/NOTIFY" in caplog.text +def test_pgmq_broker_uses_custom_partition_settings_when_provided(): + broker = PgmqBroker( + url=TEST_PGMQ_URL, + middleware=[], + archive_partition_interval="2 days", + archive_retention_interval="14 days", + ) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.enable_notify = Mock() + + broker.declare_queue("default") + + broker.client.create_partitioned_queue.assert_called_once_with( + "default", + partition_interval="2 days", + retention_interval="14 days", + conn=None, + ) + + @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) @@ -175,7 +185,7 @@ def _fake_start_listener(self): message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) payload = _expected_payload(message) - broker.client.read = Mock(side_effect=[None, Mock(msg_id=1, message=payload)]) + broker.client.read = Mock(return_value=Mock(msg_id=1, message=payload)) broker.client.read_with_poll = Mock(return_value=[]) consumer = broker.consume("default", prefetch=1, timeout=200) @@ -185,8 +195,9 @@ def _fake_start_listener(self): assert consumed is not None assert consumed.message_id == message.message_id - assert broker.client.read.call_count == 2 + broker.client.read.assert_called_once_with("default", vt=30, qty=1) broker.client.read_with_poll.assert_not_called() + assert consumer._listener_available is True consumer.close() @@ -336,7 +347,7 @@ def _fake_start_listener(self): broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) remoulade.set_broker(broker) broker.client.validate_queue_name = Mock() - broker.client.create_queue = Mock() + broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() @remoulade.actor(actor_name="typed.actor", queue_name="default") From 6983122414fa71b1ad56bfbc310ba4de44626d29 Mon Sep 17 00:00:00 2001 From: mducros Date: Tue, 2 Jun 2026 17:13:30 +0200 Subject: [PATCH 12/44] fix(StateBackend): remove postgres --- remoulade/broker.py | 6 +- remoulade/brokers/pgmq.py | 126 ++++++++++++++++++++++++---------- tests/conftest.py | 12 +--- tests/state/test_state_api.py | 84 ----------------------- tests/test_pgmq.py | 38 ++++++++-- 5 files changed, 129 insertions(+), 137 deletions(-) delete mode 100644 tests/state/test_state_api.py diff --git a/remoulade/broker.py b/remoulade/broker.py index 57b46ec1e..d8eaef526 100644 --- a/remoulade/broker.py +++ b/remoulade/broker.py @@ -421,7 +421,11 @@ def declare_queue(self, queue_name: str) -> None: # pragma: no cover raise NotImplementedError def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": - raise NotImplementedError + """If your broker doesn't support native delay, you need to override this method""" + if self.supports_native_delay: + return message + else: + raise NotImplementedError("delay is not supported natively, you need to implement it") def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": raise NotImplementedError diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 61351cb86..19d4c2948 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -1,6 +1,6 @@ # This file is a part of Remoulade. # -# Copyright (C) 2017,2018 WIREMIND SAS +# Copyright (C) 2026 WIREMIND SAS # # Remoulade is free software; you can redistribute it and/or modify it # under the terms of the GNU Lesser General Public License as published by @@ -21,13 +21,13 @@ from contextlib import contextmanager from datetime import UTC, datetime, timedelta from threading import Event, Lock, Thread, local -from typing import TYPE_CHECKING, Any, override +from typing import TYPE_CHECKING, Any, Final, override import psycopg from pgmq import SQLAlchemyPGMQueue from pgmq.messages import Message as PgmqQueueMessage from psycopg import sql as psycopg_sql -from sqlalchemy import Column, Integer, MetaData, Table, func, select +from sqlalchemy import Column, Connection, Integer, MetaData, Table, func, select from ..broker import Broker, Consumer, MessageProxy from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding @@ -37,7 +37,7 @@ from ..middleware import Middleware PgmqPayload = dict[str, Any] -LISTEN_NOTIFY_THROTTLE_MS = 250 +LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 class PgmqBroker(Broker): @@ -51,8 +51,8 @@ def __init__( group_transaction: bool = False, archive_partition_interval: str = "1 day", archive_retention_interval: str = "7 days", - visibility_timeout: int = 30, - heartbeat_interval: float = 10.0, + visibility_timeout_in_second: int = 30, + heartbeat_interval_in_second: float = 10.0, ) -> None: """Initialize a PostgreSQL-backed broker using the PGMQ extension. @@ -60,19 +60,16 @@ def __init__( url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg. middleware(list[Middleware] | None): Middleware stack applied to this broker. group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. - archive_partition_interval(int | str): Partition interval passed to PGMQ when creating partitioned queues. - archive_retention_interval(int | str): Retention interval passed to PGMQ for archive partitions. - visibility_timeout(int): Message visibility timeout in seconds after read; must be greater than 0. - heartbeat_interval(float): Heartbeat interval in seconds used to extend in-flight message visibility; - must be greater than 0 and lower than visibility_timeout. - - Raises: - ValueError: If visibility_timeout or heartbeat_interval values are invalid. + archive_partition_interval(str): Partition interval passed to PGMQ when creating partitioned queues. + archive_retention_interval(str): Retention interval passed to PGMQ for archive partitions. + visibility_timeout_in_second(int): Message visibility timeout in seconds after read; must be greater than 0. + heartbeat_interval_in_second(float): Heartbeat interval in seconds used to extend in-flight message visibility + must be greater than 0 and lower than visibility_timeout_in_second """ super().__init__(middleware=middleware) - if visibility_timeout <= 0: - raise ValueError("visibility_timeout must be greater than 0") - if heartbeat_interval <= 0 or heartbeat_interval >= visibility_timeout: + if visibility_timeout_in_second <= 0: + raise ValueError("visibility_timeout_in_second must be greater than 0") + if heartbeat_interval_in_second <= 0 or heartbeat_interval_in_second >= visibility_timeout_in_second: raise ValueError("heartbeat_interval must be greater than 0 and lower than visibility_timeout") self.url = url @@ -80,15 +77,20 @@ def __init__( self.group_transaction = group_transaction self.archive_partition_interval = archive_partition_interval self.archive_retention_interval = archive_retention_interval - self.visibility_timeout = visibility_timeout - self.heartbeat_interval = heartbeat_interval + self.visibility_timeout_in_second = visibility_timeout_in_second + self.heartbeat_interval_in_second = heartbeat_interval_in_second self.supports_native_delay = True - self.client = SQLAlchemyPGMQueue(conn_string=url, init_extension=False, vt=self.visibility_timeout) + self.client = SQLAlchemyPGMQueue(conn_string=url, init_extension=False, vt=self.visibility_timeout_in_second) @override @contextmanager def tx(self) -> Iterator[None]: + """Run broker operations inside a single SQL transaction. + + The active SQLAlchemy connection is stored on thread-local state so + queue operations can reuse it while the context is open. + """ with self.client.engine.begin() as connection: self.state.transaction_connection = connection try: @@ -97,35 +99,48 @@ def tx(self) -> Iterator[None]: self.state.transaction_connection = None @property - def _current_connection(self) -> Any | None: + def _current_connection(self) -> Connection | None: + """Return the transactional connection, if one is active.""" return getattr(self.state, "transaction_connection", None) def _try_enable_notify(self, queue_name: str) -> None: + """Try to enable LISTEN/NOTIFY for a queue. + + Failures are logged and the consumer later falls back to polling. + """ try: self.client.enable_notify( queue_name, throttle_interval_ms=LISTEN_NOTIFY_THROTTLE_MS, conn=self._current_connection, ) - except Exception: + except Exception as e: self.logger.warning( - "Failed to enable LISTEN/NOTIFY for queue %s; consumer will fall back to polling.", + "Failed to enable LISTEN/NOTIFY for queue %s; consumer will fall back to polling. Error: %s", queue_name, + e, exc_info=True, ) @override def close(self) -> None: + """Dispose the underlying PGMQ client.""" self.client.dispose() @override def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PgmqConsumer": + """Create a consumer for a declared queue. + + Raises: + QueueNotFound: If the queue has not been declared. + """ if queue_name not in self.queues: raise QueueNotFound(queue_name) return _PgmqConsumer(self, queue_name=queue_name, prefetch=prefetch, timeout=timeout) @override def declare_queue(self, queue_name: str) -> None: + """Create a partitioned PGMQ queue if it does not already exist.""" if queue_name in self.queues: return @@ -141,17 +156,19 @@ def declare_queue(self, queue_name: str) -> None: self.queues[queue_name] = None self.emit_after("declare_queue", queue_name) - @override - def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": - return message - def _encode_message(self, message: "Message") -> PgmqPayload: + """Decode a Remoulade message into a JSON object payload for PGMQ. + + Raises: + UnsupportedMessageEncoding: If the encoded payload is not valid JSON + or is not a JSON object. + """ try: payload = json.loads(message.encode().decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as e: raise UnsupportedMessageEncoding( "PgmqBroker only supports encoders that serialize messages as JSON objects." - ) from exc + ) from e if not isinstance(payload, dict): raise UnsupportedMessageEncoding("PgmqBroker requires message encoders to produce JSON objects.") @@ -160,6 +177,7 @@ def _encode_message(self, message: "Message") -> PgmqPayload: @override def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": + """Send a message to PGMQ, optionally delayed by milliseconds.""" if message.queue_name not in self.queues: raise QueueNotFound(message.queue_name) @@ -176,6 +194,7 @@ def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message" @override def flush(self, queue_name: str) -> None: + """Remove every message currently stored in a queue.""" if queue_name not in self.queues: raise QueueNotFound(queue_name) @@ -183,10 +202,12 @@ def flush(self, queue_name: str) -> None: @override def flush_all(self) -> None: + """Purge every declared queue.""" for queue_name in self.queues: self.flush(queue_name) def _count_enqueued_messages(self, queue_name: str) -> int: + """Count rows in the underlying PGMQ queue table.""" queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") count_query = select(func.count()).select_from(queue_table) @@ -257,8 +278,8 @@ def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeou self.queue_name = queue_name self.prefetch = prefetch self.timeout = timeout - self.visibility_timeout = broker.visibility_timeout - self.heartbeat_interval = broker.heartbeat_interval + self.visibility_timeout = broker.visibility_timeout_in_second + self.heartbeat_interval = broker.heartbeat_interval_in_second self.messages: deque[PgmqQueueMessage] = deque() self._notify_event = Event() self._listener_stop = Event() @@ -275,14 +296,17 @@ def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeou @property def wait_timeout_seconds(self) -> float: + """Return the listener wait timeout in seconds.""" return self.timeout / 1000 if self.timeout > 0 else 0 def _normalize_messages(self, messages: PgmqQueueMessage | list[PgmqQueueMessage] | None) -> list[PgmqQueueMessage]: + """Normalize a PGMQ read result into a list.""" if messages is None: return [] return [messages] if not isinstance(messages, list) else messages def _read_immediate(self) -> list[PgmqQueueMessage]: + """Read up to ``prefetch`` messages without polling.""" return self._normalize_messages( self.client.read( self.queue_name, @@ -292,6 +316,7 @@ def _read_immediate(self) -> list[PgmqQueueMessage]: ) def _read_with_poll(self) -> list[PgmqQueueMessage]: + """Read up to ``prefetch`` messages using PGMQ polling.""" return self._normalize_messages( self.client.read_with_poll( self.queue_name, @@ -303,6 +328,7 @@ def _read_with_poll(self) -> list[PgmqQueueMessage]: ) def _start_heartbeat(self) -> None: + """Start the background visibility-timeout renewal thread.""" self._heartbeat_thread = Thread( target=self._run_heartbeat, name=f"pgmq-heartbeat-{self.queue_name}", @@ -311,6 +337,18 @@ def _start_heartbeat(self) -> None: self._heartbeat_thread.start() def _run_heartbeat(self) -> None: + """Periodically renew the lease of in-flight messages. + + The loop sleeps for ``heartbeat_interval`` seconds, unless shutdown is + requested through ``_heartbeat_stop``. On each tick it snapshots the + current set of tracked message ids under a lock, then calls + ``set_vt`` to push their visibility timeout back to + ``self.visibility_timeout``. + + This prevents long-running handlers from becoming visible again and + being consumed twice before they are explicitly acked, nacked or + requeued. Failures are logged and retried on the next heartbeat tick. + """ while not self._heartbeat_stop.wait(self.heartbeat_interval): with self._heartbeat_msg_ids_lock: msg_ids = list(self._heartbeat_msg_ids) @@ -318,33 +356,38 @@ def _run_heartbeat(self) -> None: continue try: self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout) - except Exception: + except Exception as e: self.broker.logger.warning( - "Failed to extend visibility timeout heartbeat for queue %s.", + "Failed to extend visibility timeout heartbeat for queue %s. Error: ", self.queue_name, + str(e), exc_info=True, ) def _register_heartbeat_msg(self, message: "_PgmqMessage") -> None: + """Track a message so the heartbeat can renew it.""" with self._heartbeat_msg_ids_lock: self._heartbeat_msg_ids.add(message._pgmq_message.msg_id) def _unregister_heartbeat_msg_id(self, msg_id: int) -> None: + """Stop tracking a message for heartbeat renewal.""" with self._heartbeat_msg_ids_lock: self._heartbeat_msg_ids.discard(msg_id) def _start_listener(self) -> None: + """Start a LISTEN/NOTIFY listener for the queue when available.""" channel = f"pgmq.q_{self.queue_name}.INSERT" try: self._listener_connection = psycopg.connect(self.broker.url, autocommit=True) self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) self._listener_available = True - except Exception: + except Exception as e: self._listener_available = False self._listener_connection = None self.broker.logger.warning( - "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling.", + "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling. Error: %s", self.queue_name, + str(e), exc_info=True, ) return @@ -357,20 +400,22 @@ def _start_listener(self) -> None: self._listener_thread.start() def _run_listener(self) -> None: + """Wake the consumer when a notification arrives, or disable the listener on errors.""" while not self._listener_stop.is_set(): try: if self._listener_connection is None: break for _ in self._listener_connection.notifies(timeout=0.5, stop_after=1): self._notify_event.set() - except Exception: + except Exception as e: if self._listener_stop.is_set(): break self._listener_available = False self._notify_event.set() self.broker.logger.warning( - "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling.", + "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling. Error: %s", self.queue_name, + str(e), exc_info=True, ) break @@ -379,6 +424,7 @@ def _run_listener(self) -> None: @override def ack(self, message: "MessageProxy") -> None: + """Archive a processed message.""" if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) @@ -386,6 +432,7 @@ def ack(self, message: "MessageProxy") -> None: @override def nack(self, message: "MessageProxy") -> None: + """Archive a failed message.""" if not isinstance(message, _PgmqMessage): raise ValueError("It must be a PgmqMessage") self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) @@ -393,6 +440,7 @@ def nack(self, message: "MessageProxy") -> None: @override def requeue(self, messages: Iterable["MessageProxy"]) -> None: + """Make messages visible again immediately by resetting their visibility timeout.""" msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] if msg_ids: for msg_id in msg_ids: @@ -400,12 +448,14 @@ def requeue(self, messages: Iterable["MessageProxy"]) -> None: self.client.set_vt(self.queue_name, msg_ids, 0) def _build_message(self, pgmq_message: PgmqQueueMessage) -> "_PgmqMessage": + """Wrap a raw PGMQ row and register it for heartbeat renewal.""" message = _PgmqMessage(pgmq_message) self._register_heartbeat_msg(message) return message @override def __next__(self) -> "_PgmqMessage | None": + """Return the next available message, or ``None`` if the queue stays empty.""" if self.messages: return self._build_message(self.messages.popleft()) @@ -427,6 +477,7 @@ def __next__(self) -> "_PgmqMessage | None": @override def close(self) -> None: + """Stop background threads and close the listener connection.""" self._listener_stop.set() self._heartbeat_stop.set() self._notify_event.set() @@ -446,6 +497,7 @@ def close(self) -> None: class _PgmqMessage(MessageProxy): def __init__(self, pgmq_message: PgmqQueueMessage) -> None: + """Wrap a PGMQ message row as a Remoulade message proxy.""" payload = pgmq_message.message if not isinstance(payload, dict): raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") diff --git a/tests/conftest.py b/tests/conftest.py index 11400d613..cb6b43634 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -66,6 +66,8 @@ def check_redis(client): def check_pgmq(client): try: with client.begin() as session: + session.execute(text("CREATE SCHEMA IF NOT EXISTS partman")) + session.execute(text("CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman")) session.execute(text("CREATE EXTENSION IF NOT EXISTS pgmq")) session.execute(text("SELECT pgmq.validate_queue_name('remoulade_test_queue')")) except Exception as e: @@ -220,7 +222,7 @@ def state_backends(redis_state_backend, stub_state_backend): return {"redis": redis_state_backend, "stub": stub_state_backend} -@pytest.fixture(params=["postgres", "redis", "stub"]) +@pytest.fixture(params=["redis", "stub"]) def state_backend(request, state_backends): return state_backends[request.param] @@ -233,14 +235,6 @@ def state_middleware(state_backend): return middleware -@pytest.fixture -def postgres_state_middleware(postgres_state_backend): - broker = remoulade.get_broker() - middleware = MessageState(backend=postgres_state_backend) - broker.add_middleware(middleware) - return middleware - - @pytest.fixture def result_middleware(result_backend): broker = remoulade.get_broker() diff --git a/tests/state/test_state_api.py b/tests/state/test_state_api.py deleted file mode 100644 index 69b2a13dd..000000000 --- a/tests/state/test_state_api.py +++ /dev/null @@ -1,84 +0,0 @@ -import datetime - -import pytest -from dateutil.parser import parse - -from remoulade.api.main import app -from remoulade.state import State, StateStatusesEnum - - -@pytest.fixture -def local_api_client(): - with app.test_client() as client: - yield client - - -class TestMessageStateBackend: - def test_select_actors(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - for i in range(2): - backend.set_state(State(f"id{i}", actor_name=f"actor{i}")) - res = backend.get_states(selected_actors=["actor1"]) - assert len(res) == 1 - assert res[0].actor_name == "actor1" - - def test_select_statuses(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - for i in range(2): - backend.set_state(State(f"id{i}", StateStatusesEnum.Success if i else StateStatusesEnum.Skipped)) - res = backend.get_states(selected_statuses=["Success"]) - assert len(res) == 1 - assert res[0].status == StateStatusesEnum.Success - - def test_select_ids(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - for i in range(2): - backend.set_state(State(f"id{i}")) - res = backend.get_states(selected_message_ids=["id1"]) - assert len(res) == 1 - assert res[0].message_id == "id1" - - def test_select_datetimes(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - for i in range(2): - backend.set_state(State(f"id{i}", enqueued_datetime=parse(f"2020-08-08 1{i}:00:00"))) - - res = backend.get_states(start_datetime=parse("2020-08-08 10:30:00")) - assert len(res) == 1 - assert res[0].message_id == "id1" - - res = backend.get_states(end_datetime=parse("2020-08-08 10:30:00")) - assert len(res) == 1 - assert res[0].message_id == "id0" - - def test_clean(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - backend.set_state(State("id0")) - - assert len(backend.get_states()) == 1 - backend.clean() - assert len(backend.get_states()) == 0 - - def test_clean_max_age(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - backend.set_state(State("id0", end_datetime=datetime.datetime.now(datetime.UTC))) - backend.set_state( - State("id1", end_datetime=datetime.datetime.now(datetime.UTC) - datetime.timedelta(minutes=50)) - ) - - assert len(backend.get_states()) == 2 - backend.clean(max_age=25) - res = backend.get_states() - assert len(res) == 1 - assert res[0].message_id == "id0" - - def test_clean_not_started(self, stub_broker, postgres_state_middleware): - backend = postgres_state_middleware.backend - backend.set_state(State("id0", started_datetime=datetime.datetime.now(datetime.UTC))) - backend.set_state(State("id1")) - - assert len(backend.get_states()) == 2 - backend.clean(not_started=True) - res = backend.get_states() - assert len(res) == 1 - assert res[0].message_id == "id0" diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 98a3639f5..5703dbd66 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -147,7 +147,6 @@ def test_pgmq_broker_uses_native_visibility_delay_without_delay_queue(pgmq_broke pgmq_broker.enqueue(message, delay=250) assert pgmq_broker.client.read("default", vt=1) is None - assert not _queue_exists(pgmq_broker, "default.DQ") time.sleep(0.35) delayed = pgmq_broker.client.read("default", vt=1) @@ -237,7 +236,7 @@ def _fake_start_listener(self): monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[], visibility_timeout=17) + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[], visibility_timeout_in_second=17) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(5,), kwargs={}, options={}) @@ -305,8 +304,8 @@ def _fake_start_listener(self): broker = PgmqBroker( url=TEST_PGMQ_URL, middleware=[], - visibility_timeout=2, - heartbeat_interval=0.05, + visibility_timeout_in_second=2, + heartbeat_interval_in_second=0.05, ) broker.queues["default"] = None @@ -451,8 +450,6 @@ def do_work(value): finally: worker.stop() - assert not _queue_exists(pgmq_broker, "default.DQ") - @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_broker_join_times_out_while_processing_invisible_message(pgmq_broker): @@ -483,6 +480,35 @@ def do_work(): worker.stop() +@pytest.mark.usefixtures("pgmq_broker") +def test_pgmq_worker_processes_a_two_actor_pipeline(pgmq_broker): + seen: list[tuple[str, int]] = [] + + @remoulade.actor + def first_actor(value): + seen.append(("first", value)) + return value + 1 + + @remoulade.actor + def second_actor(value): + seen.append(("second", value)) + + pgmq_broker.declare_actor(first_actor) + pgmq_broker.declare_actor(second_actor) + + worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=1) + worker.start() + try: + remoulade.pipeline([first_actor.message(1), second_actor.message()]).run() + + pgmq_broker.join(second_actor.queue_name, timeout=10_000) + worker.join() + + assert seen == [("first", 1), ("second", 2)] + finally: + worker.stop() + + @pytest.mark.usefixtures("pgmq_broker") def test_pgmq_consumer_listener_wakes_on_enqueue_with_listen_notify(pgmq_broker): message = Message(queue_name="default", actor_name="do_work", args=(99,), kwargs={}, options={}) From 308b27ed7a2877bca6ae001f649a4fd03d64c3ce Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 3 Jun 2026 16:48:56 +0200 Subject: [PATCH 13/44] fix(Worker): remove useless code --- README.md | 13 +++++++++++-- docs/source/global.rst | 2 ++ docs/source/guide.rst | 27 +++++++++++++++++++++++++-- docs/source/index.rst | 9 +++++++-- docs/source/installation.rst | 10 ++++++++-- docs/source/reference.rst | 3 +++ remoulade/brokers/pgmq.py | 17 +++++++++++++++-- remoulade/worker.py | 14 +++++--------- tests/test_pgmq.py | 18 ++++++++++++++++++ 9 files changed, 94 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index c4a929758..76e3a8227 100644 --- a/README.md +++ b/README.md @@ -23,10 +23,17 @@ If you want to use it with [RabbitMQ] uv pip install 'remoulade[rabbitmq]' ``` -or if you want to use it with [Redis] +If you want to use it with [PostgreSQL] and [PGMQ] ```console - uv pip install 'remoulade[redis]' + uv pip install 'remoulade[postgres]' +``` + +If you want Redis-backed extras like results or cancellation, add [Redis] to the broker extra you use: + +```console + uv pip install 'remoulade[rabbitmq, redis]' + uv pip install 'remoulade[postgres, redis]' ``` ## Quickstart @@ -110,6 +117,8 @@ remoulade is licensed under the LGPL. Please see [COPYING] and [COPYING.LESSER]: https://github.com/wiremind/remoulade/blob/master/COPYING.LESSER [COPYING]: https://github.com/wiremind/remoulade/blob/master/COPYING +[PostgreSQL]: https://www.postgresql.org/ +[PGMQ]: https://pgmq.github.io/pgmq/ [RabbitMQ]: https://www.rabbitmq.com/ [Redis]: https://redis.io [user guide]: https://remoulade.readthedocs.io/en/latest/guide.html diff --git a/docs/source/global.rst b/docs/source/global.rst index a27b10ac4..416529144 100644 --- a/docs/source/global.rst +++ b/docs/source/global.rst @@ -88,5 +88,7 @@ .. _gevent: http://www.gevent.org/ .. _RabbitMQ: https://www.rabbitmq.com +.. _PostgreSQL: https://www.postgresql.org/ +.. _PGMQ: https://pgmq.github.io/pgmq/ .. _Redis: https://redis.io .. _Dramatiq: https://dramatiq.io diff --git a/docs/source/guide.rst b/docs/source/guide.rst index e786f7d63..31a5d920a 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -331,6 +331,11 @@ milliseconds):: Keep in mind that *your message broker is not a database*. Scheduled messages should represent a small subset of all your messages. +On brokers that emulate delay in worker memory, the enqueued message +will carry an ``eta`` option. ``PgmqBroker`` stores delayed messages in +PostgreSQL natively instead, so the message it returns does not include +``eta``. + Prioritizing Messages --------------------- @@ -381,7 +386,7 @@ Message Brokers --------------- Remoulade abstracts over the notion of a message broker and currently -supports RabbitMQ out of the box. +supports RabbitMQ and PostgreSQL/PGMQ out of the box. RabbitMQ Broker ^^^^^^^^^^^^^^^ @@ -398,6 +403,25 @@ execution:: remoulade.set_broker(rabbitmq_broker) +PGMQ Broker +^^^^^^^^^^^ + +To configure PostgreSQL/PGMQ, install ``remoulade[postgres]`` and +instantiate a ``PgmqBroker`` with a PostgreSQL URL as early as possible +during your program's execution:: + + import remoulade + + from remoulade.brokers.pgmq import PgmqBroker + + pgmq_broker = PgmqBroker(url="postgresql://remoulade@localhost:5432/remoulade") + remoulade.set_broker(pgmq_broker) + +PGMQ handles delayed messages natively, so ``send_with_options(delay=...)`` +does not create a worker-side delay queue or add an ``eta`` option to +the message. + + Local Broker ^^^^^^^^^^^^^^^ @@ -472,4 +496,3 @@ synchronously by calling them as you would normal functions. .. _pytest fixtures: https://docs.pytest.org/en/latest/fixture.html .. _priority documentation: https://www.rabbitmq.com/priority.html - diff --git a/docs/source/index.rst b/docs/source/index.rst index 4a1674893..a0e19d881 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -49,9 +49,14 @@ If you want to use it with RabbitMQ_:: $ pip install -U 'remoulade[rabbitmq]' -Or if you want to use it with Redis_:: +Or if you want to use it with PostgreSQL_ and PGMQ_:: - $ pip install -U 'remoulade[redis]' + $ pip install -U 'remoulade[postgres]' + +Or if you want to use Redis_ for results and cancellation:: + + $ pip install -U 'remoulade[rabbitmq, redis]' + $ pip install -U 'remoulade[postgres, redis]' Read the :doc:`guide` if you're ready to get started. diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 1057dc0c3..80333ffc7 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -14,11 +14,16 @@ To install remoulade, simply run the following command in a terminal:: $ pip install -U 'remoulade[rabbitmq]' -Remoulade use RabbitMQ_ as message broker. +Remoulade uses RabbitMQ_ as a message broker. If you want to use +PostgreSQL_ with PGMQ_ instead, install:: -If you would like to use it with Redis_ to store the results then run: + $ pip install -U 'remoulade[postgres]' + +If you would like to use Redis_-backed extras like results or +cancellation, add the ``redis`` extra to whichever broker you choose:: $ pip install -U 'remoulade[rabbitmq, redis]' + $ pip install -U 'remoulade[postgres, redis]' If you don't have `pip`_ installed, check out `this guide`_. @@ -32,6 +37,7 @@ extra requirements: Name Description ============= ======================================================================================= ``rabbitmq`` Installs the required dependencies for using Remoulade with RabbitMQ. +``postgres`` Installs the required dependencies for using Remoulade with PostgreSQL and PGMQ. ``redis`` Installs the required dependencies for using Remoulade with Redis. ============= ======================================================================================= diff --git a/docs/source/reference.rst b/docs/source/reference.rst index ab5123adb..662956cf3 100644 --- a/docs/source/reference.rst +++ b/docs/source/reference.rst @@ -62,6 +62,9 @@ Brokers .. autoclass:: remoulade.brokers.rabbitmq.RabbitmqBroker :members: :inherited-members: +.. autoclass:: remoulade.brokers.pgmq.PgmqBroker + :members: + :inherited-members: .. autoclass:: remoulade.brokers.stub.StubBroker :members: :inherited-members: diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 19d4c2948..6a0c6c8bd 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -41,7 +41,11 @@ class PgmqBroker(Broker): - """A broker backed by PostgreSQL via the PGMQ extension.""" + """A broker backed by PostgreSQL via the PGMQ extension. + + PGMQ handles delayed messages natively, so delayed sends stay in + PostgreSQL instead of being staged in worker memory. + """ def __init__( self, @@ -122,6 +126,10 @@ def _try_enable_notify(self, queue_name: str) -> None: exc_info=True, ) + def _queue_exists(self, queue_name: str) -> bool: + """Return whether the queue already exists in PostgreSQL.""" + return queue_name in [queue.queue_name for queue in self.client.list_queues()] + @override def close(self) -> None: """Dispose the underlying PGMQ client.""" @@ -145,6 +153,10 @@ def declare_queue(self, queue_name: str) -> None: return self.client.validate_queue_name(queue_name, conn=self._current_connection) + if self._queue_exists(queue_name): + self.queues[queue_name] = None + return + self.emit_before("declare_queue", queue_name) self.client.create_partitioned_queue( queue_name, @@ -508,6 +520,7 @@ def __init__(self, pgmq_message: PgmqQueueMessage) -> None: message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8")) except TypeError as exc: raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc - + if message.options.get("eta", None) is not None: + raise UnsupportedMessageEncoding("eta option isn't supported with pgmq broker") super().__init__(message) self._pgmq_message = pgmq_message diff --git a/remoulade/worker.py b/remoulade/worker.py index 9dff76a6c..a9cc4b748 100644 --- a/remoulade/worker.py +++ b/remoulade/worker.py @@ -171,8 +171,7 @@ def join(self): """ while True: for consumer in self.consumers.values(): - if consumer.uses_in_memory_delay: - consumer.delay_queue.join() + consumer.delay_queue.join() self.work_queue.join() @@ -180,7 +179,7 @@ def join(self): # joining on the work queue then it should be safe to exit. # This could still miss stuff but the chances are slim. for consumer in self.consumers.values(): - if consumer.uses_in_memory_delay and consumer.delay_queue.unfinished_tasks: + if consumer.delay_queue.unfinished_tasks: break else: if self.work_queue.unfinished_tasks: @@ -242,7 +241,6 @@ def __init__(self, *, broker, queue_name, prefetch, work_queue, worker_timeout): self.queue_name = queue_name self.work_queue = work_queue self.worker_timeout = worker_timeout - self.uses_in_memory_delay = not broker.supports_native_delay self.delay_queue = PriorityQueue() # type: PriorityQueue def run(self): @@ -272,8 +270,7 @@ def run(self): elif self.paused: break - if self.uses_in_memory_delay: - self.handle_delayed_messages() + self.handle_delayed_messages() if not self.running: break @@ -326,7 +323,7 @@ def handle_message(self, message): work queue. """ try: - if self.uses_in_memory_delay and "eta" in message.options: + if "eta" in message.options: self.logger.debug("Pushing message %r onto delay queue.", message.message_id) self.broker.emit_before("delay_message", message) self.delay_queue.put((message.options.get("eta", 0), message)) @@ -391,8 +388,7 @@ def close(self): """Close this consumer thread and its underlying connection.""" try: if self.consumer: - if self.uses_in_memory_delay: - self.requeue_messages(m for _, m in iter_queue(self.delay_queue)) + self.requeue_messages(m for _, m in iter_queue(self.delay_queue)) self.consumer.close() except ConnectionError: pass diff --git a/tests/test_pgmq.py b/tests/test_pgmq.py index 5703dbd66..06eafc442 100644 --- a/tests/test_pgmq.py +++ b/tests/test_pgmq.py @@ -66,6 +66,7 @@ def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.create_queue = Mock() + broker._queue_exists = Mock(return_value=False) broker.declare_queue("default") @@ -83,6 +84,7 @@ def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) broker.declare_queue("default") @@ -100,6 +102,7 @@ def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) + broker._queue_exists = Mock(return_value=False) broker.declare_queue("default") @@ -107,6 +110,20 @@ def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): assert "Failed to enable LISTEN/NOTIFY" in caplog.text +def test_pgmq_broker_declare_queue_is_idempotent_when_queue_already_exists(): + broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=True) + + broker.declare_queue("default") + + broker.client.create_partitioned_queue.assert_not_called() + broker.client.enable_notify.assert_not_called() + assert "default" in broker.queues + + def test_pgmq_broker_uses_custom_partition_settings_when_provided(): broker = PgmqBroker( url=TEST_PGMQ_URL, @@ -117,6 +134,7 @@ def test_pgmq_broker_uses_custom_partition_settings_when_provided(): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) broker.declare_queue("default") From 1b94275307571a4d0385566f94a3a8a9aa430159 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 5 Jun 2026 15:32:46 +0200 Subject: [PATCH 14/44] fix(postgres): code review --- CONTRIBUTORS.md | 5 +- docs/source/guide.rst | 23 +- docs/source/reference.rst | 2 +- .../2026-06-05-postgres-broker-ms-timings.md | 125 ++++ ...pgmq_broker.py => local_postgres_broker.py | 4 +- ..._consumer.py => local_postgres_consumer.py | 8 +- pyproject.toml | 2 +- remoulade/brokers/pgmq.py | 539 +---------------- remoulade/brokers/postgres.py | 541 ++++++++++++++++++ tests/conftest.py | 16 +- tests/docker-compose.yml | 2 +- ...n.sql => 01-create-postgres-extension.sql} | 0 tests/{test_pgmq.py => test_postgres.py} | 194 +++---- 13 files changed, 808 insertions(+), 653 deletions(-) create mode 100644 docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md rename local_pgmq_broker.py => local_postgres_broker.py (91%) rename local_pgmq_consumer.py => local_postgres_consumer.py (80%) create mode 100644 remoulade/brokers/postgres.py rename tests/docker/postgres/{01-create-pgmq-extension.sql => 01-create-postgres-extension.sql} (100%) rename tests/{test_pgmq.py => test_postgres.py} (67%) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index a873de1eb..34caddd4c 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -32,5 +32,6 @@ of those changes to CLEARTYPE SRL. | [@pgitips](https://github.com/pgitips) | Pierre Giraud | | [@williampollet](https://github.com/williampollet) | William Pollet | | [@mehdithez](https://github.com/mehdithez) | Zeroual Mehdi | -| [@julien-duponchelle](https://github.com/julien-duponchelle) | Julien Duponchelle | -| [@alisterd51](https://github.com/alisterd51) | Antoine Clarman | +| [@mducros-wm](https://github.com/mducros-wm) | Martin Ducros | +| [@julien-duponchelle](https://github.com/julien-duponchelle) | Juli en Duponchelle | +| [@alisterd51](https://github.com/alisterd51) | Antoine Clarman | diff --git a/docs/source/guide.rst b/docs/source/guide.rst index 31a5d920a..91c874804 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -170,7 +170,7 @@ actor an invalid URL. Let's try it:: Message Retries ---------------- +^^^^^^^^^^^^^^^ If an error occurs during message processing, it will be terminated with a failure message. Alternatively, you can add the |Retries| Middleware to the broker and set the max_retries or retry_when option to automatically retry your message on failure. @@ -203,12 +203,12 @@ If you want to use a different strategy than the default exponential backoff to The following retry options are configurable on a per-actor basis: max_retries -^^^^^^^^^^^ +--------------- The maximum number of times a message should be retried. Default to ``0``. min_backoff -^^^^^^^^^^^ +^^^^^^^^^^^^^^^ The minimum number of milliseconds of backoff to apply between retries. Must be greater than 100 milliseconds. Defaults to 15 seconds. @@ -332,7 +332,7 @@ Keep in mind that *your message broker is not a database*. Scheduled messages should represent a small subset of all your messages. On brokers that emulate delay in worker memory, the enqueued message -will carry an ``eta`` option. ``PgmqBroker`` stores delayed messages in +will carry an ``eta`` option. ``PostgresBroker`` stores delayed messages in PostgreSQL natively instead, so the message it returns does not include ``eta``. @@ -403,19 +403,20 @@ execution:: remoulade.set_broker(rabbitmq_broker) -PGMQ Broker -^^^^^^^^^^^ +Postgres Broker +^^^^^^^^^^^^^^^ To configure PostgreSQL/PGMQ, install ``remoulade[postgres]`` and -instantiate a ``PgmqBroker`` with a PostgreSQL URL as early as possible -during your program's execution:: +instantiate a ``PostgresBroker`` with a PostgreSQL URL as early as possible +during your program's execution. This broker be used with a user postgresql who can create and +delete tables:: import remoulade - from remoulade.brokers.pgmq import PgmqBroker + from remoulade.brokers.postgres import PostgresBroker - pgmq_broker = PgmqBroker(url="postgresql://remoulade@localhost:5432/remoulade") - remoulade.set_broker(pgmq_broker) + postgres_broker = PostgresBroker(url="postgresql://remoulade@localhost:5432/remoulade") + remoulade.set_broker(postgres_broker) PGMQ handles delayed messages natively, so ``send_with_options(delay=...)`` does not create a worker-side delay queue or add an ``eta`` option to diff --git a/docs/source/reference.rst b/docs/source/reference.rst index 662956cf3..43382c171 100644 --- a/docs/source/reference.rst +++ b/docs/source/reference.rst @@ -62,7 +62,7 @@ Brokers .. autoclass:: remoulade.brokers.rabbitmq.RabbitmqBroker :members: :inherited-members: -.. autoclass:: remoulade.brokers.pgmq.PgmqBroker +.. autoclass:: remoulade.brokers.postgres.PostgresBroker :members: :inherited-members: .. autoclass:: remoulade.brokers.stub.StubBroker diff --git a/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md b/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md new file mode 100644 index 000000000..64e4ca058 --- /dev/null +++ b/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md @@ -0,0 +1,125 @@ +# Postgres Broker Millisecond Timings Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make every Postgres broker timing input accept milliseconds, while keeping PGMQ and thread APIs working with their required units internally. + +**Architecture:** Rename the broker-facing timeout parameters to `*_ms`, keep public consumer timing inputs in milliseconds, and convert to seconds only at the boundary where PGMQ or `threading.Event.wait()` require it. Update the Postgres tests to assert the new API and the conversion boundary explicitly. + +**Tech Stack:** Python, pytest, psycopg, pgmq, SQLAlchemy + +--- + +### Task 1: Update timing API names + +**Files:** +- Modify: `remoulade/brokers/postgres.py` +- Modify: `tests/test_postgres.py` + +- [ ] **Step 1: Write the failing test** + +```python +broker = PostgresBroker( + url=TEST_POSTGRES_URL, + middleware=[], + visibility_timeout_ms=17_000, + heartbeat_interval_ms=500, +) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `pytest tests/test_postgres.py -k postgres_consumer_uses_broker_visibility_timeout_for_reads -v` +Expected: `TypeError` because the broker still expects `*_in_second` arguments. + +- [ ] **Step 3: Write minimal implementation** + +```python +def __init__( + self, + *, + url: str, + middleware: list["Middleware"] | None = None, + group_transaction: bool = False, + archive_partition_interval: str = "1 day", + archive_retention_interval: str = "7 days", + visibility_timeout_ms: int = 30_000, + heartbeat_interval_ms: int = 10_000, +) -> None: + if visibility_timeout_ms <= 0: + raise ValueError("visibility_timeout_ms must be greater than 0") + if heartbeat_interval_ms <= 0 or heartbeat_interval_ms >= visibility_timeout_ms: + raise ValueError("heartbeat_interval_ms must be greater than 0 and lower than visibility_timeout_ms") +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `pytest tests/test_postgres.py -k postgres_consumer_uses_broker_visibility_timeout_for_reads -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add remoulade/brokers/postgres.py tests/test_postgres.py docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md +git commit -m "feat: use milliseconds for postgres broker timing inputs" +``` + +### Task 2: Convert broker internals to seconds + +**Files:** +- Modify: `remoulade/brokers/postgres.py` +- Modify: `tests/test_postgres.py` + +- [ ] **Step 1: Write the failing test** + +```python +assert broker.client.read.call_args.kwargs["vt"] == 17 +assert broker.client.set_vt.call_args.args[2] == 2 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `pytest tests/test_postgres.py -k 'postgres_consumer_uses_broker_visibility_timeout_for_reads or postgres_consumer_heartbeat_extends_inflight_message_visibility' -v` +Expected: `vt` still reflects millisecond inputs instead of converted seconds. + +- [ ] **Step 3: Write minimal implementation** + +```python +from math import ceil + +self.visibility_timeout = max(1, ceil(self.visibility_timeout_ms / 1000)) +self.heartbeat_interval = self.heartbeat_interval_ms / 1000 +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `pytest tests/test_postgres.py -k 'postgres_consumer_uses_broker_visibility_timeout_for_reads or postgres_consumer_heartbeat_extends_inflight_message_visibility' -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add remoulade/brokers/postgres.py tests/test_postgres.py +git commit -m "fix: convert postgres broker timing inputs" +``` + +### Task 3: Verify the full Postgres test file + +**Files:** +- Test: `tests/test_postgres.py` + +- [ ] **Step 1: Run the whole file** + +Run: `pytest tests/test_postgres.py -v` +Expected: PASS + +- [ ] **Step 2: Fix any remaining unit mismatches** + +Adjust only the Postgres timing assertions or docstrings if a legacy `*_in_second` name remains. + +- [ ] **Step 3: Commit** + +```bash +git add tests/test_postgres.py +git commit -m "test: cover postgres broker millisecond timings" +``` diff --git a/local_pgmq_broker.py b/local_postgres_broker.py similarity index 91% rename from local_pgmq_broker.py rename to local_postgres_broker.py index 857a96a7e..f8046e34b 100644 --- a/local_pgmq_broker.py +++ b/local_postgres_broker.py @@ -5,7 +5,7 @@ import remoulade from remoulade import Message -from remoulade.brokers.pgmq import PgmqBroker +from remoulade.brokers.postgres import PostgresBroker # Just for local test URL = "postgresql://remoulade@localhost:5544/test" @@ -13,7 +13,7 @@ ACTOR_NAME = "demo.add" -broker = PgmqBroker( +broker = PostgresBroker( url=URL, middleware=[], ) diff --git a/local_pgmq_consumer.py b/local_postgres_consumer.py similarity index 80% rename from local_pgmq_consumer.py rename to local_postgres_consumer.py index 8ff921b94..26e019eb9 100644 --- a/local_pgmq_consumer.py +++ b/local_postgres_consumer.py @@ -2,14 +2,14 @@ import remoulade from remoulade import Worker -from remoulade.brokers.pgmq import PgmqBroker +from remoulade.brokers.postgres import PostgresBroker -# Same DSN as local_pgmq_broker.py +# Same DSN as local_postgres_broker.py URL = "postgresql://remoulade@localhost:5544/test" QUEUE_NAME = "default" ACTOR_NAME = "demo.add" -broker = PgmqBroker( +broker = PostgresBroker( url=URL, middleware=[], ) @@ -34,7 +34,7 @@ def add(x: int, y: int) -> int: worker = Worker(broker, queues={QUEUE_NAME}, worker_threads=1, worker_timeout=30_000) worker.start() - print(f"PGMQ local consumer started on queue '{QUEUE_NAME}' ({listener_mode}). Press Ctrl+C to stop.") + print(f"Postgres local consumer started on queue '{QUEUE_NAME}' ({listener_mode}). Press Ctrl+C to stop.") try: while True: time.sleep(1) diff --git a/pyproject.toml b/pyproject.toml index b78986d59..7286a6aea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ rabbitmq = ["amqpstorm>=2.6,<3"] redis = ["redis>=7.0.0"] server = ["flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec"] -postgres = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1"] +postgres = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1,<2"] pydantic = ["pydantic>=2.12", "simplejson"] limits = ["limits~=5.3.0"] tracing = ["opentelemetry-api>=1.20"] diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py index 6a0c6c8bd..787649a26 100644 --- a/remoulade/brokers/pgmq.py +++ b/remoulade/brokers/pgmq.py @@ -1,526 +1,13 @@ -# This file is a part of Remoulade. -# -# Copyright (C) 2026 WIREMIND SAS -# -# Remoulade is free software; you can redistribute it and/or modify it -# under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or (at -# your option) any later version. -# -# Remoulade is distributed in the hope that it will be useful, but WITHOUT -# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or -# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public -# License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . -import json -import time -from collections import deque -from collections.abc import Iterable, Iterator -from contextlib import contextmanager -from datetime import UTC, datetime, timedelta -from threading import Event, Lock, Thread, local -from typing import TYPE_CHECKING, Any, Final, override - -import psycopg -from pgmq import SQLAlchemyPGMQueue -from pgmq.messages import Message as PgmqQueueMessage -from psycopg import sql as psycopg_sql -from sqlalchemy import Column, Connection, Integer, MetaData, Table, func, select - -from ..broker import Broker, Consumer, MessageProxy -from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding -from ..message import Message - -if TYPE_CHECKING: - from ..middleware import Middleware - -PgmqPayload = dict[str, Any] -LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 - - -class PgmqBroker(Broker): - """A broker backed by PostgreSQL via the PGMQ extension. - - PGMQ handles delayed messages natively, so delayed sends stay in - PostgreSQL instead of being staged in worker memory. - """ - - def __init__( - self, - *, - url: str, - middleware: list["Middleware"] | None = None, - group_transaction: bool = False, - archive_partition_interval: str = "1 day", - archive_retention_interval: str = "7 days", - visibility_timeout_in_second: int = 30, - heartbeat_interval_in_second: float = 10.0, - ) -> None: - """Initialize a PostgreSQL-backed broker using the PGMQ extension. - - Parameters: - url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg. - middleware(list[Middleware] | None): Middleware stack applied to this broker. - group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. - archive_partition_interval(str): Partition interval passed to PGMQ when creating partitioned queues. - archive_retention_interval(str): Retention interval passed to PGMQ for archive partitions. - visibility_timeout_in_second(int): Message visibility timeout in seconds after read; must be greater than 0. - heartbeat_interval_in_second(float): Heartbeat interval in seconds used to extend in-flight message visibility - must be greater than 0 and lower than visibility_timeout_in_second - """ - super().__init__(middleware=middleware) - if visibility_timeout_in_second <= 0: - raise ValueError("visibility_timeout_in_second must be greater than 0") - if heartbeat_interval_in_second <= 0 or heartbeat_interval_in_second >= visibility_timeout_in_second: - raise ValueError("heartbeat_interval must be greater than 0 and lower than visibility_timeout") - - self.url = url - self.state = local() - self.group_transaction = group_transaction - self.archive_partition_interval = archive_partition_interval - self.archive_retention_interval = archive_retention_interval - self.visibility_timeout_in_second = visibility_timeout_in_second - self.heartbeat_interval_in_second = heartbeat_interval_in_second - self.supports_native_delay = True - - self.client = SQLAlchemyPGMQueue(conn_string=url, init_extension=False, vt=self.visibility_timeout_in_second) - - @override - @contextmanager - def tx(self) -> Iterator[None]: - """Run broker operations inside a single SQL transaction. - - The active SQLAlchemy connection is stored on thread-local state so - queue operations can reuse it while the context is open. - """ - with self.client.engine.begin() as connection: - self.state.transaction_connection = connection - try: - yield - finally: - self.state.transaction_connection = None - - @property - def _current_connection(self) -> Connection | None: - """Return the transactional connection, if one is active.""" - return getattr(self.state, "transaction_connection", None) - - def _try_enable_notify(self, queue_name: str) -> None: - """Try to enable LISTEN/NOTIFY for a queue. - - Failures are logged and the consumer later falls back to polling. - """ - try: - self.client.enable_notify( - queue_name, - throttle_interval_ms=LISTEN_NOTIFY_THROTTLE_MS, - conn=self._current_connection, - ) - except Exception as e: - self.logger.warning( - "Failed to enable LISTEN/NOTIFY for queue %s; consumer will fall back to polling. Error: %s", - queue_name, - e, - exc_info=True, - ) - - def _queue_exists(self, queue_name: str) -> bool: - """Return whether the queue already exists in PostgreSQL.""" - return queue_name in [queue.queue_name for queue in self.client.list_queues()] - - @override - def close(self) -> None: - """Dispose the underlying PGMQ client.""" - self.client.dispose() - - @override - def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PgmqConsumer": - """Create a consumer for a declared queue. - - Raises: - QueueNotFound: If the queue has not been declared. - """ - if queue_name not in self.queues: - raise QueueNotFound(queue_name) - return _PgmqConsumer(self, queue_name=queue_name, prefetch=prefetch, timeout=timeout) - - @override - def declare_queue(self, queue_name: str) -> None: - """Create a partitioned PGMQ queue if it does not already exist.""" - if queue_name in self.queues: - return - - self.client.validate_queue_name(queue_name, conn=self._current_connection) - if self._queue_exists(queue_name): - self.queues[queue_name] = None - return - - self.emit_before("declare_queue", queue_name) - self.client.create_partitioned_queue( - queue_name, - partition_interval=self.archive_partition_interval, - retention_interval=self.archive_retention_interval, - conn=self._current_connection, - ) - self._try_enable_notify(queue_name) - self.queues[queue_name] = None - self.emit_after("declare_queue", queue_name) - - def _encode_message(self, message: "Message") -> PgmqPayload: - """Decode a Remoulade message into a JSON object payload for PGMQ. - - Raises: - UnsupportedMessageEncoding: If the encoded payload is not valid JSON - or is not a JSON object. - """ - try: - payload = json.loads(message.encode().decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as e: - raise UnsupportedMessageEncoding( - "PgmqBroker only supports encoders that serialize messages as JSON objects." - ) from e - - if not isinstance(payload, dict): - raise UnsupportedMessageEncoding("PgmqBroker requires message encoders to produce JSON objects.") - - return payload - - @override - def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": - """Send a message to PGMQ, optionally delayed by milliseconds.""" - if message.queue_name not in self.queues: - raise QueueNotFound(message.queue_name) - - payload = self._encode_message(message) - visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) if delay is not None else None - self.client.send( - message.queue_name, - payload, - conn=self._current_connection, - delay=visible_at, - ) - - return message - - @override - def flush(self, queue_name: str) -> None: - """Remove every message currently stored in a queue.""" - if queue_name not in self.queues: - raise QueueNotFound(queue_name) - - self.client.purge(queue_name, conn=self._current_connection) - - @override - def flush_all(self) -> None: - """Purge every declared queue.""" - for queue_name in self.queues: - self.flush(queue_name) - - def _count_enqueued_messages(self, queue_name: str) -> int: - """Count rows in the underlying PGMQ queue table.""" - queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") - count_query = select(func.count()).select_from(queue_table) - - with self.client.session() as session: - return session.execute(count_query).scalar_one() - - @override - def join( - self, - queue_name: str, - min_successes: int = 10, - idle_time: int = 100, - *, - timeout: int | None = None, - ) -> None: - """Wait for all the messages on the given queue to be processed. - - This method checks the full PGMQ queue table and therefore waits for - all states: visible messages, invisible in-flight messages and native - delayed messages. - - Parameters: - queue_name(str): The queue to wait on. - min_successes(int): The minimum number of times the queue should be - observed as empty. - idle_time(int): The number of milliseconds to wait between checks. - timeout(Optional[int]): The max amount of time, in milliseconds, to - wait on this queue. - - Raises: - QueueNotFound: If the given queue was never declared. - QueueJoinTimeout: When the timeout elapses. - """ - if queue_name not in self.queues: - raise QueueNotFound(queue_name) - - deadline = time.monotonic() + timeout / 1000 if timeout is not None else None - successes = 0 - - while successes < min_successes: - if deadline and time.monotonic() >= deadline: - raise QueueJoinTimeout(queue_name) - - total_messages = self._count_enqueued_messages(queue_name) - if total_messages == 0: - successes += 1 - if successes >= min_successes: - return - else: - successes = 0 - - time.sleep(idle_time / 1000) - - -class _PgmqConsumer(Consumer): - def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int) -> None: - """Initialize a consumer for a PGMQ queue. - - Parameters: - broker(PgmqBroker): Broker instance that owns the queue and database client. - queue_name(str): Name of the declared queue to consume from. - prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1. - timeout(int): Idle wait timeout in milliseconds when polling for messages; values lower than 0 are coerced to - 0. A value of 0 performs non-blocking reads. - """ - self.broker = broker - self.client = broker.client - self.queue_name = queue_name - self.prefetch = prefetch - self.timeout = timeout - self.visibility_timeout = broker.visibility_timeout_in_second - self.heartbeat_interval = broker.heartbeat_interval_in_second - self.messages: deque[PgmqQueueMessage] = deque() - self._notify_event = Event() - self._listener_stop = Event() - self._listener_thread: Thread | None = None - self._listener_connection: psycopg.Connection[Any] | None = None - self._listener_available = False - self._heartbeat_stop = Event() - self._heartbeat_thread: Thread | None = None - self._heartbeat_msg_ids_lock = Lock() - self._heartbeat_msg_ids: set[int] = set() - - self._start_listener() - self._start_heartbeat() - - @property - def wait_timeout_seconds(self) -> float: - """Return the listener wait timeout in seconds.""" - return self.timeout / 1000 if self.timeout > 0 else 0 - - def _normalize_messages(self, messages: PgmqQueueMessage | list[PgmqQueueMessage] | None) -> list[PgmqQueueMessage]: - """Normalize a PGMQ read result into a list.""" - if messages is None: - return [] - return [messages] if not isinstance(messages, list) else messages - - def _read_immediate(self) -> list[PgmqQueueMessage]: - """Read up to ``prefetch`` messages without polling.""" - return self._normalize_messages( - self.client.read( - self.queue_name, - vt=self.visibility_timeout, - qty=self.prefetch, - ) - ) - - def _read_with_poll(self) -> list[PgmqQueueMessage]: - """Read up to ``prefetch`` messages using PGMQ polling.""" - return self._normalize_messages( - self.client.read_with_poll( - self.queue_name, - vt=self.visibility_timeout, - qty=self.prefetch, - max_poll_seconds=max((self.timeout + 999) // 1000, 1), - poll_interval_ms=max(min(self.timeout, 1000), 1), - ) - ) - - def _start_heartbeat(self) -> None: - """Start the background visibility-timeout renewal thread.""" - self._heartbeat_thread = Thread( - target=self._run_heartbeat, - name=f"pgmq-heartbeat-{self.queue_name}", - daemon=True, - ) - self._heartbeat_thread.start() - - def _run_heartbeat(self) -> None: - """Periodically renew the lease of in-flight messages. - - The loop sleeps for ``heartbeat_interval`` seconds, unless shutdown is - requested through ``_heartbeat_stop``. On each tick it snapshots the - current set of tracked message ids under a lock, then calls - ``set_vt`` to push their visibility timeout back to - ``self.visibility_timeout``. - - This prevents long-running handlers from becoming visible again and - being consumed twice before they are explicitly acked, nacked or - requeued. Failures are logged and retried on the next heartbeat tick. - """ - while not self._heartbeat_stop.wait(self.heartbeat_interval): - with self._heartbeat_msg_ids_lock: - msg_ids = list(self._heartbeat_msg_ids) - if not msg_ids: - continue - try: - self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout) - except Exception as e: - self.broker.logger.warning( - "Failed to extend visibility timeout heartbeat for queue %s. Error: ", - self.queue_name, - str(e), - exc_info=True, - ) - - def _register_heartbeat_msg(self, message: "_PgmqMessage") -> None: - """Track a message so the heartbeat can renew it.""" - with self._heartbeat_msg_ids_lock: - self._heartbeat_msg_ids.add(message._pgmq_message.msg_id) - - def _unregister_heartbeat_msg_id(self, msg_id: int) -> None: - """Stop tracking a message for heartbeat renewal.""" - with self._heartbeat_msg_ids_lock: - self._heartbeat_msg_ids.discard(msg_id) - - def _start_listener(self) -> None: - """Start a LISTEN/NOTIFY listener for the queue when available.""" - channel = f"pgmq.q_{self.queue_name}.INSERT" - try: - self._listener_connection = psycopg.connect(self.broker.url, autocommit=True) - self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) - self._listener_available = True - except Exception as e: - self._listener_available = False - self._listener_connection = None - self.broker.logger.warning( - "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling. Error: %s", - self.queue_name, - str(e), - exc_info=True, - ) - return - - self._listener_thread = Thread( - target=self._run_listener, - name=f"pgmq-listener-{self.queue_name}", - daemon=True, - ) - self._listener_thread.start() - - def _run_listener(self) -> None: - """Wake the consumer when a notification arrives, or disable the listener on errors.""" - while not self._listener_stop.is_set(): - try: - if self._listener_connection is None: - break - for _ in self._listener_connection.notifies(timeout=0.5, stop_after=1): - self._notify_event.set() - except Exception as e: - if self._listener_stop.is_set(): - break - self._listener_available = False - self._notify_event.set() - self.broker.logger.warning( - "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling. Error: %s", - self.queue_name, - str(e), - exc_info=True, - ) - break - - self._listener_available = False - - @override - def ack(self, message: "MessageProxy") -> None: - """Archive a processed message.""" - if not isinstance(message, _PgmqMessage): - raise ValueError("It must be a PgmqMessage") - self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) - self.client.archive(self.queue_name, message._pgmq_message.msg_id) - - @override - def nack(self, message: "MessageProxy") -> None: - """Archive a failed message.""" - if not isinstance(message, _PgmqMessage): - raise ValueError("It must be a PgmqMessage") - self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id) - self.client.archive(self.queue_name, message._pgmq_message.msg_id) - - @override - def requeue(self, messages: Iterable["MessageProxy"]) -> None: - """Make messages visible again immediately by resetting their visibility timeout.""" - msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)] - if msg_ids: - for msg_id in msg_ids: - self._unregister_heartbeat_msg_id(msg_id) - self.client.set_vt(self.queue_name, msg_ids, 0) - - def _build_message(self, pgmq_message: PgmqQueueMessage) -> "_PgmqMessage": - """Wrap a raw PGMQ row and register it for heartbeat renewal.""" - message = _PgmqMessage(pgmq_message) - self._register_heartbeat_msg(message) - return message - - @override - def __next__(self) -> "_PgmqMessage | None": - """Return the next available message, or ``None`` if the queue stays empty.""" - if self.messages: - return self._build_message(self.messages.popleft()) - - if self._listener_available: - self._notify_event.wait(self.wait_timeout_seconds) - messages = self._read_immediate() - if not messages: - messages = self._read_with_poll() - self._listener_available = False - self._notify_event.clear() - else: - messages = self._read_with_poll() - - if not messages: - return None - - self.messages.extend(messages) - return self._build_message(self.messages.popleft()) - - @override - def close(self) -> None: - """Stop background threads and close the listener connection.""" - self._listener_stop.set() - self._heartbeat_stop.set() - self._notify_event.set() - if self._listener_connection is not None: - try: - self._listener_connection.close() - except Exception as e: - self.broker.logger.error("Listener not joinable: %s", str(e)) - finally: - self._listener_connection = None - if self._listener_thread is not None: - self._listener_thread.join(timeout=1) - if self._heartbeat_thread is not None: - self._heartbeat_thread.join(timeout=1) - return - - -class _PgmqMessage(MessageProxy): - def __init__(self, pgmq_message: PgmqQueueMessage) -> None: - """Wrap a PGMQ message row as a Remoulade message proxy.""" - payload = pgmq_message.message - if not isinstance(payload, dict): - raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") - - try: - # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder) - # can rehydrate actor args/kwargs to their typed schemas. - message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8")) - except TypeError as exc: - raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc - if message.options.get("eta", None) is not None: - raise UnsupportedMessageEncoding("eta option isn't supported with pgmq broker") - super().__init__(message) - self._pgmq_message = pgmq_message +from .postgres import ( + PostgresBroker, + PostgresPayload, + PostgresQueueMessage, + _PostgresConsumer, + _PostgresMessage, +) + +PgmqBroker = PostgresBroker +PgmqPayload = PostgresPayload +PgmqQueueMessage = PostgresQueueMessage +_PgmqConsumer = _PostgresConsumer +_PgmqMessage = _PostgresMessage diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py new file mode 100644 index 000000000..a2c07ccd3 --- /dev/null +++ b/remoulade/brokers/postgres.py @@ -0,0 +1,541 @@ +# This file is a part of Remoulade. +# +# Copyright (C) 2026 WIREMIND SAS +# +# Remoulade is free software; you can redistribute it and/or modify it +# under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or (at +# your option) any later version. +# +# Remoulade is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public +# License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . +import json +import math +import time +from collections import deque +from collections.abc import Iterable, Iterator +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from threading import Event, Lock, Thread, local +from typing import TYPE_CHECKING, Any, Final, override +from urllib.parse import urlparse + +import psycopg +from pgmq import SQLAlchemyPGMQueue +from pgmq.messages import Message as PostgresQueueMessage +from psycopg import sql as psycopg_sql +from sqlalchemy import Column, Connection, Integer, MetaData, Table, func, select + +from ..broker import Broker, Consumer, MessageProxy +from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding +from ..message import Message + +if TYPE_CHECKING: + from ..middleware import Middleware + +PostgresPayload = dict[str, Any] +LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 + + +def _milliseconds_to_seconds(milliseconds: int) -> int: + """Convert milliseconds to whole seconds, rounding up for PGMQ inputs.""" + return max(1, math.ceil(milliseconds / 1000)) + + +class PostgresBroker(Broker): + """A broker backed by PostgreSQL via the PGMQ extension. + + PGMQ handles delayed messages natively, so delayed sends stay in + PostgreSQL instead of being staged in worker memory. + """ + + def __init__( + self, + *, + url: str, + middleware: list["Middleware"] | None = None, + group_transaction: bool = False, + archive_partition_interval: str = "1 day", + archive_retention_interval: str = "7 days", + visibility_timeout_ms: int = 30_000, + heartbeat_interval_ms: int = 10_000, + ) -> None: + """Initialize a PostgreSQL-backed broker using the PGMQ extension. + + Parameters: + url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg. + The url must be the creds for a user who can create and delete tables + middleware(list[Middleware] | None): Middleware stack applied to this broker. + group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. + archive_partition_interval(str): Partition interval passed to PGMQ when creating partitioned queues. + archive_retention_interval(str): Retention interval passed to PGMQ for archive partitions. + visibility_timeout_ms(int): Message visibility timeout in milliseconds after read; must be greater than 0. + heartbeat_interval_ms(int): Heartbeat interval in milliseconds used to extend in-flight message visibility + must be greater than 0 and lower than visibility_timeout_ms. + """ + super().__init__(middleware=middleware) + if visibility_timeout_ms <= 0: + raise ValueError("visibility_timeout_ms must be greater than 0") + if heartbeat_interval_ms <= 0 or heartbeat_interval_ms >= visibility_timeout_ms: + raise ValueError("heartbeat_interval_ms must be greater than 0 and lower than visibility_timeout_ms") + + self.url = urlparse(url).geturl() + self.state = local() + self.group_transaction = group_transaction + self.archive_partition_interval = archive_partition_interval + self.archive_retention_interval = archive_retention_interval + self.visibility_timeout_ms = visibility_timeout_ms + self.heartbeat_interval_ms = heartbeat_interval_ms + self.visibility_timeout_seconds = _milliseconds_to_seconds(visibility_timeout_ms) + self.heartbeat_interval_seconds = heartbeat_interval_ms / 1000 + self.supports_native_delay = True + + self.client = SQLAlchemyPGMQueue( + conn_string=url, + init_extension=False, + vt=self.visibility_timeout_seconds, + ) + + @override + @contextmanager + def tx(self) -> Iterator[None]: + """Run broker operations inside a single SQL transaction. + + The active SQLAlchemy connection is stored on thread-local state so + queue operations can reuse it while the context is open. + """ + with self.client.engine.begin() as connection: + self.state.transaction_connection = connection + try: + yield + finally: + self.state.transaction_connection = None + + @property + def _current_connection(self) -> Connection | None: + """Return the transactional connection, if one is active.""" + return getattr(self.state, "transaction_connection", None) + + def _try_enable_notify(self, queue_name: str) -> None: + """Try to enable LISTEN/NOTIFY for a queue. + + Failures are logged and the consumer later falls back to polling. + """ + try: + self.client.enable_notify( + queue_name, + throttle_interval_ms=LISTEN_NOTIFY_THROTTLE_MS, + conn=self._current_connection, + ) + except Exception as e: + self.logger.warning( + "Failed to enable LISTEN/NOTIFY for queue %s; consumer will fall back to polling. Error: %s", + queue_name, + e, + exc_info=True, + ) + + def _queue_exists(self, queue_name: str) -> bool: + """Return whether the queue already exists in PostgreSQL.""" + return queue_name in [queue.queue_name for queue in self.client.list_queues()] + + @override + def close(self) -> None: + """Dispose the underlying PGMQ client.""" + self.client.dispose() + + @override + def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PostgresConsumer": + """Create a consumer for a declared queue. + + Raises: + QueueNotFound: If the queue has not been declared. + """ + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + return _PostgresConsumer(self, queue_name=queue_name, prefetch=prefetch, timeout=timeout) + + @override + def declare_queue(self, queue_name: str) -> None: + """Create a partitioned PGMQ queue if it does not already exist.""" + if queue_name in self.queues: + return + + self.client.validate_queue_name(queue_name, conn=self._current_connection) + if self._queue_exists(queue_name): + self.queues[queue_name] = None + return + + self.emit_before("declare_queue", queue_name) + self.client.create_partitioned_queue( + queue_name, + partition_interval=self.archive_partition_interval, + retention_interval=self.archive_retention_interval, + conn=self._current_connection, + ) + self._try_enable_notify(queue_name) + self.queues[queue_name] = None + self.emit_after("declare_queue", queue_name) + + def _encode_message(self, message: "Message") -> PostgresPayload: + """Decode a Remoulade message into a JSON object payload for PGMQ. + + Raises: + UnsupportedMessageEncoding: If the encoded payload is not valid JSON + or is not a JSON object. + """ + try: + payload = json.loads(message.encode().decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as e: + raise UnsupportedMessageEncoding( + "PostgresBroker only supports encoders that serialize messages as JSON objects." + ) from e + + if not isinstance(payload, dict): + raise UnsupportedMessageEncoding("PostgresBroker requires message encoders to produce JSON objects.") + + return payload + + @override + def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": + """Send a message to PGMQ, optionally delayed by milliseconds.""" + if message.queue_name not in self.queues: + raise QueueNotFound(message.queue_name) + + payload = self._encode_message(message) + visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) if delay is not None else None + self.client.send( + message.queue_name, + payload, + conn=self._current_connection, + delay=visible_at, + ) + + return message + + @override + def flush(self, queue_name: str) -> None: + """Remove every message currently stored in a queue.""" + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + + self.client.purge(queue_name, conn=self._current_connection) + + @override + def flush_all(self) -> None: + """Purge every declared queue.""" + for queue_name in self.queues: + self.flush(queue_name) + + def _count_enqueued_messages(self, queue_name: str) -> int: + """Count rows in the underlying PGMQ queue table.""" + queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") + count_query = select(func.count()).select_from(queue_table) + + with self.client.session() as session: + return session.execute(count_query).scalar_one() + + @override + def join( + self, + queue_name: str, + min_successes: int = 10, + idle_time: int = 100, + *, + timeout: int | None = None, + ) -> None: + """Wait for all the messages on the given queue to be processed. + + This method checks the full PGMQ queue table and therefore waits for + all states: visible messages, invisible in-flight messages and native + delayed messages. + + Parameters: + queue_name(str): The queue to wait on. + min_successes(int): The minimum number of times the queue should be + observed as empty. + idle_time(int): The number of milliseconds to wait between checks. + timeout(Optional[int]): The max amount of time, in milliseconds, to + wait on this queue. + + Raises: + QueueNotFound: If the given queue was never declared. + QueueJoinTimeout: When the timeout elapses. + """ + if queue_name not in self.queues: + raise QueueNotFound(queue_name) + + deadline = time.monotonic() + timeout / 1000 if timeout is not None else None + successes = 0 + + while successes < min_successes: + if deadline and time.monotonic() >= deadline: + raise QueueJoinTimeout(queue_name) + + total_messages = self._count_enqueued_messages(queue_name) + if total_messages == 0: + successes += 1 + if successes >= min_successes: + return + else: + successes = 0 + + time.sleep(idle_time / 1000) + + +class _PostgresConsumer(Consumer): + def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, timeout: int) -> None: + """Initialize a consumer for a PGMQ queue. + + Parameters: + broker(PostgresBroker): Broker instance that owns the queue and database client. + queue_name(str): Name of the declared queue to consume from. + prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1. + timeout(int): Idle wait timeout in milliseconds when polling for messages; values lower than 0 are coerced to + 0. A value of 0 performs non-blocking reads. + """ + self.broker = broker + self.client = broker.client + self.queue_name = queue_name + self.prefetch = prefetch + self.timeout = timeout + self.visibility_timeout_seconds = broker.visibility_timeout_seconds + self.heartbeat_interval_seconds = broker.heartbeat_interval_seconds + self.messages: deque[PostgresQueueMessage] = deque() + self._notify_event = Event() + self._listener_stop = Event() + self._listener_thread: Thread | None = None + self._listener_connection: psycopg.Connection[Any] | None = None + self._listener_available = False + self._heartbeat_stop = Event() + self._heartbeat_thread: Thread | None = None + self._heartbeat_msg_ids_lock = Lock() + self._heartbeat_msg_ids: set[int] = set() + + self._start_listener() + self._start_heartbeat() + + @property + def wait_timeout_seconds(self) -> float: + """Return the listener wait timeout in seconds.""" + return self.timeout / 1000 if self.timeout > 0 else 0 + + def _normalize_messages( + self, messages: PostgresQueueMessage | list[PostgresQueueMessage] | None + ) -> list[PostgresQueueMessage]: + """Normalize a PGMQ read result into a list.""" + if messages is None: + return [] + return [messages] if not isinstance(messages, list) else messages + + def _read_immediate(self) -> list[PostgresQueueMessage]: + """Read up to ``prefetch`` messages without polling.""" + return self._normalize_messages( + self.client.read( + self.queue_name, + vt=self.visibility_timeout_seconds, + qty=self.prefetch, + ) + ) + + def _read_with_poll(self) -> list[PostgresQueueMessage]: + """Read up to ``prefetch`` messages using PGMQ polling.""" + return self._normalize_messages( + self.client.read_with_poll( + self.queue_name, + vt=self.visibility_timeout_seconds, + qty=self.prefetch, + max_poll_seconds=max((self.timeout + 999) // 1000, 1), + poll_interval_ms=max(min(self.timeout, 1000), 1), + ) + ) + + def _start_heartbeat(self) -> None: + """Start the background visibility-timeout renewal thread.""" + self._heartbeat_thread = Thread( + target=self._run_heartbeat, + name=f"postgres-heartbeat-{self.queue_name}", + daemon=True, + ) + self._heartbeat_thread.start() + + def _run_heartbeat(self) -> None: + """Periodically renew the lease of in-flight messages. + + The loop sleeps for ``heartbeat_interval_seconds`` seconds, derived from the + millisecond broker input, unless shutdown is requested through + ``_heartbeat_stop``. On each tick it snapshots the current set of + tracked message ids under a lock, then calls ``set_vt`` to push their + visibility timeout back to ``self.visibility_timeout_seconds`` seconds. + + This prevents long-running handlers from becoming visible again and + being consumed twice before they are explicitly acked, nacked or + requeued. Failures are logged and retried on the next heartbeat tick. + """ + while not self._heartbeat_stop.wait(self.heartbeat_interval_seconds): + with self._heartbeat_msg_ids_lock: + msg_ids = list(self._heartbeat_msg_ids) + if not msg_ids: + continue + try: + self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout_seconds) + except Exception as e: + self.broker.logger.warning( + "Failed to extend visibility timeout heartbeat for queue %s. Error: ", + self.queue_name, + str(e), + exc_info=True, + ) + + def _register_heartbeat_msg(self, message: "_PostgresMessage") -> None: + """Track a message so the heartbeat can renew it.""" + with self._heartbeat_msg_ids_lock: + self._heartbeat_msg_ids.add(message._postgres_message.msg_id) + + def _unregister_heartbeat_msg_id(self, msg_id: int) -> None: + """Stop tracking a message for heartbeat renewal.""" + with self._heartbeat_msg_ids_lock: + self._heartbeat_msg_ids.discard(msg_id) + + def _start_listener(self) -> None: + """Start a LISTEN/NOTIFY listener for the queue when available.""" + channel = f"pgmq.q_{self.queue_name}.INSERT" + try: + self._listener_connection = psycopg.connect(self.broker.url, autocommit=True) + self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) + self._listener_available = True + except Exception as e: + self._listener_available = False + self._listener_connection = None + self.broker.logger.warning( + "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling. Error: %s", + self.queue_name, + str(e), + exc_info=True, + ) + return + + self._listener_thread = Thread( + target=self._run_listener, + name=f"postgres-listener-{self.queue_name}", + daemon=True, + ) + self._listener_thread.start() + + def _run_listener(self) -> None: + """Wake the consumer when a notification arrives, or disable the listener on errors.""" + while not self._listener_stop.is_set(): + try: + if self._listener_connection is None: + break + for _ in self._listener_connection.notifies(timeout=0.5, stop_after=1): + self._notify_event.set() + except Exception as e: + if self._listener_stop.is_set(): + break + self._listener_available = False + self._notify_event.set() + self.broker.logger.warning( + "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling. Error: %s", + self.queue_name, + str(e), + exc_info=True, + ) + break + + self._listener_available = False + + @override + def ack(self, message: "MessageProxy") -> None: + """Archive a processed message.""" + if not isinstance(message, _PostgresMessage): + raise ValueError("It must be a PostgresMessage") + self._unregister_heartbeat_msg_id(message._postgres_message.msg_id) + self.client.archive(self.queue_name, message._postgres_message.msg_id) + + @override + def nack(self, message: "MessageProxy") -> None: + """Archive a failed message.""" + if not isinstance(message, _PostgresMessage): + raise ValueError("It must be a PostgresMessage") + self._unregister_heartbeat_msg_id(message._postgres_message.msg_id) + self.client.archive(self.queue_name, message._postgres_message.msg_id) + + @override + def requeue(self, messages: Iterable["MessageProxy"]) -> None: + """Make messages visible again immediately by resetting their visibility timeout.""" + msg_ids = [message._postgres_message.msg_id for message in messages if isinstance(message, _PostgresMessage)] + if msg_ids: + for msg_id in msg_ids: + self._unregister_heartbeat_msg_id(msg_id) + self.client.set_vt(self.queue_name, msg_ids, 0) + + def _build_message(self, postgres_message: PostgresQueueMessage) -> "_PostgresMessage": + """Wrap a raw PGMQ row and register it for heartbeat renewal.""" + message = _PostgresMessage(postgres_message) + self._register_heartbeat_msg(message) + return message + + @override + def __next__(self) -> "_PostgresMessage | None": + """Return the next available message, or ``None`` if the queue stays empty.""" + if self.messages: + return self._build_message(self.messages.popleft()) + + if self._listener_available: + self._notify_event.wait(self.wait_timeout_seconds) + messages = self._read_immediate() + if not messages: + messages = self._read_with_poll() + self._listener_available = False + self._notify_event.clear() + else: + messages = self._read_with_poll() + + if not messages: + return None + + self.messages.extend(messages) + return self._build_message(self.messages.popleft()) + + @override + def close(self) -> None: + """Stop background threads and close the listener connection.""" + self._listener_stop.set() + self._heartbeat_stop.set() + self._notify_event.set() + if self._listener_connection is not None: + try: + self._listener_connection.close() + except Exception as e: + self.broker.logger.error("Listener not joinable: %s", str(e)) + finally: + self._listener_connection = None + if self._listener_thread is not None: + self._listener_thread.join(timeout=1) + if self._heartbeat_thread is not None: + self._heartbeat_thread.join(timeout=1) + return + + +class _PostgresMessage(MessageProxy): + def __init__(self, postgres_message: PostgresQueueMessage) -> None: + """Wrap a PGMQ message row as a Remoulade message proxy.""" + payload = postgres_message.message + if not isinstance(payload, dict): + raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") + try: + # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder) + # can rehydrate actor args/kwargs to their typed schemas. + message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8")) + except TypeError as exc: + raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc + if message.options.get("eta", None) is not None: + raise UnsupportedMessageEncoding("eta option isn't supported with postgres broker") + super().__init__(message) + self._postgres_message = postgres_message diff --git a/tests/conftest.py b/tests/conftest.py index cb6b43634..a01b87bbf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,7 +18,7 @@ from remoulade import Worker from remoulade.api import app from remoulade.brokers.local import LocalBroker -from remoulade.brokers.pgmq import PgmqBroker +from remoulade.brokers.postgres import PostgresBroker from remoulade.brokers.rabbitmq import RabbitmqBroker from remoulade.brokers.stub import StubBroker from remoulade.cancel import backends as cl_backends @@ -63,7 +63,7 @@ def check_redis(client): client.flushall() -def check_pgmq(client): +def check_postgres(client): try: with client.begin() as session: session.execute(text("CREATE SCHEMA IF NOT EXISTS partman")) @@ -74,7 +74,7 @@ def check_pgmq(client): raise e from e if CI else pytest.skip("No connection to PostgreSQL/PGMQ server.") -def cleanup_pgmq_queues(client): +def cleanup_postgres_queues(client): with client.begin() as session: queue_names = [row[0] for row in session.execute(text("SELECT queue_name FROM pgmq.list_queues()")).all()] for queue_name in queue_names: @@ -115,18 +115,18 @@ def rabbitmq_broker(request): @pytest.fixture() -def pgmq_broker(): +def postgres_broker(): db_string = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" engine = create_engine(db_string, poolclass=NullPool) client = sessionmaker(bind=engine) - check_pgmq(client) - cleanup_pgmq_queues(client) + check_postgres(client) + cleanup_postgres_queues(client) - broker = PgmqBroker(url=db_string) + broker = PostgresBroker(url=db_string) broker.emit_after("process_boot") remoulade.set_broker(broker) yield broker - cleanup_pgmq_queues(client) + cleanup_postgres_queues(client) broker.emit_before("process_stop") broker.close() diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index 12173d68e..05707d794 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -16,4 +16,4 @@ services: POSTGRES_HOST_AUTH_METHOD: trust POSTGRES_DB: test volumes: - - ./docker/postgres/01-create-pgmq-extension.sql:/docker-entrypoint-initdb.d/01-create-pgmq-extension.sql:ro + - ./docker/postgres/01-create-postgres-extension.sql:/docker-entrypoint-initdb.d/01-create-postgres-extension.sql:ro diff --git a/tests/docker/postgres/01-create-pgmq-extension.sql b/tests/docker/postgres/01-create-postgres-extension.sql similarity index 100% rename from tests/docker/postgres/01-create-pgmq-extension.sql rename to tests/docker/postgres/01-create-postgres-extension.sql diff --git a/tests/test_pgmq.py b/tests/test_postgres.py similarity index 67% rename from tests/test_pgmq.py rename to tests/test_postgres.py index 06eafc442..1ccd99626 100644 --- a/tests/test_pgmq.py +++ b/tests/test_postgres.py @@ -11,9 +11,9 @@ import remoulade from remoulade import Message, QueueJoinTimeout, Worker -from remoulade.brokers.pgmq import PgmqBroker +from remoulade.brokers.postgres import PostgresBroker -TEST_PGMQ_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" +TEST_POSTGRES_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" def _count_messages(broker, queue_name="default"): @@ -51,16 +51,16 @@ def _expected_payload(message): return json.loads(message.encode().decode("utf-8")) -def test_pgmq_broker_uses_provided_url(): - broker_url = TEST_PGMQ_URL - broker = PgmqBroker(url=broker_url) +def test_postgres_broker_uses_provided_url(): + broker_url = TEST_POSTGRES_URL + broker = PostgresBroker(url=broker_url) assert broker.url == broker_url -def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): - broker = PgmqBroker( - url=TEST_PGMQ_URL, +def test_postgres_broker_partitions_archive_table_on_postgresql_queue_init(): + broker = PostgresBroker( + url=TEST_POSTGRES_URL, middleware=[], ) broker.client.validate_queue_name = Mock() @@ -79,8 +79,8 @@ def test_pgmq_broker_partitions_archive_table_on_postgresql_queue_init(): broker.client.create_queue.assert_not_called() -def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) +def test_postgres_broker_enables_notify_on_postgresql_queue_init(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() @@ -97,8 +97,8 @@ def test_pgmq_broker_enables_notify_on_postgresql_queue_init(): broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) -def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) +def test_postgres_broker_does_not_fail_when_enable_notify_raises(caplog): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) @@ -110,8 +110,8 @@ def test_pgmq_broker_does_not_fail_when_enable_notify_raises(caplog): assert "Failed to enable LISTEN/NOTIFY" in caplog.text -def test_pgmq_broker_declare_queue_is_idempotent_when_queue_already_exists(): - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) +def test_postgres_broker_declare_queue_is_idempotent_when_queue_already_exists(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() @@ -124,9 +124,9 @@ def test_pgmq_broker_declare_queue_is_idempotent_when_queue_already_exists(): assert "default" in broker.queues -def test_pgmq_broker_uses_custom_partition_settings_when_provided(): - broker = PgmqBroker( - url=TEST_PGMQ_URL, +def test_postgres_broker_uses_custom_partition_settings_when_provided(): + broker = PostgresBroker( + url=TEST_POSTGRES_URL, middleware=[], archive_partition_interval="2 days", archive_retention_interval="14 days", @@ -146,58 +146,58 @@ def test_pgmq_broker_uses_custom_partition_settings_when_provided(): ) -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) - pgmq_broker.declare_queue(message.queue_name) + postgres_broker.declare_queue(message.queue_name) - pgmq_broker.enqueue(message) + postgres_broker.enqueue(message) - assert _count_messages(pgmq_broker) == 1 - assert _first_payload(pgmq_broker) == _expected_payload(message) + assert _count_messages(postgres_broker) == 1 + assert _first_payload(postgres_broker) == _expected_payload(message) -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_broker_uses_native_visibility_delay_without_delay_queue(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_uses_native_visibility_delay_without_delay_queue(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) - pgmq_broker.declare_queue(message.queue_name) + postgres_broker.declare_queue(message.queue_name) - pgmq_broker.enqueue(message, delay=250) + postgres_broker.enqueue(message, delay=250) - assert pgmq_broker.client.read("default", vt=1) is None + assert postgres_broker.client.read("default", vt=1) is None time.sleep(0.35) - delayed = pgmq_broker.client.read("default", vt=1) + delayed = postgres_broker.client.read("default", vt=1) assert delayed is not None assert delayed.message == _expected_payload(message) -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_broker_transactions_commit_and_rollback_messages(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_transactions_commit_and_rollback_messages(postgres_broker): @remoulade.actor def do_work(): return 1 - pgmq_broker.declare_actor(do_work) + postgres_broker.declare_actor(do_work) - with pgmq_broker.tx(): + with postgres_broker.tx(): do_work.send() - with pytest.raises(ValueError), pgmq_broker.tx(): + with pytest.raises(ValueError), postgres_broker.tx(): do_work.send() raise ValueError("rollback") - assert _count_messages(pgmq_broker) == 1 + assert _count_messages(postgres_broker) == 1 -def test_pgmq_consumer_uses_notification_path_when_listener_is_available(monkeypatch): +def test_postgres_consumer_uses_notification_path_when_listener_is_available(monkeypatch): def _fake_start_listener(self): self._listener_available = True - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) @@ -218,13 +218,13 @@ def _fake_start_listener(self): consumer.close() -def test_pgmq_consumer_falls_back_to_polling_when_listener_is_unavailable(monkeypatch): +def test_postgres_consumer_falls_back_to_polling_when_listener_is_unavailable(monkeypatch): def _fake_start_listener(self): self._listener_available = False - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) @@ -248,13 +248,13 @@ def _fake_start_listener(self): consumer.close() -def test_pgmq_consumer_uses_broker_visibility_timeout_for_reads(monkeypatch): +def test_postgres_consumer_uses_broker_visibility_timeout_for_reads(monkeypatch): def _fake_start_listener(self): self._listener_available = False - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[], visibility_timeout_in_second=17) + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], visibility_timeout_ms=17_000) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(5,), kwargs={}, options={}) @@ -276,13 +276,13 @@ def _fake_start_listener(self): consumer.close() -def test_pgmq_consumer_falls_back_to_polling_when_listener_stops_during_wait(monkeypatch): +def test_postgres_consumer_falls_back_to_polling_when_listener_stops_during_wait(monkeypatch): def _fake_start_listener(self): self._listener_available = True - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None message = Message(queue_name="default", actor_name="do_work", args=(3,), kwargs={}, options={}) @@ -313,17 +313,17 @@ def _wait_and_stop(_timeout): consumer.close() -def test_pgmq_consumer_heartbeat_extends_inflight_message_visibility(monkeypatch): +def test_postgres_consumer_heartbeat_extends_inflight_message_visibility(monkeypatch): def _fake_start_listener(self): self._listener_available = False - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker( - url=TEST_PGMQ_URL, + broker = PostgresBroker( + url=TEST_POSTGRES_URL, middleware=[], - visibility_timeout_in_second=2, - heartbeat_interval_in_second=0.05, + visibility_timeout_ms=2_000, + heartbeat_interval_ms=50, ) broker.queues["default"] = None @@ -352,16 +352,16 @@ def _fake_start_listener(self): consumer.close() -def test_pgmq_consumer_decodes_payload_with_global_encoder(monkeypatch, pydantic_encoder): +def test_postgres_consumer_decodes_payload_with_global_encoder(monkeypatch, pydantic_encoder): class InputSchema(BaseModel): value: int def _fake_start_listener(self): self._listener_available = False - monkeypatch.setattr("remoulade.brokers.pgmq._PgmqConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PgmqBroker(url=TEST_PGMQ_URL, middleware=[]) + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) remoulade.set_broker(broker) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() @@ -393,13 +393,13 @@ def typed_actor(payload: InputSchema): consumer.close() -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_consumer_reads_messages_and_acks_with_delete(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_consumer_reads_messages_and_acks_with_delete(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(42,), kwargs={}, options={}) - pgmq_broker.declare_queue(message.queue_name) - pgmq_broker.enqueue(message) + postgres_broker.declare_queue(message.queue_name) + postgres_broker.enqueue(message) - consumer = pgmq_broker.consume("default", prefetch=2, timeout=200) + consumer = postgres_broker.consume("default", prefetch=2, timeout=200) consumed_message = next(consumer) assert consumed_message is not None assert consumed_message.message_id == message.message_id @@ -407,33 +407,33 @@ def test_pgmq_consumer_reads_messages_and_acks_with_delete(pgmq_broker): consumer.ack(consumed_message) consumer.close() - assert _count_messages(pgmq_broker) == 0 + assert _count_messages(postgres_broker) == 0 -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_consumer_nack_archives_messages(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_consumer_nack_archives_messages(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) - pgmq_broker.declare_queue(message.queue_name) - pgmq_broker.enqueue(message) + postgres_broker.declare_queue(message.queue_name) + postgres_broker.enqueue(message) - consumer = pgmq_broker.consume("default", prefetch=1, timeout=200) + consumer = postgres_broker.consume("default", prefetch=1, timeout=200) consumed_message = next(consumer) assert consumed_message is not None consumer.nack(consumed_message) consumer.close() - assert _count_messages(pgmq_broker) == 0 - assert _count_archived_messages(pgmq_broker) == 1 + assert _count_messages(postgres_broker) == 0 + assert _count_archived_messages(postgres_broker) == 1 -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_consumer_requeue_restores_visibility_with_set_vt(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_consumer_requeue_restores_visibility_with_set_vt(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) - pgmq_broker.declare_queue(message.queue_name) - pgmq_broker.enqueue(message) + postgres_broker.declare_queue(message.queue_name) + postgres_broker.enqueue(message) - consumer = pgmq_broker.consume("default", prefetch=1, timeout=200) + consumer = postgres_broker.consume("default", prefetch=1, timeout=200) consumed_message = next(consumer) assert consumed_message is not None @@ -446,31 +446,31 @@ def test_pgmq_consumer_requeue_restores_visibility_with_set_vt(pgmq_broker): consumer.ack(replayed_message) consumer.close() - assert _count_messages(pgmq_broker) == 0 + assert _count_messages(postgres_broker) == 0 -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_worker_processes_native_delayed_messages_without_delay_queue(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_processes_native_delayed_messages_without_delay_queue(postgres_broker): seen = [] @remoulade.actor def do_work(value): seen.append(value) - pgmq_broker.declare_actor(do_work) - worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=2) + postgres_broker.declare_actor(do_work) + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=2) worker.start() try: do_work.send_with_options(args=(3,), delay=150) - pgmq_broker.join(do_work.queue_name, timeout=10_000) + postgres_broker.join(do_work.queue_name, timeout=10_000) assert seen == [3] worker.join() finally: worker.stop() -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_broker_join_times_out_while_processing_invisible_message(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_join_times_out_while_processing_invisible_message(postgres_broker): started = threading.Event() release = threading.Event() @@ -479,27 +479,27 @@ def do_work(): started.set() release.wait(timeout=5) - pgmq_broker.declare_actor(do_work) + postgres_broker.declare_actor(do_work) - worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=1) + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=1) worker.start() try: do_work.send() assert started.wait(timeout=2) with pytest.raises(QueueJoinTimeout): - pgmq_broker.join(do_work.queue_name, timeout=100) + postgres_broker.join(do_work.queue_name, timeout=100) release.set() - pgmq_broker.join(do_work.queue_name, timeout=5_000) + postgres_broker.join(do_work.queue_name, timeout=5_000) worker.join() finally: release.set() worker.stop() -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_worker_processes_a_two_actor_pipeline(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_processes_a_two_actor_pipeline(postgres_broker): seen: list[tuple[str, int]] = [] @remoulade.actor @@ -511,15 +511,15 @@ def first_actor(value): def second_actor(value): seen.append(("second", value)) - pgmq_broker.declare_actor(first_actor) - pgmq_broker.declare_actor(second_actor) + postgres_broker.declare_actor(first_actor) + postgres_broker.declare_actor(second_actor) - worker = Worker(pgmq_broker, worker_timeout=100, worker_threads=1) + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=1) worker.start() try: remoulade.pipeline([first_actor.message(1), second_actor.message()]).run() - pgmq_broker.join(second_actor.queue_name, timeout=10_000) + postgres_broker.join(second_actor.queue_name, timeout=10_000) worker.join() assert seen == [("first", 1), ("second", 2)] @@ -527,11 +527,11 @@ def second_actor(value): worker.stop() -@pytest.mark.usefixtures("pgmq_broker") -def test_pgmq_consumer_listener_wakes_on_enqueue_with_listen_notify(pgmq_broker): +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_consumer_listener_wakes_on_enqueue_with_listen_notify(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(99,), kwargs={}, options={}) - pgmq_broker.declare_queue(message.queue_name) - consumer = pgmq_broker.consume(message.queue_name, prefetch=1, timeout=1500) + postgres_broker.declare_queue(message.queue_name) + consumer = postgres_broker.consume(message.queue_name, prefetch=1, timeout=1500) if not consumer._listener_available: pytest.skip("LISTEN/NOTIFY listener unavailable in this environment.") @@ -545,7 +545,7 @@ def _consume_once(): thread.start() try: time.sleep(0.15) - pgmq_broker.enqueue(message) + postgres_broker.enqueue(message) thread.join(timeout=3) assert not thread.is_alive() assert consumed_messages From 734649ffb91091c050d6d5ab8c8d561184b701e5 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 11 Jun 2026 15:54:37 +0200 Subject: [PATCH 15/44] fix(postgres): partition creation --- remoulade/brokers/postgres.py | 29 ++++++++++++++----------- tests/test_postgres.py | 41 ++++++++++++++++++++++++++++++++++- 2 files changed, 56 insertions(+), 14 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index a2c07ccd3..b2501c1db 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -40,6 +40,7 @@ PostgresPayload = dict[str, Any] LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 +FUTURE_PARTITION_HORIZON: Final[str] = "2 months" def _milliseconds_to_seconds(milliseconds: int) -> int: @@ -142,7 +143,7 @@ def _try_enable_notify(self, queue_name: str) -> None: def _queue_exists(self, queue_name: str) -> bool: """Return whether the queue already exists in PostgreSQL.""" - return queue_name in [queue.queue_name for queue in self.client.list_queues()] + return queue_name in {queue.queue_name for queue in self.client.list_queues()} @override def close(self) -> None: @@ -167,20 +168,22 @@ def declare_queue(self, queue_name: str) -> None: return self.client.validate_queue_name(queue_name, conn=self._current_connection) - if self._queue_exists(queue_name): - self.queues[queue_name] = None - return + queue_exists = self._queue_exists(queue_name) + + if not queue_exists: + self.emit_before("declare_queue", queue_name) + self.client.create_partitioned_queue( + queue_name, + partition_interval=self.archive_partition_interval, + retention_interval=self.archive_retention_interval, + conn=self._current_connection, + ) + self._try_enable_notify(queue_name) - self.emit_before("declare_queue", queue_name) - self.client.create_partitioned_queue( - queue_name, - partition_interval=self.archive_partition_interval, - retention_interval=self.archive_retention_interval, - conn=self._current_connection, - ) - self._try_enable_notify(queue_name) self.queues[queue_name] = None - self.emit_after("declare_queue", queue_name) + + if not queue_exists: + self.emit_after("declare_queue", queue_name) def _encode_message(self, message: "Message") -> PostgresPayload: """Decode a Remoulade message into a JSON object payload for PGMQ. diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 1ccd99626..b31cef075 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -58,7 +58,7 @@ def test_postgres_broker_uses_provided_url(): assert broker.url == broker_url -def test_postgres_broker_partitions_archive_table_on_postgresql_queue_init(): +def test_postgres_broker_creates_partitioned_queue_with_default_intervals(): broker = PostgresBroker( url=TEST_POSTGRES_URL, middleware=[], @@ -66,19 +66,57 @@ def test_postgres_broker_partitions_archive_table_on_postgresql_queue_init(): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.create_queue = Mock() + broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) broker.declare_queue("default") + broker.client.validate_queue_name.assert_called_once_with("default", conn=None) broker.client.create_partitioned_queue.assert_called_once_with( "default", partition_interval="1 day", retention_interval="7 days", conn=None, ) + broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) broker.client.create_queue.assert_not_called() +def test_postgres_broker_uses_current_transaction_connection_for_queue_creation(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) + + transaction_connection = object() + + class _DummyTransaction: + def __enter__(self): + return transaction_connection + + def __exit__(self, exc_type, exc, tb): + return None + + broker.client.engine.begin = Mock(return_value=_DummyTransaction()) + + with broker.tx(): + broker.declare_queue("default") + + broker.client.validate_queue_name.assert_called_once_with("default", conn=transaction_connection) + broker.client.create_partitioned_queue.assert_called_once_with( + "default", + partition_interval="1 day", + retention_interval="7 days", + conn=transaction_connection, + ) + broker.client.enable_notify.assert_called_once_with( + "default", + throttle_interval_ms=250, + conn=transaction_connection, + ) + + def test_postgres_broker_enables_notify_on_postgresql_queue_init(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() @@ -366,6 +404,7 @@ def _fake_start_listener(self): broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) @remoulade.actor(actor_name="typed.actor", queue_name="default") def typed_actor(payload: InputSchema): From e6a060873b312186b8089ef0661f77de465ea2d6 Mon Sep 17 00:00:00 2001 From: mducros Date: Mon, 15 Jun 2026 10:47:45 +0200 Subject: [PATCH 16/44] fix(api): restore main apis --- remoulade/api/main.py | 154 ++++++++++++++++- tests/state/test_state_api.py | 303 ++++++++++++++++++++++++++++++++++ tests/test_cancel.py | 36 ++++ 3 files changed, 489 insertions(+), 4 deletions(-) create mode 100644 tests/state/test_state_api.py diff --git a/remoulade/api/main.py b/remoulade/api/main.py index 3efd151f3..b0a9831d2 100644 --- a/remoulade/api/main.py +++ b/remoulade/api/main.py @@ -1,14 +1,158 @@ """This file describe the API to get the state of messages""" +import sys +from typing import Any, TypedDict + from flask import Flask -from marshmallow import ValidationError +from flask_apispec import marshal_with +from marshmallow import Schema, ValidationError, fields, validate, validates_schema from werkzeug.exceptions import HTTPException -from remoulade.errors import RemouladeError +import remoulade +from remoulade import get_broker +from remoulade.errors import NoResultBackend, NoStateBackend, RemouladeError +from remoulade.result import Result +from remoulade.results import ResultMissing -from .apispec import add_swagger +from .apispec import add_swagger, validate_schema +from .scheduler import scheduler_bp, scheduler_routes +from .state import messages_bp, messages_routes app = Flask(__name__) +app.register_blueprint(scheduler_bp) +app.register_blueprint(messages_bp) + + +class MessageSchema(Schema): + """ + Class to validate post data in /messages + """ + + actor_name = fields.Str(validate=validate.Length(min=1), required=True) + args = fields.List(fields.Raw(), allow_none=True) + kwargs = fields.Dict(allow_none=True) + options = fields.Dict(allow_none=True) + delay = fields.Float(validate=validate.Range(min=1), allow_none=True) + + @validates_schema + def validate_actor_name(self, data, **kwargs): + actor_name = data.get("actor_name") + if actor_name not in remoulade.get_broker().actors: + raise ValidationError(f"No actor named {actor_name} exists") + + +class ResponseSchema(Schema): + result = fields.Raw(allow_none=True) + error = fields.Str(allow_none=True) + + +class ArgsSchema(Schema): + name = fields.Str() + type = fields.Str() + default = fields.Str(allow_none=True) + + +class ActorSchema(Schema): + name = fields.Str() + queue_name = fields.Str() + alternative_queues = fields.List(fields.Str()) + priority = fields.Int() + args = fields.List(fields.Nested(ArgsSchema)) + + +class ActorResponseSchema(Schema): + result = fields.List(fields.Nested(ActorSchema)) + + +class OptionsResponseSchema(Schema): + options = fields.List(fields.Str()) + + +@app.route("/messages/cancel/", methods=["POST"]) +@marshal_with(ResponseSchema) +def cancel_message(message_id): + broker = get_broker() + cancel_backend = broker.get_cancel_backend() + try: + # If a single message in a composition is canceled, we cancel the whole composition + state_backend = broker.get_state_backend() + states_count = state_backend.get_states_count(selected_composition_ids=[message_id]) + state = state_backend.get_state(message_id) + if states_count == 0 and state is None: + raise ValidationError(f"This message id or composition id {message_id} does not exist.") + if state and state.composition_id: + message_id = state.composition_id + except NoStateBackend: + pass + cancel_backend.cancel([message_id]) + return {"result": "ok"} + + +@app.route("/messages/requeue/", methods=["POST"]) +@marshal_with(ResponseSchema) +def requeue_message(message_id): + broker = get_broker() + backend = broker.get_state_backend() + state = backend.get_state(message_id) + if state is None: + raise ValidationError(f"No message with id {message_id} exists") + actor = broker.get_actor(state.actor_name) + payload = {"args": state.args, "kwargs": state.kwargs} + pipe_target = state.options.get("pipe_target") + if pipe_target is None: + actor.send_with_options(**payload, **state.options) + return {"result": "ok"} + else: + return {"error": "requeue message in a pipeline not supported"}, 400 + + +@app.route("/messages/result/") +@marshal_with(ResponseSchema) +def get_results(message_id): + from ..message import get_encoder + + max_size = 1e4 + try: + result = Result[Any](message_id=message_id).get() + encoded_result = get_encoder().encode(result).decode("utf-8") + size_result = sys.getsizeof(encoded_result) + if size_result >= max_size: + encoded_result = f"The result is too big {size_result / 1e6}M" + return {"result": encoded_result} + except ResultMissing: + return {"result": "result is missing"} + except NoResultBackend: + return {"result": "no result backend"} + except (UnicodeDecodeError, TypeError): + return {"result": "non serializable result"} + + +@app.route("/messages", methods=["POST"]) +@marshal_with(ResponseSchema) +@validate_schema(MessageSchema) +def enqueue_message(**kwargs): + actor = get_broker().get_actor(kwargs.pop("actor_name")) + options = kwargs.pop("options") or {} + actor.send_with_options(**kwargs, **options) + return {"result": "ok"} + + +class GroupMessagesT(TypedDict): + group_id: str + messages: list[dict] + + +@app.route("/actors") +@marshal_with(ActorResponseSchema) +def get_actors(): + return {"result": [actor.as_dict() for actor in get_broker().actors.values()]} + + +@app.route("/options") +@marshal_with(OptionsResponseSchema) +def get_options(): + broker = get_broker() + return {"options": list(broker.actor_options)} @app.errorhandler(RemouladeError) @@ -26,4 +170,6 @@ def validation_error(e): return {"error": e.normalized_messages()}, 400 -add_swagger(app, {"main": []}) +routes = [cancel_message, requeue_message, get_results, enqueue_message, get_actors, get_options] + +add_swagger(app, {"main": routes, "scheduler": scheduler_routes, "messages": messages_routes}) diff --git a/tests/state/test_state_api.py b/tests/state/test_state_api.py new file mode 100644 index 000000000..39b319bc5 --- /dev/null +++ b/tests/state/test_state_api.py @@ -0,0 +1,303 @@ +import datetime +import json +from datetime import date +from operator import itemgetter +from random import choice +from unittest import mock +from unittest.mock import MagicMock +from zoneinfo import ZoneInfo + +import pytest + +import remoulade +from remoulade import set_scheduler +from remoulade.api.main import app +from remoulade.cancel import Cancel +from remoulade.message import Message +from remoulade.scheduler import ScheduledJob +from remoulade.state import State, StateStatusesEnum + + +class TestMessageStateAPI: + """Class Responsible to do the test of the API""" + + def test_no_messages(self, stub_broker, state_middleware, api_client): + res = api_client.post("/messages/states") + assert res.status_code == 200 + assert len(res.json["data"]) == 0 + + def test_invalid_url(self, stub_broker, state_middleware, api_client): + res = api_client.get("/invalid_url/") + assert res.status_code == 404 + + def test_invalid_state_message(self, stub_broker, api_client): + res = api_client.post( + "/messages/states", data=json.dumps({"status": "invalid_state"}), content_type="application/json" + ) + assert res.status_code == 400 + + @pytest.mark.parametrize("status", [StateStatusesEnum.Skipped, StateStatusesEnum.Success]) + def test_get_state_by_name(self, status, stub_broker, state_middleware, api_client): + state = State("1", status) + state_middleware.backend.set_state(state, ttl=1000) + args = {"selected_statuses": [status.value]} + res = api_client.post("/messages/states", data=json.dumps(args), content_type="application/json") + assert res.json == {"count": 1, "data": [state.as_dict()]} + + def test_get_by_message_id(self, stub_broker, state_middleware, api_client): + message_id = "1" + state = State(message_id, StateStatusesEnum.Pending) + state_middleware.backend.set_state(state, ttl=1000) + res = api_client.get(f"/messages/states/{message_id}") + assert res.json == state.as_dict() + + @pytest.mark.parametrize("n", [0, 10, 50, 100]) + def test_get_states_by_name(self, n, stub_broker, state_middleware, api_client): + # generate random test + random_states = [] + for i in range(n): + random_state = State(f"id{i}", choice(list(StateStatusesEnum))) # random state + random_states.append(random_state.as_dict()) + state_middleware.backend.set_state(random_state, ttl=1000) + + # check storage + res = api_client.post("/messages/states", data=json.dumps({"size": 100}), content_type="application/json") + states = res.json["data"] + assert len(states) == n + # check integrity + for state in states: + assert state in random_states + + def test_request_cancel(self, stub_broker, cancel_backend): + with app.test_client() as client: + message_id = "1" + stub_broker.add_middleware(Cancel(backend=cancel_backend)) + client.post(f"/messages/cancel/{message_id}") + assert cancel_backend.is_canceled(message_id, None) + + def test_request_cancel_without_backend(self, stub_broker, api_client): + res = api_client.post("/messages/cancel/{}".format("id")) + # if there is not cancel backend should raise an error + assert res.json["error"] == "The default broker doesn't have a cancel backend." + assert res.status_code == 500 + + def test_no_scheduler(self, stub_broker, api_client): + set_scheduler(None) + res = api_client.get("/scheduled/jobs") + assert res.status_code == 400 + + def test_scheduled_jobs(self, scheduler, api_client, do_work, frozen_datetime): + timezone = ZoneInfo("Europe/Paris") + scheduler.schedule = [ + ScheduledJob( + actor_name=do_work.actor_name, daily_time=(datetime.datetime.now(timezone)).time(), tz="Europe/Paris" + ) + ] + scheduler.sync_config() + + res = api_client.get("/scheduled/jobs") + jobs = res.json["result"] + assert jobs == [ + { + "hash": jobs[0]["hash"], + "actor_name": "do_work", + "args": [], + "daily_time": "01:00:00", + "enabled": True, + "interval": 86400, + "iso_weekday": None, + "kwargs": {}, + "last_queued": None, + "tz": "Europe/Paris", + } + ] + + def test_enqueue_message(self, stub_broker, do_work, api_client): + data = { + "actor_name": "do_work", + "args": ["1", "2"], + "kwargs": {}, + "options": {"time_limit": 1000}, + "delay": None, + } + stub_broker.enqueue = MagicMock() + stub_broker.join(do_work.queue_name) + res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") + message = Message( + queue_name="default", + actor_name="do_work", + args=("1", "2"), + kwargs={}, + options={"time_limit": 1000}, + message_id=mock.ANY, + message_timestamp=mock.ANY, + ) + assert stub_broker.enqueue.call_count == 1 + assert stub_broker.enqueue.call_args == mock.call(message, delay=None) + assert res.status_code == 200 + + @pytest.mark.parametrize( + "actor_name,error", + [ + (None, "Field may not be null."), + ("", "Shorter than minimum length 1."), + (111, "Not a valid string."), + ([], "Not a valid string."), + ], + ) + def test_invalid_actor_name_to_enqueue(self, actor_name, error, stub_broker, do_work, api_client): + data = { + "actor_name": actor_name, + } + res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") + validation_error = res.json["error"] + assert validation_error["actor_name"] == [error] + assert res.status_code == 400 + + @pytest.mark.parametrize( + "delay,error", [(-1, "Must be greater than or equal to 1."), ("str", "Not a valid number.")] + ) + def test_invalid_delay_to_enqueue(self, delay, error, stub_broker, do_work, api_client): + data = { + "actor_name": "do_work", + "delay": delay, + } + res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") + validation_error = res.json["error"] + assert validation_error["delay"] == [error] + assert res.status_code == 400 + + def test_enqueue_invalid_actor_name(self, stub_broker, api_client): + data = {"actor_name": "invalid_actor"} + res = api_client.post("/messages", data=json.dumps(data), content_type="application/json") + + assert res.status_code == 400 + + def test_get_declared_actors(self, stub_broker, do_work, api_client): + @remoulade.actor(queue_name="foo", priority=10) + def do_job(): + pass + + stub_broker.declare_actor(do_job) + res = api_client.get("/actors") + expected = [ + {"name": "do_work", "priority": 0, "queue_name": "default"}, + {"name": "do_job", "priority": 10, "queue_name": "foo"}, + ].sort(key=itemgetter("name")) + assert res.json["result"].sort(key=itemgetter("name")) == expected + + def test_filter_messages(self, stub_broker, api_client): + state = State("some_message_id") + state_backend = stub_broker.get_state_backend() + state_backend.set_state(state, ttl=1000) + data = {"sort_column": "message_id", "selected_message_ids": ["some_message_id"]} + res = api_client.post("/messages/states", data=json.dumps(data), content_type="application/json") + assert res.json == {"count": 1, "data": [state.as_dict()]} + + @pytest.mark.parametrize("offset", [0, 1, 5, 100]) + def test_get_states_offset(self, offset, stub_broker, api_client, state_middleware): + for i in range(0, 10): + state_middleware.backend.set_state(State(f"id{i}"), ttl=1000) + res = api_client.post( + "/messages/states", data=json.dumps({"offset": offset, "size": 50}), content_type="application/json" + ) + if offset >= 10: + assert res.json["data"] == [] + else: + assert len(res.json["data"]) + offset == 10 + + @pytest.mark.parametrize("size", [1, 5, 100]) + def test_get_states_page_size(self, size, stub_broker, api_client, state_middleware): + for i in range(0, 10): + state_middleware.backend.set_state(State(f"id{i}"), ttl=1000) + res = api_client.post("/messages/states", data=json.dumps({"size": size}), content_type="application/json") + if size >= 10: + assert res.json["count"] == 10 + else: + assert len(res.json["data"]) == size + + def test_not_raise_error_with_pickle_and_non_serializable( + self, pickle_encoder, stub_broker, api_client, state_middleware + ): + state = State("id1", args=["some_status", date(2020, 12, 12)]) + state_middleware.backend.set_state(state, ttl=1000) + res = api_client.post("messages/states") + assert res.json == { + "count": 1, + "data": [{"args": ["some_status", "Sat, 12 Dec 2020 00:00:00 GMT"], "message_id": "id1"}], + } + + def test_requeue_message(self, stub_broker, do_work, api_client, state_middleware): + stub_broker.enqueue = MagicMock() + state = State("id1", StateStatusesEnum.Success, actor_name="do_work", options={"time_limit": 1000}) + state_middleware.backend.set_state(state, ttl=1000) + res = api_client.post("messages/requeue/id1") + message = Message( + queue_name="default", + actor_name="do_work", + args=(), + kwargs={}, + options={"time_limit": 1000}, + message_id=mock.ANY, + message_timestamp=mock.ANY, + ) + assert stub_broker.enqueue.call_count == 1 + assert stub_broker.enqueue.call_args == mock.call(message, delay=None) + assert res.status_code == 200 + + def test_requeue_invalid_id(self, stub_broker, api_client, state_middleware): + res = api_client.post("messages/requeue/invalid_id") + assert res.status_code == 400 + + def test_requeue_message_with_pipeline(self, stub_broker, do_work, api_client, state_middleware): + stub_broker.enqueue = MagicMock() + state = State("id1", StateStatusesEnum.Success, actor_name="do_work", options={"pipe_target": "some_pipe"}) + state_middleware.backend.set_state(state, ttl=1000) + res = api_client.post("messages/requeue/id1") + assert res.json == {"error": "requeue message in a pipeline not supported"} + assert res.status_code == 400 + + def test_get_result_message(self, stub_broker, stub_worker, api_client, state_middleware, result_middleware): + @remoulade.actor(store_results=True) + def do_work(): + return 42 + + stub_broker.declare_actor(do_work) + message = do_work.send() + stub_broker.join(do_work.queue_name) + stub_worker.join() + res = api_client.get(f"/messages/result/{message.message_id}") + assert res.json["result"] == "42" + + def test_no_result_backend(self, stub_broker, stub_worker, api_client, do_work): + message = do_work.send() + res = api_client.get(f"/messages/result/{message.message_id}") + assert res.json["result"] == "no result backend" + + def test_result_not_serializable( + self, pickle_encoder, stub_broker, stub_worker, api_client, state_middleware, result_middleware + ): + @remoulade.actor(store_results=True) + def do_work(): + return {"date": date(2020, 10, 10)} + + stub_broker.declare_actor(do_work) + message = do_work.send() + stub_broker.join(do_work.queue_name) + stub_worker.join() + res = api_client.get(f"/messages/result/{message.message_id}") + assert res.json["result"] == "non serializable result" + + def test_cant_sort_by_args_kwargs_options(self, stub_broker, state_middleware, api_client): + res = api_client.post( + "/messages/states", data=json.dumps({"sort_column": "args"}), content_type="application/json" + ) + assert res.status_code == 400 + res = api_client.post( + "/messages/states", data=json.dumps({"sort_column": "kwargs"}), content_type="application/json" + ) + assert res.status_code == 400 + res = api_client.post( + "/messages/states", data=json.dumps({"sort_column": "options"}), content_type="application/json" + ) + assert res.status_code == 400 diff --git a/tests/test_cancel.py b/tests/test_cancel.py index fec9c9b95..af20f745e 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -89,6 +89,32 @@ def do_fail(arg): assert cancel_backend.is_canceled("", "mocked_composition_id") == cancel +def test_compositions_are_canceled_on_message_cancel(stub_broker, cancel_backend, state_middleware, api_client): + # Given a cancel backend + # And a broker with the cancel middleware + stub_broker.add_middleware(Cancel(backend=cancel_backend)) + + # And an actor + @remoulade.actor + def do_work(arg=None): + return 1 + + # And those actors are declared + stub_broker.declare_actor(do_work) + + message_to_cancel = do_work.message() + + # And a group that I enqueue + g = group([message_to_cancel | do_work.message(), do_work.message()], cancel_on_error=True) + g.run() + + # When I cancel a message of this group + api_client.post("messages/cancel/" + message_to_cancel.message_id) + + # The whole composition should be canceled + assert cancel_backend.is_canceled("", g.group_id) + + def test_cannot_cancel_on_error_if_no_cancel(stub_broker): # Given an actor @remoulade.actor() @@ -162,3 +188,13 @@ def do_work(): # It messages should not have runMessageSchema assert calls_count == 0 + + +def test_raise_error_if_unknown_id(stub_broker, cancel_backend, api_client): + # Given a cancel middleware + stub_broker.add_middleware(Cancel(backend=cancel_backend)) + + # If I try to cancel a id that is not a id of a message or composition + res = api_client.post("messages/cancel/invalid_id") + + assert res.status_code == 400 From da23afc52b193205fbb950f76be13a299173cc61 Mon Sep 17 00:00:00 2001 From: mducros Date: Mon, 15 Jun 2026 11:51:36 +0200 Subject: [PATCH 17/44] fix(pgmessage): compilation --- docs/source/changelog.rst | 17 +++++++++++++++++ remoulade/brokers/postgres.py | 2 +- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 6f0446b84..6a37435db 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -5,6 +5,23 @@ Changelog All notable changes to this project will be documented in this file. +`7.0.0`_ -- 2026-06-15 +------------ +Breaking changes +^^^^^^^^^^^^^^^^ +* Remove the legacy PostgreSQL state backend. +* Rework the broker API around the new PostgreSQL/PGMQ implementation. + +Feat +^^^^ +* Add a PostgreSQL/PGMQ broker with partitioned queues, native delayed messages, ``LISTEN/NOTIFY`` wakeups, and queue join support. +* Add the ``PgmqBroker`` alias for the PostgreSQL-backed broker. + +Changed +^^^^^^^ +* Restore the main APIs after the broker refactor. +* Update the documentation, examples, CI, and test suite for the new stack. + ======= `6.2.0`_ -- 2026-05-18 ------------- diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index b2501c1db..b0d28c3d1 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -535,7 +535,7 @@ def __init__(self, postgres_message: PostgresQueueMessage) -> None: try: # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder) # can rehydrate actor args/kwargs to their typed schemas. - message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8")) + message = Message.decode(json.dumps(payload).encode("utf-8")) except TypeError as exc: raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc if message.options.get("eta", None) is not None: From c74adbfbff74879553f5509b473b1dbc2c6d31e8 Mon Sep 17 00:00:00 2001 From: mducros Date: Tue, 16 Jun 2026 16:12:00 +0200 Subject: [PATCH 18/44] fix(Encoder): less useless parsing --- CONTRIBUTORS.md | 5 +- docs/source/changelog.rst | 2 - remoulade/api/main.py | 2 +- remoulade/brokers/pgmq.py | 13 --- remoulade/brokers/postgres.py | 17 ++-- remoulade/brokers/rabbitmq.py | 4 +- remoulade/brokers/stub.py | 4 +- remoulade/encoder.py | 120 +++++++++++++++++++--------- remoulade/message.py | 17 +++- remoulade/results/backends/redis.py | 18 ++--- remoulade/results/backends/stub.py | 6 +- remoulade/scheduler/scheduler.py | 8 +- remoulade/state/backend.py | 8 +- tests/test_actors.py | 2 +- tests/test_encoders.py | 78 ++++++++++++------ tests/test_postgres.py | 52 +++++++++++- tests/test_results.py | 6 +- 17 files changed, 240 insertions(+), 122 deletions(-) delete mode 100644 remoulade/brokers/pgmq.py diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 34caddd4c..0d6510bcc 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -30,8 +30,11 @@ of those changes to CLEARTYPE SRL. | [@thomasLeMeur](https://github.com/thomasLeMeur) | Thomas Le Meur | | [@fregogui](https://github.com/fregogui) | Guillaume Fregosi | | [@pgitips](https://github.com/pgitips) | Pierre Giraud | + + + | [@williampollet](https://github.com/williampollet) | William Pollet | | [@mehdithez](https://github.com/mehdithez) | Zeroual Mehdi | | [@mducros-wm](https://github.com/mducros-wm) | Martin Ducros | -| [@julien-duponchelle](https://github.com/julien-duponchelle) | Juli en Duponchelle | +| [@julien-duponchelle](https://github.com/julien-duponchelle) | Julien Duponchelle | | [@alisterd51](https://github.com/alisterd51) | Antoine Clarman | diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 6a37435db..3d5a59564 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -15,8 +15,6 @@ Breaking changes Feat ^^^^ * Add a PostgreSQL/PGMQ broker with partitioned queues, native delayed messages, ``LISTEN/NOTIFY`` wakeups, and queue join support. -* Add the ``PgmqBroker`` alias for the PostgreSQL-backed broker. - Changed ^^^^^^^ * Restore the main APIs after the broker refactor. diff --git a/remoulade/api/main.py b/remoulade/api/main.py index b0a9831d2..7008d80ec 100644 --- a/remoulade/api/main.py +++ b/remoulade/api/main.py @@ -114,7 +114,7 @@ def get_results(message_id): max_size = 1e4 try: result = Result[Any](message_id=message_id).get() - encoded_result = get_encoder().encode(result).decode("utf-8") + encoded_result = get_encoder().encode_in_bytes(result).decode("utf-8") size_result = sys.getsizeof(encoded_result) if size_result >= max_size: encoded_result = f"The result is too big {size_result / 1e6}M" diff --git a/remoulade/brokers/pgmq.py b/remoulade/brokers/pgmq.py deleted file mode 100644 index 787649a26..000000000 --- a/remoulade/brokers/pgmq.py +++ /dev/null @@ -1,13 +0,0 @@ -from .postgres import ( - PostgresBroker, - PostgresPayload, - PostgresQueueMessage, - _PostgresConsumer, - _PostgresMessage, -) - -PgmqBroker = PostgresBroker -PgmqPayload = PostgresPayload -PgmqQueueMessage = PostgresQueueMessage -_PgmqConsumer = _PostgresConsumer -_PgmqMessage = _PostgresMessage diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index b0d28c3d1..bdf5d014f 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -193,14 +193,17 @@ def _encode_message(self, message: "Message") -> PostgresPayload: or is not a JSON object. """ try: - payload = json.loads(message.encode().decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as e: - raise UnsupportedMessageEncoding( - "PostgresBroker only supports encoders that serialize messages as JSON objects." - ) from e + payload = message.encode_in_json() + except (TypeError, ValueError, UnsupportedMessageEncoding) as exc: + raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") from exc if not isinstance(payload, dict): - raise UnsupportedMessageEncoding("PostgresBroker requires message encoders to produce JSON objects.") + raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") + + try: + json.dumps(payload) + except (TypeError, ValueError) as exc: + raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") from exc return payload @@ -535,7 +538,7 @@ def __init__(self, postgres_message: PostgresQueueMessage) -> None: try: # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder) # can rehydrate actor args/kwargs to their typed schemas. - message = Message.decode(json.dumps(payload).encode("utf-8")) + message = Message.decode_json(payload) except TypeError as exc: raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc if message.options.get("eta", None) is not None: diff --git a/remoulade/brokers/rabbitmq.py b/remoulade/brokers/rabbitmq.py index 7b5003aa9..844c4a180 100644 --- a/remoulade/brokers/rabbitmq.py +++ b/remoulade/brokers/rabbitmq.py @@ -350,7 +350,7 @@ def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message" self.logger.debug("Enqueueing message %r on queue %r.", message.message_id, queue_name) with self._get_channel(confirm_delivery) as channel: confirmation = channel.basic.publish( - exchange="", routing_key=queue_name, body=message.encode(), properties=properties + exchange="", routing_key=queue_name, body=message.encode_in_bytes(), properties=properties ) if confirm_delivery and not confirmation: raise MessageNotDelivered("Message could not be delivered") @@ -515,7 +515,7 @@ def close(self): class _RabbitmqMessage(MessageProxy): def __init__(self, rabbitmq_message): - super().__init__(Message.decode(rabbitmq_message.body)) + super().__init__(Message.decode_bytes(rabbitmq_message.body)) self._rabbitmq_message = rabbitmq_message diff --git a/remoulade/brokers/stub.py b/remoulade/brokers/stub.py index 968151db3..7dd137f4f 100644 --- a/remoulade/brokers/stub.py +++ b/remoulade/brokers/stub.py @@ -99,7 +99,7 @@ def _enqueue(self, message, *, delay=None): if queue_name not in self.queues: raise QueueNotFound(queue_name) - self.queues[queue_name].put(message.encode()) + self.queues[queue_name].put(message.encode_in_bytes()) return message def flush(self, queue_name): @@ -165,7 +165,7 @@ def nack(self, message): def __next__(self): try: data = self.queue.get(timeout=self.timeout / 1000) - message = Message.decode(data) + message = Message.decode_bytes(data) return MessageProxy(message) except Empty: return None diff --git a/remoulade/encoder.py b/remoulade/encoder.py index 3c6f52653..e92c022ea 100644 --- a/remoulade/encoder.py +++ b/remoulade/encoder.py @@ -19,15 +19,15 @@ import json import pickle import warnings -from typing import Annotated, Any, get_type_hints +from typing import Annotated, Any, get_type_hints, override + +from remoulade.errors import UnsupportedMessageEncoding try: - from pydantic import BaseModel, TypeAdapter, WithJsonSchema - from simplejson.decoder import JSONDecoder - from simplejson.encoder import JSONEncoder as _JSONEncoder + from pydantic import TypeAdapter, WithJsonSchema except ImportError: # pragma: no cover warnings.warn( - "Pydantic and simplejson are not available. Run `pip install remoulade[pydantic]`", + "Pydantic is not available. Run `pip install remoulade[pydantic]`", ImportWarning, stacklevel=2, ) @@ -35,30 +35,57 @@ #: Represents the contents of a Message object as a dict. MessageData = dict[str, Any] +JsonData = dict[str, Any] class Encoder(abc.ABC): """Base class for message encoders.""" @abc.abstractmethod - def encode(self, data: MessageData) -> bytes: # pragma: no cover - """Convert message metadata into a bytestring.""" + def encode_in_bytes(self, data: MessageData) -> bytes: raise NotImplementedError @abc.abstractmethod - def decode(self, data: bytes) -> MessageData: # pragma: no cover - """Convert a bytestring into message metadata.""" + def decode_bytes(self, data: bytes) -> MessageData: + raise NotImplementedError + + def encode_in_json(self, data: MessageData) -> JsonData: + encoded = self._encode_in_json(data) + try: + json.dumps(encoded) + except (TypeError, ValueError) as e: + raise UnsupportedMessageEncoding("This is not a valid json") from e + return encoded + + @abc.abstractmethod + def _encode_in_json(self, data: MessageData) -> JsonData: + return data + + @abc.abstractmethod + def decode_json(self, data: JsonData) -> MessageData: raise NotImplementedError class JSONEncoder(Encoder): """Encodes messages as JSON. This is the default encoder.""" - def encode(self, data: MessageData) -> bytes: - return json.dumps(data, separators=(",", ":")).encode("utf-8") + @override + def encode_in_bytes(self, data: MessageData) -> bytes: # pragma: no cover + """Convert message metadata into a bytestring.""" + return json.dumps(self.encode_in_json(data), separators=(",", ":")).encode("utf-8") - def decode(self, data: bytes) -> MessageData: - return json.loads(data.decode("utf-8")) + @override + def decode_bytes(self, data: bytes) -> MessageData: # pragma: no cover + """Convert a bytestring into message metadata.""" + return self.decode_json(json.loads(data.decode("utf-8"))) + + @override + def _encode_in_json(self, data: MessageData) -> JsonData: + return data + + @override + def decode_json(self, data: JsonData) -> MessageData: + return data class PickleEncoder(Encoder): @@ -69,8 +96,21 @@ class PickleEncoder(Encoder): Use it at your own risk. """ - encode = pickle.dumps # type: ignore - decode = pickle.loads # type: ignore + @override + def encode_in_bytes(self, data: MessageData) -> bytes: + return pickle.dumps(data) + + @override + def decode_bytes(self, data: bytes) -> MessageData: + return pickle.loads(data) # noqa: S301 + + @override + def _encode_in_json(self, data: MessageData) -> JsonData: + raise TypeError("PickleEncoder does not support JSON encoding.") + + @override + def decode_json(self, data: JsonData) -> MessageData: + raise TypeError("PickleEncoder does not support JSON decoding.") class PydanticEncoder(Encoder): @@ -91,31 +131,36 @@ def my_actor(input_1: MyActorInputSchema, input_2: MyActorInputSchema | None = N def __init__(self, fallback_encoder: Encoder | None = None): self.fallback_encoder = fallback_encoder - self.json_encoder = _JSONEncoder(default=self.default) - self.json_decoder = JSONDecoder() - - @staticmethod - def default(o): - if isinstance(o, BaseModel): - # keep dict otherwise it will be serialized as a string (see Pydantic .json()) - return json.loads(o.model_dump_json()) - raise TypeError(f"Object of type {o.__class__.__name__} is not JSON serializable") + self.json_adapter = TypeAdapter(object) - def encode(self, data: MessageData) -> bytes: + @override + def encode_in_bytes(self, data: MessageData) -> bytes: try: - return self.json_encoder.encode(data).encode("utf-8") + return json.dumps(self.encode_in_json(data)).encode("utf-8") except Exception as e: if self.fallback_encoder is not None: - return self.fallback_encoder.encode(data) - else: - raise e + return self.fallback_encoder.encode_in_bytes(data) + raise e - def decode(self, data: bytes) -> MessageData: + @override + def decode_bytes(self, data: bytes) -> MessageData: + try: + return self.decode_json(json.loads(data.decode("utf-8"))) + except Exception: + if self.fallback_encoder is not None: + return self.fallback_encoder.decode_bytes(data) + raise + + @override + def _encode_in_json(self, data: MessageData) -> JsonData: + return self.json_adapter.dump_python(data, mode="json") + + @override + def decode_json(self, data: JsonData) -> MessageData: from remoulade import get_broker try: - raw_message = self.json_decoder.decode(data.decode("utf-8")) - actor_name = raw_message["actor_name"] + actor_name = data["actor_name"] actor_fn = get_broker().get_actor(actor_name).fn # Retrieve the Pydantic schemas from typing @@ -130,9 +175,9 @@ def decode(self, data: bytes) -> MessageData: ] ) - # Override message_data with Pydantic schema when it matches + # Override message_data with Pydantic schema when it matches. parsed_message: dict[str, Any] = {} - for key, values in raw_message.items(): + for key, values in data.items(): if key == "kwargs": if not isinstance(values, dict): raise TypeError(f"Expected `values` to be a dict, got {type(values).__name__}") @@ -156,8 +201,7 @@ def decode(self, data: bytes) -> MessageData: parsed_message[key] = values return parsed_message - except Exception as e: + except Exception: if self.fallback_encoder is not None: - return self.fallback_encoder.decode(data) - else: - raise e + return self.fallback_encoder.decode_json(data) + raise diff --git a/remoulade/message.py b/remoulade/message.py index 516203cc0..be8e37be7 100644 --- a/remoulade/message.py +++ b/remoulade/message.py @@ -89,13 +89,22 @@ def asdict(self): return attr.asdict(self, recurse=False) @classmethod - def decode(cls, data): + def decode_bytes(cls, data: bytes): """Convert a bytestring to a message.""" - return cls(**global_encoder.decode(data)) + return cls(**global_encoder.decode_bytes(data)) - def encode(self): + def encode_in_bytes(self): """Convert this message to a bytestring.""" - return global_encoder.encode(self.asdict()) + return global_encoder.encode_in_bytes(self.asdict()) + + @classmethod + def decode_json(cls, data: dict[str, Any]): + """Convert a bytestring to a message.""" + return cls(**global_encoder.decode_json(data)) + + def encode_in_json(self): + """Convert this message to a bytestring.""" + return global_encoder.encode_in_json(self.asdict()) def copy(self, **attributes): """Create a copy of this message.""" diff --git a/remoulade/results/backends/redis.py b/remoulade/results/backends/redis.py index 9bc606ee4..f3895b73c 100644 --- a/remoulade/results/backends/redis.py +++ b/remoulade/results/backends/redis.py @@ -88,7 +88,7 @@ def get_results( for message_id in message_ids: message_key = self.build_message_key(message_id) if forget: - pipe.rpushx(message_key, self.encoder.encode(ForgottenResult.asdict())) + pipe.rpushx(message_key, self.encoder.encode_in_bytes(ForgottenResult.asdict())) pipe.lpop(message_key) else: pipe.rpoplpush(message_key, message_key) @@ -98,7 +98,7 @@ def get_results( continue # skip one row in two if forget as there is two commands if row is None: raise ResultMissing(message_id) - yield self.process_result(BackendResult(**self.encoder.decode(row)), raise_on_error) + yield self.process_result(BackendResult(**self.encoder.decode_bytes(row)), raise_on_error) def _brpoplpush_with_timeout(self, src, dst, timeout: int): """ @@ -150,7 +150,7 @@ async def async_get_result( if timeout > 0: if forget: async with self.async_client.pipeline() as pipe: - await pipe.rpushx(message_key, self.encoder.encode(ForgottenResult.asdict())) + await pipe.rpushx(message_key, self.encoder.encode_in_bytes(ForgottenResult.asdict())) await pipe.lpop(message_key) pipe_exec = await pipe.execute() data = pipe_exec[1] @@ -176,7 +176,7 @@ async def async_get_result( timeout = int(deadline - time.monotonic()) if timeout <= 0: # do not retry is timeout is expired raise - result = BackendResult(**self.encoder.decode(data)) + result = BackendResult(**self.encoder.decode_bytes(data)) return self.process_result(result, raise_on_error) def get_result(self, message_id: str, *, block=False, timeout=None, forget=False, raise_on_error=True): @@ -214,14 +214,14 @@ def get_result(self, message_id: str, *, block=False, timeout=None, forget=False data = self._brpoplpush_with_timeout(message_key, message_key, timeout=timeout) if forget and data is not None: with self.client.pipeline() as pipe: - pipe.lpushx(message_key, self.encoder.encode(ForgottenResult.asdict())) + pipe.lpushx(message_key, self.encoder.encode_in_bytes(ForgottenResult.asdict())) pipe.ltrim(message_key, 0, 0) pipe.execute() else: if forget: with self.client.pipeline() as pipe: - pipe.rpushx(message_key, self.encoder.encode(ForgottenResult.asdict())) + pipe.rpushx(message_key, self.encoder.encode_in_bytes(ForgottenResult.asdict())) pipe.lpop(message_key) data = pipe.execute()[1] else: @@ -251,21 +251,21 @@ def get_result(self, message_id: str, *, block=False, timeout=None, forget=False if block and timeout <= 0: # do not retry is timeout is expired raise - result = BackendResult(**self.encoder.decode(data)) + result = BackendResult(**self.encoder.decode_bytes(data)) return self.process_result(result, raise_on_error) def _store(self, message_keys, results, ttl): with self.client.pipeline() as pipe: for message_key, result in zip(message_keys, results, strict=False): pipe.delete(message_key) - pipe.lpush(message_key, self.encoder.encode(result)) + pipe.lpush(message_key, self.encoder.encode_in_bytes(result)) pipe.pexpire(message_key, ttl) pipe.execute() def _get(self, key, forget=False): data = self.client.rpop(key) if forget else self.client.rpoplpush(key, key) if data: - return self.encoder.decode(data) + return self.encoder.decode_bytes(data) return Missing def _delete(self, key): diff --git a/remoulade/results/backends/stub.py b/remoulade/results/backends/stub.py index 21ab3147a..3d613a47b 100644 --- a/remoulade/results/backends/stub.py +++ b/remoulade/results/backends/stub.py @@ -36,17 +36,17 @@ def _get(self, message_key: str, forget: bool = False): if forget: data, expiration = self.results.get(message_key, (None, None)) if data is not None: - self.results[message_key] = self.encoder.encode(ForgottenResult.asdict()), expiration + self.results[message_key] = self.encoder.encode_in_bytes(ForgottenResult.asdict()), expiration else: data, expiration = self.results.get(message_key, (None, None)) if data is not None and time.monotonic() < expiration: - return self.encoder.decode(data) + return self.encoder.decode_bytes(data) return Missing def _store(self, message_keys, results, ttl): for message_key, result in zip(message_keys, results, strict=False): - result_data = self.encoder.encode(result) + result_data = self.encoder.encode_in_bytes(result) expiration = time.monotonic() + int(ttl / 1000) self.results[message_key] = (result_data, expiration) diff --git a/remoulade/scheduler/scheduler.py b/remoulade/scheduler/scheduler.py index 6332621d2..45095500d 100644 --- a/remoulade/scheduler/scheduler.py +++ b/remoulade/scheduler/scheduler.py @@ -72,9 +72,9 @@ def get_hash(self) -> str: str(self.iso_weekday), str(self.enabled), self.tz, - *(encoder.encode(arg).decode() for arg in self.args), + *(encoder.encode_in_bytes(arg).decode() for arg in self.args), *( - f"{name}: {encoder.encode(arg).decode()}" + f"{name}: {encoder.encode_in_bytes(arg).decode()}" for name, arg in sorted(self.kwargs.items(), key=itemgetter(0)) ), ] @@ -104,11 +104,11 @@ def as_dict(self, encode: bool = False) -> dict: return job_dict def encode(self) -> bytes: - return get_encoder().encode(self.as_dict(encode=True)) + return get_encoder().encode_in_bytes(self.as_dict(encode=True)) @classmethod def decode(cls, data: bytes) -> "ScheduledJob": - data = get_encoder().decode(data) + data = get_encoder().decode_bytes(data) return ScheduledJob( actor_name=data["actor_name"], interval=data["interval"], diff --git a/remoulade/state/backend.py b/remoulade/state/backend.py index 783265917..b044dd6ec 100644 --- a/remoulade/state/backend.py +++ b/remoulade/state/backend.py @@ -99,7 +99,7 @@ def as_dict(self, exclude_keys=(), encode_args=False): for key in (item for item in ["args", "kwargs", "options"] if item in as_dict): try: - as_dict[key] = get_encoder().encode(as_dict[key]).decode("utf-8") + as_dict[key] = get_encoder().encode_in_bytes(as_dict[key]).decode("utf-8") except (UnicodeDecodeError, TypeError): as_dict[key] = "encoded_data" return as_dict @@ -201,16 +201,16 @@ def _encode_dict(self, data): """Return the (keys, values) of a dictionary encoded""" encoded_data = {} for key, value in data.items(): - encoded_value = self.encoder.encode(value) + encoded_value = self.encoder.encode_in_bytes(value) if sys.getsizeof(encoded_value) <= self.max_size: - encoded_data[self.encoder.encode(key)] = self.encoder.encode(value) + encoded_data[self.encoder.encode_in_bytes(key)] = self.encoder.encode_in_bytes(value) return encoded_data def _decode_dict(self, data): """Return the (keys, values) of a dictionary decoded""" decoded_data = {} for key, value in data.items(): - decoded_data[self.encoder.decode(key)] = self.encoder.decode(value) + decoded_data[self.encoder.decode_bytes(key)] = self.encoder.decode_bytes(value) return decoded_data def clean(self, max_age: int | None = None, not_started: bool = False): diff --git a/tests/test_actors.py b/tests/test_actors.py index 6d40acec2..57baef068 100644 --- a/tests/test_actors.py +++ b/tests/test_actors.py @@ -103,7 +103,7 @@ def add(x, y): # I expect it to enqueue a message enqueued_message = add.send(1, 2) enqueued_message_data = stub_broker.queues["default"].get(timeout=1) - assert enqueued_message == Message.decode(enqueued_message_data) + assert enqueued_message == Message.decode_bytes(enqueued_message_data) def test_actors_no_broker(): diff --git a/tests/test_encoders.py b/tests/test_encoders.py index 76d0ba817..67035b4b9 100644 --- a/tests/test_encoders.py +++ b/tests/test_encoders.py @@ -45,6 +45,15 @@ def add_value(x): assert db == [1] +def test_json_encoder_json_round_trip(): + encoder = JSONEncoder() + data = {"queue_name": "default", "args": [1, 2], "kwargs": {"debug": True}} + + assert encoder.encode_in_json(data) == data + assert encoder.decode_json(data) == data + assert encoder.decode_bytes(encoder.encode_in_bytes(data)) == data + + class MyEnum(Enum): val = "val" other = "other" @@ -155,29 +164,46 @@ def encoder_with_fallback(stub_broker, stub_worker, result_backend) -> PydanticE def test_encoder_message(encoder: PydanticEncoder, message_data_decoded: MessageData, message_data_encoded: bytes): - encoded_result = encoder.encode(message_data_decoded) + encoded_result = encoder.encode_in_bytes(message_data_decoded) assert encoded_result == message_data_encoded - decoded_result = encoder.decode(message_data_encoded) + decoded_result = encoder.decode_bytes(message_data_encoded) # Args tuple are assumed to become list in remoulade assert decoded_result == tuple_to_list(message_data_decoded, "args") - assert encoder.decode(encoder.encode(message_data_decoded)) == tuple_to_list(message_data_decoded, "args") - assert encoder.encode(encoder.decode(message_data_encoded)) == message_data_encoded + assert encoder.decode_bytes(encoder.encode_in_bytes(message_data_decoded)) == tuple_to_list( + message_data_decoded, "args" + ) + assert encoder.encode_in_bytes(encoder.decode_bytes(message_data_encoded)) == message_data_encoded + + +def test_encoder_json_message_round_trip(pydantic_encoder, encoder: PydanticEncoder, message_data_decoded: MessageData): + message = Message(**message_data_decoded) + + encoded_json = message.encode_in_json() + + assert encoded_json["args"] == [json.loads(input_1.model_dump_json()) for input_1 in message.args] + assert encoded_json["kwargs"] == {"input_2": json.loads(message.kwargs["input_2"].model_dump_json())} + + decoded_message = Message.decode_json(encoded_json) + + assert decoded_message == message + assert isinstance(decoded_message.args[0], MyFirstSchema) + assert isinstance(decoded_message.kwargs["input_2"], MySecondSchema) def test_message_unknown_actor(encoder: PydanticEncoder, message_data_encoded: bytes): message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) message_json_decoded["actor_name"] = "titi" with pytest.raises(ActorNotFound): - encoder.decode(json.dumps(message_json_decoded).encode("utf-8")) + encoder.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) def test_message_fallback_no_actor_name(encoder_with_fallback: PydanticEncoder, message_data_encoded: bytes): message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) message_json_decoded["actor_name"] = "titi" - decoded_result = encoder_with_fallback.decode(json.dumps(message_json_decoded).encode("utf-8")) + decoded_result = encoder_with_fallback.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) # Do not raise and keep dict instead of schema assert decoded_result == tuple_to_list(message_json_decoded, "args") @@ -187,7 +213,7 @@ def test_message_schema_not_matching(encoder: PydanticEncoder, message_data_enco message_json_decoded["args"] = [{"toto": "a"}] message_json_decoded["kwargs"]["input_2"] = {"val": "aaa"} with pytest.raises(ValidationError): - encoder.decode(json.dumps(message_json_decoded).encode("utf-8")) + encoder.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) @pytest.fixture @@ -236,76 +262,76 @@ def backend_result_encoded_none() -> bytes: def test_encoder_result(encoder: PydanticEncoder, backend_result_decoded: MessageData, backend_result_encoded: bytes): - encoded_value = encoder.encode(backend_result_decoded) + encoded_value = encoder.encode_in_bytes(backend_result_decoded) assert encoded_value == backend_result_encoded - decoded_result = encoder.decode(backend_result_encoded) + decoded_result = encoder.decode_bytes(backend_result_encoded) assert decoded_result == backend_result_decoded - assert encoder.decode(encoder.encode(backend_result_decoded)) == backend_result_decoded - assert encoder.encode(encoder.decode(backend_result_encoded)) == backend_result_encoded + assert encoder.decode_bytes(encoder.encode_in_bytes(backend_result_decoded)) == backend_result_decoded + assert encoder.encode_in_bytes(encoder.decode_bytes(backend_result_encoded)) == backend_result_encoded def test_encoder_result_when_raise( encoder: PydanticEncoder, backend_result_decoded_raise: MessageData, backend_result_encoded_raise: bytes ): - encoded_value = encoder.encode(backend_result_decoded_raise) + encoded_value = encoder.encode_in_bytes(backend_result_decoded_raise) assert encoded_value == backend_result_encoded_raise - decoded_result = encoder.decode(backend_result_encoded_raise) + decoded_result = encoder.decode_bytes(backend_result_encoded_raise) assert decoded_result == backend_result_decoded_raise - assert encoder.decode(encoder.encode(backend_result_decoded_raise)) == backend_result_decoded_raise - assert encoder.encode(encoder.decode(backend_result_encoded_raise)) == backend_result_encoded_raise + assert encoder.decode_bytes(encoder.encode_in_bytes(backend_result_decoded_raise)) == backend_result_decoded_raise + assert encoder.encode_in_bytes(encoder.decode_bytes(backend_result_encoded_raise)) == backend_result_encoded_raise def test_encoder_result_tuple( encoder: PydanticEncoder, backend_result_decoded_tuple: MessageData, backend_result_encoded_tuple: bytes ): - encoded_value = encoder.encode(backend_result_decoded_tuple) + encoded_value = encoder.encode_in_bytes(backend_result_decoded_tuple) assert encoded_value == backend_result_encoded_tuple - decoded_result = encoder.decode(backend_result_encoded_tuple) + decoded_result = encoder.decode_bytes(backend_result_encoded_tuple) assert decoded_result == backend_result_decoded_tuple - assert encoder.decode(encoder.encode(backend_result_decoded_tuple)) == backend_result_decoded_tuple - assert encoder.encode(encoder.decode(backend_result_encoded_tuple)) == backend_result_encoded_tuple + assert encoder.decode_bytes(encoder.encode_in_bytes(backend_result_decoded_tuple)) == backend_result_decoded_tuple + assert encoder.encode_in_bytes(encoder.decode_bytes(backend_result_encoded_tuple)) == backend_result_encoded_tuple def test_encoder_result_with_none( encoder: PydanticEncoder, backend_result_decoded_none: MessageData, backend_result_encoded_none: bytes ): - encoded_value = encoder.encode(backend_result_decoded_none) + encoded_value = encoder.encode_in_bytes(backend_result_decoded_none) assert encoded_value == backend_result_encoded_none - decoded_result = encoder.decode(backend_result_encoded_none) + decoded_result = encoder.decode_bytes(backend_result_encoded_none) assert decoded_result == backend_result_decoded_none - assert encoder.decode(encoder.encode(backend_result_decoded_none)) == backend_result_decoded_none - assert encoder.encode(encoder.decode(backend_result_encoded_none)) == backend_result_encoded_none + assert encoder.decode_bytes(encoder.encode_in_bytes(backend_result_decoded_none)) == backend_result_decoded_none + assert encoder.encode_in_bytes(encoder.decode_bytes(backend_result_encoded_none)) == backend_result_encoded_none def test_backend_result_unknown_actor(encoder: PydanticEncoder, backend_result_encoded_tuple: bytes): backend_result_json_decoded = json.loads(backend_result_encoded_tuple.decode("utf-8")) backend_result_json_decoded["actor_name"] = "titi" with pytest.raises(ActorNotFound): - encoder.decode(json.dumps(backend_result_json_decoded).encode("utf-8")) + encoder.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) def test_backend_result_schema_not_matching(encoder: PydanticEncoder, backend_result_encoded_tuple: bytes): backend_result_json_decoded = json.loads(backend_result_encoded_tuple.decode("utf-8")) backend_result_json_decoded["result"] = [{"val": "titi"}] with pytest.raises(ValidationError): - encoder.decode(json.dumps(backend_result_json_decoded).encode("utf-8")) + encoder.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) def test_fallback_no_schema(encoder_with_fallback: PydanticEncoder, backend_result_encoded_tuple: bytes): backend_result_json_decoded = json.loads(backend_result_encoded_tuple.decode("utf-8")) backend_result_json_decoded["result"] = [{"val": "titi"}] - decoded_backend_result = encoder_with_fallback.decode(json.dumps(backend_result_json_decoded).encode("utf-8")) + decoded_backend_result = encoder_with_fallback.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) # Do not raise and keep dict instead of schema assert decoded_backend_result == backend_result_json_decoded diff --git a/tests/test_postgres.py b/tests/test_postgres.py index b31cef075..6920ce6ee 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -10,8 +10,9 @@ from sqlalchemy.sql import select, text import remoulade -from remoulade import Message, QueueJoinTimeout, Worker +from remoulade import Message, QueueJoinTimeout, UnsupportedMessageEncoding, Worker from remoulade.brokers.postgres import PostgresBroker +from remoulade.encoder import Encoder, MessageData TEST_POSTGRES_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" @@ -48,7 +49,7 @@ def _queue_exists(broker, queue_name): def _expected_payload(message): - return json.loads(message.encode().decode("utf-8")) + return json.loads(message.encode_in_bytes().decode("utf-8")) def test_postgres_broker_uses_provided_url(): @@ -184,6 +185,53 @@ def test_postgres_broker_uses_custom_partition_settings_when_provided(): ) +def test_postgres_broker_rejects_non_json_message_encoders(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + + class _MessageWithInvalidJson(Encoder): + def _encode_in_json(self, data): + raise TypeError("not json") + + def decode_json(self, data): + return data + + def encode_in_bytes(self, data: MessageData) -> bytes: + return b"" + + def decode_bytes(self, data: bytes) -> MessageData: + return {} + + with pytest.raises(UnsupportedMessageEncoding): + broker._encode_message(_MessageWithInvalidJson()) + + +def test_postgres_broker_rejects_nested_non_json_safe_payloads(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + old_encoder = remoulade.get_encoder() + + class _NestedInvalidJsonEncoder(Encoder): + def _encode_in_json(self, data): + return {**data, "options": {"nested": {1}}} + + def decode_json(self, data): + return data + + def encode_in_bytes(self, data: MessageData) -> bytes: + return b"" + + def decode_bytes(self, data: bytes) -> MessageData: + return {} + + remoulade.set_encoder(_NestedInvalidJsonEncoder()) + try: + message = Message(queue_name="default", actor_name="do_work", args=(), kwargs={}, options={}) + + with pytest.raises(UnsupportedMessageEncoding): + broker._encode_message(message) + finally: + remoulade.set_encoder(old_encoder) + + @pytest.mark.usefixtures("postgres_broker") def test_postgres_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) diff --git a/tests/test_results.py b/tests/test_results.py index f45cac153..0d82030f3 100644 --- a/tests/test_results.py +++ b/tests/test_results.py @@ -534,7 +534,7 @@ def test_redis_get_result_with_block_timeout_larger_than_socket_timeout(redis_re redis.TimeoutError(), redis.TimeoutError(), redis.TimeoutError(), - redis_result_backend.encoder.encode(ForgottenResult.asdict()), + redis_result_backend.encoder.encode_in_bytes(ForgottenResult.asdict()), ] redis_result_backend.max_retries = 1 redis_result_backend.get_result("message-id", block=True, forget=False, timeout=60 * 1000) @@ -545,7 +545,7 @@ def test_redis_get_result_with_block_timeout_larger_than_socket_timeout(redis_re @mock.patch("remoulade.results.backends.redis.compute_backoff", fast_backoff) def test_redis_get_result_still_return_result_if_forget_fails(redis_result_backend): with patch.object(redis_result_backend, "client") as mock_client: - mock_client.brpoplpush.return_value = redis_result_backend.encoder.encode(ForgottenResult.asdict()) + mock_client.brpoplpush.return_value = redis_result_backend.encoder.encode_in_bytes(ForgottenResult.asdict()) mock_client.pipeline.side_effect = redis.ConnectionError() assert redis_result_backend.get_result("message-id", block=True, forget=True) is None assert mock_client.brpoplpush.call_count == 1 @@ -613,7 +613,7 @@ async def test_redis_async_get_result_with_block_timeout_larger_than_socket_time redis.TimeoutError(), redis.TimeoutError(), redis.TimeoutError(), - redis_result_backend.encoder.encode(ForgottenResult.asdict()), + redis_result_backend.encoder.encode_in_bytes(ForgottenResult.asdict()), ] redis_result_backend.max_retries = 1 await redis_result_backend.async_get_result("message-id", forget=False, timeout=60 * 1000) From 01b3511d3bfffe6182a3f59c76a40da3f121d543 Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 17 Jun 2026 10:34:48 +0200 Subject: [PATCH 19/44] fix(broker): interval unit --- remoulade/broker.py | 3 ++- remoulade/brokers/postgres.py | 23 ++++++++++++++++------- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/remoulade/broker.py b/remoulade/broker.py index d8eaef526..8732b6d7e 100644 --- a/remoulade/broker.py +++ b/remoulade/broker.py @@ -194,12 +194,13 @@ class Broker: overwrite when they are declared. """ + supports_native_delay = False + def __init__(self, middleware: "Iterable[Middleware] | None" = None): self.logger = get_logger(__name__, type(self)) self.actors: dict[str, Actor] = {} self.queues: dict[str, Queue | None] = {} self.delay_queues: set[str] = set() - self.supports_native_delay = False self.actor_options: set[str] = set() self.middleware: list[Middleware] = [] diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index bdf5d014f..2632f56a6 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -55,14 +55,16 @@ class PostgresBroker(Broker): PostgreSQL instead of being staged in worker memory. """ + supports_native_delay = True + def __init__( self, *, url: str, middleware: list["Middleware"] | None = None, group_transaction: bool = False, - archive_partition_interval: str = "1 day", - archive_retention_interval: str = "7 days", + archive_partition_interval_in_days: int = 1, + archive_retention_interval_in_days: int = 7, visibility_timeout_ms: int = 30_000, heartbeat_interval_ms: int = 10_000, ) -> None: @@ -73,8 +75,8 @@ def __init__( The url must be the creds for a user who can create and delete tables middleware(list[Middleware] | None): Middleware stack applied to this broker. group_transaction(bool): If True, wraps group and pipeline operations in a single transaction. - archive_partition_interval(str): Partition interval passed to PGMQ when creating partitioned queues. - archive_retention_interval(str): Retention interval passed to PGMQ for archive partitions. + archive_partition_interval_in_days(int): Partition interval passed to PGMQ when creating partitioned queues. + archive_retention_interval_in_days(int): Retention interval passed to PGMQ for archive partitions. visibility_timeout_ms(int): Message visibility timeout in milliseconds after read; must be greater than 0. heartbeat_interval_ms(int): Heartbeat interval in milliseconds used to extend in-flight message visibility must be greater than 0 and lower than visibility_timeout_ms. @@ -88,13 +90,12 @@ def __init__( self.url = urlparse(url).geturl() self.state = local() self.group_transaction = group_transaction - self.archive_partition_interval = archive_partition_interval - self.archive_retention_interval = archive_retention_interval + self.archive_partition_interval = self.convert_days_in_partman_syntax(archive_partition_interval_in_days) + self.archive_retention_interval = self.convert_days_in_partman_syntax(archive_retention_interval_in_days) self.visibility_timeout_ms = visibility_timeout_ms self.heartbeat_interval_ms = heartbeat_interval_ms self.visibility_timeout_seconds = _milliseconds_to_seconds(visibility_timeout_ms) self.heartbeat_interval_seconds = heartbeat_interval_ms / 1000 - self.supports_native_delay = True self.client = SQLAlchemyPGMQueue( conn_string=url, @@ -102,6 +103,14 @@ def __init__( vt=self.visibility_timeout_seconds, ) + def convert_days_in_partman_syntax(self, interval_in_day: int) -> str: + """Convert int into partman syntax""" + if interval_in_day <= 0: + raise ValueError("interval_in_day must be greater than 0") + if interval_in_day == 1: + return "1 day" + return f"{interval_in_day} days" + @override @contextmanager def tx(self) -> Iterator[None]: From d85986eb13bd0b9013a2b4d382654e9b398ee24a Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 17 Jun 2026 14:33:11 +0200 Subject: [PATCH 20/44] fix(postgres-broker): close properly --- remoulade/brokers/postgres.py | 61 +++++++++++++++-------------- tests/test_postgres.py | 74 ++++++++++++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 31 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 2632f56a6..0ac24ce1b 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -109,7 +109,7 @@ def convert_days_in_partman_syntax(self, interval_in_day: int) -> str: raise ValueError("interval_in_day must be greater than 0") if interval_in_day == 1: return "1 day" - return f"{interval_in_day} days" + return f"{interval_in_day} days" @override @contextmanager @@ -329,8 +329,8 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti self._listener_available = False self._heartbeat_stop = Event() self._heartbeat_thread: Thread | None = None - self._heartbeat_msg_ids_lock = Lock() - self._heartbeat_msg_ids: set[int] = set() + self._heartbeat_message_ids_lock = Lock() + self._heartbeat_message_ids: set[int] = set() self._start_listener() self._start_heartbeat() @@ -393,12 +393,12 @@ def _run_heartbeat(self) -> None: requeued. Failures are logged and retried on the next heartbeat tick. """ while not self._heartbeat_stop.wait(self.heartbeat_interval_seconds): - with self._heartbeat_msg_ids_lock: - msg_ids = list(self._heartbeat_msg_ids) - if not msg_ids: + with self._heartbeat_message_ids_lock: + message_ids = list(self._heartbeat_message_ids) + if not message_ids: continue try: - self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout_seconds) + self.client.set_vt(self.queue_name, message_ids, self.visibility_timeout_seconds) except Exception as e: self.broker.logger.warning( "Failed to extend visibility timeout heartbeat for queue %s. Error: ", @@ -407,15 +407,17 @@ def _run_heartbeat(self) -> None: exc_info=True, ) - def _register_heartbeat_msg(self, message: "_PostgresMessage") -> None: - """Track a message so the heartbeat can renew it.""" - with self._heartbeat_msg_ids_lock: - self._heartbeat_msg_ids.add(message._postgres_message.msg_id) - - def _unregister_heartbeat_msg_id(self, msg_id: int) -> None: + def _unregister_heartbeat_message_id(self, message_id: int) -> None: """Stop tracking a message for heartbeat renewal.""" - with self._heartbeat_msg_ids_lock: - self._heartbeat_msg_ids.discard(msg_id) + with self._heartbeat_message_ids_lock: + self._heartbeat_message_ids.discard(message_id) + + def _requeue_message_ids(self, message_ids: list[int]) -> None: + """Make a batch of message ids visible again immediately.""" + for message_id in message_ids: + self._unregister_heartbeat_message_id(message_id) + if len(message_ids) > 0: + self.client.set_vt(self.queue_name, message_ids, 0) def _start_listener(self) -> None: """Start a LISTEN/NOTIFY listener for the queue when available.""" @@ -470,7 +472,7 @@ def ack(self, message: "MessageProxy") -> None: """Archive a processed message.""" if not isinstance(message, _PostgresMessage): raise ValueError("It must be a PostgresMessage") - self._unregister_heartbeat_msg_id(message._postgres_message.msg_id) + self._unregister_heartbeat_message_id(message._postgres_message.msg_id) self.client.archive(self.queue_name, message._postgres_message.msg_id) @override @@ -478,23 +480,20 @@ def nack(self, message: "MessageProxy") -> None: """Archive a failed message.""" if not isinstance(message, _PostgresMessage): raise ValueError("It must be a PostgresMessage") - self._unregister_heartbeat_msg_id(message._postgres_message.msg_id) + self._unregister_heartbeat_message_id(message._postgres_message.msg_id) self.client.archive(self.queue_name, message._postgres_message.msg_id) @override def requeue(self, messages: Iterable["MessageProxy"]) -> None: """Make messages visible again immediately by resetting their visibility timeout.""" - msg_ids = [message._postgres_message.msg_id for message in messages if isinstance(message, _PostgresMessage)] - if msg_ids: - for msg_id in msg_ids: - self._unregister_heartbeat_msg_id(msg_id) - self.client.set_vt(self.queue_name, msg_ids, 0) + message_ids = [ + message._postgres_message.msg_id for message in messages if isinstance(message, _PostgresMessage) + ] + self._requeue_message_ids(message_ids) def _build_message(self, postgres_message: PostgresQueueMessage) -> "_PostgresMessage": - """Wrap a raw PGMQ row and register it for heartbeat renewal.""" - message = _PostgresMessage(postgres_message) - self._register_heartbeat_msg(message) - return message + """Wrap a raw PGMQ row as a Remoulade message proxy.""" + return _PostgresMessage(postgres_message) @override def __next__(self) -> "_PostgresMessage | None": @@ -514,7 +513,9 @@ def __next__(self) -> "_PostgresMessage | None": if not messages: return None - + message_ids = [message.msg_id for message in messages] + with self._heartbeat_message_ids_lock: + self._heartbeat_message_ids.update(message_ids) self.messages.extend(messages) return self._build_message(self.messages.popleft()) @@ -531,10 +532,12 @@ def close(self) -> None: self.broker.logger.error("Listener not joinable: %s", str(e)) finally: self._listener_connection = None - if self._listener_thread is not None: - self._listener_thread.join(timeout=1) if self._heartbeat_thread is not None: self._heartbeat_thread.join(timeout=1) + self._requeue_message_ids([message.msg_id for message in self.messages]) + self.messages.clear() + if self._listener_thread is not None: + self._listener_thread.join(timeout=1) return diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 6920ce6ee..6bfac4c62 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -167,8 +167,8 @@ def test_postgres_broker_uses_custom_partition_settings_when_provided(): broker = PostgresBroker( url=TEST_POSTGRES_URL, middleware=[], - archive_partition_interval="2 days", - archive_retention_interval="14 days", + archive_partition_interval_in_days=2, + archive_retention_interval_in_days=14, ) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() @@ -362,6 +362,76 @@ def _fake_start_listener(self): consumer.close() +def test_postgres_consumer_tracks_all_prefetched_messages_for_heartbeat(monkeypatch): + def _fake_start_listener(self): + self._listener_available = False + + def _fake_start_heartbeat(self): + self._heartbeat_thread = None + + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) + + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["default"] = None + + first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock( + return_value=[ + Mock(msg_id=1, message=_expected_payload(first_message)), + Mock(msg_id=2, message=_expected_payload(second_message)), + ] + ) + broker.client.set_vt = Mock() + + consumer = broker.consume("default", prefetch=2, timeout=200) + consumed = next(consumer) + + assert consumed is not None + assert consumed.message_id == first_message.message_id + with consumer._heartbeat_message_ids_lock: + assert consumer._heartbeat_message_ids == {1, 2} + + consumer.close() + + +def test_postgres_consumer_close_requeues_buffered_prefetched_messages(monkeypatch): + def _fake_start_listener(self): + self._listener_available = False + + def _fake_start_heartbeat(self): + self._heartbeat_thread = None + + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) + + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["default"] = None + + first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock( + return_value=[ + Mock(msg_id=1, message=_expected_payload(first_message)), + Mock(msg_id=2, message=_expected_payload(second_message)), + ] + ) + broker.client.set_vt = Mock() + + consumer = broker.consume("default", prefetch=2, timeout=200) + consumed = next(consumer) + + assert consumed is not None + assert consumed.message_id == first_message.message_id + + consumer.close() + + broker.client.set_vt.assert_called_once_with("default", [2], 0) + + def test_postgres_consumer_falls_back_to_polling_when_listener_stops_during_wait(monkeypatch): def _fake_start_listener(self): self._listener_available = True From cfc3f2a991185ee47f58512f44a16917cc15da78 Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 17 Jun 2026 14:55:38 +0200 Subject: [PATCH 21/44] fix(listen-notify): code review --- remoulade/brokers/postgres.py | 24 +++++++++++++++--------- tests/test_postgres.py | 28 +++++++++++++++++++++++++++- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 0ac24ce1b..ea277cbab 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -370,6 +370,20 @@ def _read_with_poll(self) -> list[PostgresQueueMessage]: ) ) + def _read_next_batch(self) -> list[PostgresQueueMessage]: + """Read the next batch, favoring LISTEN/NOTIFY when available.""" + if not self._listener_available: + return self._read_with_poll() + + messages = self._read_immediate() + if messages: + self._notify_event.clear() + return messages + + self._notify_event.wait(self.wait_timeout_seconds) + self._notify_event.clear() + return self._read_immediate() if self._listener_available else self._read_with_poll() + def _start_heartbeat(self) -> None: """Start the background visibility-timeout renewal thread.""" self._heartbeat_thread = Thread( @@ -501,15 +515,7 @@ def __next__(self) -> "_PostgresMessage | None": if self.messages: return self._build_message(self.messages.popleft()) - if self._listener_available: - self._notify_event.wait(self.wait_timeout_seconds) - messages = self._read_immediate() - if not messages: - messages = self._read_with_poll() - self._listener_available = False - self._notify_event.clear() - else: - messages = self._read_with_poll() + messages = self._read_next_batch() if not messages: return None diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 6bfac4c62..4e555bed7 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -292,7 +292,7 @@ def _fake_start_listener(self): broker.client.read_with_poll = Mock(return_value=[]) consumer = broker.consume("default", prefetch=1, timeout=200) - consumer._notify_event.set() + consumer._notify_event.wait = Mock(return_value=False) consumed = next(consumer) @@ -300,6 +300,7 @@ def _fake_start_listener(self): assert consumed.message_id == message.message_id broker.client.read.assert_called_once_with("default", vt=30, qty=1) broker.client.read_with_poll.assert_not_called() + consumer._notify_event.wait.assert_not_called() assert consumer._listener_available is True consumer.close() @@ -334,6 +335,31 @@ def _fake_start_listener(self): consumer.close() +def test_postgres_consumer_keeps_listener_path_after_an_empty_cycle(monkeypatch): + def _fake_start_listener(self): + self._listener_available = True + + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) + + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["default"] = None + + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumer._notify_event.wait = Mock(return_value=False) + + consumed = next(consumer) + + assert consumed is None + assert consumer._listener_available is True + assert broker.client.read.call_count == 2 + broker.client.read_with_poll.assert_not_called() + consumer._notify_event.wait.assert_called_once_with(0.2) + consumer.close() + + def test_postgres_consumer_uses_broker_visibility_timeout_for_reads(monkeypatch): def _fake_start_listener(self): self._listener_available = False From 457a4adbb80bf890cc8369b6e0a40ef2625319c9 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 18 Jun 2026 16:42:05 +0200 Subject: [PATCH 22/44] fix(postgres): improve code --- .../2026-06-05-postgres-broker-ms-timings.md | 125 --------- pyproject.toml | 7 +- remoulade/brokers/postgres.py | 252 ++++++++++++------ remoulade/encoder.py | 2 +- tests/test_postgres.py | 153 +++++++---- 5 files changed, 281 insertions(+), 258 deletions(-) delete mode 100644 docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md diff --git a/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md b/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md deleted file mode 100644 index 64e4ca058..000000000 --- a/docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md +++ /dev/null @@ -1,125 +0,0 @@ -# Postgres Broker Millisecond Timings Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Make every Postgres broker timing input accept milliseconds, while keeping PGMQ and thread APIs working with their required units internally. - -**Architecture:** Rename the broker-facing timeout parameters to `*_ms`, keep public consumer timing inputs in milliseconds, and convert to seconds only at the boundary where PGMQ or `threading.Event.wait()` require it. Update the Postgres tests to assert the new API and the conversion boundary explicitly. - -**Tech Stack:** Python, pytest, psycopg, pgmq, SQLAlchemy - ---- - -### Task 1: Update timing API names - -**Files:** -- Modify: `remoulade/brokers/postgres.py` -- Modify: `tests/test_postgres.py` - -- [ ] **Step 1: Write the failing test** - -```python -broker = PostgresBroker( - url=TEST_POSTGRES_URL, - middleware=[], - visibility_timeout_ms=17_000, - heartbeat_interval_ms=500, -) -``` - -- [ ] **Step 2: Run test to verify it fails** - -Run: `pytest tests/test_postgres.py -k postgres_consumer_uses_broker_visibility_timeout_for_reads -v` -Expected: `TypeError` because the broker still expects `*_in_second` arguments. - -- [ ] **Step 3: Write minimal implementation** - -```python -def __init__( - self, - *, - url: str, - middleware: list["Middleware"] | None = None, - group_transaction: bool = False, - archive_partition_interval: str = "1 day", - archive_retention_interval: str = "7 days", - visibility_timeout_ms: int = 30_000, - heartbeat_interval_ms: int = 10_000, -) -> None: - if visibility_timeout_ms <= 0: - raise ValueError("visibility_timeout_ms must be greater than 0") - if heartbeat_interval_ms <= 0 or heartbeat_interval_ms >= visibility_timeout_ms: - raise ValueError("heartbeat_interval_ms must be greater than 0 and lower than visibility_timeout_ms") -``` - -- [ ] **Step 4: Run test to verify it passes** - -Run: `pytest tests/test_postgres.py -k postgres_consumer_uses_broker_visibility_timeout_for_reads -v` -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add remoulade/brokers/postgres.py tests/test_postgres.py docs/superpowers/plans/2026-06-05-postgres-broker-ms-timings.md -git commit -m "feat: use milliseconds for postgres broker timing inputs" -``` - -### Task 2: Convert broker internals to seconds - -**Files:** -- Modify: `remoulade/brokers/postgres.py` -- Modify: `tests/test_postgres.py` - -- [ ] **Step 1: Write the failing test** - -```python -assert broker.client.read.call_args.kwargs["vt"] == 17 -assert broker.client.set_vt.call_args.args[2] == 2 -``` - -- [ ] **Step 2: Run test to verify it fails** - -Run: `pytest tests/test_postgres.py -k 'postgres_consumer_uses_broker_visibility_timeout_for_reads or postgres_consumer_heartbeat_extends_inflight_message_visibility' -v` -Expected: `vt` still reflects millisecond inputs instead of converted seconds. - -- [ ] **Step 3: Write minimal implementation** - -```python -from math import ceil - -self.visibility_timeout = max(1, ceil(self.visibility_timeout_ms / 1000)) -self.heartbeat_interval = self.heartbeat_interval_ms / 1000 -``` - -- [ ] **Step 4: Run test to verify it passes** - -Run: `pytest tests/test_postgres.py -k 'postgres_consumer_uses_broker_visibility_timeout_for_reads or postgres_consumer_heartbeat_extends_inflight_message_visibility' -v` -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add remoulade/brokers/postgres.py tests/test_postgres.py -git commit -m "fix: convert postgres broker timing inputs" -``` - -### Task 3: Verify the full Postgres test file - -**Files:** -- Test: `tests/test_postgres.py` - -- [ ] **Step 1: Run the whole file** - -Run: `pytest tests/test_postgres.py -v` -Expected: PASS - -- [ ] **Step 2: Fix any remaining unit mismatches** - -Adjust only the Postgres timing assertions or docstrings if a legacy `*_in_second` name remains. - -- [ ] **Step 3: Commit** - -```bash -git add tests/test_postgres.py -git commit -m "test: cover postgres broker millisecond timings" -``` diff --git a/pyproject.toml b/pyproject.toml index 7286a6aea..bf279d855 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ rabbitmq = ["amqpstorm>=2.6,<3"] redis = ["redis>=7.0.0"] server = ["flask~=2.3.3", "marshmallow>=3,<4", "flask-apispec"] postgres = ["sqlalchemy>=2.0,<3", "psycopg>=3.2", "pgmq[sqlalchemy]>=1.1.1,<2"] -pydantic = ["pydantic>=2.12", "simplejson"] +pydantic = ["pydantic>=2.12"] limits = ["limits~=5.3.0"] tracing = ["opentelemetry-api>=1.20"] @@ -35,11 +35,9 @@ dev = [ "marshmallow>=3,<4", "flask-apispec", "sqlalchemy>=2.0,<3", - "psycopg2>=2.9.11", "pgmq[sqlalchemy]>=1.1.1", "psycopg>=3.2", "pydantic>=2.12", - "simplejson", "limits~=5.3.0", # Docs "alabaster", @@ -53,7 +51,6 @@ dev = [ "sqlalchemy[mypy]>=2.0,<3", "types-redis", "types-python-dateutil", - "types-simplejson", "types-requests", # Misc "pre-commit", @@ -99,8 +96,6 @@ include = ["remoulade*"] testpaths = ["tests"] asyncio_mode = "auto" markers = ["confirm_delivery", "group_transaction"] -filterwarnings = [ -] [tool.ruff] target-version = "py312" diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index ea277cbab..e3347647c 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -14,7 +14,7 @@ # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . -import json +import logging import math import time from collections import deque @@ -36,6 +36,8 @@ from ..message import Message if TYPE_CHECKING: + from sqlalchemy import Engine + from ..middleware import Middleware PostgresPayload = dict[str, Any] @@ -53,6 +55,17 @@ class PostgresBroker(Broker): PGMQ handles delayed messages natively, so delayed sends stay in PostgreSQL instead of being staged in worker memory. + + Connection budget (per worker process): + * the shared SQLAlchemy pool (``pool_size``), used by reads, acks and + the heartbeat, bounded regardless of the number of consumed queues; + * a single shared LISTEN/NOTIFY connection (one per process, not one per + queue) when ``enable_listen_notify`` is True. + + Set ``enable_listen_notify=False`` for a poll-only mode that opens no + dedicated listener connection. This is required behind a connection pooler + in transaction pooling mode (e.g. pgbouncer), where LISTEN/NOTIFY is not + supported, and useful to cap connections on very large fan-outs. """ supports_native_delay = True @@ -67,6 +80,9 @@ def __init__( archive_retention_interval_in_days: int = 7, visibility_timeout_ms: int = 30_000, heartbeat_interval_ms: int = 10_000, + enable_listen_notify: bool = True, + pool_size: int = 10, + engine: "Engine | None" = None, ) -> None: """Initialize a PostgreSQL-backed broker using the PGMQ extension. @@ -80,6 +96,12 @@ def __init__( visibility_timeout_ms(int): Message visibility timeout in milliseconds after read; must be greater than 0. heartbeat_interval_ms(int): Heartbeat interval in milliseconds used to extend in-flight message visibility must be greater than 0 and lower than visibility_timeout_ms. + enable_listen_notify(bool): If True (default), consumers are woken by a single process-wide LISTEN/NOTIFY + connection. If False, consumers poll only and no dedicated listener connection is opened (required behind + a transaction-pooling connection pooler such as pgbouncer). + pool_size(int): Size of the shared SQLAlchemy connection pool. Ignored when ``engine`` is provided. + engine(Engine | None): A pre-configured SQLAlchemy engine to reuse instead of letting PGMQ build one, so + the pool can be sized and shared by the caller. """ super().__init__(middleware=middleware) if visibility_timeout_ms <= 0: @@ -96,13 +118,18 @@ def __init__( self.heartbeat_interval_ms = heartbeat_interval_ms self.visibility_timeout_seconds = _milliseconds_to_seconds(visibility_timeout_ms) self.heartbeat_interval_seconds = heartbeat_interval_ms / 1000 + self.enable_listen_notify = enable_listen_notify self.client = SQLAlchemyPGMQueue( conn_string=url, init_extension=False, vt=self.visibility_timeout_seconds, + pool_size=pool_size, + engine=engine, ) + self._listener = _PostgresListener(self.url, self.logger) if enable_listen_notify else None + def convert_days_in_partman_syntax(self, interval_in_day: int) -> str: """Convert int into partman syntax""" if interval_in_day <= 0: @@ -156,7 +183,9 @@ def _queue_exists(self, queue_name: str) -> bool: @override def close(self) -> None: - """Dispose the underlying PGMQ client.""" + """Stop the shared listener and dispose the underlying PGMQ client.""" + if self._listener is not None: + self._listener.close() self.client.dispose() @override @@ -187,7 +216,8 @@ def declare_queue(self, queue_name: str) -> None: retention_interval=self.archive_retention_interval, conn=self._current_connection, ) - self._try_enable_notify(queue_name) + if self.enable_listen_notify: + self._try_enable_notify(queue_name) self.queues[queue_name] = None @@ -195,7 +225,11 @@ def declare_queue(self, queue_name: str) -> None: self.emit_after("declare_queue", queue_name) def _encode_message(self, message: "Message") -> PostgresPayload: - """Decode a Remoulade message into a JSON object payload for PGMQ. + """Encode a Remoulade message into a JSON object payload for PGMQ. + + The encoder is responsible for validating that the payload is valid + JSON; here we only enforce the PGMQ-specific requirement that it be a + JSON object. Raises: UnsupportedMessageEncoding: If the encoded payload is not valid JSON @@ -209,11 +243,6 @@ def _encode_message(self, message: "Message") -> PostgresPayload: if not isinstance(payload, dict): raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") - try: - json.dumps(payload) - except (TypeError, ValueError) as exc: - raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.") from exc - return payload @override @@ -303,6 +332,134 @@ def join( time.sleep(idle_time / 1000) +class _PostgresListener: + """Process-wide LISTEN/NOTIFY dispatcher shared by all consumers of a broker. + + A single psycopg connection LISTENs on every consumed queue's + ``pgmq.q_.INSERT`` channel and a single background thread routes each + notification to the matching consumer's wake event. This keeps the number + of dedicated listener connections at one per process instead of one per + consumed queue. + + The connection is opened lazily on the first registration. If it cannot be + opened, or if the listener thread later crashes, ``available`` flips to + False and consumers transparently fall back to polling. + """ + + def __init__(self, url: str, logger: logging.Logger) -> None: + self._url = url + self._logger = logger + self._lock = Lock() + self._stop = Event() + self._connection: psycopg.Connection[Any] | None = None + self._thread: Thread | None = None + self._started = False + self.available = False + self._events: dict[str, Event] = {} + self._channel_to_queue: dict[str, str] = {} + self._pending_listen: set[str] = set() + + @staticmethod + def _channel_for(queue_name: str) -> str: + return f"pgmq.q_{queue_name}.INSERT" + + def register(self, queue_name: str, event: Event) -> None: + """Register a consumer's wake event and ensure its channel is listened to.""" + with self._lock: + self._events[queue_name] = event + self._channel_to_queue[self._channel_for(queue_name)] = queue_name + self._pending_listen.add(queue_name) + if not self._started: + self._start_locked() + + def unregister(self, queue_name: str) -> None: + """Stop routing notifications to a consumer that is shutting down.""" + with self._lock: + self._events.pop(queue_name, None) + self._channel_to_queue.pop(self._channel_for(queue_name), None) + self._pending_listen.discard(queue_name) + + def _start_locked(self) -> None: + """Open the shared connection and start the dispatch thread (under lock).""" + self._started = True + try: + self._connection = psycopg.connect(self._url, autocommit=True) + except Exception as e: + self.available = False + self._connection = None + self._logger.warning( + "Failed to start shared LISTEN/NOTIFY listener; consumers will fall back to polling. Error: %s", + str(e), + exc_info=True, + ) + return + self.available = True + self._thread = Thread(target=self._run, name="postgres-listener", daemon=True) + self._thread.start() + + def _drain_pending_listen(self) -> None: + """Issue LISTEN for queues registered since the last loop (listener thread only).""" + with self._lock: + pending = list(self._pending_listen) + self._pending_listen.clear() + if self._connection is None: + return + for queue_name in pending: + channel = self._channel_for(queue_name) + self._connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) + + def _wake_channel(self, channel: str) -> None: + """Wake the single consumer registered for a notification channel.""" + with self._lock: + queue_name = self._channel_to_queue.get(channel) + event = self._events.get(queue_name) if queue_name is not None else None + if event is not None: + event.set() + + def _wake_all(self) -> None: + """Wake every registered consumer (used on shutdown or listener failure).""" + with self._lock: + events = list(self._events.values()) + for event in events: + event.set() + + def _run(self) -> None: + """Route notifications to consumer events, falling back to polling on errors.""" + while not self._stop.is_set(): + try: + if self._connection is None: + break + self._drain_pending_listen() + for notify in self._connection.notifies(timeout=0.5, stop_after=1): + self._wake_channel(notify.channel) + except Exception as e: + if self._stop.is_set(): + break + self.available = False + self._wake_all() + self._logger.warning( + "Shared LISTEN/NOTIFY listener crashed; consumers will fall back to polling. Error: %s", + str(e), + exc_info=True, + ) + break + + self.available = False + + def close(self) -> None: + """Stop the dispatch thread and close the shared connection.""" + self._stop.set() + if self._connection is not None: + try: + self._connection.close() + except Exception as e: + self._logger.error("Failed to close shared listener connection: %s", str(e)) + self._wake_all() + if self._thread is not None: + self._thread.join(timeout=1) + self.available = False + + class _PostgresConsumer(Consumer): def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, timeout: int) -> None: """Initialize a consumer for a PGMQ queue. @@ -323,16 +480,13 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti self.heartbeat_interval_seconds = broker.heartbeat_interval_seconds self.messages: deque[PostgresQueueMessage] = deque() self._notify_event = Event() - self._listener_stop = Event() - self._listener_thread: Thread | None = None - self._listener_connection: psycopg.Connection[Any] | None = None - self._listener_available = False self._heartbeat_stop = Event() self._heartbeat_thread: Thread | None = None self._heartbeat_message_ids_lock = Lock() self._heartbeat_message_ids: set[int] = set() - self._start_listener() + if broker._listener is not None: + broker._listener.register(queue_name, self._notify_event) self._start_heartbeat() @property @@ -340,6 +494,11 @@ def wait_timeout_seconds(self) -> float: """Return the listener wait timeout in seconds.""" return self.timeout / 1000 if self.timeout > 0 else 0 + @property + def _listener_available(self) -> bool: + """Whether the broker's shared LISTEN/NOTIFY listener is currently usable.""" + return self.broker._listener is not None and self.broker._listener.available + def _normalize_messages( self, messages: PostgresQueueMessage | list[PostgresQueueMessage] | None ) -> list[PostgresQueueMessage]: @@ -415,7 +574,7 @@ def _run_heartbeat(self) -> None: self.client.set_vt(self.queue_name, message_ids, self.visibility_timeout_seconds) except Exception as e: self.broker.logger.warning( - "Failed to extend visibility timeout heartbeat for queue %s. Error: ", + "Failed to extend visibility timeout heartbeat for queue %s. Error: %s", self.queue_name, str(e), exc_info=True, @@ -433,54 +592,6 @@ def _requeue_message_ids(self, message_ids: list[int]) -> None: if len(message_ids) > 0: self.client.set_vt(self.queue_name, message_ids, 0) - def _start_listener(self) -> None: - """Start a LISTEN/NOTIFY listener for the queue when available.""" - channel = f"pgmq.q_{self.queue_name}.INSERT" - try: - self._listener_connection = psycopg.connect(self.broker.url, autocommit=True) - self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) - self._listener_available = True - except Exception as e: - self._listener_available = False - self._listener_connection = None - self.broker.logger.warning( - "Failed to start LISTEN/NOTIFY listener for queue %s; falling back to polling. Error: %s", - self.queue_name, - str(e), - exc_info=True, - ) - return - - self._listener_thread = Thread( - target=self._run_listener, - name=f"postgres-listener-{self.queue_name}", - daemon=True, - ) - self._listener_thread.start() - - def _run_listener(self) -> None: - """Wake the consumer when a notification arrives, or disable the listener on errors.""" - while not self._listener_stop.is_set(): - try: - if self._listener_connection is None: - break - for _ in self._listener_connection.notifies(timeout=0.5, stop_after=1): - self._notify_event.set() - except Exception as e: - if self._listener_stop.is_set(): - break - self._listener_available = False - self._notify_event.set() - self.broker.logger.warning( - "LISTEN/NOTIFY listener crashed for queue %s; falling back to polling. Error: %s", - self.queue_name, - str(e), - exc_info=True, - ) - break - - self._listener_available = False - @override def ack(self, message: "MessageProxy") -> None: """Archive a processed message.""" @@ -527,24 +638,15 @@ def __next__(self) -> "_PostgresMessage | None": @override def close(self) -> None: - """Stop background threads and close the listener connection.""" - self._listener_stop.set() + """Stop the heartbeat, unregister from the shared listener and requeue buffered messages.""" self._heartbeat_stop.set() self._notify_event.set() - if self._listener_connection is not None: - try: - self._listener_connection.close() - except Exception as e: - self.broker.logger.error("Listener not joinable: %s", str(e)) - finally: - self._listener_connection = None + if self.broker._listener is not None: + self.broker._listener.unregister(self.queue_name) if self._heartbeat_thread is not None: self._heartbeat_thread.join(timeout=1) self._requeue_message_ids([message.msg_id for message in self.messages]) self.messages.clear() - if self._listener_thread is not None: - self._listener_thread.join(timeout=1) - return class _PostgresMessage(MessageProxy): diff --git a/remoulade/encoder.py b/remoulade/encoder.py index e92c022ea..ba4d08f46 100644 --- a/remoulade/encoder.py +++ b/remoulade/encoder.py @@ -59,7 +59,7 @@ def encode_in_json(self, data: MessageData) -> JsonData: @abc.abstractmethod def _encode_in_json(self, data: MessageData) -> JsonData: - return data + raise NotImplementedError @abc.abstractmethod def decode_json(self, data: JsonData) -> MessageData: diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 4e555bed7..505c556c2 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -52,6 +52,33 @@ def _expected_payload(message): return json.loads(message.encode_in_bytes().decode("utf-8")) +class _FakeListener: + """Stand-in for the broker's shared LISTEN/NOTIFY listener. + + Lets tests control listener availability deterministically without opening + a real psycopg connection or starting the dispatch thread. + """ + + def __init__(self, available): + self.available = available + + def register(self, queue_name, event): + pass + + def unregister(self, queue_name): + pass + + def close(self): + pass + + +def _install_listener(broker, *, available): + """Replace the broker's shared listener with a controllable fake.""" + listener = _FakeListener(available) + broker._listener = listener + return listener + + def test_postgres_broker_uses_provided_url(): broker_url = TEST_POSTGRES_URL broker = PostgresBroker(url=broker_url) @@ -163,6 +190,64 @@ def test_postgres_broker_declare_queue_is_idempotent_when_queue_already_exists() assert "default" in broker.queues +def test_postgres_broker_poll_only_mode_opens_no_listener_and_skips_enable_notify(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) + + broker.declare_queue("default") + + assert broker._listener is None + broker.client.enable_notify.assert_not_called() + + +def test_postgres_broker_poll_only_consumer_never_reports_listener_available(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) + broker.queues["default"] = None + + message = Message(queue_name="default", actor_name="do_work", args=(9,), kwargs={}, options={}) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=[Mock(msg_id=9, message=_expected_payload(message))]) + + consumer = broker.consume("default", prefetch=1, timeout=200) + consumed = next(consumer) + + assert consumed is not None + assert consumer._listener_available is False + broker.client.read.assert_not_called() + broker.client.read_with_poll.assert_called_once() + consumer.close() + + +def test_postgres_broker_shares_a_single_listener_across_consumers(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["first"] = None + broker.queues["second"] = None + listener = _install_listener(broker, available=True) + listener.register = Mock() + listener.unregister = Mock() + + first = broker.consume("first", prefetch=1, timeout=0) + second = broker.consume("second", prefetch=1, timeout=0) + + assert first.broker._listener is second.broker._listener + assert listener.register.call_count == 2 + registered_queues = {call.args[0] for call in listener.register.call_args_list} + assert registered_queues == {"first", "second"} + + first.close() + second.close() + assert listener.unregister.call_count == 2 + + +def test_postgres_broker_forwards_pool_size_to_client(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], pool_size=3) + + assert broker.client.engine.pool.size() == 3 + + def test_postgres_broker_uses_custom_partition_settings_when_provided(): broker = PostgresBroker( url=TEST_POSTGRES_URL, @@ -277,14 +362,10 @@ def do_work(): assert _count_messages(postgres_broker) == 1 -def test_postgres_consumer_uses_notification_path_when_listener_is_available(monkeypatch): - def _fake_start_listener(self): - self._listener_available = True - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_uses_notification_path_when_listener_is_available(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + _install_listener(broker, available=True) message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) payload = _expected_payload(message) @@ -305,14 +386,10 @@ def _fake_start_listener(self): consumer.close() -def test_postgres_consumer_falls_back_to_polling_when_listener_is_unavailable(monkeypatch): - def _fake_start_listener(self): - self._listener_available = False - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_falls_back_to_polling_when_listener_is_unavailable(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + _install_listener(broker, available=False) message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) payload = _expected_payload(message) @@ -335,14 +412,10 @@ def _fake_start_listener(self): consumer.close() -def test_postgres_consumer_keeps_listener_path_after_an_empty_cycle(monkeypatch): - def _fake_start_listener(self): - self._listener_available = True - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_keeps_listener_path_after_an_empty_cycle(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + _install_listener(broker, available=True) broker.client.read = Mock(return_value=None) broker.client.read_with_poll = Mock(return_value=[]) @@ -360,14 +433,10 @@ def _fake_start_listener(self): consumer.close() -def test_postgres_consumer_uses_broker_visibility_timeout_for_reads(monkeypatch): - def _fake_start_listener(self): - self._listener_available = False - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_uses_broker_visibility_timeout_for_reads(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], visibility_timeout_ms=17_000) broker.queues["default"] = None + _install_listener(broker, available=False) message = Message(queue_name="default", actor_name="do_work", args=(5,), kwargs={}, options={}) payload = _expected_payload(message) @@ -389,17 +458,14 @@ def _fake_start_listener(self): def test_postgres_consumer_tracks_all_prefetched_messages_for_heartbeat(monkeypatch): - def _fake_start_listener(self): - self._listener_available = False - def _fake_start_heartbeat(self): self._heartbeat_thread = None - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + _install_listener(broker, available=False) first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) @@ -424,17 +490,14 @@ def _fake_start_heartbeat(self): def test_postgres_consumer_close_requeues_buffered_prefetched_messages(monkeypatch): - def _fake_start_listener(self): - self._listener_available = False - def _fake_start_heartbeat(self): self._heartbeat_thread = None - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + _install_listener(broker, available=False) first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) @@ -458,14 +521,10 @@ def _fake_start_heartbeat(self): broker.client.set_vt.assert_called_once_with("default", [2], 0) -def test_postgres_consumer_falls_back_to_polling_when_listener_stops_during_wait(monkeypatch): - def _fake_start_listener(self): - self._listener_available = True - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_falls_back_to_polling_when_listener_stops_during_wait(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None + listener = _install_listener(broker, available=True) message = Message(queue_name="default", actor_name="do_work", args=(3,), kwargs={}, options={}) payload = _expected_payload(message) @@ -475,7 +534,7 @@ def _fake_start_listener(self): consumer = broker.consume("default", prefetch=1, timeout=200) def _wait_and_stop(_timeout): - consumer._listener_available = False + listener.available = False return False consumer._notify_event.wait = Mock(side_effect=_wait_and_stop) @@ -495,12 +554,7 @@ def _wait_and_stop(_timeout): consumer.close() -def test_postgres_consumer_heartbeat_extends_inflight_message_visibility(monkeypatch): - def _fake_start_listener(self): - self._listener_available = False - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - +def test_postgres_consumer_heartbeat_extends_inflight_message_visibility(): broker = PostgresBroker( url=TEST_POSTGRES_URL, middleware=[], @@ -508,6 +562,7 @@ def _fake_start_listener(self): heartbeat_interval_ms=50, ) broker.queues["default"] = None + _install_listener(broker, available=False) message = Message(queue_name="default", actor_name="do_work", args=(7,), kwargs={}, options={}) payload = _expected_payload(message) @@ -534,21 +589,17 @@ def _fake_start_listener(self): consumer.close() -def test_postgres_consumer_decodes_payload_with_global_encoder(monkeypatch, pydantic_encoder): +def test_postgres_consumer_decodes_payload_with_global_encoder(pydantic_encoder): class InputSchema(BaseModel): value: int - def _fake_start_listener(self): - self._listener_available = False - - monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_listener", _fake_start_listener) - broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) remoulade.set_broker(broker) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + _install_listener(broker, available=False) @remoulade.actor(actor_name="typed.actor", queue_name="default") def typed_actor(payload: InputSchema): From ce01e4c3304a8eef77cb35135f3801c669f3587d Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 10:57:02 +0200 Subject: [PATCH 23/44] fix(postgres): self-healing LISTEN/NOTIFY listener Reconnect the shared listener with capped backoff instead of degrading to polling permanently after the first connection failure. On any connection error the listener drops the connection, wakes consumers so they fall back to polling immediately, then reopens and re-LISTENs on every registered channel; `available` flips back to True so consumers resume LISTEN/NOTIFY automatically. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 122 ++++++++++++++++++++++++++-------- tests/test_postgres.py | 110 +++++++++++++++++++++++++++++- 2 files changed, 202 insertions(+), 30 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index e3347647c..ef2bf079b 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -43,6 +43,8 @@ PostgresPayload = dict[str, Any] LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 FUTURE_PARTITION_HORIZON: Final[str] = "2 months" +LISTENER_RECONNECT_BACKOFF_MIN_S: Final[float] = 0.5 +LISTENER_RECONNECT_BACKOFF_MAX_S: Final[float] = 30.0 def _milliseconds_to_seconds(milliseconds: int) -> int: @@ -341,9 +343,13 @@ class _PostgresListener: of dedicated listener connections at one per process instead of one per consumed queue. - The connection is opened lazily on the first registration. If it cannot be - opened, or if the listener thread later crashes, ``available`` flips to - False and consumers transparently fall back to polling. + The connection is opened synchronously when the first consumer registers. + On any connection failure ``available`` flips to False and + consumers transparently fall back to polling, while the thread keeps + retrying (with capped backoff) to reopen the connection and re-LISTEN on + every registered channel. ``available`` flips back to True as soon as a + reconnection succeeds, so consumers resume LISTEN/NOTIFY automatically + instead of polling for the rest of the process lifetime. """ def __init__(self, url: str, logger: logging.Logger) -> None: @@ -369,8 +375,13 @@ def register(self, queue_name: str, event: Event) -> None: self._events[queue_name] = event self._channel_to_queue[self._channel_for(queue_name)] = queue_name self._pending_listen.add(queue_name) - if not self._started: - self._start_locked() + starting = not self._started + if starting: + self._started = True + # Start outside the lock: the initial connection attempt acquires the + # same lock to snapshot channels, and would otherwise deadlock. + if starting: + self._start() def unregister(self, queue_name: str) -> None: """Stop routing notifications to a consumer that is shutting down.""" @@ -379,23 +390,71 @@ def unregister(self, queue_name: str) -> None: self._channel_to_queue.pop(self._channel_for(queue_name), None) self._pending_listen.discard(queue_name) - def _start_locked(self) -> None: - """Open the shared connection and start the dispatch thread (under lock).""" - self._started = True + def _start(self) -> None: + """Open the shared connection (best effort) and start the dispatch thread. + + The initial connection attempt is synchronous so ``available`` reflects + a real outcome by the time the first consumer starts reading. If it + fails, the dispatch thread keeps retrying with backoff, reopening the + connection and re-LISTENing on every registered channel, so a transient + database outage degrades consumers to polling instead of disabling + LISTEN/NOTIFY for the rest of the process lifetime. + """ + self._open_connection() + self._thread = Thread(target=self._run, name="postgres-listener", daemon=True) + self._thread.start() + + def _open_connection(self) -> bool: + """Open the shared connection and LISTEN on every registered channel. + + Called once synchronously when the listener starts and then from the + dispatch thread on every reconnection. Must not be called while holding + ``self._lock``, which it acquires to snapshot the registered channels. + Returns True when the connection is ready and all channels are listened + to, False when the connection could not be opened (the caller then backs + off and retries). + """ + connection = None try: - self._connection = psycopg.connect(self._url, autocommit=True) + connection = psycopg.connect(self._url, autocommit=True) + with self._lock: + # A fresh connection listens to nothing, so re-LISTEN every + # channel currently registered rather than only those pending + # since the last loop. + channels = list(self._channel_to_queue) + self._pending_listen.clear() + for channel in channels: + connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) except Exception as e: - self.available = False - self._connection = None self._logger.warning( - "Failed to start shared LISTEN/NOTIFY listener; consumers will fall back to polling. Error: %s", + "Failed to open shared LISTEN/NOTIFY connection; consumers will fall back to polling. Error: %s", str(e), exc_info=True, ) - return + if connection is not None: + try: + connection.close() + except Exception: + self._logger.debug("Failed to close partially-opened listener connection", exc_info=True) + return False + self._connection = connection self.available = True - self._thread = Thread(target=self._run, name="postgres-listener", daemon=True) - self._thread.start() + return True + + def _drop_connection(self) -> None: + """Mark the listener unavailable and discard the current connection. + + Waiting consumers are woken so they fall back to polling immediately + instead of blocking on their wake event for the full timeout. + """ + self.available = False + self._wake_all() + if self._connection is not None: + try: + self._connection.close() + except Exception as e: + self._logger.error("Failed to close shared listener connection: %s", str(e)) + self._connection = None def _drain_pending_listen(self) -> None: """Issue LISTEN for queues registered since the last loop (listener thread only).""" @@ -424,40 +483,45 @@ def _wake_all(self) -> None: event.set() def _run(self) -> None: - """Route notifications to consumer events, falling back to polling on errors.""" + """Route notifications to consumer events, reconnecting on errors. + + When the connection is down it is reopened with an exponential backoff + capped at ``LISTENER_RECONNECT_BACKOFF_MAX_S``; the backoff resets on + every successful reconnection. Consumers poll while the connection is + unavailable and resume LISTEN/NOTIFY once it is restored. + """ + backoff = LISTENER_RECONNECT_BACKOFF_MIN_S while not self._stop.is_set(): + if self._connection is None: + if not self._open_connection(): + self._stop.wait(backoff) + backoff = min(backoff * 2, LISTENER_RECONNECT_BACKOFF_MAX_S) + continue + backoff = LISTENER_RECONNECT_BACKOFF_MIN_S try: - if self._connection is None: - break self._drain_pending_listen() for notify in self._connection.notifies(timeout=0.5, stop_after=1): self._wake_channel(notify.channel) except Exception as e: if self._stop.is_set(): break - self.available = False - self._wake_all() self._logger.warning( - "Shared LISTEN/NOTIFY listener crashed; consumers will fall back to polling. Error: %s", + "Shared LISTEN/NOTIFY listener error; consumers will fall back to polling while it " + "reconnects. Error: %s", str(e), exc_info=True, ) - break + self._drop_connection() - self.available = False + self._drop_connection() def close(self) -> None: """Stop the dispatch thread and close the shared connection.""" self._stop.set() - if self._connection is not None: - try: - self._connection.close() - except Exception as e: - self._logger.error("Failed to close shared listener connection: %s", str(e)) self._wake_all() if self._thread is not None: self._thread.join(timeout=1) - self.available = False + self._drop_connection() class _PostgresConsumer(Consumer): diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 505c556c2..b975a0d48 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -1,4 +1,5 @@ import json +import logging import os import threading import time @@ -11,7 +12,7 @@ import remoulade from remoulade import Message, QueueJoinTimeout, UnsupportedMessageEncoding, Worker -from remoulade.brokers.postgres import PostgresBroker +from remoulade.brokers.postgres import PostgresBroker, _PostgresListener from remoulade.encoder import Encoder, MessageData TEST_POSTGRES_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" @@ -789,3 +790,110 @@ def _consume_once(): consumer.ack(consumed) finally: consumer.close() + + +class _FakeConnection: + """Minimal psycopg connection stand-in for listener reconnection tests.""" + + def __init__(self, *, raise_on_notify=False, on_notify=None): + self.raise_on_notify = raise_on_notify + self.on_notify = on_notify + self.listened = [] + self.closed = False + + def execute(self, statement): + self.listened.append(statement) + + def notifies(self, timeout=0.5, stop_after=1): + if self.raise_on_notify: + self.raise_on_notify = False + raise RuntimeError("connection lost") + if self.on_notify is not None: + self.on_notify() + return [] + + def close(self): + self.closed = True + + +def test_listener_open_connection_relistens_every_registered_channel(monkeypatch): + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + listener._channel_to_queue = { + listener._channel_for("first"): "first", + listener._channel_for("second"): "second", + } + connection = _FakeConnection() + monkeypatch.setattr("remoulade.brokers.postgres.psycopg.connect", Mock(return_value=connection)) + + assert listener._open_connection() is True + assert listener.available is True + assert listener._connection is connection + assert len(connection.listened) == 2 + + +def test_listener_open_connection_failure_keeps_listener_unavailable(monkeypatch, caplog): + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + monkeypatch.setattr("remoulade.brokers.postgres.psycopg.connect", Mock(side_effect=OSError("database is down"))) + + with caplog.at_level(logging.WARNING): + assert listener._open_connection() is False + + assert listener.available is False + assert listener._connection is None + assert "Failed to open shared LISTEN/NOTIFY connection" in caplog.text + + +def test_listener_reconnects_after_a_connection_drop(monkeypatch): + monkeypatch.setattr("remoulade.brokers.postgres.LISTENER_RECONNECT_BACKOFF_MIN_S", 0.01) + monkeypatch.setattr("remoulade.brokers.postgres.LISTENER_RECONNECT_BACKOFF_MAX_S", 0.01) + + recovered = threading.Event() + healthy_connection = _FakeConnection(on_notify=recovered.set) + connections = [_FakeConnection(raise_on_notify=True), healthy_connection] + connect = Mock(side_effect=connections) + monkeypatch.setattr("remoulade.brokers.postgres.psycopg.connect", connect) + + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + listener.register("default", threading.Event()) + try: + assert recovered.wait(timeout=2) + assert listener.available is True + assert connect.call_count >= 2 + assert connections[0].closed is True + assert len(healthy_connection.listened) == 1 + finally: + listener.close() + + +def test_listener_recovers_when_initial_connection_fails(monkeypatch): + monkeypatch.setattr("remoulade.brokers.postgres.LISTENER_RECONNECT_BACKOFF_MIN_S", 0.01) + monkeypatch.setattr("remoulade.brokers.postgres.LISTENER_RECONNECT_BACKOFF_MAX_S", 0.01) + + recovered = threading.Event() + healthy_connection = _FakeConnection(on_notify=recovered.set) + connect = Mock(side_effect=[OSError("database is down"), healthy_connection]) + monkeypatch.setattr("remoulade.brokers.postgres.psycopg.connect", connect) + + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + listener.register("default", threading.Event()) + try: + assert recovered.wait(timeout=2) + assert listener.available is True + assert connect.call_count >= 2 + finally: + listener.close() + + +def test_listener_close_stops_the_dispatch_thread(monkeypatch): + connection = _FakeConnection() + monkeypatch.setattr("remoulade.brokers.postgres.psycopg.connect", Mock(return_value=connection)) + + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + listener.register("default", threading.Event()) + + listener.close() + + assert listener.available is False + assert connection.closed is True + assert listener._thread is not None + assert listener._thread.is_alive() is False From cadd4a8910fcfb80a016145b29629eb8cd03ab00 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 11:19:04 +0200 Subject: [PATCH 24/44] docs(api): drop SuperBowl references and prune dead API surface Remove all SuperBowl dashboard mentions from the docs (getting_started, README). Also remove API artifacts left unused after the legacy state backend removal and unused by the SuperBowl frontend: - DeleteSchema, orphaned by the earlier clean_states route removal - update_job (PUT /scheduled/jobs/), never called by SuperBowl - GroupMessagesT, a defined-but-unused TypedDict Document the removed DELETE /messages/states and PUT /scheduled/jobs/ routes in the 7.0.0 changelog breaking changes. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 6 ---- docs/source/changelog.rst | 2 ++ docs/source/getting_started.rst | 50 --------------------------------- remoulade/api/main.py | 7 +---- remoulade/api/scheduler.py | 12 +------- remoulade/api/state.py | 9 ------ 6 files changed, 4 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index 76e3a8227..66f0c555f 100644 --- a/README.md +++ b/README.md @@ -91,12 +91,6 @@ If you want to contribute to the project. First make a Pull request and get appr This will trigger a CI/CD pipeline that publish the package -## Dashboard - -Check out [SuperBowl](https://github.com/wiremind/super-bowl) a dashboard for real-time monitoring and administrating all your Remoulade tasks. -***See the current progress, enqueue, requeue, cancel and more ...*** -Super easy to use !. - ## Kubernetes Remoulade is tailored to run transparently in containers on [Kubernetes](https://kubernetes.io/) and to make the most of their features. This does not mean it cannot run outside of Kubernetes ;) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 3d5a59564..1e3073f7c 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -10,6 +10,8 @@ All notable changes to this project will be documented in this file. Breaking changes ^^^^^^^^^^^^^^^^ * Remove the legacy PostgreSQL state backend. +* Remove the ``DELETE /messages/states`` API route, which only worked with the removed PostgreSQL state backend. +* Remove the ``PUT /scheduled/jobs/`` API route (single-job update); use ``PUT /scheduled/jobs`` instead. * Rework the broker API around the new PostgreSQL/PGMQ implementation. Feat diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 23303cf4e..c3405fb98 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -12,7 +12,6 @@ By the end of this tutorial, you will be able to do the following: * :ref:`create a pipeline of tasks that will sequentially process data` * :ref:`use the result middleware to wait and get actor results` * :ref:`use the remoulade scheduler to periodically run tasks` -* :ref:`use SuperBowl to monitor and manage your tasks` Prerequisites ------------- @@ -341,55 +340,6 @@ To set up the scheduler, we instantiate it, set it as the global scheduler, and If you run this script and get back to the worker terminal, you will see ``get_weather`` being executed every 10 seconds. -Monitoring and Managing your tasks ----------------------------------- - -To monitor and manage your tasks, you can use the Superbowl_ dashboard. - -.. _Superbowl: https://github.com/wiremind/super-bowl - -First, you will need to install Node.js_. Then, clone Superbowl_ in another directory, install its dependencies and run it:: - - $ cd .. - $ git clone https://github.com/wiremind/super-bowl.git - $ npm install - $ npm run serve - -.. _Node.js: https://nodejs.org/en/download/ - -Now, if you open ``localhost:8080`` in your browser, you will see the SuperBowl dashboard, but you will not see your messages yet. In order to see and manage them, you will have to modify the ``get_weather.py`` script to serve the remoulade api. - -.. code-block:: python - :caption: get_weather.py - :emphasize-lines: 4, 22, 23 - - import requests - import remoulade - from remoulade.brokers.rabbitmq import RabbitmqBroker - from remoulade.api.main import app - - - @remoulade.actor - def get_weather(city): - url = f"https://goweather.herokuapp.com/weather/{city}" - - response = requests.get(url).json() - - url_endpoint = - text = f'{city}: {response["description"]}' - requests.post(url_endpoint, json=text) - - - rabbitmq_broker = RabbitmqBroker() - remoulade.set_broker(rabbitmq_broker) - remoulade.declare_actors([get_weather]) - - if __name__ == "__main__": - app.run(host="localhost", port=5005) - -Now you can use the Enqueue tab to enqueue messages with custom arguments, and then see their progress in the messages tab. -Additionally, if you run groups or scheduled jobs in your script, you will be able to see them in their respective tabs. - Next Steps ---------- diff --git a/remoulade/api/main.py b/remoulade/api/main.py index 7008d80ec..c5d216a4e 100644 --- a/remoulade/api/main.py +++ b/remoulade/api/main.py @@ -1,7 +1,7 @@ """This file describe the API to get the state of messages""" import sys -from typing import Any, TypedDict +from typing import Any from flask import Flask from flask_apispec import marshal_with @@ -137,11 +137,6 @@ def enqueue_message(**kwargs): return {"result": "ok"} -class GroupMessagesT(TypedDict): - group_id: str - messages: list[dict] - - @app.route("/actors") @marshal_with(ActorResponseSchema) def get_actors(): diff --git a/remoulade/api/scheduler.py b/remoulade/api/scheduler.py index d3dfbd6c1..9f17c03fd 100644 --- a/remoulade/api/scheduler.py +++ b/remoulade/api/scheduler.py @@ -108,16 +108,6 @@ def delete_job(scheduler, job_hash): scheduler.delete_job(job_hash) -@scheduler_bp.route("/jobs/", methods=["PUT"]) -@doc(tags=["scheduler"]) -@marshal_with(ScheduleResponseSchema) -@validate_schema(ScheduledJobSchema) -@with_scheduler -def update_job(scheduler, job_hash, **kwargs): - scheduler.delete_job(job_hash) - scheduler.add_job(ScheduledJob(**kwargs)) - - @scheduler_bp.route("/jobs", methods=["PUT"]) @doc(tags=["scheduler"]) @marshal_with(ScheduleResponseSchema) @@ -129,4 +119,4 @@ def update_jobs(scheduler, **kwargs): scheduler.add_job(ScheduledJob(**job_dict)) -scheduler_routes = [get_jobs, add_job, delete_job, update_job, update_jobs] +scheduler_routes = [get_jobs, add_job, delete_job, update_jobs] diff --git a/remoulade/api/state.py b/remoulade/api/state.py index 91bb868a0..4982e5d1f 100644 --- a/remoulade/api/state.py +++ b/remoulade/api/state.py @@ -11,15 +11,6 @@ messages_bp = Blueprint("messages", __name__, url_prefix="/messages") -class DeleteSchema(Schema): - """ - Class to validate delete body data in /messages/states - """ - - max_age = fields.Int(allow_none=True) - not_started = fields.Bool(load_default=False) - - class StatesParamsSchema(Schema): """ Class to validate the state search parameters From 09a2e2dc49d05a14211c19920e8967018ea04dd1 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 11:24:29 +0200 Subject: [PATCH 25/44] docs(changelog): list all removed/renamed public APIs for 7.0.0 Complete the 7.0.0 breaking changes with the Encoder and Message method renames (encode/decode -> encode_in_bytes/decode_bytes, plus the new JSON encoder hooks), alongside the already-listed state backend and removed HTTP routes. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/changelog.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 1e3073f7c..9f359f195 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -13,6 +13,8 @@ Breaking changes * Remove the ``DELETE /messages/states`` API route, which only worked with the removed PostgreSQL state backend. * Remove the ``PUT /scheduled/jobs/`` API route (single-job update); use ``PUT /scheduled/jobs`` instead. * Rework the broker API around the new PostgreSQL/PGMQ implementation. +* Rename ``Encoder.encode``/``Encoder.decode`` to ``Encoder.encode_in_bytes``/``Encoder.decode_bytes``; custom encoders must now also implement ``Encoder._encode_in_json`` and ``Encoder.decode_json``. +* Rename ``Message.encode``/``Message.decode`` to ``Message.encode_in_bytes``/``Message.decode_bytes``. Feat ^^^^ From 0b577e54f539312bddd205e821837b5e827d39d6 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 11:25:38 +0200 Subject: [PATCH 26/44] docs(changelog): note the repurposed postgres extra as breaking The postgres extra now pulls the PGMQ broker dependencies instead of the removed legacy state backend dependencies. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/changelog.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 9f359f195..5fa041c14 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -15,6 +15,7 @@ Breaking changes * Rework the broker API around the new PostgreSQL/PGMQ implementation. * Rename ``Encoder.encode``/``Encoder.decode`` to ``Encoder.encode_in_bytes``/``Encoder.decode_bytes``; custom encoders must now also implement ``Encoder._encode_in_json`` and ``Encoder.decode_json``. * Rename ``Message.encode``/``Message.decode`` to ``Message.encode_in_bytes``/``Message.decode_bytes``. +* Repurpose the ``postgres`` extra: it now installs the PGMQ broker dependencies (``sqlalchemy>=2``, ``psycopg>=3``, ``pgmq``) instead of the legacy state backend dependencies (``sqlalchemy<2``, ``psycopg2``). Feat ^^^^ From ec138eee0d619e4361775171a0fd58c9ed909966 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 12:03:13 +0200 Subject: [PATCH 27/44] test(postgres): use psycopg v3 driver for the fixture's SQLAlchemy engine SQLAlchemy defaults `postgresql://` to the psycopg2 dialect, which the project does not depend on. CI installs only psycopg v3, so the raw engine built in the postgres_broker fixture failed at setup with ModuleNotFoundError: No module named 'psycopg2'. Swap the scheme to postgresql+psycopg://, mirroring what PGMQ does internally for the broker. The plain url is kept for PostgresBroker, whose listener opens a raw psycopg.connect() connection that does not understand +psycopg. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/conftest.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index a01b87bbf..e07b44846 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -117,7 +117,11 @@ def rabbitmq_broker(request): @pytest.fixture() def postgres_broker(): db_string = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" - engine = create_engine(db_string, poolclass=NullPool) + # Force the psycopg (v3) driver for the raw SQLAlchemy engine: SQLAlchemy defaults + # `postgresql://` to psycopg2, which the project does not depend on. This mirrors the driver + # swap PGMQ does internally for the broker. The plain `db_string` is kept for PostgresBroker, + # whose listener opens a raw psycopg.connect() connection that does not understand `+psycopg`. + engine = create_engine(db_string.replace("postgresql://", "postgresql+psycopg://", 1), poolclass=NullPool) client = sessionmaker(bind=engine) check_postgres(client) cleanup_postgres_queues(client) From f1398902eb357d9cc7a4135abef244385ce143d0 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 12:22:45 +0200 Subject: [PATCH 28/44] docs(guide): keep max_retries heading consistent with its siblings Revert the unjustified ^^^ -> --- underline change on the max_retries retry option so it matches the other retry options (min_backoff, max_backoff, retry_when, backoff_strategy). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/guide.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/guide.rst b/docs/source/guide.rst index 91c874804..2a7a4ebad 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -203,7 +203,7 @@ If you want to use a different strategy than the default exponential backoff to The following retry options are configurable on a per-actor basis: max_retries ---------------- +^^^^^^^^^^^ The maximum number of times a message should be retried. Default to ``0``. From e873695270c0ec81380c632df298e9ed33f59118 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 13:51:19 +0200 Subject: [PATCH 29/44] docs(guide): document PostgresBroker partitioning and retention Explain that queues are partitioned PGMQ queues backed by pg_partman, what archive_partition_interval_in_days / archive_retention_interval_in_days control, and that pg_partman maintenance must run periodically. Also fix the broken sentence about the required PostgreSQL user and remove the unused FUTURE_PARTITION_HORIZON constant. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/guide.rst | 29 +++++++++++++++++++++++++++-- remoulade/brokers/postgres.py | 1 - 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/docs/source/guide.rst b/docs/source/guide.rst index 2a7a4ebad..0f41ed9c8 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -408,8 +408,8 @@ Postgres Broker To configure PostgreSQL/PGMQ, install ``remoulade[postgres]`` and instantiate a ``PostgresBroker`` with a PostgreSQL URL as early as possible -during your program's execution. This broker be used with a user postgresql who can create and -delete tables:: +during your program's execution. This broker must be used with a PostgreSQL +user allowed to create and delete tables:: import remoulade @@ -422,6 +422,31 @@ PGMQ handles delayed messages natively, so ``send_with_options(delay=...)`` does not create a worker-side delay queue or add an ``eta`` option to the message. +Each queue is created as a **partitioned** PGMQ queue (through +``pgmq.create_partitioned``), which relies on the PostgreSQL ``pg_partman`` +extension. Messages are stored in time-based partitions of the queue table, +and once a message is acked or nacked it is moved to the queue's archive +table, which is partitioned the same way. Two broker parameters control this: + +``archive_partition_interval_in_days`` (default ``1``) + The time span covered by a single partition. Smaller values create more, + smaller partitions; larger values create fewer, larger ones. + +``archive_retention_interval_in_days`` (default ``7``) + How long a partition is kept before ``pg_partman`` drops it. Archived + messages older than this window are removed together with their partition, + so set it comfortably above your longest expected processing and retry + window. + +.. note:: + + Partitioning is maintained by ``pg_partman``: its maintenance routine + (``partman.run_maintenance_proc()``) must run periodically — through the + ``pg_partman`` background worker or a ``pg_cron`` job — to create upcoming + partitions ahead of time and drop expired ones. The official PGMQ Docker + image (``ghcr.io/pgmq/pg18-pgmq``) ships ``pg_partman`` and schedules this + for you; on a self-managed or hosted PostgreSQL you must enable it yourself. + Local Broker ^^^^^^^^^^^^^^^ diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index ef2bf079b..09c8d4b60 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -42,7 +42,6 @@ PostgresPayload = dict[str, Any] LISTEN_NOTIFY_THROTTLE_MS: Final[int] = 250 -FUTURE_PARTITION_HORIZON: Final[str] = "2 months" LISTENER_RECONNECT_BACKOFF_MIN_S: Final[float] = 0.5 LISTENER_RECONNECT_BACKOFF_MAX_S: Final[float] = 30.0 From da0cc5a2d003b7d5dc344a9f54c7bea1ab47e514 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 14:09:28 +0200 Subject: [PATCH 30/44] fix(postgres): never let a failed ack/nack kill the worker thread archive() runs in process_message's finally, outside the worker loop's except Empty, so a transient DB error during ack/nack used to propagate and kill the worker thread (which has no restart). Swallow and log it like the RabbitMQ consumer does; the message is redelivered after its visibility timeout, honoring at-least-once. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 09c8d4b60..4bfb8fafe 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -655,21 +655,37 @@ def _requeue_message_ids(self, message_ids: list[int]) -> None: if len(message_ids) > 0: self.client.set_vt(self.queue_name, message_ids, 0) - @override - def ack(self, message: "MessageProxy") -> None: - """Archive a processed message.""" + def _archive_message(self, message: "MessageProxy") -> None: + """Stop tracking a message and archive it, tolerating transient failures. + + A failed archive (connection blip, pool exhaustion, ...) is logged and + swallowed rather than propagated: letting it bubble up would kill the + worker thread, which has no restart logic. The message simply becomes + visible again once its visibility timeout expires and is redelivered, + which is the broker's at-least-once guarantee. + """ if not isinstance(message, _PostgresMessage): raise ValueError("It must be a PostgresMessage") self._unregister_heartbeat_message_id(message._postgres_message.msg_id) - self.client.archive(self.queue_name, message._postgres_message.msg_id) + try: + self.client.archive(self.queue_name, message._postgres_message.msg_id) + except Exception: + self.broker.logger.error( + "Failed to archive message %s on queue %s; it will be redelivered after its visibility timeout.", + message._postgres_message.msg_id, + self.queue_name, + exc_info=True, + ) + + @override + def ack(self, message: "MessageProxy") -> None: + """Archive a processed message.""" + self._archive_message(message) @override def nack(self, message: "MessageProxy") -> None: """Archive a failed message.""" - if not isinstance(message, _PostgresMessage): - raise ValueError("It must be a PostgresMessage") - self._unregister_heartbeat_message_id(message._postgres_message.msg_id) - self.client.archive(self.queue_name, message._postgres_message.msg_id) + self._archive_message(message) @override def requeue(self, messages: Iterable["MessageProxy"]) -> None: From 7d72884ce62596c3c64e9e311fac80898d7401a1 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 14:25:05 +0200 Subject: [PATCH 31/44] perf(encoder): serialize once on the bytes path; drop committed dev scripts encode_in_bytes routed through encode_in_json, whose throwaway json.dumps validation pass ran on top of the real serialization: 2 passes for the default JSONEncoder (RabbitMQ/Stub) and 3 for PydanticEncoder. Serialize directly from the data / _encode_in_json output instead, bringing the default path back to a single pass and Pydantic down to two. Output bytes are unchanged. Also drop the misleading `# pragma: no cover` on the default codec, which the Stub broker actually exercises. Remove local_postgres_broker.py / local_postgres_consumer.py: local dev scripts with hard-coded DSNs and prints that should never have been committed. Co-Authored-By: Claude Opus 4.8 (1M context) --- local_postgres_broker.py | 39 --------------------------------- local_postgres_consumer.py | 45 -------------------------------------- remoulade/encoder.py | 10 +++++---- 3 files changed, 6 insertions(+), 88 deletions(-) delete mode 100644 local_postgres_broker.py delete mode 100644 local_postgres_consumer.py diff --git a/local_postgres_broker.py b/local_postgres_broker.py deleted file mode 100644 index f8046e34b..000000000 --- a/local_postgres_broker.py +++ /dev/null @@ -1,39 +0,0 @@ -from collections.abc import Iterable -from typing import Any, cast - -from sqlalchemy import text - -import remoulade -from remoulade import Message -from remoulade.brokers.postgres import PostgresBroker - -# Just for local test -URL = "postgresql://remoulade@localhost:5544/test" -QUEUE_NAME = "default" -ACTOR_NAME = "demo.add" - - -broker = PostgresBroker( - url=URL, - middleware=[], -) -remoulade.set_broker(broker) -broker.declare_queue(QUEUE_NAME) - -msg_args: Iterable[Any] = cast(Iterable[Any], (1, 2)) -msg: Message = Message( - queue_name=QUEUE_NAME, - actor_name=ACTOR_NAME, - args=msg_args, - kwargs={}, - options={}, -) -broker.enqueue(msg) - -with broker.client.session() as session: - row = session.execute(text("SELECT msg_id, message FROM pgmq.q_default ORDER BY msg_id DESC LIMIT 1")).one() - print("msg_id:", row.msg_id) - print("payload:", row.message) - print("listen_notify_channel:", f"pgmq.q_{QUEUE_NAME}.INSERT") - -broker.close() diff --git a/local_postgres_consumer.py b/local_postgres_consumer.py deleted file mode 100644 index 26e019eb9..000000000 --- a/local_postgres_consumer.py +++ /dev/null @@ -1,45 +0,0 @@ -import time - -import remoulade -from remoulade import Worker -from remoulade.brokers.postgres import PostgresBroker - -# Same DSN as local_postgres_broker.py -URL = "postgresql://remoulade@localhost:5544/test" -QUEUE_NAME = "default" -ACTOR_NAME = "demo.add" - -broker = PostgresBroker( - url=URL, - middleware=[], -) -remoulade.set_broker(broker) -broker.declare_queue(QUEUE_NAME) - - -@remoulade.actor(actor_name=ACTOR_NAME, queue_name=QUEUE_NAME) -def add(x: int, y: int) -> int: - result = x + y - print(f"[demo.add] {x} + {y} = {result}") - return result - - -remoulade.declare_actors([add]) - - -if __name__ == "__main__": - probe = broker.consume(QUEUE_NAME, prefetch=1, timeout=30_000) - listener_mode = "LISTEN/NOTIFY" if probe._listener_available else "polling fallback" - probe.close() - - worker = Worker(broker, queues={QUEUE_NAME}, worker_threads=1, worker_timeout=30_000) - worker.start() - print(f"Postgres local consumer started on queue '{QUEUE_NAME}' ({listener_mode}). Press Ctrl+C to stop.") - try: - while True: - time.sleep(1) - except KeyboardInterrupt: - pass - finally: - worker.stop() - broker.close() diff --git a/remoulade/encoder.py b/remoulade/encoder.py index ba4d08f46..1ab8daea2 100644 --- a/remoulade/encoder.py +++ b/remoulade/encoder.py @@ -70,12 +70,14 @@ class JSONEncoder(Encoder): """Encodes messages as JSON. This is the default encoder.""" @override - def encode_in_bytes(self, data: MessageData) -> bytes: # pragma: no cover + def encode_in_bytes(self, data: MessageData) -> bytes: """Convert message metadata into a bytestring.""" - return json.dumps(self.encode_in_json(data), separators=(",", ":")).encode("utf-8") + # Serialize directly: routing through encode_in_json would add a + # throwaway json.dumps validation pass on top of this one. + return json.dumps(data, separators=(",", ":")).encode("utf-8") @override - def decode_bytes(self, data: bytes) -> MessageData: # pragma: no cover + def decode_bytes(self, data: bytes) -> MessageData: """Convert a bytestring into message metadata.""" return self.decode_json(json.loads(data.decode("utf-8"))) @@ -136,7 +138,7 @@ def __init__(self, fallback_encoder: Encoder | None = None): @override def encode_in_bytes(self, data: MessageData) -> bytes: try: - return json.dumps(self.encode_in_json(data)).encode("utf-8") + return json.dumps(self._encode_in_json(data)).encode("utf-8") except Exception as e: if self.fallback_encoder is not None: return self.fallback_encoder.encode_in_bytes(data) From eac24c42a2c81b43ff88f9fd0961bcceba986c36 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 15:14:02 +0200 Subject: [PATCH 32/44] test(postgres): cover retries, results and groups end-to-end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing suite tests the broker's mechanics largely in isolation; the middleware-driven behaviours were only exercised against the stub and RabbitMQ brokers. Add real-Worker tests running through PGMQ: - a flaky actor that fails then succeeds on retry, - an always-failing actor whose exhausted retries end archived (proving the nack does not leave the message to be redelivered forever — no DLQ), - result storage/retrieval via a Results middleware, - a group aggregating results. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_postgres.py | 113 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 112 insertions(+), 1 deletion(-) diff --git a/tests/test_postgres.py b/tests/test_postgres.py index b975a0d48..886aa2ee9 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -11,9 +11,11 @@ from sqlalchemy.sql import select, text import remoulade -from remoulade import Message, QueueJoinTimeout, UnsupportedMessageEncoding, Worker +from remoulade import Message, QueueJoinTimeout, UnsupportedMessageEncoding, Worker, group from remoulade.brokers.postgres import PostgresBroker, _PostgresListener from remoulade.encoder import Encoder, MessageData +from remoulade.results import Results +from remoulade.results.backends import StubBackend TEST_POSTGRES_URL = os.getenv("REMOULADE_TEST_DB_URL") or "postgresql://remoulade@localhost:5544/test" @@ -897,3 +899,112 @@ def test_listener_close_stops_the_dispatch_thread(monkeypatch): assert connection.closed is True assert listener._thread is not None assert listener._thread.is_alive() is False + + +# End-to-end tests running a real Worker against PGMQ, exercising the +# middleware-driven behaviours (retries, results, groups) that the unit tests +# above stub out. These validate the broker's ack/nack/requeue/heartbeat +# interplay through the full message lifecycle, not just in isolation. + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_retries_failed_message_then_succeeds(postgres_broker): + attempts = [] + + @remoulade.actor(max_retries=3, min_backoff=50, max_backoff=50, jitter=False) + def flaky(): + attempts.append(1) + if len(attempts) < 2: + raise RuntimeError("boom") + + postgres_broker.declare_actor(flaky) + + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=1) + worker.start() + try: + flaky.send() + postgres_broker.join(flaky.queue_name, timeout=10_000) + worker.join() + + # The actor ran twice: the failed attempt, then the retry that succeeded. + assert len(attempts) == 2 + # The failed attempt is archived (acked) and a delayed retry enqueued; + # the successful retry is archived too. Nothing is left in the queue. + assert _count_messages(postgres_broker) == 0 + assert _count_archived_messages(postgres_broker) == 2 + finally: + worker.stop() + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_archives_message_when_retries_exhausted(postgres_broker): + attempts = [] + + @remoulade.actor(max_retries=1, min_backoff=50, max_backoff=50, jitter=False) + def always_fails(): + attempts.append(1) + raise RuntimeError("boom") + + postgres_broker.declare_actor(always_fails) + + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=1) + worker.start() + try: + always_fails.send() + postgres_broker.join(always_fails.queue_name, timeout=10_000) + worker.join() + + # max_retries=1 -> one initial attempt plus one retry, then it is failed. + assert len(attempts) == 2 + # Both the acked retry-source and the nacked exhausted message end up + # archived; an empty queue proves the nack does not leave the message + # invisible to be redelivered forever (PostgresBroker has no DLQ). + assert _count_messages(postgres_broker) == 0 + assert _count_archived_messages(postgres_broker) == 2 + finally: + worker.stop() + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_stores_and_retrieves_actor_result(postgres_broker): + postgres_broker.add_middleware(Results(backend=StubBackend())) + + @remoulade.actor(store_results=True) + def do_work(): + return 42 + + postgres_broker.declare_actor(do_work) + + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=1) + worker.start() + try: + message = do_work.send() + postgres_broker.join(do_work.queue_name, timeout=10_000) + worker.join() + + assert message.result.get(block=True) == 42 + finally: + worker.stop() + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_worker_runs_group_with_results(postgres_broker): + postgres_broker.add_middleware(Results(backend=StubBackend())) + + @remoulade.actor(store_results=True) + def square(value): + return value * value + + postgres_broker.declare_actor(square) + + worker = Worker(postgres_broker, worker_timeout=100, worker_threads=4) + worker.start() + try: + g = group([square.message(value) for value in range(4)]) + g.run() + postgres_broker.join(square.queue_name, timeout=10_000) + worker.join() + + assert sorted(g.results.get(block=True)) == [0, 1, 4, 9] + finally: + worker.stop() From 9c75b1217e75d79e052a0851473a17afff3743f0 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 15:20:35 +0200 Subject: [PATCH 33/44] refactor(postgres): count queue length via PGMQ metrics, not the raw table join() counted messages by hand-building a Table("q_", schema="pgmq") and running count(*), coupling the broker to PGMQ's internal table layout. Use the public metrics() API instead: queue_length covers visible, invisible in-flight and native delayed messages alike, matching join()'s contract. Prune the now-unused SQLAlchemy imports. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 4bfb8fafe..f6adf2f24 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -29,7 +29,7 @@ from pgmq import SQLAlchemyPGMQueue from pgmq.messages import Message as PostgresQueueMessage from psycopg import sql as psycopg_sql -from sqlalchemy import Column, Connection, Integer, MetaData, Table, func, select +from sqlalchemy import Connection from ..broker import Broker, Consumer, MessageProxy from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding @@ -278,12 +278,8 @@ def flush_all(self) -> None: self.flush(queue_name) def _count_enqueued_messages(self, queue_name: str) -> int: - """Count rows in the underlying PGMQ queue table.""" - queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq") - count_query = select(func.count()).select_from(queue_table) - - with self.client.session() as session: - return session.execute(count_query).scalar_one() + """Count every message stored in the queue.""" + return self.client.metrics(queue_name).queue_length @override def join( From 125976e275192e45f02dca1e400bd393ad244ec7 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 15:38:49 +0200 Subject: [PATCH 34/44] fix(postgres): honor timeout=0 as a non-blocking read; reject negative timeout The polling fallback rounded timeout up to a one-second minimum, so a consumer created with timeout=0 blocked ~1s per read instead of being non-blocking as documented. Do a single immediate read when timeout is 0, and validate in the consumer's __init__ that timeout is not negative (matching the broker's other input validation) instead of silently pretending to coerce it. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 18 +++++++++++++++--- tests/test_postgres.py | 26 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index f6adf2f24..02e8bb291 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -527,9 +527,15 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti broker(PostgresBroker): Broker instance that owns the queue and database client. queue_name(str): Name of the declared queue to consume from. prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1. - timeout(int): Idle wait timeout in milliseconds when polling for messages; values lower than 0 are coerced to - 0. A value of 0 performs non-blocking reads. + timeout(int): Idle wait timeout in milliseconds when polling for messages; must be greater than or equal to 0. + A value of 0 performs non-blocking reads. + + Raises: + ValueError: If ``timeout`` is negative. """ + if timeout < 0: + raise ValueError("timeout must be greater than or equal to 0") + self.broker = broker self.client = broker.client self.queue_name = queue_name @@ -577,7 +583,13 @@ def _read_immediate(self) -> list[PostgresQueueMessage]: ) def _read_with_poll(self) -> list[PostgresQueueMessage]: - """Read up to ``prefetch`` messages using PGMQ polling.""" + """Read up to ``prefetch`` messages using PGMQ polling. + + A zero timeout means non-blocking: do a single immediate read instead + of polling for the rounded-up one-second minimum. + """ + if self.timeout == 0: + return self._read_immediate() return self._normalize_messages( self.client.read_with_poll( self.queue_name, diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 886aa2ee9..9a8b0848d 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -224,6 +224,32 @@ def test_postgres_broker_poll_only_consumer_never_reports_listener_available(): consumer.close() +def test_postgres_consumer_rejects_negative_timeout(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) + broker.queues["default"] = None + + with pytest.raises(ValueError, match="timeout must be greater than or equal to 0"): + broker.consume("default", prefetch=1, timeout=-1) + + broker.close() + + +def test_postgres_poll_only_consumer_reads_immediately_when_timeout_is_zero(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) + broker.queues["default"] = None + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock(return_value=None) + + consumer = broker.consume("default", prefetch=1, timeout=0) + + # A zero timeout is non-blocking: the polling fallback must do a single + # immediate read and return None at once instead of polling for ~1s. + assert next(consumer) is None + broker.client.read.assert_called_once() + broker.client.read_with_poll.assert_not_called() + consumer.close() + + def test_postgres_broker_shares_a_single_listener_across_consumers(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["first"] = None From 96dd9b3f53392ce662fd55d9a624a9898b76624c Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 15:50:29 +0200 Subject: [PATCH 35/44] fix(encoder): drop PydanticEncoder fallback, raise on decode failure The fallback_encoder swallowed every decode error and returned the raw, un-rehydrated payload through the fallback encoder, silently handing actors the wrong data shape and hiding the original failure. Remove the argument so decoding raises explicitly (ActorNotFound, ValidationError, ...). Drop the two tests that asserted the silent-fallback behaviour; the raising paths are already covered. Document both this and the simplejson->Pydantic serialization change (Decimal now encoded as a string) as breaking. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/changelog.rst | 2 + remoulade/encoder.py | 98 ++++++++++++++++----------------------- tests/test_encoders.py | 23 --------- 3 files changed, 42 insertions(+), 81 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 5fa041c14..53544dbd2 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -15,6 +15,8 @@ Breaking changes * Rework the broker API around the new PostgreSQL/PGMQ implementation. * Rename ``Encoder.encode``/``Encoder.decode`` to ``Encoder.encode_in_bytes``/``Encoder.decode_bytes``; custom encoders must now also implement ``Encoder._encode_in_json`` and ``Encoder.decode_json``. * Rename ``Message.encode``/``Message.decode`` to ``Message.encode_in_bytes``/``Message.decode_bytes``. +* Remove ``PydanticEncoder``'s ``fallback_encoder`` argument: decoding now raises explicitly instead of silently falling back to another encoder and returning un-rehydrated payloads. +* ``PydanticEncoder`` no longer depends on ``simplejson``; serialization now goes through Pydantic, so ``Decimal`` values are encoded as JSON strings instead of numbers. * Repurpose the ``postgres`` extra: it now installs the PGMQ broker dependencies (``sqlalchemy>=2``, ``psycopg>=3``, ``pgmq``) instead of the legacy state backend dependencies (``sqlalchemy<2``, ``psycopg2``). Feat diff --git a/remoulade/encoder.py b/remoulade/encoder.py index 1ab8daea2..629bc4a31 100644 --- a/remoulade/encoder.py +++ b/remoulade/encoder.py @@ -131,27 +131,16 @@ def my_actor(input_1: MyActorInputSchema, input_2: MyActorInputSchema | None = N return MyActorOutputSchema() """ - def __init__(self, fallback_encoder: Encoder | None = None): - self.fallback_encoder = fallback_encoder + def __init__(self): self.json_adapter = TypeAdapter(object) @override def encode_in_bytes(self, data: MessageData) -> bytes: - try: - return json.dumps(self._encode_in_json(data)).encode("utf-8") - except Exception as e: - if self.fallback_encoder is not None: - return self.fallback_encoder.encode_in_bytes(data) - raise e + return json.dumps(self._encode_in_json(data)).encode("utf-8") @override def decode_bytes(self, data: bytes) -> MessageData: - try: - return self.decode_json(json.loads(data.decode("utf-8"))) - except Exception: - if self.fallback_encoder is not None: - return self.fallback_encoder.decode_bytes(data) - raise + return self.decode_json(json.loads(data.decode("utf-8"))) @override def _encode_in_json(self, data: MessageData) -> JsonData: @@ -161,49 +150,42 @@ def _encode_in_json(self, data: MessageData) -> JsonData: def decode_json(self, data: JsonData) -> MessageData: from remoulade import get_broker - try: - actor_name = data["actor_name"] - actor_fn = get_broker().get_actor(actor_name).fn - - # Retrieve the Pydantic schemas from typing - schemas_by_param_name: dict[str, TypeAdapter] = {} - for param_name, type_hint in get_type_hints(actor_fn).items(): - schemas_by_param_name[param_name] = TypeAdapter( - Annotated[ - type_hint, - WithJsonSchema( - {"type": type_hint, "description": f"{param_name}_schema"}, mode="serialization" - ), - ] - ) - - # Override message_data with Pydantic schema when it matches. - parsed_message: dict[str, Any] = {} - for key, values in data.items(): - if key == "kwargs": - if not isinstance(values, dict): - raise TypeError(f"Expected `values` to be a dict, got {type(values).__name__}") - parsed_message[key] = { - param_name: schemas_by_param_name[param_name].validate_python(raw_value) - for param_name, raw_value in values.items() - } - elif key == "args": - if not isinstance(values, list): - raise TypeError(f"Expected `values` to be a list, got {type(values).__name__}") - schemas = list(schemas_by_param_name.values()) - parsed_message[key] = [ - schemas[order].validate_python(raw_value) for order, raw_value in enumerate(values) - ] - elif key == "result": - if values is None: - parsed_message[key] = None - else: - parsed_message[key] = schemas_by_param_name["return"].validate_python(values) + actor_name = data["actor_name"] + actor_fn = get_broker().get_actor(actor_name).fn + + # Retrieve the Pydantic schemas from typing + schemas_by_param_name: dict[str, TypeAdapter] = {} + for param_name, type_hint in get_type_hints(actor_fn).items(): + schemas_by_param_name[param_name] = TypeAdapter( + Annotated[ + type_hint, + WithJsonSchema({"type": type_hint, "description": f"{param_name}_schema"}, mode="serialization"), + ] + ) + + # Override message_data with Pydantic schema when it matches. + parsed_message: dict[str, Any] = {} + for key, values in data.items(): + if key == "kwargs": + if not isinstance(values, dict): + raise TypeError(f"Expected `values` to be a dict, got {type(values).__name__}") + parsed_message[key] = { + param_name: schemas_by_param_name[param_name].validate_python(raw_value) + for param_name, raw_value in values.items() + } + elif key == "args": + if not isinstance(values, list): + raise TypeError(f"Expected `values` to be a list, got {type(values).__name__}") + schemas = list(schemas_by_param_name.values()) + parsed_message[key] = [ + schemas[order].validate_python(raw_value) for order, raw_value in enumerate(values) + ] + elif key == "result": + if values is None: + parsed_message[key] = None else: - parsed_message[key] = values + parsed_message[key] = schemas_by_param_name["return"].validate_python(values) + else: + parsed_message[key] = values - return parsed_message - except Exception: - if self.fallback_encoder is not None: - return self.fallback_encoder.decode_json(data) - raise + return parsed_message diff --git a/tests/test_encoders.py b/tests/test_encoders.py index 67035b4b9..0b2f9555f 100644 --- a/tests/test_encoders.py +++ b/tests/test_encoders.py @@ -156,13 +156,6 @@ def encoder(stub_broker, stub_worker, result_backend) -> PydanticEncoder: return PydanticEncoder() -@pytest.fixture -def encoder_with_fallback(stub_broker, stub_worker, result_backend) -> PydanticEncoder: - stub_broker.add_middleware(Results(backend=result_backend)) - stub_broker.declare_actor(my_actor_tuple) - return PydanticEncoder(fallback_encoder=JSONEncoder()) - - def test_encoder_message(encoder: PydanticEncoder, message_data_decoded: MessageData, message_data_encoded: bytes): encoded_result = encoder.encode_in_bytes(message_data_decoded) assert encoded_result == message_data_encoded @@ -200,14 +193,6 @@ def test_message_unknown_actor(encoder: PydanticEncoder, message_data_encoded: b encoder.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) -def test_message_fallback_no_actor_name(encoder_with_fallback: PydanticEncoder, message_data_encoded: bytes): - message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) - message_json_decoded["actor_name"] = "titi" - decoded_result = encoder_with_fallback.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) - # Do not raise and keep dict instead of schema - assert decoded_result == tuple_to_list(message_json_decoded, "args") - - def test_message_schema_not_matching(encoder: PydanticEncoder, message_data_encoded: bytes): message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) message_json_decoded["args"] = [{"toto": "a"}] @@ -327,11 +312,3 @@ def test_backend_result_schema_not_matching(encoder: PydanticEncoder, backend_re backend_result_json_decoded["result"] = [{"val": "titi"}] with pytest.raises(ValidationError): encoder.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) - - -def test_fallback_no_schema(encoder_with_fallback: PydanticEncoder, backend_result_encoded_tuple: bytes): - backend_result_json_decoded = json.loads(backend_result_encoded_tuple.decode("utf-8")) - backend_result_json_decoded["result"] = [{"val": "titi"}] - decoded_backend_result = encoder_with_fallback.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) - # Do not raise and keep dict instead of schema - assert decoded_backend_result == backend_result_json_decoded From c0cb0fff6200535eff9f274b8d568311ff67a471 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 16:01:27 +0200 Subject: [PATCH 36/44] refactor(postgres): minor review cleanups - Fix copy-pasted Message.decode_json/encode_in_json docstrings that said "bytestring" while they handle JSON objects. - Make convert_days_in_partman_syntax a staticmethod; it never used self. - Reject prefetch < 1 in the consumer with a ValueError (matching the timeout validation) instead of claiming a coercion that never happened. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 9 +++++---- remoulade/message.py | 4 ++-- tests/test_postgres.py | 10 ++++++++++ 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 02e8bb291..27f538aa0 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -131,7 +131,8 @@ def __init__( self._listener = _PostgresListener(self.url, self.logger) if enable_listen_notify else None - def convert_days_in_partman_syntax(self, interval_in_day: int) -> str: + @staticmethod + def convert_days_in_partman_syntax(interval_in_day: int) -> str: """Convert int into partman syntax""" if interval_in_day <= 0: raise ValueError("interval_in_day must be greater than 0") @@ -526,13 +527,13 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti Parameters: broker(PostgresBroker): Broker instance that owns the queue and database client. queue_name(str): Name of the declared queue to consume from. - prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1. + prefetch(int): Maximum number of messages fetched per read call; must be greater than or equal to 1. timeout(int): Idle wait timeout in milliseconds when polling for messages; must be greater than or equal to 0. A value of 0 performs non-blocking reads. - Raises: - ValueError: If ``timeout`` is negative. """ + if prefetch < 1: + raise ValueError("prefetch must be greater than or equal to 1") if timeout < 0: raise ValueError("timeout must be greater than or equal to 0") diff --git a/remoulade/message.py b/remoulade/message.py index be8e37be7..5b3836c74 100644 --- a/remoulade/message.py +++ b/remoulade/message.py @@ -99,11 +99,11 @@ def encode_in_bytes(self): @classmethod def decode_json(cls, data: dict[str, Any]): - """Convert a bytestring to a message.""" + """Convert a JSON object to a message.""" return cls(**global_encoder.decode_json(data)) def encode_in_json(self): - """Convert this message to a bytestring.""" + """Convert this message to a JSON object.""" return global_encoder.encode_in_json(self.asdict()) def copy(self, **attributes): diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 9a8b0848d..fd5fbf203 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -234,6 +234,16 @@ def test_postgres_consumer_rejects_negative_timeout(): broker.close() +def test_postgres_consumer_rejects_prefetch_below_one(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) + broker.queues["default"] = None + + with pytest.raises(ValueError, match="prefetch must be greater than or equal to 1"): + broker.consume("default", prefetch=0, timeout=200) + + broker.close() + + def test_postgres_poll_only_consumer_reads_immediately_when_timeout_is_zero(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) broker.queues["default"] = None From 4d98f7b6be955adbbc05d98c5ca702ab67b93a17 Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 16:17:27 +0200 Subject: [PATCH 37/44] fix(postgres): support several consumers per queue on the shared listener The shared LISTEN/NOTIFY listener keyed a single wake event per queue, so a second consumer on the same queue overwrote the first's event and a closing consumer's unregister() silently stopped notifications for a still-active sibling. Track a set of events per queue, wake them all, and only drop a queue's channel routing once its last consumer leaves. Also back `available` with a threading.Event instead of a bare bool so the dispatch thread's writes and the consumer threads' reads are synchronized rather than relying on CPython GIL atomicity; expose it as a read-only bool property so callers and tests are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) --- remoulade/brokers/postgres.py | 50 ++++++++++++++++++++++++----------- tests/test_postgres.py | 31 +++++++++++++++++++++- 2 files changed, 65 insertions(+), 16 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 27f538aa0..f1e3a7a24 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -356,11 +356,21 @@ def __init__(self, url: str, logger: logging.Logger) -> None: self._connection: psycopg.Connection[Any] | None = None self._thread: Thread | None = None self._started = False - self.available = False - self._events: dict[str, Event] = {} + # An Event rather than a bare bool so reads from consumer threads and + # writes from the dispatch thread are synchronized without relying on + # CPython's GIL making attribute access atomic. + self._available = Event() + # Several consumers may register on the same queue, so map each queue to + # the set of their wake events instead of a single one. + self._events: dict[str, set[Event]] = {} self._channel_to_queue: dict[str, str] = {} self._pending_listen: set[str] = set() + @property + def available(self) -> bool: + """Whether the shared LISTEN/NOTIFY connection is currently usable.""" + return self._available.is_set() + @staticmethod def _channel_for(queue_name: str) -> str: return f"pgmq.q_{queue_name}.INSERT" @@ -368,7 +378,7 @@ def _channel_for(queue_name: str) -> str: def register(self, queue_name: str, event: Event) -> None: """Register a consumer's wake event and ensure its channel is listened to.""" with self._lock: - self._events[queue_name] = event + self._events.setdefault(queue_name, set()).add(event) self._channel_to_queue[self._channel_for(queue_name)] = queue_name self._pending_listen.add(queue_name) starting = not self._started @@ -379,12 +389,22 @@ def register(self, queue_name: str, event: Event) -> None: if starting: self._start() - def unregister(self, queue_name: str) -> None: - """Stop routing notifications to a consumer that is shutting down.""" + def unregister(self, queue_name: str, event: Event) -> None: + """Stop routing notifications to a consumer that is shutting down. + + Only the queue's channel routing is dropped once its last consumer + leaves, so a closing consumer never stops notifications for a sibling + still consuming the same queue. + """ with self._lock: - self._events.pop(queue_name, None) - self._channel_to_queue.pop(self._channel_for(queue_name), None) - self._pending_listen.discard(queue_name) + events = self._events.get(queue_name) + if events is None: + return + events.discard(event) + if not events: + self._events.pop(queue_name, None) + self._channel_to_queue.pop(self._channel_for(queue_name), None) + self._pending_listen.discard(queue_name) def _start(self) -> None: """Open the shared connection (best effort) and start the dispatch thread. @@ -434,7 +454,7 @@ def _open_connection(self) -> bool: self._logger.debug("Failed to close partially-opened listener connection", exc_info=True) return False self._connection = connection - self.available = True + self._available.set() return True def _drop_connection(self) -> None: @@ -443,7 +463,7 @@ def _drop_connection(self) -> None: Waiting consumers are woken so they fall back to polling immediately instead of blocking on their wake event for the full timeout. """ - self.available = False + self._available.clear() self._wake_all() if self._connection is not None: try: @@ -464,17 +484,17 @@ def _drain_pending_listen(self) -> None: self._connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel))) def _wake_channel(self, channel: str) -> None: - """Wake the single consumer registered for a notification channel.""" + """Wake every consumer registered for a notification channel.""" with self._lock: queue_name = self._channel_to_queue.get(channel) - event = self._events.get(queue_name) if queue_name is not None else None - if event is not None: + events = list(self._events.get(queue_name, ())) if queue_name is not None else [] + for event in events: event.set() def _wake_all(self) -> None: """Wake every registered consumer (used on shutdown or listener failure).""" with self._lock: - events = list(self._events.values()) + events = [event for queue_events in self._events.values() for event in queue_events] for event in events: event.set() @@ -730,7 +750,7 @@ def close(self) -> None: self._heartbeat_stop.set() self._notify_event.set() if self.broker._listener is not None: - self.broker._listener.unregister(self.queue_name) + self.broker._listener.unregister(self.queue_name, self._notify_event) if self._heartbeat_thread is not None: self._heartbeat_thread.join(timeout=1) self._requeue_message_ids([message.msg_id for message in self.messages]) diff --git a/tests/test_postgres.py b/tests/test_postgres.py index fd5fbf203..66501a507 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -68,7 +68,7 @@ def __init__(self, available): def register(self, queue_name, event): pass - def unregister(self, queue_name): + def unregister(self, queue_name, event): pass def close(self): @@ -1044,3 +1044,32 @@ def square(value): assert sorted(g.results.get(block=True)) == [0, 1, 4, 9] finally: worker.stop() + + +def test_listener_wakes_all_consumers_on_a_queue_and_unregister_spares_siblings(): + listener = _PostgresListener("postgresql://localhost/test", logging.getLogger("test")) + # Skip starting the dispatch thread: this test only exercises wake routing. + listener._started = True + + first, second = threading.Event(), threading.Event() + listener.register("default", first) + listener.register("default", second) + channel = listener._channel_for("default") + + # A notification on the shared queue wakes every registered consumer. + listener._wake_channel(channel) + assert first.is_set() + assert second.is_set() + first.clear() + second.clear() + + # One consumer leaving must not stop notifications for its sibling. + listener.unregister("default", first) + listener._wake_channel(channel) + assert not first.is_set() + assert second.is_set() + + # The channel routing is dropped only once the last consumer leaves. + listener.unregister("default", second) + assert channel not in listener._channel_to_queue + assert "default" not in listener._events From 8e5b6562ed7a5ad888b35d1b045cab5cf040410e Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 19 Jun 2026 16:29:50 +0200 Subject: [PATCH 38/44] fix(doc): reformulation --- docs/source/installation.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 80333ffc7..ac7c77bea 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -9,13 +9,13 @@ Remoulade supports Python versions 3.12 and up and is installable via Via pip ------- +remoulade can be used with a RabbbitMQ_ or a PostgreSQL_ broker. -To install remoulade, simply run the following command in a terminal:: +If you want to use it with RabbitMQ_, simply run the following command in a terminal:: $ pip install -U 'remoulade[rabbitmq]' -Remoulade uses RabbitMQ_ as a message broker. If you want to use -PostgreSQL_ with PGMQ_ instead, install:: +If you want to use PostgreSQL_ with PGMQ_ instead, install:: $ pip install -U 'remoulade[postgres]' From bc2bdc6b77f703213c569a0ab2b0bff0b9abc9b0 Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 24 Jun 2026 16:23:43 +0200 Subject: [PATCH 39/44] feat(encoder): restore PydanticEncoder fallback_encoder support Reintroduce the opt-in fallback_encoder argument so decoding can fall back to another encoder instead of raising. Defaults to None, keeping strict decoding as the default behaviour. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/changelog.rst | 1 - remoulade/encoder.py | 98 +++++++++++++++++++++++---------------- tests/test_encoders.py | 23 +++++++++ 3 files changed, 81 insertions(+), 41 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 53544dbd2..5c4dd4e4e 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -15,7 +15,6 @@ Breaking changes * Rework the broker API around the new PostgreSQL/PGMQ implementation. * Rename ``Encoder.encode``/``Encoder.decode`` to ``Encoder.encode_in_bytes``/``Encoder.decode_bytes``; custom encoders must now also implement ``Encoder._encode_in_json`` and ``Encoder.decode_json``. * Rename ``Message.encode``/``Message.decode`` to ``Message.encode_in_bytes``/``Message.decode_bytes``. -* Remove ``PydanticEncoder``'s ``fallback_encoder`` argument: decoding now raises explicitly instead of silently falling back to another encoder and returning un-rehydrated payloads. * ``PydanticEncoder`` no longer depends on ``simplejson``; serialization now goes through Pydantic, so ``Decimal`` values are encoded as JSON strings instead of numbers. * Repurpose the ``postgres`` extra: it now installs the PGMQ broker dependencies (``sqlalchemy>=2``, ``psycopg>=3``, ``pgmq``) instead of the legacy state backend dependencies (``sqlalchemy<2``, ``psycopg2``). diff --git a/remoulade/encoder.py b/remoulade/encoder.py index 629bc4a31..126bfede1 100644 --- a/remoulade/encoder.py +++ b/remoulade/encoder.py @@ -131,16 +131,27 @@ def my_actor(input_1: MyActorInputSchema, input_2: MyActorInputSchema | None = N return MyActorOutputSchema() """ - def __init__(self): + def __init__(self, fallback_encoder: Encoder | None = None): + self.fallback_encoder = fallback_encoder self.json_adapter = TypeAdapter(object) @override def encode_in_bytes(self, data: MessageData) -> bytes: - return json.dumps(self._encode_in_json(data)).encode("utf-8") + try: + return json.dumps(self._encode_in_json(data)).encode("utf-8") + except Exception: + if self.fallback_encoder is not None: + return self.fallback_encoder.encode_in_bytes(data) + raise @override def decode_bytes(self, data: bytes) -> MessageData: - return self.decode_json(json.loads(data.decode("utf-8"))) + try: + return self.decode_json(json.loads(data.decode("utf-8"))) + except Exception: + if self.fallback_encoder is not None: + return self.fallback_encoder.decode_bytes(data) + raise @override def _encode_in_json(self, data: MessageData) -> JsonData: @@ -150,42 +161,49 @@ def _encode_in_json(self, data: MessageData) -> JsonData: def decode_json(self, data: JsonData) -> MessageData: from remoulade import get_broker - actor_name = data["actor_name"] - actor_fn = get_broker().get_actor(actor_name).fn - - # Retrieve the Pydantic schemas from typing - schemas_by_param_name: dict[str, TypeAdapter] = {} - for param_name, type_hint in get_type_hints(actor_fn).items(): - schemas_by_param_name[param_name] = TypeAdapter( - Annotated[ - type_hint, - WithJsonSchema({"type": type_hint, "description": f"{param_name}_schema"}, mode="serialization"), - ] - ) - - # Override message_data with Pydantic schema when it matches. - parsed_message: dict[str, Any] = {} - for key, values in data.items(): - if key == "kwargs": - if not isinstance(values, dict): - raise TypeError(f"Expected `values` to be a dict, got {type(values).__name__}") - parsed_message[key] = { - param_name: schemas_by_param_name[param_name].validate_python(raw_value) - for param_name, raw_value in values.items() - } - elif key == "args": - if not isinstance(values, list): - raise TypeError(f"Expected `values` to be a list, got {type(values).__name__}") - schemas = list(schemas_by_param_name.values()) - parsed_message[key] = [ - schemas[order].validate_python(raw_value) for order, raw_value in enumerate(values) - ] - elif key == "result": - if values is None: - parsed_message[key] = None + try: + actor_name = data["actor_name"] + actor_fn = get_broker().get_actor(actor_name).fn + + # Retrieve the Pydantic schemas from typing + schemas_by_param_name: dict[str, TypeAdapter] = {} + for param_name, type_hint in get_type_hints(actor_fn).items(): + schemas_by_param_name[param_name] = TypeAdapter( + Annotated[ + type_hint, + WithJsonSchema( + {"type": type_hint, "description": f"{param_name}_schema"}, mode="serialization" + ), + ] + ) + + # Override message_data with Pydantic schema when it matches. + parsed_message: dict[str, Any] = {} + for key, values in data.items(): + if key == "kwargs": + if not isinstance(values, dict): + raise TypeError(f"Expected `values` to be a dict, got {type(values).__name__}") + parsed_message[key] = { + param_name: schemas_by_param_name[param_name].validate_python(raw_value) + for param_name, raw_value in values.items() + } + elif key == "args": + if not isinstance(values, list): + raise TypeError(f"Expected `values` to be a list, got {type(values).__name__}") + schemas = list(schemas_by_param_name.values()) + parsed_message[key] = [ + schemas[order].validate_python(raw_value) for order, raw_value in enumerate(values) + ] + elif key == "result": + if values is None: + parsed_message[key] = None + else: + parsed_message[key] = schemas_by_param_name["return"].validate_python(values) else: - parsed_message[key] = schemas_by_param_name["return"].validate_python(values) - else: - parsed_message[key] = values + parsed_message[key] = values - return parsed_message + return parsed_message + except Exception: + if self.fallback_encoder is not None: + return self.fallback_encoder.decode_json(data) + raise diff --git a/tests/test_encoders.py b/tests/test_encoders.py index 0b2f9555f..67035b4b9 100644 --- a/tests/test_encoders.py +++ b/tests/test_encoders.py @@ -156,6 +156,13 @@ def encoder(stub_broker, stub_worker, result_backend) -> PydanticEncoder: return PydanticEncoder() +@pytest.fixture +def encoder_with_fallback(stub_broker, stub_worker, result_backend) -> PydanticEncoder: + stub_broker.add_middleware(Results(backend=result_backend)) + stub_broker.declare_actor(my_actor_tuple) + return PydanticEncoder(fallback_encoder=JSONEncoder()) + + def test_encoder_message(encoder: PydanticEncoder, message_data_decoded: MessageData, message_data_encoded: bytes): encoded_result = encoder.encode_in_bytes(message_data_decoded) assert encoded_result == message_data_encoded @@ -193,6 +200,14 @@ def test_message_unknown_actor(encoder: PydanticEncoder, message_data_encoded: b encoder.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) +def test_message_fallback_no_actor_name(encoder_with_fallback: PydanticEncoder, message_data_encoded: bytes): + message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) + message_json_decoded["actor_name"] = "titi" + decoded_result = encoder_with_fallback.decode_bytes(json.dumps(message_json_decoded).encode("utf-8")) + # Do not raise and keep dict instead of schema + assert decoded_result == tuple_to_list(message_json_decoded, "args") + + def test_message_schema_not_matching(encoder: PydanticEncoder, message_data_encoded: bytes): message_json_decoded = json.loads(message_data_encoded.decode("utf-8")) message_json_decoded["args"] = [{"toto": "a"}] @@ -312,3 +327,11 @@ def test_backend_result_schema_not_matching(encoder: PydanticEncoder, backend_re backend_result_json_decoded["result"] = [{"val": "titi"}] with pytest.raises(ValidationError): encoder.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) + + +def test_fallback_no_schema(encoder_with_fallback: PydanticEncoder, backend_result_encoded_tuple: bytes): + backend_result_json_decoded = json.loads(backend_result_encoded_tuple.decode("utf-8")) + backend_result_json_decoded["result"] = [{"val": "titi"}] + decoded_backend_result = encoder_with_fallback.decode_bytes(json.dumps(backend_result_json_decoded).encode("utf-8")) + # Do not raise and keep dict instead of schema + assert decoded_backend_result == backend_result_json_decoded From 0dc3aa8c42e999096ab7149ee0e81ea60363c001 Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 25 Jun 2026 17:14:09 +0200 Subject: [PATCH 40/44] fix(broker): declare queue concurrency --- remoulade/brokers/postgres.py | 35 +++++++++++++----------- tests/test_postgres.py | 50 ++++++++++++++++++++++++----------- 2 files changed, 55 insertions(+), 30 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index f1e3a7a24..572fe76cc 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -29,7 +29,7 @@ from pgmq import SQLAlchemyPGMQueue from pgmq.messages import Message as PostgresQueueMessage from psycopg import sql as psycopg_sql -from sqlalchemy import Connection +from sqlalchemy import Connection, text from ..broker import Broker, Consumer, MessageProxy from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding @@ -206,20 +206,25 @@ def declare_queue(self, queue_name: str) -> None: """Create a partitioned PGMQ queue if it does not already exist.""" if queue_name in self.queues: return - - self.client.validate_queue_name(queue_name, conn=self._current_connection) - queue_exists = self._queue_exists(queue_name) - - if not queue_exists: - self.emit_before("declare_queue", queue_name) - self.client.create_partitioned_queue( - queue_name, - partition_interval=self.archive_partition_interval, - retention_interval=self.archive_retention_interval, - conn=self._current_connection, - ) - if self.enable_listen_notify: - self._try_enable_notify(queue_name) + with self.tx(): + if self._current_connection is None: + raise ValueError("cannot be None we are inside a tx") + self._current_connection.execute( + text("SELECT pg_advisory_xact_lock(hashtext(:k))"), {"k": f"remoulade.declare_queue.{queue_name}"} + ) # Concurrency issues if broker try to create the same queues + self.client.validate_queue_name(queue_name, conn=self._current_connection) + queue_exists = self._queue_exists(queue_name) + + if not queue_exists: + self.emit_before("declare_queue", queue_name) + self.client.create_partitioned_queue( + queue_name, + partition_interval=self.archive_partition_interval, + retention_interval=self.archive_retention_interval, + conn=self._current_connection, + ) + if self.enable_listen_notify: + self._try_enable_notify(queue_name) self.queues[queue_name] = None diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 66501a507..e6de40a36 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -55,6 +55,24 @@ def _expected_payload(message): return json.loads(message.encode_in_bytes().decode("utf-8")) +class _StubTransaction: + """Stand-in for ``engine.begin()`` that yields a controllable connection. + + ``declare_queue`` always opens its own transaction and runs queue + operations against that connection, so tests inject a fake connection to + assert on the ``conn`` argument forwarded to the PGMQ client. + """ + + def __init__(self, connection): + self._connection = connection + + def __enter__(self): + return self._connection + + def __exit__(self, exc_type, exc, tb): + return None + + class _FakeListener: """Stand-in for the broker's shared LISTEN/NOTIFY listener. @@ -100,16 +118,19 @@ def test_postgres_broker_creates_partitioned_queue_with_default_intervals(): broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + conn = Mock() + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + broker.declare_queue("default") - broker.client.validate_queue_name.assert_called_once_with("default", conn=None) + broker.client.validate_queue_name.assert_called_once_with("default", conn=conn) broker.client.create_partitioned_queue.assert_called_once_with( "default", partition_interval="1 day", retention_interval="7 days", - conn=None, + conn=conn, ) - broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) + broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=conn) broker.client.create_queue.assert_not_called() @@ -120,16 +141,9 @@ def test_postgres_broker_uses_current_transaction_connection_for_queue_creation( broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) - transaction_connection = object() - - class _DummyTransaction: - def __enter__(self): - return transaction_connection + transaction_connection = Mock() - def __exit__(self, exc_type, exc, tb): - return None - - broker.client.engine.begin = Mock(return_value=_DummyTransaction()) + broker.client.engine.begin = Mock(return_value=_StubTransaction(transaction_connection)) with broker.tx(): broker.declare_queue("default") @@ -155,15 +169,18 @@ def test_postgres_broker_enables_notify_on_postgresql_queue_init(): broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + conn = Mock() + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + broker.declare_queue("default") broker.client.create_partitioned_queue.assert_called_once_with( "default", partition_interval="1 day", retention_interval="7 days", - conn=None, + conn=conn, ) - broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=None) + broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=conn) def test_postgres_broker_does_not_fail_when_enable_notify_raises(caplog): @@ -299,13 +316,16 @@ def test_postgres_broker_uses_custom_partition_settings_when_provided(): broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + conn = Mock() + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + broker.declare_queue("default") broker.client.create_partitioned_queue.assert_called_once_with( "default", partition_interval="2 days", retention_interval="14 days", - conn=None, + conn=conn, ) From aed735c4d3031947fbe44a5957bb102146c769c7 Mon Sep 17 00:00:00 2001 From: mducros Date: Wed, 1 Jul 2026 14:51:16 +0200 Subject: [PATCH 41/44] chore(pgmq): add partition cli --- docs/source/guide.rst | 10 ++++++ pyproject.toml | 1 + remoulade/brokers/postgres.py | 34 ++++++++++++++++++ remoulade/cli/remoulade_partitions.py | 52 +++++++++++++++++++++++++++ tests/test_postgres.py | 27 ++++++++++++++ 5 files changed, 124 insertions(+) create mode 100644 remoulade/cli/remoulade_partitions.py diff --git a/docs/source/guide.rst b/docs/source/guide.rst index 0f41ed9c8..48b87df73 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -447,6 +447,16 @@ table, which is partitioned the same way. Two broker parameters control this: image (``ghcr.io/pgmq/pg18-pgmq``) ships ``pg_partman`` and schedules this for you; on a self-managed or hosted PostgreSQL you must enable it yourself. +The ``remoulade-partitions`` CLI sets ``infinite_time_partitions = true`` on every queue and +archive partition set (so ``pg_partman`` keeps maintaining partitions up to the current time even +after a gap in activity), then runs ``pg_partman`` maintenance immediately:: + + REMOULADE_POSTGRES_URL=postgresql://remoulade@localhost:5432/remoulade remoulade-partitions + +It reads the database URL from the ``REMOULADE_POSTGRES_URL`` environment variable and creates a +minimal ``PostgresBroker`` on its own (no middleware, no LISTEN/NOTIFY connection), so it does not +need to import your application code. + Local Broker ^^^^^^^^^^^^^^^ diff --git a/pyproject.toml b/pyproject.toml index bf279d855..fa0166903 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,7 @@ remoulade = "remoulade.__main__:main" remoulade-ls = "remoulade.cli.remoulade_ls:main" remoulade-run = "remoulade.cli.remoulade_run:main" remoulade-scheduler = "remoulade.cli.remoulade_scheduler:main" +remoulade-partitions = "remoulade.cli.remoulade_partitions:main" [build-system] requires = ["setuptools>=80.9", "wheel", "setuptools-scm[simple]>=8"] diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 572fe76cc..8d1b00d1c 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -231,6 +231,40 @@ def declare_queue(self, queue_name: str) -> None: if not queue_exists: self.emit_after("declare_queue", queue_name) + def enable_infinite_time_partitions(self) -> list[str]: + """Set ``infinite_time_partitions = true`` on every PGMQ queue and archive partition set. + + Both the queue table (``pgmq.q_``) and its archive (``pgmq.a_``) are + partitioned time-based sets registered in ``partman.part_config``. Enabling + ``infinite_time_partitions`` forces ``pg_partman`` to keep maintaining partitions up to the + current time even when no recent rows exist, avoiding failed inserts after a gap in activity. + + Returns: + The list of parent tables that were updated. + """ + with self.client.engine.begin() as conn: + rows = conn.execute( + text( + "UPDATE part_config SET infinite_time_partitions = true " + "WHERE parent_table LIKE 'pgmq.q\\_%' OR parent_table LIKE 'pgmq.a\\_%' " + "RETURNING parent_table" + ) + ).all() + tables = [row[0] for row in rows] + self.logger.info("Enabled infinite_time_partitions on %d partition sets: %s", len(tables), tables) + return tables + + def run_partition_maintenance(self) -> None: + """Run ``pg_partman`` maintenance to create upcoming partitions and drop expired ones. + + ``partman.run_maintenance_proc()`` is a procedure that ``COMMIT``s internally, which + PostgreSQL forbids inside an already-open transaction. The ``CALL`` therefore runs in + autocommit (no surrounding ``BEGIN``) rather than through ``engine.begin()``. + """ + with self.client.engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: + conn.execute(text("CALL run_maintenance_proc()")) + self.logger.info("Ran pg_partman maintenance (partman.run_maintenance_proc())") + def _encode_message(self, message: "Message") -> PostgresPayload: """Encode a Remoulade message into a JSON object payload for PGMQ. diff --git a/remoulade/cli/remoulade_partitions.py b/remoulade/cli/remoulade_partitions.py new file mode 100644 index 000000000..12d278402 --- /dev/null +++ b/remoulade/cli/remoulade_partitions.py @@ -0,0 +1,52 @@ +"""``remoulade-partitions`` CLI. + +One-shot maintenance command for the PostgreSQL broker. It creates a minimal +``PostgresBroker`` from the ``REMOULADE_POSTGRES_URL`` environment variable (no +middleware, no LISTEN/NOTIFY connection), enables ``infinite_time_partitions`` +on every PGMQ queue and archive partition set so ``pg_partman`` keeps +maintaining partitions up to the current time even after a gap in activity, and +then runs ``pg_partman`` maintenance to apply it immediately. +""" + +import argparse +import os + +from remoulade.brokers.postgres import PostgresBroker + +POSTGRES_URL_ENV_VAR = "REMOULADE_POSTGRES_URL" + + +def parse_arguments(): + parser = argparse.ArgumentParser( + prog="remoulade-partitions", + description=( + "Enable infinite_time_partitions on every PostgreSQL queue and archive partition set, " + f"then run pg_partman maintenance. The database URL is read from ${POSTGRES_URL_ENV_VAR}." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + return parser.parse_args() + + +def main(): + """Enable infinite_time_partitions on all queue/archive partition sets, then run maintenance.""" + parse_arguments() + + url = os.getenv(POSTGRES_URL_ENV_VAR) + if not url: + print(f"{POSTGRES_URL_ENV_VAR} environment variable is not set") + return 1 + + broker = PostgresBroker(url=url, middleware=[], enable_listen_notify=False) + try: + tables = broker.enable_infinite_time_partitions() + print(f"Enabled infinite_time_partitions on {len(tables)} partition sets:") + for table in tables: + print(f" {table}") + + broker.run_partition_maintenance() + print("Ran pg_partman maintenance.") + finally: + broker.close() + + return 0 diff --git a/tests/test_postgres.py b/tests/test_postgres.py index e6de40a36..d12e7cfc0 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -183,6 +183,33 @@ def test_postgres_broker_enables_notify_on_postgresql_queue_init(): broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=conn) +def test_postgres_broker_enable_infinite_time_partitions_runs_update_and_returns_tables(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + + conn = Mock() + conn.execute.return_value.all.return_value = [("pgmq.q_default",), ("pgmq.a_default",)] + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + + tables = broker.enable_infinite_time_partitions() + + assert tables == ["pgmq.q_default", "pgmq.a_default"] + sql = str(conn.execute.call_args.args[0]) + assert "infinite_time_partitions = true" in sql + assert "part_config" in sql + + +def test_postgres_broker_run_partition_maintenance_calls_run_maintenance_proc(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + + conn = Mock() + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + + broker.run_partition_maintenance() + + sql = str(conn.execute.call_args.args[0]) + assert "run_maintenance_proc()" in sql + + def test_postgres_broker_does_not_fail_when_enable_notify_raises(caplog): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() From 2a71baba6888b55adb459561b1a7b984bd03532b Mon Sep 17 00:00:00 2001 From: mducros Date: Thu, 2 Jul 2026 14:08:52 +0200 Subject: [PATCH 42/44] fix(broker): remove partition cli --- docs/source/guide.rst | 23 ++++-------- pyproject.toml | 1 - remoulade/brokers/postgres.py | 34 ------------------ remoulade/cli/remoulade_partitions.py | 52 --------------------------- tests/test_postgres.py | 27 -------------- 5 files changed, 7 insertions(+), 130 deletions(-) delete mode 100644 remoulade/cli/remoulade_partitions.py diff --git a/docs/source/guide.rst b/docs/source/guide.rst index 48b87df73..d5aa0cefe 100644 --- a/docs/source/guide.rst +++ b/docs/source/guide.rst @@ -440,22 +440,13 @@ table, which is partitioned the same way. Two broker parameters control this: .. note:: - Partitioning is maintained by ``pg_partman``: its maintenance routine - (``partman.run_maintenance_proc()``) must run periodically — through the - ``pg_partman`` background worker or a ``pg_cron`` job — to create upcoming - partitions ahead of time and drop expired ones. The official PGMQ Docker - image (``ghcr.io/pgmq/pg18-pgmq``) ships ``pg_partman`` and schedules this - for you; on a self-managed or hosted PostgreSQL you must enable it yourself. - -The ``remoulade-partitions`` CLI sets ``infinite_time_partitions = true`` on every queue and -archive partition set (so ``pg_partman`` keeps maintaining partitions up to the current time even -after a gap in activity), then runs ``pg_partman`` maintenance immediately:: - - REMOULADE_POSTGRES_URL=postgresql://remoulade@localhost:5432/remoulade remoulade-partitions - -It reads the database URL from the ``REMOULADE_POSTGRES_URL`` environment variable and creates a -minimal ``PostgresBroker`` on its own (no middleware, no LISTEN/NOTIFY connection), so it does not -need to import your application code. + Partitioning is maintained by ``pg_partman``: its maintenance routine must run periodically — + through the ``pg_partman`` background worker, a ``pg_cron`` job, or an application-level scheduled + task — to create upcoming partitions ahead of time and drop expired ones. It should also set + ``infinite_time_partitions = true`` on every queue and archive partition set so ``pg_partman`` + keeps maintaining partitions up to the current time even after a gap in activity. The official + PGMQ Docker image (``ghcr.io/pgmq/pg18-pgmq``) ships ``pg_partman`` and schedules this for you; on + a self-managed or hosted PostgreSQL you must enable it yourself. Local Broker diff --git a/pyproject.toml b/pyproject.toml index fa0166903..bf279d855 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,6 @@ remoulade = "remoulade.__main__:main" remoulade-ls = "remoulade.cli.remoulade_ls:main" remoulade-run = "remoulade.cli.remoulade_run:main" remoulade-scheduler = "remoulade.cli.remoulade_scheduler:main" -remoulade-partitions = "remoulade.cli.remoulade_partitions:main" [build-system] requires = ["setuptools>=80.9", "wheel", "setuptools-scm[simple]>=8"] diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 8d1b00d1c..572fe76cc 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -231,40 +231,6 @@ def declare_queue(self, queue_name: str) -> None: if not queue_exists: self.emit_after("declare_queue", queue_name) - def enable_infinite_time_partitions(self) -> list[str]: - """Set ``infinite_time_partitions = true`` on every PGMQ queue and archive partition set. - - Both the queue table (``pgmq.q_``) and its archive (``pgmq.a_``) are - partitioned time-based sets registered in ``partman.part_config``. Enabling - ``infinite_time_partitions`` forces ``pg_partman`` to keep maintaining partitions up to the - current time even when no recent rows exist, avoiding failed inserts after a gap in activity. - - Returns: - The list of parent tables that were updated. - """ - with self.client.engine.begin() as conn: - rows = conn.execute( - text( - "UPDATE part_config SET infinite_time_partitions = true " - "WHERE parent_table LIKE 'pgmq.q\\_%' OR parent_table LIKE 'pgmq.a\\_%' " - "RETURNING parent_table" - ) - ).all() - tables = [row[0] for row in rows] - self.logger.info("Enabled infinite_time_partitions on %d partition sets: %s", len(tables), tables) - return tables - - def run_partition_maintenance(self) -> None: - """Run ``pg_partman`` maintenance to create upcoming partitions and drop expired ones. - - ``partman.run_maintenance_proc()`` is a procedure that ``COMMIT``s internally, which - PostgreSQL forbids inside an already-open transaction. The ``CALL`` therefore runs in - autocommit (no surrounding ``BEGIN``) rather than through ``engine.begin()``. - """ - with self.client.engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: - conn.execute(text("CALL run_maintenance_proc()")) - self.logger.info("Ran pg_partman maintenance (partman.run_maintenance_proc())") - def _encode_message(self, message: "Message") -> PostgresPayload: """Encode a Remoulade message into a JSON object payload for PGMQ. diff --git a/remoulade/cli/remoulade_partitions.py b/remoulade/cli/remoulade_partitions.py deleted file mode 100644 index 12d278402..000000000 --- a/remoulade/cli/remoulade_partitions.py +++ /dev/null @@ -1,52 +0,0 @@ -"""``remoulade-partitions`` CLI. - -One-shot maintenance command for the PostgreSQL broker. It creates a minimal -``PostgresBroker`` from the ``REMOULADE_POSTGRES_URL`` environment variable (no -middleware, no LISTEN/NOTIFY connection), enables ``infinite_time_partitions`` -on every PGMQ queue and archive partition set so ``pg_partman`` keeps -maintaining partitions up to the current time even after a gap in activity, and -then runs ``pg_partman`` maintenance to apply it immediately. -""" - -import argparse -import os - -from remoulade.brokers.postgres import PostgresBroker - -POSTGRES_URL_ENV_VAR = "REMOULADE_POSTGRES_URL" - - -def parse_arguments(): - parser = argparse.ArgumentParser( - prog="remoulade-partitions", - description=( - "Enable infinite_time_partitions on every PostgreSQL queue and archive partition set, " - f"then run pg_partman maintenance. The database URL is read from ${POSTGRES_URL_ENV_VAR}." - ), - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - return parser.parse_args() - - -def main(): - """Enable infinite_time_partitions on all queue/archive partition sets, then run maintenance.""" - parse_arguments() - - url = os.getenv(POSTGRES_URL_ENV_VAR) - if not url: - print(f"{POSTGRES_URL_ENV_VAR} environment variable is not set") - return 1 - - broker = PostgresBroker(url=url, middleware=[], enable_listen_notify=False) - try: - tables = broker.enable_infinite_time_partitions() - print(f"Enabled infinite_time_partitions on {len(tables)} partition sets:") - for table in tables: - print(f" {table}") - - broker.run_partition_maintenance() - print("Ran pg_partman maintenance.") - finally: - broker.close() - - return 0 diff --git a/tests/test_postgres.py b/tests/test_postgres.py index d12e7cfc0..e6de40a36 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -183,33 +183,6 @@ def test_postgres_broker_enables_notify_on_postgresql_queue_init(): broker.client.enable_notify.assert_called_once_with("default", throttle_interval_ms=250, conn=conn) -def test_postgres_broker_enable_infinite_time_partitions_runs_update_and_returns_tables(): - broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) - - conn = Mock() - conn.execute.return_value.all.return_value = [("pgmq.q_default",), ("pgmq.a_default",)] - broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) - - tables = broker.enable_infinite_time_partitions() - - assert tables == ["pgmq.q_default", "pgmq.a_default"] - sql = str(conn.execute.call_args.args[0]) - assert "infinite_time_partitions = true" in sql - assert "part_config" in sql - - -def test_postgres_broker_run_partition_maintenance_calls_run_maintenance_proc(): - broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) - - conn = Mock() - broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) - - broker.run_partition_maintenance() - - sql = str(conn.execute.call_args.args[0]) - assert "run_maintenance_proc()" in sql - - def test_postgres_broker_does_not_fail_when_enable_notify_raises(caplog): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.client.validate_queue_name = Mock() From e1532ad3c7db1266664764fbebf1fb8c1e1295c8 Mon Sep 17 00:00:00 2001 From: mducros-wm Date: Fri, 3 Jul 2026 16:52:07 +0200 Subject: [PATCH 43/44] feat(postgres): batch group enqueue with a single PGMQ send_batch (#400) * feat(postgres): batch group enqueue with a single PGMQ send_batch Enqueuing a group of N messages on the PostgreSQL/PGMQ broker previously performed N separate INSERTs (and, without group_transaction, N commits). Add a batch-enqueue path so a group fans out in a single send_batch per queue. - Broker.enqueue_many(): runs the enqueue middleware per message (before_enqueue grouped before the write, after_enqueue after it) and delegates the write to _enqueue_many. On failure, after_enqueue(exception) is emitted only for the messages whose before_enqueue already ran, so no span/state leaks. - Broker._enqueue_many(): default loops _enqueue, so rabbitmq/local/stub are unchanged. - PostgresBroker._enqueue_many(): one client.send_batch per queue, reusing the current transaction connection (1 INSERT, 1 connection, 1 commit with group_transaction). - group.run() now uses broker.enqueue_many(self.build()). Co-Authored-By: Claude Opus 4.8 (1M context) * feat(postgres): make group enqueue send_batch size configurable Add an optional `enqueue_batch_size` to PostgresBroker. When set, _enqueue_many splits each queue's messages into chunks of that size and sends one send_batch per chunk (all within the same transaction), so a very large fan-out does not build a single oversized INSERT. None (default) keeps the current behavior of a single send_batch per queue. Values <= 0 are rejected. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- remoulade/broker.py | 46 ++++++++++++++++++++++ remoulade/brokers/local.py | 51 ++++++++++++++++++------ remoulade/brokers/postgres.py | 40 +++++++++++++++++++ remoulade/brokers/rabbitmq.py | 4 ++ remoulade/brokers/stub.py | 3 ++ remoulade/composition.py | 3 +- tests/test_broker.py | 62 +++++++++++++++++++++++++++++ tests/test_postgres.py | 74 ++++++++++++++++++++++++++++++++++- 8 files changed, 267 insertions(+), 16 deletions(-) diff --git a/remoulade/broker.py b/remoulade/broker.py index 8732b6d7e..1f3c1fbee 100644 --- a/remoulade/broker.py +++ b/remoulade/broker.py @@ -431,6 +431,16 @@ def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": raise NotImplementedError + def _enqueue_many(self, messages: list["Message[Any]"], *, delay: int | None = None) -> list["Message[Any]"]: + """Raw batch-enqueue hook, called by :meth:`enqueue_many` after ``emit_before``. + + The default implementation simply loops over :meth:`_enqueue`, so brokers that do not + override it behave exactly as if each message had been enqueued on its own. Brokers whose + backend supports a native batch insert (e.g. PGMQ ``send_batch``) should override this to + collapse the N writes into a single round-trip. + """ + raise NotImplementedError + def enqueue(self, message: "Message[Any]", *, delay: int | None = None) -> "Message[Any]": # pragma: no cover """Enqueue a message on this broker. @@ -454,6 +464,42 @@ def enqueue(self, message: "Message[Any]", *, delay: int | None = None) -> "Mess self.emit_after("enqueue", message, delay, exception=e) raise e from e + def enqueue_many(self, messages: list["Message[Any]"], *, delay: int | None = None) -> list["Message[Any]"]: + """Enqueue several messages on this broker, batching the backend write when possible. + + The enqueue middleware still runs once per message, but the ``emit_before`` calls are grouped + before the (batched) write and the ``emit_after`` calls after it, instead of being interleaved + per message as in :meth:`enqueue`. This ordering is what lets a broker collapse the N backend + writes into one (see :meth:`_enqueue_many`). ``group.run()`` uses this to fan out a group in a + single backend round-trip. + + On failure, ``emit_after`` is emitted (with the exception) only for the messages whose + ``emit_before`` already ran, so tracing spans / message state opened in ``before_enqueue`` are + always closed. + + Parameters: + messages(list[Message]): The messages to enqueue. + delay(int): The number of milliseconds to delay every message for. + + Returns: + list[Message]: The enqueued messages. + """ + messages = [self._apply_delay(message, delay) for message in messages] + emitted_before: list[Message[Any]] = [] + try: + for message in messages: + self.emit_before("enqueue", message, delay) + emitted_before.append(message) + messages = self._enqueue_many(messages, delay=delay) + for message in messages: + self.emit_after("enqueue", message, delay) + return messages + + except BaseException as e: + for message in emitted_before: + self.emit_after("enqueue", message, delay, exception=e) + raise e from e + def get_actor(self, actor_name: str) -> "Actor": # pragma: no cover """Look up an actor by its name. diff --git a/remoulade/brokers/local.py b/remoulade/brokers/local.py index f7091714c..38cbb57d7 100644 --- a/remoulade/brokers/local.py +++ b/remoulade/brokers/local.py @@ -15,13 +15,20 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . -from ..broker import Broker, MessageProxy +from typing import TYPE_CHECKING, Any, override + +from ..broker import Broker, Consumer, MessageProxy from ..common import current_millis from ..message import Message from ..middleware import SkipMessage from ..results import Results from ..results.backends import LocalBackend +if TYPE_CHECKING: + from collections.abc import Iterable + + from ..middleware import Middleware + class LocalBroker(Broker): """Broker that calculate the message result immediately @@ -29,41 +36,51 @@ class LocalBroker(Broker): It can only be used with LocalBackend as a result backend """ - def __init__(self, middleware=None): + def __init__(self, middleware: "Iterable[Middleware] | None" = None) -> None: super().__init__(middleware) @property - def local(self): + @override + def local(self) -> bool: return True - def add_middleware(self, middleware, *, before=None, after=None): + @override + def add_middleware( + self, middleware: "Middleware", *, before: "Middleware | None" = None, after: "Middleware | None" = None + ) -> None: if isinstance(middleware, Results) and not isinstance(middleware.backend, LocalBackend): raise RuntimeError("LocalBroker can only be used with LocalBackend.") super().add_middleware(middleware) - def emit_before(self, signal, *args, **kwargs): + @override + def emit_before(self, signal: str, *args: Any, **kwargs: Any) -> None: # A local broker should not catch any exception because we are not in a worker but in the main thread for middleware in self.middleware: getattr(middleware, "before_" + signal)(self, *args, **kwargs) - def emit_after(self, signal, *args, **kwargs): + @override + def emit_after(self, signal: str, *args: Any, **kwargs: Any) -> None: for middleware in reversed(self.middleware): getattr(middleware, "after_" + signal)(self, *args, **kwargs) - def consume(self, queue_name, prefetch=1, timeout=100): + @override + def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 100) -> "Consumer": raise ValueError("LocalBroker is not destined to use with a Worker") - def declare_queue(self, queue_name): + @override + def declare_queue(self, queue_name: str) -> None: self.emit_before("declare_queue", queue_name) self.queues[queue_name] = None self.emit_after("declare_queue", queue_name) - def _apply_delay(self, message, delay: int | None = None): + @override + def _apply_delay(self, message: "Message", delay: int | None = None) -> "Message": if delay is not None: message_eta = current_millis() + delay message = message.copy(options={"eta": message_eta}) return message + @override def enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": # pragma: no cover """Enqueue a message on this broker. @@ -82,7 +99,8 @@ def enqueue(self, message: "Message", *, delay: int | None = None) -> "Message": message = self._enqueue(message, delay=delay) return message - def _enqueue(self, message, *, delay=None): + @override + def _enqueue(self, message: "Message", *, delay: int | None = None): """Enqueue and compute a message. Parameters: @@ -115,11 +133,18 @@ def _enqueue(self, message, *, delay=None): return message_proxy - def flush(self, _): + @override + def _enqueue_many(self, messages: list["Message[Any]"], *, delay: int | None = None) -> list["Message[Any]"]: + return [self._enqueue(message, delay=delay) for message in messages] + + @override + def flush(self, _: str) -> None: pass - def flush_all(self): + @override + def flush_all(self) -> None: pass - def join(self, queue_name: str, *, timeout: int | None = None): + @override + def join(self, queue_name: str, *, timeout: int | None = None) -> None: return diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 572fe76cc..106faba37 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -83,6 +83,7 @@ def __init__( heartbeat_interval_ms: int = 10_000, enable_listen_notify: bool = True, pool_size: int = 10, + enqueue_batch_size: int | None = None, engine: "Engine | None" = None, ) -> None: """Initialize a PostgreSQL-backed broker using the PGMQ extension. @@ -101,6 +102,10 @@ def __init__( connection. If False, consumers poll only and no dedicated listener connection is opened (required behind a transaction-pooling connection pooler such as pgbouncer). pool_size(int): Size of the shared SQLAlchemy connection pool. Ignored when ``engine`` is provided. + enqueue_batch_size(int | None): Maximum number of messages sent per ``send_batch`` call in + :meth:`_enqueue_many`. When None (default), a queue's messages are all sent in a single + ``send_batch``; set it to cap the size of each ``INSERT`` on very large fan-outs. All chunks + still run inside the same transaction. Must be greater than 0 when provided. engine(Engine | None): A pre-configured SQLAlchemy engine to reuse instead of letting PGMQ build one, so the pool can be sized and shared by the caller. """ @@ -109,6 +114,8 @@ def __init__( raise ValueError("visibility_timeout_ms must be greater than 0") if heartbeat_interval_ms <= 0 or heartbeat_interval_ms >= visibility_timeout_ms: raise ValueError("heartbeat_interval_ms must be greater than 0 and lower than visibility_timeout_ms") + if enqueue_batch_size is not None and enqueue_batch_size <= 0: + raise ValueError("enqueue_batch_size must be greater than 0") self.url = urlparse(url).geturl() self.state = local() @@ -120,6 +127,7 @@ def __init__( self.visibility_timeout_seconds = _milliseconds_to_seconds(visibility_timeout_ms) self.heartbeat_interval_seconds = heartbeat_interval_ms / 1000 self.enable_listen_notify = enable_listen_notify + self.enqueue_batch_size = enqueue_batch_size self.client = SQLAlchemyPGMQueue( conn_string=url, @@ -269,6 +277,38 @@ def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message" return message + @override + def _enqueue_many(self, messages: list["Message"], *, delay: int | None = None) -> list["Message"]: + """Send several messages to PGMQ with a ``send_batch`` per queue (optionally chunked). + + For large fan-outs (e.g. ``group.run()``) this collapses N ``INSERT`` round-trips into one + statement per queue. Combined with ``group_transaction`` it means one connection, one + transaction and one insert for the whole group instead of N of each. When + ``enqueue_batch_size`` is set, each queue's messages are split into chunks of that size and + sent with one ``send_batch`` per chunk, so a huge fan-out does not build a single oversized + ``INSERT``; every chunk still runs inside the same transaction. + """ + visible_at = datetime.now(UTC) + timedelta(milliseconds=delay) if delay is not None else None + messages_by_queue: dict[str, list[Message]] = {} + for message in messages: + if message.queue_name not in self.queues: + raise QueueNotFound(message.queue_name) + messages_by_queue.setdefault(message.queue_name, []).append(message) + + for queue_name, queue_messages in messages_by_queue.items(): + chunk_size = self.enqueue_batch_size or len(queue_messages) + for chunk_start in range(0, len(queue_messages), chunk_size): + chunk = queue_messages[chunk_start : chunk_start + chunk_size] + payloads = [self._encode_message(message) for message in chunk] + self.client.send_batch( + queue_name, + payloads, + conn=self._current_connection, + delay=visible_at, + ) + + return messages + @override def flush(self, queue_name: str) -> None: """Remove every message currently stored in a queue.""" diff --git a/remoulade/brokers/rabbitmq.py b/remoulade/brokers/rabbitmq.py index 844c4a180..8348b8a2f 100644 --- a/remoulade/brokers/rabbitmq.py +++ b/remoulade/brokers/rabbitmq.py @@ -377,6 +377,10 @@ def _enqueue(self, message: "Message", *, delay: int | None = None) -> "Message" self.logger.info("Retrying enqueue due to closed connection. [%d/%d]", attempts, MAX_ENQUEUE_ATTEMPTS) + @override + def _enqueue_many(self, messages: list["Message[Any]"], *, delay: int | None = None) -> list["Message[Any]"]: + return [self._enqueue(message, delay=delay) for message in messages] + def get_queue_message_counts(self, queue_name: str): """Get the number of messages in a queue. This method is only meant to be used in unit and integration tests. diff --git a/remoulade/brokers/stub.py b/remoulade/brokers/stub.py index 7dd137f4f..dc85df2aa 100644 --- a/remoulade/brokers/stub.py +++ b/remoulade/brokers/stub.py @@ -102,6 +102,9 @@ def _enqueue(self, message, *, delay=None): self.queues[queue_name].put(message.encode_in_bytes()) return message + def _enqueue_many(self, messages, *, delay=None): + return [self._enqueue(message, delay=delay) for message in messages] + def flush(self, queue_name): """Drop all the messages from a queue. diff --git a/remoulade/composition.py b/remoulade/composition.py index 00b519b1b..ab7309f79 100644 --- a/remoulade/composition.py +++ b/remoulade/composition.py @@ -305,8 +305,7 @@ def run(self, *, delay: int | None = None, transaction: bool | None = None) -> S """ transaction = transaction if transaction is not None else self.broker.group_transaction with self.broker.tx() if transaction else nullcontext(): - for message in self.build(): - self.broker.enqueue(message, delay=delay) + self.broker.enqueue_many(self.build(), delay=delay) return self diff --git a/tests/test_broker.py b/tests/test_broker.py index d067be420..c6e77a0e9 100644 --- a/tests/test_broker.py +++ b/tests/test_broker.py @@ -104,6 +104,68 @@ class TestMiddleware(Middleware): assert isinstance(stub_broker.middleware[middleware_count], TestMiddleware) +class _EnqueueRecorder(Middleware): + def __init__(self): + self.events = [] + + def before_enqueue(self, broker, message, delay): + self.events.append(("before", message.args[0])) + + def after_enqueue(self, broker, message, delay, exception=None): + self.events.append(("after", message.args[0], exception is not None)) + + +def test_enqueue_many_groups_before_calls_then_after_calls(stub_broker): + @remoulade.actor + def do_work(index): + return index + + stub_broker.declare_actor(do_work) + recorder = _EnqueueRecorder() + stub_broker.add_middleware(recorder) + + messages = [do_work.message(0), do_work.message(1)] + stub_broker.enqueue_many(messages) + + # Every message goes through the enqueue middleware, but all before_enqueue calls happen before the + # (batched) write and all after_enqueue calls after it, instead of being interleaved per message. + assert recorder.events == [ + ("before", 0), + ("before", 1), + ("after", 0, False), + ("after", 1, False), + ] + + +def test_enqueue_many_emits_after_with_exception_only_for_started_messages(stub_broker): + @remoulade.actor + def do_work(index): + return index + + stub_broker.declare_actor(do_work) + recorder = _EnqueueRecorder() + stub_broker.add_middleware(recorder) + + boom = RuntimeError("backend down") + + def failing_enqueue_many(messages, *, delay=None): + raise boom + + stub_broker._enqueue_many = failing_enqueue_many + + messages = [do_work.message(0), do_work.message(1)] + with pytest.raises(RuntimeError): + stub_broker.enqueue_many(messages) + + # before ran for both, and after (with exception) is emitted for both so nothing leaks a dangling span/state + assert recorder.events == [ + ("before", 0), + ("before", 1), + ("after", 0, True), + ("after", 1, True), + ] + + def test_can_instantiate_brokers_without_middleware(): # Given that I have an empty list of middleware # When I pass that to the RMQ Broker diff --git a/tests/test_postgres.py b/tests/test_postgres.py index e6de40a36..5c0f75d47 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -3,7 +3,7 @@ import os import threading import time -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest from pydantic import BaseModel @@ -14,6 +14,7 @@ from remoulade import Message, QueueJoinTimeout, UnsupportedMessageEncoding, Worker, group from remoulade.brokers.postgres import PostgresBroker, _PostgresListener from remoulade.encoder import Encoder, MessageData +from remoulade.errors import QueueNotFound from remoulade.results import Results from remoulade.results.backends import StubBackend @@ -376,6 +377,77 @@ def decode_bytes(self, data: bytes) -> MessageData: remoulade.set_encoder(old_encoder) +def test_postgres_broker_enqueue_many_sends_one_send_batch_per_queue(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.client = Mock() + broker.queues = {"default": None, "other": None} + + messages = [ + Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}), + Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}), + Message(queue_name="other", actor_name="do_work", args=(3,), kwargs={}, options={}), + ] + + returned = broker._enqueue_many(messages) + + assert returned == messages + # one send_batch per queue, no per-message send + broker.client.send.assert_not_called() + assert broker.client.send_batch.call_count == 2 + calls_by_queue = {call.args[0]: call.args[1] for call in broker.client.send_batch.call_args_list} + assert [payload["args"] for payload in calls_by_queue["default"]] == [(1,), (2,)] + assert [payload["args"] for payload in calls_by_queue["other"]] == [(3,)] + + +def test_postgres_broker_enqueue_many_raises_for_unknown_queue(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.client = Mock() + broker.queues = {"default": None} + + with pytest.raises(QueueNotFound): + broker._enqueue_many([Message(queue_name="missing", actor_name="do_work", args=(), kwargs={}, options={})]) + + broker.client.send_batch.assert_not_called() + + +def test_postgres_broker_rejects_non_positive_enqueue_batch_size(): + with pytest.raises(ValueError, match="enqueue_batch_size must be greater than 0"): + PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enqueue_batch_size=0) + + +def test_postgres_broker_enqueue_many_chunks_send_batch_by_enqueue_batch_size(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enqueue_batch_size=2) + broker.client = Mock() + broker.queues = {"default": None} + + messages = [ + Message(queue_name="default", actor_name="do_work", args=(index,), kwargs={}, options={}) for index in range(5) + ] + + broker._enqueue_many(messages) + + # 5 messages, batch size 2 -> chunks of 2, 2, 1 + assert broker.client.send_batch.call_count == 3 + chunk_args = [[payload["args"] for payload in call.args[1]] for call in broker.client.send_batch.call_args_list] + assert chunk_args == [[(0,), (1,)], [(2,), (3,)], [(4,)]] + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_group_run_enqueues_every_message_in_a_single_batch(postgres_broker): + postgres_broker.declare_queue("default") + + messages = [Message(queue_name="default", actor_name="do_work", args=(i,), kwargs={}, options={}) for i in range(5)] + with ( + patch.object(postgres_broker.client, "send_batch", wraps=postgres_broker.client.send_batch) as spy_batch, + patch.object(postgres_broker.client, "send", wraps=postgres_broker.client.send) as spy_send, + ): + group(messages).run() + + assert _count_messages(postgres_broker) == 5 + spy_batch.assert_called_once() # a single INSERT for the whole group + spy_send.assert_not_called() + + @pytest.mark.usefixtures("postgres_broker") def test_postgres_broker_enqueue_stores_a_standard_remoulade_payload_as_jsonb(postgres_broker): message = Message(queue_name="default", actor_name="do_work", args=(1, 2), kwargs={"debug": True}, options={}) From 2197ef012d31ba410107a4e84f697aa0534c65ad Mon Sep 17 00:00:00 2001 From: mducros Date: Fri, 3 Jul 2026 17:35:50 +0200 Subject: [PATCH 44/44] fix(postgres): cap in-flight messages at prefetch like RabbitMQ QoS PGMQ has no server-side equivalent of RabbitMQ's basic_qos, so the consumer kept reading regardless of how many messages were unacked: with any backlog it drained the queue into worker memory, and the heartbeat renewed an ever-growing set of ids every 10s. The resulting giant set_vt updates bloated the queue table, serialized acks behind row locks and collapsed worker throughput until the process restarted. Enforce prefetch as the in-flight cap in the consumer: only read when unacked count is below prefetch, only request the free slots, and wait for an ack to free a slot instead of busy-looping when at capacity. Co-Authored-By: Claude Fable 5 fix(postgres-broker): add msg_id index on queue tables Time-partitioned PGMQ queue tables carry no index on msg_id, so every archive (ack/nack) and set_vt (heartbeat, requeue) seq scans all partitions and throughput degrades with backlog size. declare_queue now ensures a btree index on the partitioned parent, idempotently and also for pre-existing queues, so it propagates to existing and future partitions. Co-Authored-By: Claude Fable 5 --- remoulade/brokers/postgres.py | 86 ++++++++++++++++----- tests/test_postgres.py | 141 ++++++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+), 19 deletions(-) diff --git a/remoulade/brokers/postgres.py b/remoulade/brokers/postgres.py index 106faba37..dbeec7f06 100644 --- a/remoulade/brokers/postgres.py +++ b/remoulade/brokers/postgres.py @@ -202,6 +202,10 @@ def close(self) -> None: def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "_PostgresConsumer": """Create a consumer for a declared queue. + ``prefetch`` caps the number of in-flight (unacked) messages the + consumer may hold at any time — the same contract RabbitMQ enforces + server-side with ``basic_qos`` — not just the batch size per read. + Raises: QueueNotFound: If the queue has not been declared. """ @@ -211,7 +215,13 @@ def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> " @override def declare_queue(self, queue_name: str) -> None: - """Create a partitioned PGMQ queue if it does not already exist.""" + """Create a partitioned PGMQ queue if it does not already exist. + + Also ensures the queue table has a btree index on ``msg_id`` — even + for pre-existing queues, so queues created before remoulade added the + index pick it up on the next declaration. On a large existing queue + the initial index build locks the table for its duration. + """ if queue_name in self.queues: return with self.tx(): @@ -234,11 +244,29 @@ def declare_queue(self, queue_name: str) -> None: if self.enable_listen_notify: self._try_enable_notify(queue_name) + self._create_msg_id_index(queue_name, self._current_connection) + self.queues[queue_name] = None if not queue_exists: self.emit_after("declare_queue", queue_name) + def _create_msg_id_index(self, queue_name: str, connection: "Connection") -> None: + """Ensure the queue table has a btree index on ``msg_id``. + + PGMQ's time-partitioned queue tables ship without one, so every + ``archive`` (ack/nack) and ``set_vt`` (heartbeat, requeue) lookup seq + scans all partitions. Created on the partitioned parent, the index + propagates to existing and future partitions. ``msg_id`` never + changes, so the index does not defeat HOT updates of ``vt``/``read_ct``. + + The queue name must already be validated (``validate_queue_name``) + since it is interpolated as an identifier. + """ + connection.execute( + text(f'CREATE INDEX IF NOT EXISTS "q_{queue_name}_msg_id_idx" ON pgmq."q_{queue_name}" (msg_id)') + ) + def _encode_message(self, message: "Message") -> PostgresPayload: """Encode a Remoulade message into a JSON object payload for PGMQ. @@ -592,9 +620,10 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti Parameters: broker(PostgresBroker): Broker instance that owns the queue and database client. queue_name(str): Name of the declared queue to consume from. - prefetch(int): Maximum number of messages fetched per read call; must be greater than or equal to 1. - timeout(int): Idle wait timeout in milliseconds when polling for messages; must be greater than or equal to 0. - A value of 0 performs non-blocking reads. + prefetch(int): Maximum number of in-flight (unacked) messages the consumer may hold, and therefore + also the upper bound of messages fetched per read call; must be greater than or equal to 1. + timeout(int): Idle wait timeout in milliseconds when polling for messages or waiting for in-flight + capacity; must be greater than or equal to 0. A value of 0 performs non-blocking reads. """ if prefetch < 1: @@ -611,6 +640,7 @@ def __init__(self, broker: PostgresBroker, *, queue_name: str, prefetch: int, ti self.heartbeat_interval_seconds = broker.heartbeat_interval_seconds self.messages: deque[PostgresQueueMessage] = deque() self._notify_event = Event() + self._capacity_event = Event() self._heartbeat_stop = Event() self._heartbeat_thread: Thread | None = None self._heartbeat_message_ids_lock = Lock() @@ -638,47 +668,47 @@ def _normalize_messages( return [] return [messages] if not isinstance(messages, list) else messages - def _read_immediate(self) -> list[PostgresQueueMessage]: - """Read up to ``prefetch`` messages without polling.""" + def _read_immediate(self, qty: int) -> list[PostgresQueueMessage]: + """Read up to ``qty`` messages without polling.""" return self._normalize_messages( self.client.read( self.queue_name, vt=self.visibility_timeout_seconds, - qty=self.prefetch, + qty=qty, ) ) - def _read_with_poll(self) -> list[PostgresQueueMessage]: - """Read up to ``prefetch`` messages using PGMQ polling. + def _read_with_poll(self, qty: int) -> list[PostgresQueueMessage]: + """Read up to ``qty`` messages using PGMQ polling. A zero timeout means non-blocking: do a single immediate read instead of polling for the rounded-up one-second minimum. """ if self.timeout == 0: - return self._read_immediate() + return self._read_immediate(qty) return self._normalize_messages( self.client.read_with_poll( self.queue_name, vt=self.visibility_timeout_seconds, - qty=self.prefetch, + qty=qty, max_poll_seconds=max((self.timeout + 999) // 1000, 1), poll_interval_ms=max(min(self.timeout, 1000), 1), ) ) - def _read_next_batch(self) -> list[PostgresQueueMessage]: - """Read the next batch, favoring LISTEN/NOTIFY when available.""" + def _read_next_batch(self, qty: int) -> list[PostgresQueueMessage]: + """Read the next batch of at most ``qty`` messages, favoring LISTEN/NOTIFY when available.""" if not self._listener_available: - return self._read_with_poll() + return self._read_with_poll(qty) - messages = self._read_immediate() + messages = self._read_immediate(qty) if messages: self._notify_event.clear() return messages self._notify_event.wait(self.wait_timeout_seconds) self._notify_event.clear() - return self._read_immediate() if self._listener_available else self._read_with_poll() + return self._read_immediate(qty) if self._listener_available else self._read_with_poll(qty) def _start_heartbeat(self) -> None: """Start the background visibility-timeout renewal thread.""" @@ -717,10 +747,16 @@ def _run_heartbeat(self) -> None: exc_info=True, ) + def _pending_ack_count(self) -> int: + """Count messages read but not yet acked, nacked or requeued.""" + with self._heartbeat_message_ids_lock: + return len(self._heartbeat_message_ids) + def _unregister_heartbeat_message_id(self, message_id: int) -> None: - """Stop tracking a message for heartbeat renewal.""" + """Stop tracking a message for heartbeat renewal and free an in-flight slot.""" with self._heartbeat_message_ids_lock: self._heartbeat_message_ids.discard(message_id) + self._capacity_event.set() def _requeue_message_ids(self, message_ids: list[int]) -> None: """Make a batch of message ids visible again immediately.""" @@ -775,11 +811,22 @@ def _build_message(self, postgres_message: PostgresQueueMessage) -> "_PostgresMe @override def __next__(self) -> "_PostgresMessage | None": - """Return the next available message, or ``None`` if the queue stays empty.""" + """Return the next available message, or ``None`` if the queue stays empty + or the consumer is at its in-flight capacity.""" if self.messages: return self._build_message(self.messages.popleft()) - messages = self._read_next_batch() + # PGMQ has no server-side equivalent of RabbitMQ's basic_qos, so the + # in-flight cap must be enforced here: reading regardless of unacked + # count would drain the whole queue into worker memory and make the + # heartbeat renew an ever-growing set of ids. + capacity = self.prefetch - self._pending_ack_count() + if capacity <= 0: + self._capacity_event.wait(self.wait_timeout_seconds) + self._capacity_event.clear() + return None + + messages = self._read_next_batch(capacity) if not messages: return None @@ -794,6 +841,7 @@ def close(self) -> None: """Stop the heartbeat, unregister from the shared listener and requeue buffered messages.""" self._heartbeat_stop.set() self._notify_event.set() + self._capacity_event.set() if self.broker._listener is not None: self.broker._listener.unregister(self.queue_name, self._notify_event) if self._heartbeat_thread is not None: diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 5c0f75d47..440d2a573 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -190,6 +190,7 @@ def test_postgres_broker_does_not_fail_when_enable_notify_raises(caplog): broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock(side_effect=RuntimeError("notify unavailable")) broker._queue_exists = Mock(return_value=False) + broker._create_msg_id_index = Mock() broker.declare_queue("default") @@ -203,20 +204,41 @@ def test_postgres_broker_declare_queue_is_idempotent_when_queue_already_exists() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=True) + broker._create_msg_id_index = Mock() broker.declare_queue("default") broker.client.create_partitioned_queue.assert_not_called() broker.client.enable_notify.assert_not_called() + # The index is still ensured on pre-existing queues, so queues created + # before remoulade added it pick it up on the next declaration. + broker._create_msg_id_index.assert_called_once() assert "default" in broker.queues +def test_postgres_broker_declare_queue_creates_msg_id_index(): + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.client.validate_queue_name = Mock() + broker.client.create_partitioned_queue = Mock() + broker.client.enable_notify = Mock() + broker._queue_exists = Mock(return_value=False) + + conn = Mock() + broker.client.engine.begin = Mock(return_value=_StubTransaction(conn)) + + broker.declare_queue("default") + + executed = [str(call.args[0]) for call in conn.execute.call_args_list] + assert 'CREATE INDEX IF NOT EXISTS "q_default_msg_id_idx" ON pgmq."q_default" (msg_id)' in executed + + def test_postgres_broker_poll_only_mode_opens_no_listener_and_skips_enable_notify(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[], enable_listen_notify=False) broker.client.validate_queue_name = Mock() broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + broker._create_msg_id_index = Mock() broker.declare_queue("default") @@ -652,6 +674,82 @@ def _fake_start_heartbeat(self): broker.client.set_vt.assert_called_once_with("default", [2], 0) +def test_postgres_consumer_caps_in_flight_messages_at_prefetch(monkeypatch): + def _fake_start_heartbeat(self): + self._heartbeat_thread = None + + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) + + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["default"] = None + _install_listener(broker, available=False) + + first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock( + return_value=[ + Mock(msg_id=1, message=_expected_payload(first_message)), + Mock(msg_id=2, message=_expected_payload(second_message)), + ] + ) + broker.client.set_vt = Mock() + + consumer = broker.consume("default", prefetch=2, timeout=10) + + assert next(consumer) is not None + assert next(consumer) is not None + + # Both messages are unacked: the consumer is at capacity and must not + # read again until a slot is freed. + assert next(consumer) is None + broker.client.read_with_poll.assert_called_once() + + consumer.close() + + +def test_postgres_consumer_resumes_reading_when_an_ack_frees_capacity(monkeypatch): + def _fake_start_heartbeat(self): + self._heartbeat_thread = None + + monkeypatch.setattr("remoulade.brokers.postgres._PostgresConsumer._start_heartbeat", _fake_start_heartbeat) + + broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) + broker.queues["default"] = None + _install_listener(broker, available=False) + + first_message = Message(queue_name="default", actor_name="do_work", args=(1,), kwargs={}, options={}) + second_message = Message(queue_name="default", actor_name="do_work", args=(2,), kwargs={}, options={}) + broker.client.read = Mock(return_value=None) + broker.client.read_with_poll = Mock( + side_effect=[ + [ + Mock(msg_id=1, message=_expected_payload(first_message)), + Mock(msg_id=2, message=_expected_payload(second_message)), + ], + [], + ] + ) + broker.client.set_vt = Mock() + broker.client.archive = Mock() + + consumer = broker.consume("default", prefetch=2, timeout=10) + + first = next(consumer) + assert first is not None + assert next(consumer) is not None + assert next(consumer) is None # at capacity, no read issued + + consumer.ack(first) + + assert next(consumer) is None # one slot free again: reads, queue is empty + assert broker.client.read_with_poll.call_count == 2 + # Only the freed slot may be requested, keeping in-flight <= prefetch. + assert broker.client.read_with_poll.call_args.kwargs["qty"] == 1 + + consumer.close() + + def test_postgres_consumer_falls_back_to_polling_when_listener_stops_during_wait(): broker = PostgresBroker(url=TEST_POSTGRES_URL, middleware=[]) broker.queues["default"] = None @@ -730,6 +828,7 @@ class InputSchema(BaseModel): broker.client.create_partitioned_queue = Mock() broker.client.enable_notify = Mock() broker._queue_exists = Mock(return_value=False) + broker._create_msg_id_index = Mock() _install_listener(broker, available=False) @remoulade.actor(actor_name="typed.actor", queue_name="default") @@ -814,6 +913,48 @@ def test_postgres_consumer_requeue_restores_visibility_with_set_vt(postgres_brok assert _count_messages(postgres_broker) == 0 +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_declare_queue_creates_msg_id_index_on_all_partitions(postgres_broker): + postgres_broker.declare_queue("default") + + with postgres_broker.client.session() as session: + tables = ( + session.execute( + text( + "SELECT tablename FROM pg_indexes " + "WHERE schemaname = 'pgmq' AND indexdef LIKE '%USING btree (msg_id)' " + "AND (tablename = 'q_default' OR tablename LIKE 'q\\_default\\_p%')" + ) + ) + .scalars() + .all() + ) + + assert "q_default" in tables # the partitioned parent + assert len(tables) > 1 # propagated to at least one partition + + +@pytest.mark.usefixtures("postgres_broker") +def test_postgres_broker_declare_queue_restores_missing_msg_id_index(postgres_broker): + postgres_broker.declare_queue("default") + + # Simulate a queue created before remoulade started ensuring the index. + with postgres_broker.client.engine.begin() as connection: + connection.execute(text('DROP INDEX pgmq."q_default_msg_id_idx"')) + + postgres_broker.queues.pop("default") + postgres_broker.declare_queue("default") + + with postgres_broker.client.session() as session: + index_exists = session.execute( + text( + "SELECT EXISTS(SELECT 1 FROM pg_indexes " + "WHERE schemaname = 'pgmq' AND indexname = 'q_default_msg_id_idx')" + ) + ).scalar_one() + assert index_exists + + @pytest.mark.usefixtures("postgres_broker") def test_postgres_worker_processes_native_delayed_messages_without_delay_queue(postgres_broker): seen = []