|
| 1 | +"""The BabelQueue runtime: produce and consume polyglot messages. |
| 2 | +
|
| 3 | + from babelqueue import BabelQueue |
| 4 | +
|
| 5 | + app = BabelQueue("redis://localhost:6379/0", queue="orders") |
| 6 | +
|
| 7 | + @app.handler("urn:babel:orders:created") |
| 8 | + def on_order_created(data, meta): |
| 9 | + ... # AI/ML, data processing, anything |
| 10 | +
|
| 11 | + app.publish("urn:babel:orders:created", {"order_id": 1042}) |
| 12 | + app.run() # consume forever |
| 13 | +
|
| 14 | +Routing is by URN; the wire format is the canonical envelope (shared core codec), |
| 15 | +so this interoperates with the PHP/Laravel, Symfony, Go, ... SDKs. Retry uses the |
| 16 | +top-level ``attempts`` counter; failures past ``max_attempts`` go to a dead-letter |
| 17 | +queue when enabled. |
| 18 | +""" |
| 19 | + |
| 20 | +from __future__ import annotations |
| 21 | + |
| 22 | +import inspect |
| 23 | +from typing import Any, Callable, Dict, Mapping, Optional |
| 24 | + |
| 25 | +from . import dead_letter |
| 26 | +from .codec import EnvelopeCodec |
| 27 | +from .exceptions import UnknownUrnError |
| 28 | +from .routing import UnknownUrnStrategy |
| 29 | +from .transport import ReceivedMessage, Transport, make_transport |
| 30 | + |
| 31 | +Handler = Callable[..., None] |
| 32 | + |
| 33 | + |
| 34 | +class BabelQueue: |
| 35 | + def __init__( |
| 36 | + self, |
| 37 | + broker_url: str = "memory://", |
| 38 | + *, |
| 39 | + transport: Optional[Transport] = None, |
| 40 | + queue: str = "default", |
| 41 | + on_unknown_urn: str = UnknownUrnStrategy.FAIL, |
| 42 | + max_attempts: int = 3, |
| 43 | + dead_letter: bool = False, |
| 44 | + dead_letter_queue: Optional[str] = None, |
| 45 | + dead_letter_suffix: str = ".dlq", |
| 46 | + ) -> None: |
| 47 | + self.transport = transport if transport is not None else make_transport(broker_url) |
| 48 | + self.queue = queue |
| 49 | + self.on_unknown_urn = on_unknown_urn |
| 50 | + self.max_attempts = max_attempts |
| 51 | + self.dead_letter_enabled = bool(dead_letter) |
| 52 | + self.dead_letter_queue = dead_letter_queue |
| 53 | + self.dead_letter_suffix = dead_letter_suffix |
| 54 | + self._handlers: Dict[str, Handler] = {} |
| 55 | + |
| 56 | + # -- Produce ------------------------------------------------------------ |
| 57 | + |
| 58 | + def publish( |
| 59 | + self, |
| 60 | + urn: str, |
| 61 | + data: Mapping[str, Any], |
| 62 | + *, |
| 63 | + queue: Optional[str] = None, |
| 64 | + trace_id: Optional[str] = None, |
| 65 | + ) -> str: |
| 66 | + """Publish a message; returns its id (``meta.id``).""" |
| 67 | + target = queue or self.queue |
| 68 | + envelope = EnvelopeCodec.make(urn, data, queue=target, trace_id=trace_id) |
| 69 | + self.transport.publish(target, EnvelopeCodec.encode(envelope)) |
| 70 | + return envelope["meta"]["id"] |
| 71 | + |
| 72 | + # -- Register handlers -------------------------------------------------- |
| 73 | + |
| 74 | + def handler(self, urn: str) -> Callable[[Handler], Handler]: |
| 75 | + """Decorator: register ``fn`` as the handler for ``urn``.""" |
| 76 | + |
| 77 | + def decorator(fn: Handler) -> Handler: |
| 78 | + self._handlers[urn] = fn |
| 79 | + return fn |
| 80 | + |
| 81 | + return decorator |
| 82 | + |
| 83 | + def register(self, urn: str, fn: Handler) -> None: |
| 84 | + self._handlers[urn] = fn |
| 85 | + |
| 86 | + # -- Consume ------------------------------------------------------------ |
| 87 | + |
| 88 | + def consume( |
| 89 | + self, |
| 90 | + queue: Optional[str] = None, |
| 91 | + *, |
| 92 | + max_messages: Optional[int] = None, |
| 93 | + timeout: float = 1.0, |
| 94 | + ) -> int: |
| 95 | + """Consume messages until interrupted (or ``max_messages`` processed). |
| 96 | +
|
| 97 | + Returns the number of messages processed. With ``max_messages`` set, the |
| 98 | + loop stops once that many are handled or the queue drains within ``timeout``. |
| 99 | + """ |
| 100 | + target = queue or self.queue |
| 101 | + processed = 0 |
| 102 | + try: |
| 103 | + while max_messages is None or processed < max_messages: |
| 104 | + received = self.transport.pop(target, timeout=timeout) |
| 105 | + if received is None: |
| 106 | + if max_messages is not None: |
| 107 | + break |
| 108 | + continue |
| 109 | + self.dispatch(received) |
| 110 | + processed += 1 |
| 111 | + except KeyboardInterrupt: # pragma: no cover - graceful Ctrl-C |
| 112 | + pass |
| 113 | + return processed |
| 114 | + |
| 115 | + run = consume |
| 116 | + |
| 117 | + def dispatch(self, received: ReceivedMessage) -> None: |
| 118 | + """Route one reserved message to its handler and acknowledge it.""" |
| 119 | + envelope = EnvelopeCodec.decode(received.body) |
| 120 | + urn = str(envelope.get("job") or envelope.get("urn") or "") |
| 121 | + handler = self._handlers.get(urn) if urn else None |
| 122 | + |
| 123 | + try: |
| 124 | + if handler is None: |
| 125 | + self._route_unknown(urn, received, envelope) |
| 126 | + return |
| 127 | + self._invoke(handler, envelope) |
| 128 | + self.transport.ack(received) |
| 129 | + except Exception as exc: # noqa: BLE001 - one bad message must not kill the loop |
| 130 | + self._retry_or_dead_letter(received, envelope, exc) |
| 131 | + |
| 132 | + # -- Internals ---------------------------------------------------------- |
| 133 | + |
| 134 | + def _invoke(self, handler: Handler, envelope: Mapping[str, Any]) -> None: |
| 135 | + data = dict(envelope.get("data") or {}) |
| 136 | + meta = dict(envelope.get("meta") or {}) |
| 137 | + if _handler_wants_envelope(handler): |
| 138 | + handler(data, meta, dict(envelope)) |
| 139 | + else: |
| 140 | + handler(data, meta) |
| 141 | + |
| 142 | + def _route_unknown(self, urn: str, received: ReceivedMessage, envelope: Mapping[str, Any]) -> None: |
| 143 | + strategy = self.on_unknown_urn |
| 144 | + if strategy == UnknownUrnStrategy.DELETE: |
| 145 | + self.transport.ack(received) |
| 146 | + return |
| 147 | + if strategy == UnknownUrnStrategy.RELEASE: |
| 148 | + self.transport.publish(received.queue, received.body) |
| 149 | + self.transport.ack(received) |
| 150 | + return |
| 151 | + if strategy == UnknownUrnStrategy.DEAD_LETTER: |
| 152 | + self._dead_letter(received, dict(envelope), "unknown_urn", None) |
| 153 | + return |
| 154 | + # FAIL — surfaced through the retry/dead-letter path (never kills the loop). |
| 155 | + raise UnknownUrnError( |
| 156 | + f"No handler mapped for URN [{urn or '(empty)'}]." |
| 157 | + ) |
| 158 | + |
| 159 | + def _retry_or_dead_letter( |
| 160 | + self, received: ReceivedMessage, envelope: Dict[str, Any], exc: BaseException |
| 161 | + ) -> None: |
| 162 | + attempts = int(envelope.get("attempts", 0)) + 1 |
| 163 | + envelope["attempts"] = attempts |
| 164 | + |
| 165 | + if attempts < self.max_attempts: |
| 166 | + self.transport.publish(received.queue, EnvelopeCodec.encode(envelope)) |
| 167 | + self.transport.ack(received) |
| 168 | + return |
| 169 | + |
| 170 | + if self.dead_letter_enabled: |
| 171 | + reason = "unknown_urn" if isinstance(exc, UnknownUrnError) else "failed" |
| 172 | + self._dead_letter(received, envelope, reason, exc) |
| 173 | + return |
| 174 | + |
| 175 | + # Retries exhausted, no DLQ configured — drop it (ack so it leaves the queue). |
| 176 | + self.transport.ack(received) |
| 177 | + |
| 178 | + def _dead_letter( |
| 179 | + self, |
| 180 | + received: ReceivedMessage, |
| 181 | + envelope: Dict[str, Any], |
| 182 | + reason: str, |
| 183 | + exc: Optional[BaseException], |
| 184 | + ) -> None: |
| 185 | + original_queue = str((envelope.get("meta") or {}).get("queue") or received.queue) |
| 186 | + annotated = dead_letter.annotate( |
| 187 | + envelope, |
| 188 | + reason, |
| 189 | + original_queue, |
| 190 | + int(envelope.get("attempts", 0)), |
| 191 | + error=(str(exc) if exc is not None else None), |
| 192 | + exception=(type(exc).__name__ if exc is not None else None), |
| 193 | + ) |
| 194 | + target = self.dead_letter_queue or (received.queue + self.dead_letter_suffix) |
| 195 | + self.transport.publish(target, EnvelopeCodec.encode(annotated)) |
| 196 | + self.transport.ack(received) |
| 197 | + |
| 198 | + |
| 199 | +def _handler_wants_envelope(fn: Handler) -> bool: |
| 200 | + """True if the handler takes a 3rd positional arg (the full envelope).""" |
| 201 | + try: |
| 202 | + params = list(inspect.signature(fn).parameters.values()) |
| 203 | + except (TypeError, ValueError): # pragma: no cover - builtins/C callables |
| 204 | + return False |
| 205 | + positional = [ |
| 206 | + p for p in params |
| 207 | + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) |
| 208 | + ] |
| 209 | + has_varargs = any(p.kind == p.VAR_POSITIONAL for p in params) |
| 210 | + return has_varargs or len(positional) >= 3 |
0 commit comments