feat(state): add PostgresBackend, a state backend stored in the pgmq …

docs/source/changelog.rst

@@ -5,6 +5,43 @@ Changelog
 
 All notable changes to this project will be documented in this file.
 
+Unreleased
+----------
+Upgrading
+^^^^^^^^^
+* **PGMQ queues created before this version must be backfilled with their new indexes.** The PGMQ broker now
+  creates its indexes together with the queue, and declaring an already existing queue deliberately does
+  nothing. An existing queue therefore keeps only the ``msg_id`` index it was created with, and every state
+  lookup on it seq scans all of its partitions — the same failure mode that made throughput collapse when the
+  ``msg_id`` index itself was missing. Nothing raises, so this degrades silently. Run once, after upgrading::
+
+    for queue in broker.get_declared_queues():
+        broker.client.create_indexes(queue)
+
+  On a large queue the index build locks the table for its duration, so plan it like any other index creation.
+
+Feat
+^^^^
+* Add ``PostgresBackend``, a state backend for the PostgreSQL/PGMQ broker that stores state in the message itself
+  instead of a dedicated table. ``Pending``/``Started``, the timestamps and the attempt count are derived from
+  PGMQ's own columns, and the terminal status is merged into the archive the broker already writes on ack/nack,
+  so tracking a message's lifecycle costs no additional statement. Unlike the Redis and stub state backends,
+  ``get_states``/``get_states_count`` honour their filters, sorting and pagination in SQL. ``Message.set_progress``
+  is not supported by this backend and raises: storing a progress would mean an ``UPDATE`` on the broker's
+  queue table per call.
+* The PGMQ broker now also indexes ``(message->>'message_id')`` on queue and archive tables, and
+  ``(headers->>'status')`` on archive tables, so states can be looked up and failures found without seq scans.
+  Index creation moved into ``RemouladePostgresClient.create_partitioned_queue``, so a queue gets them the moment
+  it is created — see *Upgrading* for existing queues.
+
+Changed
+^^^^^^^
+* ``StateBackend.set_state`` accepts an optional ``message`` keyword argument, letting a backend persist state
+  through the broker's own writes rather than a statement of its own. Backends that do not need it ignore it,
+  but a state backend defined outside remoulade must accept it: the middleware now always passes it, so an
+  override still declared as ``set_state(self, state, ttl=3600)`` raises ``TypeError``.
+* Move the PGMQ broker's hand-written SQL into ``RemouladePostgresClient`` (``remoulade.helpers.postgres_client``).
+
 `7.0.0`_ -- 2026-06-15
 ------------
 Breaking changes

docs/source/guide.rst

@@ -449,6 +449,79 @@ table, which is partitioned the same way. Two broker parameters control this:
    a self-managed or hosted PostgreSQL you must enable it yourself.
 
 
+Postgres State Backend
+^^^^^^^^^^^^^^^^^^^^^^
+
+With the Postgres broker, message state can be tracked without a table of its
+own: a message's state *is* its PGMQ row. ``PostgresBackend`` reads most of a state
+straight off PGMQ's own columns and the message payload, and stores only what
+they cannot express — the terminal status — in the message's ``headers``
+column::
+
+  import remoulade
+
+  from remoulade.brokers.postgres import PostgresBroker
+  from remoulade.state import MessageState
+  from remoulade.state.backends import PostgresBackend
+
+  broker = PostgresBroker(url="postgresql://remoulade@localhost:5432/remoulade")
+  broker.add_middleware(MessageState(PostgresBackend(broker)))
+  remoulade.set_broker(broker)
+
+``Pending`` and ``Started``, the enqueued/started/end timestamps and the attempt
+count are *derived* at read time from ``read_ct``, ``last_read_at``,
+``enqueued_at``, ``archived_at`` and from whether the row still sits in the
+queue table. Nothing is written for them. The terminal status is handed to the
+broker, which folds it into the archive it performs on ack or nack anyway, so
+tracking a message's whole lifecycle costs **no additional statement**.
+
+``Message.set_progress`` is **not supported** by this backend and raises. It is
+called while the actor runs and has to be visible before the message finishes, so
+it cannot ride along with the ack: every call would be an ``UPDATE`` on the
+broker's queue table, which is its throughput-critical one. Report the progress
+of long-running work through your metrics instead.
+
+Because a retried message keeps its ``message_id`` and is re-enqueued as a new
+row, the archive keeps one row per attempt — a ready-made audit trail::
+
+  SELECT message->>'message_id', message->>'actor_name', read_ct, archived_at
+  FROM pgmq.a_default
+  WHERE headers->>'status' = 'Failure'
+    AND archived_at > now() - interval '1 day';
+
+Reads always resolve to the current attempt, so ``get_state`` and the dashboard
+filters see a message's present status, not a past failure.
+
+Two differences from the Redis state backend are worth planning for:
+
+* **Retention is the archive's.** ``state_ttl`` no longer bounds how long a
+  state is kept; ``archive_retention_interval_in_days`` does, through
+  ``pg_partman``. A state disappears when its message's archive partition is
+  dropped, and purging or dropping a queue destroys its states with it.
+* **A state cannot outlive its message.** The backend stores nothing of its own,
+  so ``set_state`` for a message that was never enqueued is a no-op.
+* **No progress.** ``Message.set_progress`` raises with this backend, and
+  ``State.progress`` is always ``None``.
+
+In exchange, ``get_states`` and ``get_states_count`` actually honour their
+filters, sorting and pagination — the Redis and stub backends ignore them and
+paginate in Python after reading everything.
+
+The indexes those lookups need are created with the queue itself, including one
+on ``(message->>'message_id')`` on the queue table. That index is on the broker's
+hot insert path, so it is worth measuring on a high-throughput queue.
+
+.. note::
+
+   Only queues *created* by this version get those indexes: declaring a queue
+   that already exists does nothing. A queue created by an earlier version of
+   remoulade will keep seq scanning every partition on state lookups until you
+   backfill it once::
+
+     for queue in broker.get_declared_queues():
+         broker.client.create_indexes(queue)
+
+
 Local Broker
 ^^^^^^^^^^^^^^^
 

remoulade/brokers/postgres.py

@@ -26,13 +26,13 @@
 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 Connection, text
 
 from ..broker import Broker, Consumer, MessageProxy
 from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding
+from ..helpers.postgres_client import RemouladePostgresClient
 from ..message import Message
 
 if TYPE_CHECKING:
@@ -129,7 +129,7 @@ def __init__(
         self.enable_listen_notify = enable_listen_notify
         self.enqueue_batch_size = enqueue_batch_size
 
-        self.client = SQLAlchemyPGMQueue(
+        self.client = RemouladePostgresClient(
             conn_string=url,
             init_extension=False,
             vt=self.visibility_timeout_seconds,
@@ -217,10 +217,11 @@ def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> "
     def declare_queue(self, queue_name: str) -> None:
         """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.
+        Creating the queue brings the indexes remoulade needs with it, through
+        :meth:`RemouladePostgresClient.create_partitioned_queue`. Nothing is done for
+        a queue that already exists, so a queue created by a version of remoulade
+        that did not yet declare one of those indexes will not gain it here; call
+        :meth:`RemouladePostgresClient.create_indexes` once to backfill it.
         """
         if queue_name in self.queues:
             return
@@ -244,29 +245,11 @@ 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.
 
@@ -768,17 +751,24 @@ def _requeue_message_ids(self, message_ids: list[int]) -> None:
     def _archive_message(self, message: "MessageProxy") -> None:
         """Stop tracking a message and archive it, tolerating transient failures.
 
+        Any header metadata a middleware staged on the message (the state backend
+        recording its outcome, typically) is handed to the client, which merges it
+        as part of the archive rather than by a statement of its own — so
+        recording that a message succeeded or failed costs nothing on top of the
+        ack.
+
         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.
+        which is the broker's at-least-once guarantee. A staged patch is lost
+        with it, and is rebuilt when the message is processed again.
         """
         if not isinstance(message, _PostgresMessage):
             raise ValueError("It must be a PostgresMessage")
         self._unregister_heartbeat_message_id(message._postgres_message.msg_id)
         try:
-            self.client.archive(self.queue_name, message._postgres_message.msg_id)
+            self.client.archive(self.queue_name, message._postgres_message.msg_id, headers=message._header_patch)
         except Exception:
             self.broker.logger.error(
                 "Failed to archive message %s on queue %s; it will be redelivered after its visibility timeout.",
@@ -866,3 +856,22 @@ def __init__(self, postgres_message: PostgresQueueMessage) -> None:
             raise UnsupportedMessageEncoding("eta option isn't supported with postgres broker")
         super().__init__(message)
         self._postgres_message = postgres_message
+        self._header_patch: dict[str, Any] = {}
+
+    def stage_headers(self, patch: dict[str, Any]) -> bool:
+        """Merge a jsonb patch into the headers this message will be archived with.
+
+        Lets a state backend record a message's outcome without a write of its
+        own: the patch rides along with the archive that ack/nack performs
+        anyway.
+
+        Purely in memory: the patch is flushed by ``ack``/``nack``, folded into
+        the archive statement. Successive calls merge, so a progress update
+        followed by a terminal status both land.
+
+        Returns:
+          bool: Always True. The return value exists so a caller holding a proxy
+          of an unknown type can tell whether it still owes a write of its own.
+        """
+        self._header_patch.update(patch)
+        return True