acemq_amqp.patterns

53class Claim(Enum):
54    """What a store found when it was asked to claim a message."""
55
56    #: Nobody had it, or the last holder's lease had run out: it is ours to run.
57    CLAIMED = "claimed"
58    #: It was handled and confirmed. Accept it without running anything.
59    DUPLICATE = "duplicate"
60    #: Somebody holds a live claim and has not confirmed. Neither run it nor
61    #: accept it: put it back and look again later.
62    IN_PROGRESS = "in_progress"

303class ClaimCheckCodec:
304    """Puts a large payload aside and sends the key instead::
305
306        from acemq_amqp.codec import JsonCodec
307        from acemq_amqp.patterns import ClaimCheckCodec, FilesystemClaimCheckStore
308
309        store = FilesystemClaimCheckStore("/mnt/payloads")
310        codec = ClaimCheckCodec(JsonCodec(), store)
311
312        mq = await Connection.open(url, codec=codec)
313
314    A payload smaller than ``threshold`` goes on the wire inline and a larger
315    one goes to the store, and the three bytes at the front of the body say
316    which. A body this codec did not write is handed to the delegate untouched,
317    which is what makes it safe to introduce on a queue that already has
318    messages in it.
319
320    :param delegate: the codec that turns payloads into bytes. What gets stored
321        or inlined is its output
322    :param store: where large payloads go
323    :param threshold: payloads of at least this many bytes are offloaded. Zero
324        offloads everything, which is occasionally what a store-backed audit
325        trail wants
326    """
327
328    def __init__(
329        self,
330        delegate: Codec,
331        store: ClaimCheckStore,
332        *,
333        threshold: int = DEFAULT_THRESHOLD,
334    ) -> None:
335        if threshold < 0:
336            raise ValueError(
337                f"acemq: a claim-check threshold cannot be negative, got {threshold}"
338            )
339        self._delegate = delegate
340        self._store = store
341        self._threshold = threshold
342
343    @property
344    def delegate(self) -> Codec:
345        """The wrapped codec, whose output is what gets stored or inlined."""
346        return self._delegate
347
348    @property
349    def store(self) -> ClaimCheckStore:
350        """Where large payloads go."""
351        return self._store
352
353    @property
354    def threshold(self) -> int:
355        """Payloads of at least this many bytes are offloaded."""
356        return self._threshold
357
358    @property
359    def content_type(self) -> str:
360        """The delegate's. A claim-checked document is still a document."""
361        return self._delegate.content_type
362
363    def encode(self, payload: Any) -> bytes:
364        encoded = self._delegate.encode(payload)
365        if len(encoded) < self._threshold:
366            return _frame(_INLINE, encoded)
367        return _frame(_CHECKED, self._store.put(encoded).encode("utf-8"))
368
369    def decode(self, body: bytes, content_type: str | None = None) -> Any:
370        if not _is_framed(body):
371            # Written before this codec was introduced, or by a publisher that
372            # does not use it. Reading it as the delegate would is the only
373            # useful answer, and it is what makes adding a claim check to a live
374            # queue safe.
375            return self._delegate.decode(body, content_type)
376
377        rest = body[_HEADER:]
378        if body[2] == _INLINE:
379            return self._delegate.decode(rest, content_type)
380
381        key = rest.decode("utf-8")
382        content = self._store.get(key)
383        if content is None:
384            # Fatal, not retryable: the payload is not coming back, and a
385            # message that keeps being redelivered for it only holds a queue
386            # open until it ages out.
387            raise FatalError(
388                f"acemq: the claim check {key!r} is not in the store, so this message "
389                "cannot be read. The payload was removed while a message referring to it "
390                "was still deliverable — the store's retention has to outlast every queue, "
391                "every dead-letter queue, and any replay somebody might do by hand."
392            )
393        return self._delegate.decode(content, content_type)
394
395    def can_decode(self, content_type: str | None) -> bool:
396        """Whatever the delegate accepts. A claim check does not change what the
397        message is."""
398        return self._delegate.can_decode(content_type)
399
400    def __repr__(self) -> str:
401        return f"ClaimCheckCodec({self._delegate!r}, above={self._threshold} bytes)"

 88@runtime_checkable
 89class ClaimCheckStore(Protocol):
 90    """Where a payload too large for a broker actually goes.
 91
 92    Three methods, so a store backed by S3, Azure Blob Storage, a filesystem or
 93    a database table is a small class. Nothing here knows about messaging: the
 94    store holds bytes under a key and hands them back, and the codec is what
 95    turns that into a claim check on the wire.
 96
 97    Synchronous, unlike the idempotency and outbox seams, because
 98    :class:`~acemq_amqp.codec.Codec` is synchronous — a codec is called from
 99    inside the publish and consume paths of both the async and the blocking API,
100    and a store that had to be awaited could not serve the second. A store
101    talking to a network should hold a small pool and block on it.
102
103    **Retention is the part that goes wrong.** The store and the queue have
104    different lifetimes and nothing enforces a relationship between them. A
105    message replayed a month later carries a key, and if the store expired that
106    key the replay produces a message nobody can read — worse than a lost
107    message, because it looks like a message and fails deep inside a consumer
108    rather than visibly. So the store's retention must exceed every retention
109    that could bring a message back: queue TTLs, dead-letter queues, and however
110    long somebody might sit on a message before replaying it by hand. When in
111    doubt, longer.
112
113    Implementations must be safe to use from several threads, because one is
114    shared by every publisher and consumer on a connection.
115    """
116
117    def put(self, content: bytes) -> str:
118        """Stores a payload.
119
120        :param content: the bytes
121        :returns: the key the message will carry, which must be unique for the
122            life of the store
123        """
124
125    def get(self, key: str) -> bytes | None:
126        """Redeems a claim check.
127
128        :param key: what the message carried
129        :returns: the payload, or ``None`` when the store no longer holds it —
130            which is retention having expired underneath a message that outlived
131            it
132        """
133
134    def delete(self, key: str) -> None:
135        """Removes a payload.
136
137        Not called by the codec. Deleting on read would break the second
138        consumer of the same message, and deleting on acknowledgement would
139        break a replay — so when a payload may be removed is a retention
140        decision, and retention decisions belong to whoever owns the data.
141
142        :param key: what to remove
143        """

 41class ConsumerGroup:
 42    """Consumers reading one queue, closed as one thing.
 43
 44    Built by :meth:`start`, because a group that has not started is a group
 45    holding nothing::
 46
 47        async with await ConsumerGroup.start(mq, "orders", 4, handle):
 48            ...
 49    """
 50
 51    def __init__(self, queue: str, consumers: list[Consumer]) -> None:
 52        self._queue = queue
 53        self._consumers = consumers
 54
 55    @classmethod
 56    async def start(
 57        cls,
 58        connection: Connection,
 59        queue: str,
 60        size: int,
 61        handler: Handler,
 62        *,
 63        codec: Codec | None = None,
 64        retry: RetryPolicy | None = None,
 65        prefetch: int | None = None,
 66        concurrency: int = 1,
 67        tag: str = "",
 68        args: Any = None,
 69        declare: bool = True,
 70    ) -> ConsumerGroup:
 71        """Starts ``size`` consumers over one queue.
 72
 73        If one fails to start, the ones already running are closed before the
 74        failure is raised. A half-started group is worse than none: it holds
 75        messages that nothing is going to finish handling, and the caller that
 76        saw the exception has no handle to close it with.
 77
 78        :param connection: whose consumers these are
 79        :param queue: what they read
 80        :param size: how many
 81        :param handler: what each does with a message
 82        :param codec: a codec other than the connection's
 83        :param retry: a policy other than the connection's
 84        :param prefetch: how many unacknowledged messages each holds
 85        :param concurrency: how many messages each works on at once
 86        :param tag: what to call them to the broker, numbered from 1. Each gets
 87            its own name so the management interface shows which consumer is
 88            holding a message rather than four identical rows
 89        :param args: broker-specific consumer arguments
 90        :param declare: declare the queues a failed message goes to before
 91            subscribing. On by default; see
 92            :meth:`acemq_amqp.Connection.consume`. Every consumer in the group
 93            declares the same things, which is a few extra round trips at
 94            start-up and nothing else — the declarations are identical, so all
 95            but the first are no-ops
 96        :returns: the running group
 97        """
 98        if size < 1:
 99            raise ValueError(f"acemq: a consumer group needs at least one consumer, got {size}")
100
101        base = tag or f"acemq-{queue}"
102        started: list[Consumer] = []
103        group = cls(queue, started)
104        try:
105            for number in range(1, size + 1):
106                started.append(
107                    await connection.consume(
108                        queue,
109                        handler,
110                        codec=codec,
111                        retry=retry,
112                        prefetch=prefetch,
113                        concurrency=concurrency,
114                        tag=f"{base}-{number}",
115                        args=args,
116                        declare=declare,
117                    )
118                )
119        except BaseException:
120            await group.close()
121            raise
122        return group
123
124    @property
125    def queue(self) -> str:
126        """What the group is reading."""
127        return self._queue
128
129    @property
130    def size(self) -> int:
131        """How many consumers are running."""
132        return len(self._consumers)
133
134    async def close(self) -> None:
135        """Stops every consumer and waits for the handlers already running.
136
137        All of them are closed even when one fails, and the first failure is
138        raised afterwards. Stopping at the first one would leave the rest
139        running after a shutdown the caller believes happened, which is worse
140        than the failure that started it.
141        """
142        stopping = list(self._consumers)
143        self._consumers.clear()
144
145        first: BaseException | None = None
146        for consumer in stopping:
147            try:
148                await consumer.close()
149            except BaseException as failure:
150                first = first or failure
151        if first is not None:
152            raise first
153
154    async def __aenter__(self) -> ConsumerGroup:
155        return self
156
157    async def __aexit__(self, *_: object) -> None:
158        await self.close()

143@runtime_checkable
144class DbConnection(Protocol):
145    """As much of a DB-API 2.0 connection as this module uses.
146
147    Deliberately small. Anything following PEP 249 has these four, so nothing
148    here is tied to a driver — and a caller's connection, whatever framework it
149    came out of, is one of these.
150    """
151
152    def cursor(self) -> DbCursor:
153        """A cursor on this connection."""
154
155    def commit(self) -> None:
156        """Commits the open transaction."""
157
158    def rollback(self) -> None:
159        """Abandons the open transaction."""
160
161    def close(self) -> None:
162        """Releases the connection."""

123class DbCursor(Protocol):
124    """As much of a DB-API 2.0 cursor as this module uses."""
125
126    @property
127    def rowcount(self) -> int:
128        """How many rows the last statement affected."""
129
130    def execute(self, operation: str, parameters: Any = ..., /) -> Any:
131        """Runs one statement."""
132
133    def fetchone(self) -> Any:
134        """The next row, or ``None``."""
135
136    def fetchall(self) -> Any:
137        """Every remaining row."""
138
139    def close(self) -> None:
140        """Releases the cursor."""

200class FilesystemClaimCheckStore:
201    """A claim-check store on a filesystem.
202
203    Useful where the filesystem is shared and durable — an NFS mount, a
204    persistent volume — and the honest middle ground between a dictionary and
205    object storage. On a container's local disk it is the in-memory store with
206    extra steps: the consumer is on another host and finds nothing.
207
208    Object storage is the usual right answer, and a store in front of S3 or
209    Azure Blob Storage is three short methods. This one exists because "write it
210    to the mount everything already has" is a real deployment and not a bad one.
211
212    **Writes are atomic.** The payload is written to a temporary file and moved
213    into place. Without that, a consumer fast enough to read the key before the
214    writer finished gets a truncated payload and a parse error somewhere
215    unhelpful — and messaging is exactly the arrangement that makes a consumer
216    that fast normal rather than unlikely.
217
218    :param directory: where payloads are written; created if it is not there
219    """
220
221    #: A key reaches the filesystem as a path segment, so it is checked rather
222    #: than trusted. Every key this store issues is a UUID; one arriving from a
223    #: message is whatever a publisher put there, and ``../../etc/passwd`` is a
224    #: key too. The same expression as Java's, so the two stores accept and
225    #: refuse the same keys over one shared mount.
226    _SAFE_KEY = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}")
227
228    def __init__(self, directory: str | os.PathLike[str]) -> None:
229        self._directory = Path(directory)
230        try:
231            self._directory.mkdir(parents=True, exist_ok=True)
232        except OSError as failure:
233            raise AceMQError(
234                f"acemq: could not create the claim-check directory "
235                f"{self._directory}: {failure}"
236            ) from failure
237
238    @property
239    def directory(self) -> Path:
240        """Where the payloads are."""
241        return self._directory
242
243    def put(self, content: bytes) -> str:
244        key = str(uuid.uuid4())
245        target = self._directory / key
246        try:
247            handle, staging = tempfile.mkstemp(
248                prefix=key, suffix=".partial", dir=self._directory
249            )
250            try:
251                with os.fdopen(handle, "wb") as writing:
252                    writing.write(content)
253                # Moved into place, so a reader sees the whole payload or no
254                # payload. os.replace is atomic within a filesystem, which is
255                # the case here because the staging file was made in the same
256                # directory on purpose.
257                os.replace(staging, target)
258            except BaseException:
259                # The staging file is ours and nothing refers to it, so a failed
260                # write leaves nothing behind rather than a .partial somebody
261                # has to explain later.
262                Path(staging).unlink(missing_ok=True)
263                raise
264        except OSError as failure:
265            raise AceMQError(
266                f"acemq: could not write the claim-check payload {key}: {failure}"
267            ) from failure
268        return key
269
270    def get(self, key: str) -> bytes | None:
271        try:
272            return self._path_for(key).read_bytes()
273        except FileNotFoundError:
274            # Absent rather than failed: a key the store no longer holds is a
275            # retention answer, and the codec turns it into a message that
276            # explains itself.
277            return None
278        except OSError as failure:
279            raise AceMQError(
280                f"acemq: could not read the claim-check payload {key}: {failure}"
281            ) from failure
282
283    def delete(self, key: str) -> None:
284        try:
285            self._path_for(key).unlink(missing_ok=True)
286        except OSError as failure:
287            raise AceMQError(
288                f"acemq: could not delete the claim-check payload {key}: {failure}"
289            ) from failure
290
291    def _path_for(self, key: str) -> Path:
292        if not self._SAFE_KEY.fullmatch(key):
293            raise AceMQError(
294                f"acemq: {key!r} is not a key this store issued. A key becomes a path "
295                "segment, so one arriving from a message is checked rather than trusted."
296            )
297        return self._directory / key
298
299    def __repr__(self) -> str:
300        return f"FilesystemClaimCheckStore({self._directory})"

65@runtime_checkable
66class IdempotencyStore(Protocol):
67    """Remembers which messages have been handled."""
68
69    async def first_time(self, key: str) -> bool:
70        """Records a key and says whether it had not been seen before.
71
72        It has to be atomic. Two consumers handling the same message at the same
73        moment must not both be told they are first, which is the entire job:
74        a check followed by a write is not this, however short the gap looks.
75
76        :param key: what identifies the message
77        :returns: whether this is the first time
78        """
79
80    async def forget(self, key: str) -> None:
81        """Removes a key, so a message that failed can be tried again.
82
83        :param key: what identifies the message
84        """