|
| 1 | +"""Azure Service Bus transport. Requires the ``azureservicebus`` extra: |
| 2 | +
|
| 3 | + pip install "babelqueue[azureservicebus]" |
| 4 | +
|
| 5 | +Producing sends the canonical envelope as the message body and projects the contract |
| 6 | +envelope fields onto native Service Bus fields — ``Subject`` = URN, ``CorrelationId`` = |
| 7 | +trace_id, ``MessageId`` = meta.id, plus the ``bq-`` application properties |
| 8 | +(``bq-schema-version`` / ``bq-source-lang`` / ``bq-created-at``) — so a .NET/Java/... peer |
| 9 | +can route on ``Subject`` and correlate on ``CorrelationId`` without parsing the body. |
| 10 | +Consuming uses the PeekLock reservation model (``receive_messages`` -> process -> |
| 11 | +``complete_message``); the authoritative attempt count is the broker's native |
| 12 | +``DeliveryCount`` (1-based), reconciled onto the envelope as ``attempts = DeliveryCount - 1``. |
| 13 | +A message left un-acked has its lock expire and is redelivered (at-least-once). |
| 14 | +
|
| 15 | +This implements §4 of the broker-bindings contract. The envelope is unchanged |
| 16 | +(``schema_version`` stays 1); Azure Service Bus is purely additive. |
| 17 | +
|
| 18 | +URL form: ``sb://<namespace>.servicebus.windows.net`` (Azure AD via |
| 19 | +``DefaultAzureCredential``). For connection-string auth or a custom client, build the |
| 20 | +transport directly and pass it via ``BabelQueue(transport=...)`` or |
| 21 | +``AsbTransport(connection_string=...)``. |
| 22 | +""" |
| 23 | + |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +from typing import Any, Dict, Optional |
| 27 | +from urllib.parse import urlsplit |
| 28 | + |
| 29 | +from .codec import EnvelopeCodec |
| 30 | +from .transport import ReceivedMessage, Transport |
| 31 | + |
| 32 | + |
| 33 | +class AsbTransport(Transport): |
| 34 | + def __init__( |
| 35 | + self, |
| 36 | + url: str = "sb://", |
| 37 | + *, |
| 38 | + client: Any = None, |
| 39 | + connection_string: Optional[str] = None, |
| 40 | + credential: Any = None, |
| 41 | + max_wait_time: Optional[float] = None, |
| 42 | + ) -> None: |
| 43 | + parts = urlsplit(url) if url else urlsplit("sb://") |
| 44 | + self._namespace = parts.hostname or None |
| 45 | + self._connection_string = connection_string |
| 46 | + self._credential = credential |
| 47 | + self._max_wait_time = max_wait_time |
| 48 | + self._senders: Dict[str, Any] = {} |
| 49 | + self._receivers: Dict[str, Any] = {} |
| 50 | + |
| 51 | + if client is not None: |
| 52 | + self._client = client |
| 53 | + return |
| 54 | + self._client = self._build_client() # pragma: no cover - needs Azure / network |
| 55 | + |
| 56 | + def _build_client(self) -> Any: # pragma: no cover - needs Azure / network |
| 57 | + try: |
| 58 | + from azure.servicebus import ServiceBusClient |
| 59 | + except ImportError as exc: |
| 60 | + raise ImportError( |
| 61 | + "AsbTransport requires the 'azure-servicebus' package. Install with " |
| 62 | + 'pip install "babelqueue[azureservicebus]".' |
| 63 | + ) from exc |
| 64 | + import os |
| 65 | + |
| 66 | + cs = self._connection_string or os.environ.get("AZURE_SERVICEBUS_CONNECTION_STRING") |
| 67 | + if cs: |
| 68 | + return ServiceBusClient.from_connection_string(cs) |
| 69 | + if self._namespace: |
| 70 | + credential = self._credential |
| 71 | + if credential is None: |
| 72 | + from azure.identity import DefaultAzureCredential |
| 73 | + |
| 74 | + credential = DefaultAzureCredential() |
| 75 | + return ServiceBusClient(self._namespace, credential) |
| 76 | + raise ValueError( |
| 77 | + "AsbTransport needs a connection string, a namespace + credential, or an injected client." |
| 78 | + ) |
| 79 | + |
| 80 | + # -- helpers ------------------------------------------------------------ |
| 81 | + |
| 82 | + def _sender(self, queue: str) -> Any: |
| 83 | + sender = self._senders.get(queue) |
| 84 | + if sender is None: |
| 85 | + sender = self._client.get_queue_sender(queue) |
| 86 | + self._senders[queue] = sender |
| 87 | + return sender |
| 88 | + |
| 89 | + def _receiver(self, queue: str) -> Any: |
| 90 | + receiver = self._receivers.get(queue) |
| 91 | + if receiver is None: |
| 92 | + receiver = self._client.get_queue_receiver(queue) |
| 93 | + self._receivers[queue] = receiver |
| 94 | + return receiver |
| 95 | + |
| 96 | + @staticmethod |
| 97 | + def _projection(body: str) -> Dict[str, Any]: |
| 98 | + """Native ServiceBusMessage kwargs — Subject/CorrelationId/MessageId + the bq- |
| 99 | + application properties (a redundant, routable view of the body). §4.2–§4.3.""" |
| 100 | + try: |
| 101 | + env: Dict[str, Any] = EnvelopeCodec.decode(body) |
| 102 | + except (ValueError, TypeError): # pragma: no cover - defensive |
| 103 | + return {"content_type": "application/json"} |
| 104 | + meta = env.get("meta") or {} |
| 105 | + |
| 106 | + props: Dict[str, Any] = {} |
| 107 | + if meta.get("schema_version") is not None: |
| 108 | + props["bq-schema-version"] = meta["schema_version"] |
| 109 | + if meta.get("lang"): |
| 110 | + props["bq-source-lang"] = meta["lang"] |
| 111 | + if meta.get("created_at") is not None: |
| 112 | + props["bq-created-at"] = meta["created_at"] |
| 113 | + |
| 114 | + kwargs: Dict[str, Any] = {"content_type": "application/json"} |
| 115 | + if env.get("job"): |
| 116 | + kwargs["subject"] = env["job"] |
| 117 | + if env.get("trace_id"): |
| 118 | + kwargs["correlation_id"] = env["trace_id"] |
| 119 | + if meta.get("id"): |
| 120 | + kwargs["message_id"] = meta["id"] |
| 121 | + if props: |
| 122 | + kwargs["application_properties"] = props |
| 123 | + return kwargs |
| 124 | + |
| 125 | + @staticmethod |
| 126 | + def _reconcile(body: str, delivery_count: Any) -> str: |
| 127 | + """Set attempts to max(current, DeliveryCount - 1) — DeliveryCount (1-based) is the |
| 128 | + broker's native redelivery floor, but the runtime retries by republishing with |
| 129 | + attempts+1 in the body, so a republished message (DeliveryCount back to 1) must not |
| 130 | + have its higher body count lowered. First delivery (DeliveryCount 1) reads 0.""" |
| 131 | + try: |
| 132 | + dc = int(delivery_count) |
| 133 | + except (ValueError, TypeError): |
| 134 | + return body |
| 135 | + if dc <= 1: |
| 136 | + return body |
| 137 | + native = dc - 1 |
| 138 | + env = EnvelopeCodec.decode(body) |
| 139 | + if not env or native <= int(env.get("attempts", 0)): |
| 140 | + return body |
| 141 | + env["attempts"] = native |
| 142 | + return EnvelopeCodec.encode(env) |
| 143 | + |
| 144 | + # -- Transport ---------------------------------------------------------- |
| 145 | + |
| 146 | + def publish(self, queue: str, body: str) -> None: |
| 147 | + from azure.servicebus import ServiceBusMessage |
| 148 | + |
| 149 | + message = ServiceBusMessage(body, **self._projection(body)) |
| 150 | + self._sender(queue).send_messages(message) |
| 151 | + |
| 152 | + def pop(self, queue: str, timeout: float = 1.0) -> Optional[ReceivedMessage]: |
| 153 | + wait = self._max_wait_time |
| 154 | + if wait is None and timeout and timeout > 0: |
| 155 | + wait = timeout |
| 156 | + messages = self._receiver(queue).receive_messages(max_message_count=1, max_wait_time=wait) |
| 157 | + if not messages: |
| 158 | + return None |
| 159 | + message = messages[0] |
| 160 | + body = self._reconcile(str(message), getattr(message, "delivery_count", None)) |
| 161 | + return ReceivedMessage(body=body, queue=queue, handle=message) |
| 162 | + |
| 163 | + def ack(self, message: ReceivedMessage) -> None: |
| 164 | + if message.handle is None: |
| 165 | + return |
| 166 | + self._receiver(message.queue).complete_message(message.handle) |
| 167 | + |
| 168 | + def close(self) -> None: # pragma: no cover - resource cleanup |
| 169 | + for resource in (*self._senders.values(), *self._receivers.values(), self._client): |
| 170 | + try: |
| 171 | + resource.close() |
| 172 | + except Exception: # noqa: BLE001 - best-effort cleanup |
| 173 | + pass |
0 commit comments