This document is for Kombu's development version, which can be significantly different from previous releases. Get the stable docs here: 5.3.
Source code for kombu.transport.azureservicebus
"""Azure Service Bus Message Queue transport module for kombu.
Note that the Shared Access Policy used to connect to Azure Service Bus
requires Manage, Send and Listen claims since the broker will create new
queues and delete old queues as required.
Notes when using with Celery if you are experiencing issues with programs not
terminating properly. The Azure Service Bus SDK uses the Azure uAMQP library
which in turn creates some threads. If the AzureServiceBus Connection is
closed (e.g. via ``Connection.release()``), said threads will be closed
properly, but it seems there are times when Celery does not do this so these
threads will be left running. As the uAMQP threads are not marked as Daemon
threads, they will not be killed when the main thread exits. Setting the
``uamqp_keep_alive_interval`` transport option to 0 will prevent the
keep_alive thread from starting
More information about Azure Service Bus:
https://azure.microsoft.com/en-us/services/service-bus/
Features
========
* Type: Virtual
* Supports Direct: *Unreviewed*
* Supports Topic: *Unreviewed*
* Supports Fanout: *Unreviewed*
* Supports Priority: *Unreviewed*
* Supports TTL: *Unreviewed*
Connection String
=================
Connection string has the following formats:
.. code-block::
azureservicebus://SAS_POLICY_NAME:SAS_KEY@SERVICE_BUSNAMESPACE
azureservicebus://DefaultAzureCredential@SERVICE_BUSNAMESPACE
azureservicebus://ManagedIdentityCredential@SERVICE_BUSNAMESPACE
Transport Options
=================
* ``queue_name_prefix`` - String prefix to prepend to queue names in a
service bus namespace.
* ``wait_time_seconds`` - Number of seconds to wait to receive messages.
Default ``5``
* ``peek_lock_seconds`` - Number of seconds the message is visible for before
it is requeued and sent to another consumer. Default ``60``
* ``uamqp_keep_alive_interval`` - Interval in seconds the Azure uAMQP library
should send keepalive messages. Default ``30``
* ``retry_total`` - Azure SDK retry total. Default ``3``
* ``retry_backoff_factor`` - Azure SDK exponential backoff factor.
Default ``0.8``
* ``retry_backoff_max`` - Azure SDK retry total time. Default ``120``
* ``use_lock_renewal`` - Enable Azure SDK ``AutoLockRenewer`` to keep
message locks alive while a worker is processing. Only effective when
receive mode is ``PEEK_LOCK`` (the default). Default ``False``.
* ``max_lock_renewal_duration`` - Time in seconds that locks registered
to the renewer should be maintained for. Default ``3600`` (1 hour).
.. versionadded:: 5.7.0
``use_lock_renewal`` and ``max_lock_renewal_duration`` transport options.
"""
from __future__ import annotations
import string
from queue import Empty
from typing import Any
import azure.core.exceptions
import azure.servicebus.exceptions
import isodate
from azure.servicebus import (AutoLockRenewer, ServiceBusClient,
ServiceBusMessage, ServiceBusReceiveMode,
ServiceBusReceiver, ServiceBusSender)
from azure.servicebus.management import ServiceBusAdministrationClient
try:
from azure.identity import (DefaultAzureCredential,
ManagedIdentityCredential)
except ImportError:
DefaultAzureCredential = None
ManagedIdentityCredential = None
from kombu.log import get_logger
from kombu.utils.encoding import bytes_to_str, safe_str
from kombu.utils.json import dumps, loads
from kombu.utils.objects import cached_property
from . import virtual
logger = get_logger(__name__)
# dots are replaced by dash, all other punctuation replaced by underscore.
PUNCTUATIONS_TO_REPLACE = set(string.punctuation) - {'_', '.', '-'}
CHARS_REPLACE_TABLE = {
**{ord(c): ord('_') for c in PUNCTUATIONS_TO_REPLACE}
}
class SendReceive:
"""Container for Sender and Receiver."""
def __init__(self,
receiver: ServiceBusReceiver | None = None,
sender: ServiceBusSender | None = None):
self.receiver: ServiceBusReceiver = receiver
self.sender: ServiceBusSender = sender
def close(self) -> None:
if self.receiver:
self.receiver.close()
self.receiver = None
if self.sender:
self.sender.close()
self.sender = None
[docs]
class Channel(virtual.Channel):
"""Azure Service Bus channel."""
default_wait_time_seconds: int = 5 # in seconds
default_peek_lock_seconds: int = 60 # in seconds (default 60, max 300)
# in seconds (is the default from service bus repo)
default_uamqp_keep_alive_interval: int = 30
# number of retries (is the default from service bus repo)
default_retry_total: int = 3
# exponential backoff factor (is the default from service bus repo)
default_retry_backoff_factor: float = 0.8
# Max time to backoff (is the default from service bus repo)
default_retry_backoff_max: int = 120
default_use_lock_renewal: bool = False
default_max_lock_renewal_duration: float = 3600 # in seconds (1 hour)
domain_format: str = 'kombu%(vhost)s'
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._namespace = None
self._policy = None
self._sas_key = None
self._connection_string = None
self._try_parse_connection_string()
self.qos.restore_at_shutdown = False
def _try_parse_connection_string(self) -> None:
self._namespace, self._credential = Transport.parse_uri(
self.conninfo.hostname)
if (
DefaultAzureCredential is not None
and isinstance(self._credential, DefaultAzureCredential)
) or (
ManagedIdentityCredential is not None
and isinstance(self._credential, ManagedIdentityCredential)
):
return None
if ":" in self._credential:
self._policy, self._sas_key = self._credential.split(':', 1)
conn_dict = {
'Endpoint': 'sb://' + self._namespace,
'SharedAccessKeyName': self._policy,
'SharedAccessKey': self._sas_key,
}
self._connection_string = ';'.join(
[key + '=' + value for key, value in conn_dict.items()])
[docs]
def basic_consume(self, queue, no_ack, callback, consumer_tag,
*args, **kwargs):
if no_ack:
self.connection._noack_consumer_tags.add(consumer_tag)
return super().basic_consume(
queue, no_ack, callback, consumer_tag, *args, **kwargs
)
[docs]
def basic_cancel(self, consumer_tag):
self.connection._noack_consumer_tags.discard(consumer_tag)
return super().basic_cancel(consumer_tag)
def _add_queue_to_cache(
self, name: str,
receiver: ServiceBusReceiver | None = None,
sender: ServiceBusSender | None = None
) -> SendReceive:
if name in self._queue_cache:
obj = self._queue_cache[name]
obj.sender = obj.sender or sender
obj.receiver = obj.receiver or receiver
else:
obj = SendReceive(receiver, sender)
self._queue_cache[name] = obj
return obj
def _get_asb_sender(self, queue: str) -> SendReceive:
queue_obj = self._queue_cache.get(queue, None)
if queue_obj is None or queue_obj.sender is None:
sender = self.queue_service.get_queue_sender(
queue, keep_alive=self.uamqp_keep_alive_interval)
queue_obj = self._add_queue_to_cache(queue, sender=sender)
return queue_obj
def _get_asb_receiver(
self, queue: str,
recv_mode: ServiceBusReceiveMode = ServiceBusReceiveMode.PEEK_LOCK,
queue_cache_key: str | None = None) -> SendReceive:
cache_key = queue_cache_key or f"{queue}::{recv_mode.name}"
queue_obj = self._queue_cache.get(cache_key, None)
if queue_obj is None or queue_obj.receiver is None:
auto_lock_renewer = None
if (self.use_lock_renewal
and recv_mode == ServiceBusReceiveMode.PEEK_LOCK):
if self.connection._renewer is None:
self.connection._renewer = AutoLockRenewer(
max_lock_renewal_duration=(
self.max_lock_renewal_duration)
)
auto_lock_renewer = self.connection._renewer
receiver = self.queue_service.get_queue_receiver(
queue_name=queue, receive_mode=recv_mode,
keep_alive=self.uamqp_keep_alive_interval,
auto_lock_renewer=auto_lock_renewer)
queue_obj = self._add_queue_to_cache(cache_key, receiver=receiver)
return queue_obj
[docs]
def entity_name(
self, name: str, table: dict[int, int] | None = None) -> str:
"""Format AMQP queue name into a valid ServiceBus queue name."""
return str(safe_str(name)).translate(table or CHARS_REPLACE_TABLE)
def _restore(self, message: virtual.base.Message) -> None:
# Not be needed as ASB handles unacked messages
# Remove 'azure_message' as its not JSON serializable
# message.delivery_info.pop('azure_message', None)
# super()._restore(message)
pass
def _new_queue(self, queue: str, **kwargs) -> SendReceive:
"""Ensure a queue exists in ServiceBus."""
queue = self.entity_name(self.queue_name_prefix + queue)
try:
return self._queue_cache[queue]
except KeyError:
# Converts seconds into ISO8601 duration format
# ie 66seconds = P1M6S
lock_duration = isodate.duration_isoformat(
isodate.Duration(seconds=self.peek_lock_seconds))
try:
self.queue_mgmt_service.create_queue(
queue_name=queue, lock_duration=lock_duration)
except azure.core.exceptions.ResourceExistsError:
pass
return self._add_queue_to_cache(queue)
def _delete(self, queue: str, *args, **kwargs) -> None:
"""Delete queue by name."""
queue = self.entity_name(self.queue_name_prefix + queue)
self.queue_mgmt_service.delete_queue(queue)
keys = [
k for k in self._queue_cache
if k == queue or k.startswith(f"{queue}::")
]
for k in keys:
obj = self._queue_cache.pop(k, None)
if obj is None:
continue
try:
obj.close()
except Exception:
logger.exception(
"Failed to close cached SendReceive for %r; continuing",
k)
def _put(self, queue: str, message, **kwargs) -> None:
"""Put message onto queue."""
queue = self.entity_name(self.queue_name_prefix + queue)
msg = ServiceBusMessage(dumps(message))
queue_obj = self._get_asb_sender(queue)
queue_obj.sender.send_messages(msg)
def _get(
self, queue: str,
timeout: float | int | None = None
) -> dict[str, Any]:
"""Try to retrieve a single message off ``queue``."""
# If we're not ack'ing for this queue, just change receive_mode
recv_mode = ServiceBusReceiveMode.RECEIVE_AND_DELETE \
if queue in self._noack_queues else ServiceBusReceiveMode.PEEK_LOCK
queue = self.entity_name(self.queue_name_prefix + queue)
queue_obj = self._get_asb_receiver(queue, recv_mode)
messages = queue_obj.receiver.receive_messages(
max_message_count=1,
max_wait_time=timeout or self.wait_time_seconds)
if not messages:
raise Empty()
# message.body is either byte or generator[bytes]
message = messages[0]
if not isinstance(message.body, bytes):
body = b''.join(message.body)
else:
body = message.body
msg = loads(bytes_to_str(body))
msg['properties']['delivery_info']['azure_message'] = message
msg['properties']['delivery_info']['azure_queue_name'] = queue
return msg
[docs]
def basic_ack(self, delivery_tag: str, multiple: bool = False) -> None:
try:
delivery_info = self.qos.get(delivery_tag).delivery_info
except KeyError:
super().basic_ack(delivery_tag)
else:
queue = delivery_info['azure_queue_name']
# recv_mode is PEEK_LOCK when ack'ing messages
queue_obj = self._get_asb_receiver(queue)
try:
queue_obj.receiver.complete_message(
delivery_info['azure_message'])
except azure.servicebus.exceptions.MessageAlreadySettled:
super().basic_ack(delivery_tag)
except azure.servicebus.exceptions.MessageLockLostError:
logger.warning(
"complete_message on queue %r failed with "
"MessageLockLostError; the broker may have redelivered "
"this message. Consider shorter handlers or enabling "
"lock renewal.",
queue,
)
super().basic_reject(delivery_tag)
except Exception:
logger.exception(
"complete_message on queue %r failed", queue)
super().basic_reject(delivery_tag)
else:
super().basic_ack(delivery_tag)
def _size(self, queue: str) -> int:
"""Return the number of messages in a queue."""
queue = self.entity_name(self.queue_name_prefix + queue)
props = self.queue_mgmt_service.get_queue_runtime_properties(queue)
return props.total_message_count
def _purge(self, queue) -> int:
"""Delete all current messages in a queue."""
# Azure has no broker-side purge API. Drain via an ephemeral
# RECEIVE_AND_DELETE receiver scoped to this call so we do not
# leak the receiver into _queue_cache.
n = 0
max_purge_count = 10
queue = self.entity_name(self.queue_name_prefix + queue)
with self.queue_service.get_queue_receiver(
queue_name=queue,
receive_mode=ServiceBusReceiveMode.RECEIVE_AND_DELETE,
keep_alive=self.uamqp_keep_alive_interval,
) as receiver:
while True:
messages = receiver.receive_messages(
max_message_count=max_purge_count,
max_wait_time=0.2,
)
n += len(messages)
if len(messages) < max_purge_count:
break
return n
[docs]
def close(self) -> None:
# Cache and noack set live on Transport; see Transport.close_connection.
super().close()
@cached_property
def queue_service(self) -> ServiceBusClient:
if self._connection_string:
return ServiceBusClient.from_connection_string(
self._connection_string,
retry_total=self.retry_total,
retry_backoff_factor=self.retry_backoff_factor,
retry_backoff_max=self.retry_backoff_max
)
return ServiceBusClient(
self._namespace,
self._credential,
retry_total=self.retry_total,
retry_backoff_factor=self.retry_backoff_factor,
retry_backoff_max=self.retry_backoff_max
)
@cached_property
def queue_mgmt_service(self) -> ServiceBusAdministrationClient:
if self._connection_string:
return ServiceBusAdministrationClient.from_connection_string(
self._connection_string
)
return ServiceBusAdministrationClient(
self._namespace, self._credential
)
@property
def conninfo(self):
return self.connection.client
@property
def transport_options(self):
return self.connection.client.transport_options
@property
def _queue_cache(self) -> dict[str, SendReceive]:
return self.connection._queue_cache
@property
def _noack_queues(self) -> set[str]:
tag_to_queue = self._tag_to_queue
return {tag_to_queue[t]
for t in self.connection._noack_consumer_tags
if t in tag_to_queue}
@cached_property
def queue_name_prefix(self) -> str:
return self.transport_options.get('queue_name_prefix', '')
@cached_property
def wait_time_seconds(self) -> int:
return self.transport_options.get('wait_time_seconds',
self.default_wait_time_seconds)
@cached_property
def peek_lock_seconds(self) -> int:
return min(self.transport_options.get('peek_lock_seconds',
self.default_peek_lock_seconds),
300) # Limit upper bounds to 300
@cached_property
def uamqp_keep_alive_interval(self) -> int:
return self.transport_options.get(
'uamqp_keep_alive_interval',
self.default_uamqp_keep_alive_interval
)
@cached_property
def retry_total(self) -> int:
return self.transport_options.get(
'retry_total', self.default_retry_total)
@cached_property
def retry_backoff_factor(self) -> float:
return self.transport_options.get(
'retry_backoff_factor', self.default_retry_backoff_factor)
@cached_property
def retry_backoff_max(self) -> int:
return self.transport_options.get(
'retry_backoff_max', self.default_retry_backoff_max)
@cached_property
def use_lock_renewal(self) -> bool:
return self.transport_options.get(
'use_lock_renewal', self.default_use_lock_renewal)
@cached_property
def max_lock_renewal_duration(self) -> float:
return self.transport_options.get(
'max_lock_renewal_duration',
self.default_max_lock_renewal_duration)
[docs]
class Transport(virtual.Transport):
"""Azure Service Bus transport."""
Channel = Channel
polling_interval = 1
default_port = None
can_parse_url = True
def __init__(self, client, **kwargs):
super().__init__(client, **kwargs)
self._queue_cache: dict[str, SendReceive] = {}
self._noack_consumer_tags: set[str] = set()
self._renewer: AutoLockRenewer | None = None
[docs]
def close_connection(self, connection) -> None:
try:
super().close_connection(connection)
finally:
if self._renewer is not None:
try:
self._renewer.close()
except Exception:
logger.exception(
"Failed to close AutoLockRenewer; continuing")
self._renewer = None
for queue_obj in self._queue_cache.values():
try:
queue_obj.close()
except Exception:
logger.exception(
"Failed to close cached SendReceive; continuing")
self._queue_cache.clear()
self._noack_consumer_tags.clear()
[docs]
@staticmethod
def parse_uri(uri: str) -> tuple[str, str | DefaultAzureCredential |
ManagedIdentityCredential]:
# URL like:
# azureservicebus://{SAS policy name}:{SAS key}@{ServiceBus Namespace}
# urllib parse does not work as the sas key could contain a slash
# e.g.: azureservicebus://rootpolicy:some/key@somenamespace
# > 'rootpolicy:some/key@somenamespace'
uri = uri.replace('azureservicebus://', '')
# > 'rootpolicy:some/key', 'somenamespace'
credential, namespace = uri.rsplit('@', 1)
if not namespace.endswith('.net'):
namespace += '.servicebus.windows.net'
if "DefaultAzureCredential".lower() == credential.lower():
if DefaultAzureCredential is None:
raise ImportError('Azure Service Bus transport with a '
'DefaultAzureCredential requires the '
'azure-identity library')
credential = DefaultAzureCredential()
elif "ManagedIdentityCredential".lower() == credential.lower():
if ManagedIdentityCredential is None:
raise ImportError('Azure Service Bus transport with a '
'ManagedIdentityCredential requires the '
'azure-identity library')
credential = ManagedIdentityCredential()
else:
# > 'rootpolicy', 'some/key'
policy, sas_key = credential.split(':', 1)
credential = f"{policy}:{sas_key}"
# Validate ASB connection string
if not all([namespace, credential]):
raise ValueError(
'Need a URI like '
'azureservicebus://{SAS policy name}:{SAS key}@{ServiceBus Namespace} ' # noqa
'or the azure Endpoint connection string'
)
return namespace, credential
[docs]
@classmethod
def as_uri(cls, uri: str, include_password=False, mask='**') -> str:
namespace, credential = cls.parse_uri(uri)
if isinstance(credential, str) and ":" in credential:
policy, sas_key = credential.split(':', 1)
return 'azureservicebus://{}:{}@{}'.format(
policy,
sas_key if include_password else mask,
namespace
)
return 'azureservicebus://{}@{}'.format(
credential.__class__.__name__,
namespace
)
