acemq_amqp.patterns.replay

  1# Copyright 2026 AceMQ.
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     https://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""Putting dead letters back, when the thing that killed them is fixed.
 16
 17This is the pattern somebody actually needs at three in the morning: a
 18dead-letter queue has two thousand messages in it, the bug is fixed, and they
 19have to go back through — but not all of them, not silently, and not in a way
 20that cannot be stopped half way.
 21
 22So a replay takes a filter, a limit and a deadline, and reports what it did and
 23why it stopped. "Moved 500" means something quite different when the limit was
 24500, which is why the reason comes back with the count.
 25"""
 26
 27from __future__ import annotations
 28
 29import asyncio
 30import logging
 31from collections.abc import Callable
 32from dataclasses import dataclass
 33from datetime import datetime, timedelta, timezone
 34from typing import Any, TypeAlias
 35
 36from ..connection import Connection
 37from ..envelope import Envelope
 38from ..errors import AceMQError, PublishError
 39from ..transport import Delivery, Outbound
 40
 41log = logging.getLogger("acemq")
 42
 43#: The queue a message was replayed out of.
 44HEADER_REPLAYED_FROM = "acemq-replayed-from"
 45
 46#: When it was replayed, as RFC 3339.
 47HEADER_REPLAYED_AT = "acemq-replayed-at"
 48
 49#: How many times it has been replayed, from 1.
 50HEADER_REPLAY_COUNT = "acemq-replay-count"
 51
 52#: Decides whether a message goes back. Returning ``False`` leaves it where it
 53#: is, which is what makes a replay something that can be done in stages.
 54ReplayFilter: TypeAlias = Callable[[Envelope, bytes], bool]
 55
 56
 57@dataclass(frozen=True, slots=True)
 58class ReplayResult:
 59    """What a replay did, and why it stopped.
 60
 61    :param moved: how many messages were republished
 62    :param skipped: how many the filter declined, and so left where they were
 63    :param reason: ``drained``, ``limit`` or ``deadline``
 64    """
 65
 66    moved: int
 67    skipped: int
 68    reason: str
 69
 70    def __str__(self) -> str:
 71        return f"moved {self.moved}, skipped {self.skipped} ({self.reason})"
 72
 73
 74class ReplayError(AceMQError):
 75    """A replay stopped because something failed.
 76
 77    It carries the :class:`ReplayResult` it had reached, because "it failed" and
 78    "it failed after moving four hundred" are different things to be told at
 79    three in the morning.
 80    """
 81
 82    def __init__(self, result: ReplayResult, what: str) -> None:
 83        self.result = result
 84        super().__init__(f"acemq: {what} (after {result})")
 85
 86
 87async def replay(
 88    connection: Connection,
 89    queue: str,
 90    *,
 91    exchange: str = "",
 92    routing_key: str | None = None,
 93    limit: int = 0,
 94    deadline: timedelta | None = None,
 95    only: ReplayFilter | None = None,
 96    restart: bool = True,
 97) -> ReplayResult:
 98    """Moves messages from a queue back onto an exchange::
 99
100        result = await replay(
101            mq,
102            dead_letter_queue("orders.new"),
103            limit=500,
104            only=lambda envelope, body: "timeout" in envelope.error,
105        )
106        print(result)   # moved 137, skipped 363 (drained)
107
108    Each message is stamped with :data:`HEADER_REPLAYED_FROM`,
109    :data:`HEADER_REPLAYED_AT` and :data:`HEADER_REPLAY_COUNT`, so a consumer
110    that needs to treat a replayed message differently can, and one that does
111    not is unaffected.
112
113    :param connection: where to read from and publish to
114    :param queue: where the messages are now, usually a dead-letter queue
115    :param exchange: where they go back to, empty to publish to a queue by name
116    :param routing_key: what to publish them under. ``None`` keeps each
117        message's own, so a message goes back where it came from rather than
118        everywhere
119    :param limit: stop after this many. Zero means no limit, which against a
120        queue something is still writing to may mean it does not stop at all —
121        prefer a limit
122    :param deadline: stop after this long. ``None`` means no deadline
123    :param only: which messages to move. ``None`` moves all of them
124    :param restart: give each message a fresh set of attempts, clearing the
125        counter and the error that dead-lettered it. On by default because it is
126        what a replay is for: a message that arrives back on attempt five of a
127        five-attempt policy is dead-lettered again before a handler sees it, and
128        the operator who has just fixed the bug has moved two thousand messages
129        from one queue to the same queue. Turn it off to put back exactly what
130        was there — for an audit, or for a queue read by something that counts
131        attempts itself
132    :returns: what it did and why it stopped
133    :raises ReplayError: when reading or republishing failed. Whatever had
134        already been moved has been moved; the rest is still on the queue
135    """
136    if not queue:
137        raise ValueError("acemq: a replay needs a queue to read from")
138
139    clock = asyncio.get_running_loop()
140    expires = None if deadline is None else clock.time() + deadline.total_seconds()
141    moved = 0
142    skipped = 0
143
144    # Messages the filter declines are held unsettled until the pass is over
145    # rather than returned one at a time. Returning one immediately does not
146    # work: the broker puts it back where it was, at the head of the queue, so
147    # the next pull hands over the same message and nothing behind it is ever
148    # looked at. Holding them takes them out of the way for the length of the
149    # pass, and because the broker is the one holding them, a crash returns them
150    # rather than losing them.
151    declined: list[Delivery] = []
152
153    try:
154        while True:
155            if limit > 0 and moved >= limit:
156                return ReplayResult(moved, skipped, "limit")
157            if expires is not None and clock.time() >= expires:
158                return ReplayResult(moved, skipped, "deadline")
159
160            try:
161                message = await connection.pull(queue)
162            except Exception as failure:
163                raise ReplayError(
164                    ReplayResult(moved, skipped, "failed"),
165                    f"cannot read {queue!r} to replay it: {failure}",
166                ) from failure
167
168            if message is None:
169                # Everything available has been offered. What is left is what
170                # this pass declined, and it goes back as the pass ends.
171                return ReplayResult(moved, skipped, "drained")
172
173            envelope = Envelope.from_headers(message.headers, message.routing_key)
174            if only is not None and not only(envelope, message.body):
175                skipped += 1
176                declined.append(message)
177                continue
178
179            try:
180                await _move(
181                    connection, message, envelope, queue, exchange, routing_key, restart
182                )
183            except Exception as failure:
184                raise ReplayError(
185                    ReplayResult(moved, skipped, "failed"),
186                    f"cannot republish a message from {queue!r}: {failure}",
187                ) from failure
188            moved += 1
189    finally:
190        for held in declined:
191            # Nothing to report if this fails: the pass is over, and a message
192            # that cannot be returned is one the broker returns itself when this
193            # connection goes.
194            await _return_quietly(held)
195
196
197async def _move(
198    connection: Connection,
199    message: Delivery,
200    envelope: Envelope,
201    queue: str,
202    exchange: str,
203    routing_key: str | None,
204    restart: bool,
205) -> None:
206    """Republishes one message and then acknowledges the original.
207
208    That order, and not the other one. Acknowledging first would lose the
209    message whenever the publish failed; this way a failure in the gap replays
210    it twice, which is the right way round for a queue somebody is putting back
211    by hand.
212    """
213    going_back = envelope.with_(attempt=1, error="") if restart else envelope
214    key = message.routing_key if routing_key is None else routing_key
215
216    headers: dict[str, Any] = dict(going_back.to_headers(routing_key=key))
217    headers[HEADER_REPLAYED_FROM] = queue
218    headers[HEADER_REPLAYED_AT] = datetime.now(timezone.utc).isoformat()
219    headers[HEADER_REPLAY_COUNT] = _replays_so_far(envelope) + 1
220
221    try:
222        # Mandatory, because the original is acknowledged as soon as this
223        # returns: a destination nothing is bound to is confirmed and dropped by
224        # the broker, and without the return the replay would delete the last
225        # copy of a message it was asked to recover.
226        result = await connection.publish_raw(
227            exchange,
228            key,
229            Outbound(
230                body=message.body,
231                content_type=message.content_type or "",
232                message_id=going_back.id,
233                headers=headers,
234                persistent=True,
235                mandatory=True,
236            ),
237        )
238        if not result.routed:
239            raise PublishError(
240                going_back.id,
241                exchange,
242                key,
243                result.return_reason or "no queue is bound to receive it",
244                unroutable=True,
245            )
246    except Exception:
247        # Returned rather than dropped, and the replay stops. A replay that
248        # loses messages is worse than one that stops early.
249        await _return_quietly(message)
250        raise
251
252    await message.ack()
253
254
255async def _return_quietly(message: Delivery) -> None:
256    """Puts a message back without letting that failure hide the first one.
257
258    Logged rather than raised: by the time this is called the replay is already
259    stopping, and the message is one the broker returns itself when this
260    connection goes.
261    """
262    try:
263        await message.nack(True)
264    except Exception:
265        log.warning("acemq: could not return a message to the queue it was replayed from")
266
267
268def _replays_so_far(envelope: Envelope) -> int:
269    """How many times this message has been replayed already.
270
271    Read leniently, because the header may have been written by another
272    language's client or by a broker that renders integers as strings.
273    """
274    written = envelope.headers.get(HEADER_REPLAY_COUNT)
275    if isinstance(written, bool):
276        return 0
277    if isinstance(written, int):
278        return written
279    try:
280        return int(str(written))
281    except (TypeError, ValueError):
282        return 0