diff --git a/eng/ci_tools.txt b/eng/ci_tools.txt index 1b8c5e85b47b..817b97cece72 100644 --- a/eng/ci_tools.txt +++ b/eng/ci_tools.txt @@ -8,7 +8,8 @@ doc-warden==0.7.2 beautifulsoup4==4.13.3 pkginfo==1.12.1.2 pip==26.2.2 -typing-extensions==4.12.2 +typing-extensions==4.12.2; python_version < "3.15" +typing-extensions==4.14.0; python_version >= "3.15" pyproject-api==1.8.0 cibuildwheel importlib-metadata==8.5.0 diff --git a/sdk/servicebus/azure-servicebus/CHANGELOG.md b/sdk/servicebus/azure-servicebus/CHANGELOG.md index 80d69ae208df..cbe072410574 100644 --- a/sdk/servicebus/azure-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-servicebus/CHANGELOG.md @@ -1,5 +1,11 @@ # Release History +## 7.16.0b1 (Unreleased) + +### Features Added + +- Added sync and async `ServiceBusReceiver.delete_messages()` and `purge_messages()` methods. Basic and Standard support up to 500 messages per request, Premium supports up to 4,000, and purge handles smaller batches caused by large messages. + ## 7.15.1 (Unreleased) ### Features Added diff --git a/sdk/servicebus/azure-servicebus/api.md b/sdk/servicebus/azure-servicebus/api.md index 05f648785159..491d8c5cb9bf 100644 --- a/sdk/servicebus/azure-servicebus/api.md +++ b/sdk/servicebus/azure-servicebus/api.md @@ -25,6 +25,18 @@ namespace azure.servicebus ) -> None: ... + class azure.servicebus.DeleteMessagesResult: + property deleted_message_count: int # Read-only + + def __init__(self, deleted_message_count: int) -> None: ... + + + class azure.servicebus.PurgeMessagesResult(DeleteMessagesResult): + property deleted_message_count: int # Read-only + + def __init__(self, deleted_message_count: int) -> None: ... + + class azure.servicebus.ServiceBusClient: implements ContextManager fully_qualified_namespace: str @@ -366,6 +378,14 @@ namespace azure.servicebus def defer_message(self, message: ServiceBusReceivedMessage) -> None: ... + def delete_messages( + self, + message_count: int, + *, + before_enqueued_time: Optional[datetime] = ..., + timeout: Optional[float] = ... + ) -> DeleteMessagesResult: ... + def peek_messages( self, max_message_count: int = 1, @@ -375,6 +395,14 @@ namespace azure.servicebus **kwargs: Any ) -> List[ServiceBusReceivedMessage]: ... + def purge_messages( + self, + *, + before_enqueued_time: Optional[datetime] = ..., + max_message_count_per_batch: int = 500, + timeout: Optional[float] = ... + ) -> PurgeMessagesResult: ... + def receive_deferred_messages( self, sequence_numbers: Union[int, List[int]], @@ -679,6 +707,14 @@ namespace azure.servicebus.aio async def defer_message(self, message: ServiceBusReceivedMessage) -> None: ... + async def delete_messages( + self, + message_count: int, + *, + before_enqueued_time: Optional[datetime] = ..., + timeout: Optional[float] = ... + ) -> DeleteMessagesResult: ... + async def peek_messages( self, max_message_count: int = 1, @@ -688,6 +724,14 @@ namespace azure.servicebus.aio **kwargs: Any ) -> List[ServiceBusReceivedMessage]: ... + async def purge_messages( + self, + *, + before_enqueued_time: Optional[datetime] = ..., + max_message_count_per_batch: int = 500, + timeout: Optional[float] = ... + ) -> PurgeMessagesResult: ... + async def receive_deferred_messages( self, sequence_numbers: Union[int, List[int]], diff --git a/sdk/servicebus/azure-servicebus/api.metadata.yml b/sdk/servicebus/azure-servicebus/api.metadata.yml index 114718312c3a..f98fc9a56495 100644 --- a/sdk/servicebus/azure-servicebus/api.metadata.yml +++ b/sdk/servicebus/azure-servicebus/api.metadata.yml @@ -1,4 +1,4 @@ -apiMdSha256: a5f9da8d7980799f054d782f99eabf9b190925718f2bc4bdd86f01752a5c1bf8 -packageVersion: 7.15.0 +apiMdSha256: 46c52fca7868fd19d0ec06b18161ba148cd8496df1ea6bdf1600296afa83a7c7 +packageVersion: 7.16.0b1 parserVersion: 0.3.31 -pythonVersion: 3.14.0 +pythonVersion: 3.11.14 diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py b/sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py index 2ab2fa09dfd4..c3939c772350 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py @@ -11,6 +11,7 @@ from ._servicebus_sender import ServiceBusSender from ._servicebus_receiver import ServiceBusReceiver from ._servicebus_session import ServiceBusSession +from ._models import DeleteMessagesResult, PurgeMessagesResult from ._common.message import ( ServiceBusMessage, ServiceBusMessageBatch, @@ -43,6 +44,8 @@ "ServiceBusReceiver", "ServiceBusSession", "ServiceBusSender", + "DeleteMessagesResult", + "PurgeMessagesResult", "TransportType", "AutoLockRenewer", "parse_connection_string", diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py index 4419bd701c6a..357483782ffe 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_base_handler.py @@ -384,6 +384,9 @@ def _do_retryable_operation( # pylint: disable=inconsistent-return-statements ) -> Any: require_last_exception = kwargs.pop("require_last_exception", False) operation_requires_timeout = kwargs.pop("operation_requires_timeout", False) + suppress_next_session_timeout_message = kwargs.pop( + "suppress_next_session_timeout_message", False + ) # Opt-in per call site so long-poll operations stay bounded only by the caller. apply_try_timeout = kwargs.pop("apply_try_timeout", False) try_timeout = self._config.try_timeout if apply_try_timeout else None @@ -420,7 +423,10 @@ def _do_retryable_operation( # pylint: disable=inconsistent-return-statements # Keep the original message: it identifies which phase timed out. # The session hint only applies to NEXT_AVAILABLE_SESSION receivers. description = str(last_exception) - if getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION: + if ( + getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION + and not suppress_next_session_timeout_message + ): description += ( " If trying to receive from NEXT_AVAILABLE_SESSION, use max_wait_time" " on the ServiceBusReceiver to control the timeout." @@ -434,6 +440,9 @@ def _do_retryable_operation( # pylint: disable=inconsistent-return-statements retried_times=retried_times, last_exception=last_exception, abs_timeout_time=abs_timeout_time, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), ) def _backoff( @@ -442,6 +451,7 @@ def _backoff( last_exception: Exception, abs_timeout_time: Optional[float] = None, entity_name: Optional[str] = None, + suppress_next_session_timeout_message: bool = False, ) -> None: entity_name = entity_name or self._container_id backoff = _get_backoff_time( @@ -468,7 +478,10 @@ def _backoff( if isinstance(last_exception, OperationTimeoutError): # Keep the original message: it identifies which phase timed out. description = str(last_exception) - if getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION: + if ( + getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION + and not suppress_next_session_timeout_message + ): description += ( " If trying to receive from NEXT_AVAILABLE_SESSION, use max_wait_time" " on the ServiceBusReceiver to control the timeout." @@ -570,8 +583,64 @@ def _mgmt_request_response_with_retry( def _open(self, timeout: Optional[float] = None): raise ValueError("Subclass should override the method.") - def _open_with_retry(self): - return self._do_retryable_operation(self._open, operation_requires_timeout=True, apply_try_timeout=True) + def _open_with_timeout(self, timeout: float): + return self._open(timeout=timeout) + + def _open_with_retry( + self, + timeout: Optional[float] = None, + *, + suppress_next_session_timeout_message: bool = False, + ): + def open_with_timeout(timeout: Optional[float] = None): + if timeout is not None and timeout <= 0: + raise OperationTimeoutError() + if timeout is None: + return self._open() + return self._open_with_timeout(timeout) + + return self._do_retryable_operation( + open_with_timeout, + timeout=timeout, + operation_requires_timeout=True, + apply_try_timeout=True, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), + ) + + def _open_mgmt_link_with_retry( + self, + timeout: Optional[float] = None, + *, + suppress_next_session_timeout_message: bool = False, + ): + def open_mgmt_link(timeout: Optional[float] = None): + if timeout is not None and timeout <= 0: + raise OperationTimeoutError() + start_time = time.monotonic() + if timeout is None: + self._open() + else: + self._open_with_timeout(timeout) + timeout -= time.monotonic() - start_time + if timeout <= 0: + raise OperationTimeoutError() + return self._amqp_transport.mgmt_client_setup( + self._handler, + node=self._mgmt_target.encode(self._config.encoding), + timeout=timeout, + ) + + return self._do_retryable_operation( + open_mgmt_link, + timeout=timeout, + operation_requires_timeout=True, + apply_try_timeout=True, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), + ) def _close_handler(self): if self._handler: diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py index c41a3d6c510b..350eacd59159 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py @@ -52,6 +52,7 @@ REQUEST_RESPONSE_ADD_RULE_OPERATION = VENDOR + b":add-rule" REQUEST_RESPONSE_REMOVE_RULE_OPERATION = VENDOR + b":remove-rule" REQUEST_RESPONSE_GET_RULES_OPERATION = VENDOR + b":enumerate-rules" +REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION = VENDOR + b":batch-delete-messages" CONTAINER_PREFIX = "servicebus.pysdk-" JWT_TOKEN_SCOPE = "https://servicebus.azure.net//.default" @@ -84,6 +85,8 @@ MGMT_REQUEST_RECEIVER_SETTLE_MODE = "receiver-settle-mode" MGMT_REQUEST_FROM_SEQUENCE_NUMBER = "from-sequence-number" MGMT_REQUEST_MAX_MESSAGE_COUNT = "message-count" +MGMT_REQUEST_MESSAGE_COUNT = MGMT_REQUEST_MAX_MESSAGE_COUNT +MGMT_REQUEST_ENQUEUED_TIME_UTC = "enqueued-time-utc" MGMT_REQUEST_MESSAGE = "message" MGMT_REQUEST_MESSAGES = "messages" MGMT_REQUEST_MESSAGE_ID = "message-id" diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/mgmt_handlers.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/mgmt_handlers.py index 2dc9e87f38f5..7ec99714d7eb 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/mgmt_handlers.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/mgmt_handlers.py @@ -69,6 +69,25 @@ def peek_op( # pylint: disable=inconsistent-return-statements ) +def batch_delete_op( # pylint: disable=inconsistent-return-statements + status_code, message, description, amqp_transport, max_message_count +): + condition = message.application_properties.get(MGMT_RESPONSE_MESSAGE_ERROR_CONDITION) + if status_code == 200 or (status_code == 404 and condition == ERROR_CODE_MESSAGE_NOT_FOUND): + deleted_count = message.value.get(b"message-count") if isinstance(message.value, dict) else None + if ( + isinstance(deleted_count, bool) + or not isinstance(deleted_count, int) + or deleted_count < 0 + or deleted_count > max_message_count + ): + raise ValueError("Batch delete response did not contain a valid message-count.") + return deleted_count + amqp_transport.handle_amqp_mgmt_error( + _LOGGER, "Batch delete messages failed.", condition, description, status_code + ) + + def list_sessions_op( # pylint: disable=inconsistent-return-statements status_code, message, description, amqp_transport ): diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/utils.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/utils.py index 9a98dab464e1..75789a4265b7 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/utils.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/utils.py @@ -74,6 +74,16 @@ # Number of seconds between the Unix epoch (1/1/1970) and year 1 CE. # This is the lowest value that can be represented by an AMQP timestamp. CE_ZERO_SECONDS: int = -62_135_596_800 +EPOCH_UTC: datetime.datetime = datetime.datetime(1970, 1, 1, tzinfo=TZ_UTC) + + +def datetime_to_timestamp_ms(value: datetime.datetime) -> int: + if value.tzinfo is None: + normalized = value.replace(tzinfo=TZ_UTC) + else: + normalized = value.astimezone(TZ_UTC) + return (normalized - EPOCH_UTC) // datetime.timedelta(milliseconds=1) + def utc_from_timestamp(timestamp: float) -> datetime.datetime: """ diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_models.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_models.py new file mode 100644 index 000000000000..e85a248759d7 --- /dev/null +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_models.py @@ -0,0 +1,24 @@ +# ------------------------------------------------------------------------ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + + +class DeleteMessagesResult: + """The result of a batch delete operation.""" + + def __init__(self, deleted_message_count: int) -> None: + self._deleted_message_count = deleted_message_count + + @property + def deleted_message_count(self) -> int: + """The number of messages deleted by the service. + + :rtype: int + """ + return self._deleted_message_count + + +class PurgeMessagesResult(DeleteMessagesResult): + """The result of a purge operation, reporting the total number of messages deleted.""" diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py index 75fa29046e45..63f63f8aef89 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py @@ -5,6 +5,7 @@ # -------------------------------------------------------------------------- # TODO: Check types of kwargs (issue exists for this) import asyncio # pylint:disable=do-not-import-asyncio +import inspect import logging import time import queue @@ -35,6 +36,7 @@ from ..error import AMQPError, ErrorCondition, AMQPException, MessageException _logger = logging.getLogger(__name__) +_MGMT_LINK_CLEANUP_TIMEOUT = 1.0 class AMQPClientAsync(AMQPClientSync): @@ -369,29 +371,12 @@ async def mgmt_request_async( :rtype: ~pyamqp.message.Message """ - # The method also takes "status_code_field" and "status_description_field" - # keyword arguments as alternate names for the status code and description - # in the response body. Those two keyword arguments are used in Azure services only. started = time.monotonic() - async with self._mgmt_link_lock_async: - try: - mgmt_link = self._mgmt_links[node] - except KeyError: - mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) - await mgmt_link.open() - self._mgmt_links[node] = mgmt_link - - while True: - _get_mgmt_request_remaining_timeout(timeout, started) - authenticated = await self.auth_complete_async() - remaining = _get_mgmt_request_remaining_timeout(timeout, started) - if authenticated: - break - await asyncio.sleep(min(0.05, remaining) if remaining else 0.05) - - while not await mgmt_link.ready(): - _get_mgmt_request_remaining_timeout(timeout, started) - await self._connection.listen(wait=False) + mgmt_link = await self.open_mgmt_link_async( + node=node, + timeout=_get_mgmt_request_remaining_timeout(timeout, started), + **kwargs + ) operation_type = operation_type or b"empty" status, description, response = await mgmt_link.execute( @@ -402,6 +387,53 @@ async def mgmt_request_async( ) return status, description, response + async def open_mgmt_link_async(self, node: Union[str, bytes] = "$management", timeout: float = 0, **kwargs): + """Open and wait for a management link without dispatching a request. + + Failure cleanup has a separate one-second budget, even after the setup timeout expires. + + :param str or bytes node: Management target. + :param float timeout: Timeout in seconds. + :returns: The opened management link. + :rtype: ~pyamqp.aio.management_link_async.ManagementOperation + """ + start_time = time.monotonic() + mgmt_link = None + created_link = False + try: + async with self._mgmt_link_lock_async: + try: + mgmt_link = self._mgmt_links[node] + except KeyError: + mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) + created_link = True + self._mgmt_links[node] = mgmt_link + await mgmt_link.open() + while True: + _get_mgmt_request_remaining_timeout(timeout, start_time) + authenticated = await self.auth_complete_async() + remaining = _get_mgmt_request_remaining_timeout(timeout, start_time) + if authenticated: + break + await asyncio.sleep(min(0.05, remaining) if remaining else 0.05) + + while not await mgmt_link.ready(): + _get_mgmt_request_remaining_timeout(timeout, start_time) + await self._connection.listen(wait=False) + return mgmt_link + except BaseException: + if created_link and mgmt_link is not None: + # Evict without yielding so another stalled setup cannot block error cleanup. + if self._mgmt_links.get(node) is mgmt_link: + self._mgmt_links.pop(node, None) + try: + close_result = mgmt_link.close() + if inspect.isawaitable(close_result): + await asyncio.wait_for(close_result, timeout=_MGMT_LINK_CLEANUP_TIMEOUT) + except BaseException: # pylint: disable=broad-except + _logger.warning("Management link cleanup failed; preserving the original error.", exc_info=True) + raise + class SendClientAsync(SendClientSync, AMQPClientAsync): """An asynchronous AMQP client. diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py index c48d2f1e14f8..7a83f2aa8dbe 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py @@ -467,29 +467,12 @@ def mgmt_request( :rtype: ~pyamqp.message.Message """ - # The method also takes "status_code_field" and "status_description_field" - # keyword arguments as alternate names for the status code and description - # in the response body. Those two keyword arguments are used in Azure services only. started = time.monotonic() - with self._mgmt_link_lock: - try: - mgmt_link = self._mgmt_links[node] - except KeyError: - mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) - self._mgmt_links[node] = mgmt_link - mgmt_link.open() - - while True: - _get_mgmt_request_remaining_timeout(timeout, started) - authenticated = self.auth_complete() - remaining = _get_mgmt_request_remaining_timeout(timeout, started) - if authenticated: - break - time.sleep(min(0.05, remaining) if remaining else 0.05) - - while not mgmt_link.ready(): - _get_mgmt_request_remaining_timeout(timeout, started) - self._connection.listen(wait=False) + mgmt_link = self.open_mgmt_link( + node=node, + timeout=_get_mgmt_request_remaining_timeout(timeout, started), + **kwargs + ) operation_type = operation_type or b"empty" status, description, response = mgmt_link.execute( message, @@ -500,6 +483,49 @@ def mgmt_request( _get_mgmt_request_remaining_timeout(timeout, started) return status, description, response + def open_mgmt_link(self, node: Union[str, bytes] = "$management", timeout: float = 0, **kwargs): + """Open and wait for a management link without dispatching a request. + + :param str or bytes node: Management target. + :param float timeout: Timeout in seconds. + :returns: The opened management link. + :rtype: ~pyamqp.management_link.ManagementOperation + """ + start_time = time.monotonic() + mgmt_link = None + created_link = False + try: + with self._mgmt_link_lock: + try: + mgmt_link = self._mgmt_links[node] + except KeyError: + mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs) + created_link = True + self._mgmt_links[node] = mgmt_link + mgmt_link.open() + while True: + _get_mgmt_request_remaining_timeout(timeout, start_time) + authenticated = self.auth_complete() + remaining = _get_mgmt_request_remaining_timeout(timeout, start_time) + if authenticated: + break + time.sleep(min(0.05, remaining) if remaining else 0.05) + + while not mgmt_link.ready(): + _get_mgmt_request_remaining_timeout(timeout, start_time) + self._connection.listen(wait=False) + return mgmt_link + except Exception: + if created_link and mgmt_link is not None: + with self._mgmt_link_lock: + if self._mgmt_links.get(node) is mgmt_link: + self._mgmt_links.pop(node, None) + try: + mgmt_link.close() + except Exception: # pylint: disable=broad-except + _logger.warning("Management link cleanup failed; preserving the original error.", exc_info=True) + raise + class SendClient(AMQPClient): """ diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py index f0443946652a..aa97e680c6d4 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py @@ -12,11 +12,13 @@ from enum import Enum from typing import Any, List, Optional, Dict, Iterator, Union, TYPE_CHECKING, cast -from .exceptions import MessageLockLostError +from .exceptions import MessageLockLostError, OperationTimeoutError from ._base_handler import BaseHandler from ._common.message import ServiceBusReceivedMessage +from ._models import DeleteMessagesResult, PurgeMessagesResult from ._common.utils import ( create_authentication, + datetime_to_timestamp_ms, get_attempt_timeout, get_link_ready_deadline, check_link_ready_deadline, @@ -37,6 +39,7 @@ REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION, REQUEST_RESPONSE_RENEWLOCK_OPERATION, REQUEST_RESPONSE_PEEK_OPERATION, + REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION, ServiceBusReceiveMode, MGMT_REQUEST_DISPOSITION_STATUS, MGMT_REQUEST_LOCK_TOKENS, @@ -44,6 +47,8 @@ MGMT_REQUEST_RECEIVER_SETTLE_MODE, MGMT_REQUEST_FROM_SEQUENCE_NUMBER, MGMT_REQUEST_MAX_MESSAGE_COUNT, + MGMT_REQUEST_MESSAGE_COUNT, + MGMT_REQUEST_ENQUEUED_TIME_UTC, MESSAGE_COMPLETE, MESSAGE_ABANDON, MESSAGE_DEFER, @@ -61,7 +66,10 @@ if TYPE_CHECKING: try: - from uamqp import ReceiveClient as uamqp_ReceiveClientSync, Message as uamqp_Message + from uamqp import ( + ReceiveClient as uamqp_ReceiveClientSync, + Message as uamqp_Message, + ) from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth except ImportError: pass @@ -79,7 +87,9 @@ _LOGGER = logging.getLogger(__name__) -class ServiceBusReceiver(BaseHandler, ReceiverMixin): # pylint: disable=too-many-instance-attributes +class ServiceBusReceiver( + BaseHandler, ReceiverMixin +): # pylint: disable=too-many-instance-attributes """The ServiceBusReceiver class defines a high level interface for receiving messages from the Azure Service Bus Queue or Topic Subscription. @@ -156,12 +166,16 @@ class ServiceBusReceiver(BaseHandler, ReceiverMixin): # pylint: disable=too-many def __init__( self, fully_qualified_namespace: str, - credential: Union["TokenCredential", "AzureSasCredential", "AzureNamedKeyCredential"], + credential: Union[ + "TokenCredential", "AzureSasCredential", "AzureNamedKeyCredential" + ], *, queue_name: Optional[str] = None, topic_name: Optional[str] = None, subscription_name: Optional[str] = None, - receive_mode: Union[ServiceBusReceiveMode, str] = ServiceBusReceiveMode.PEEK_LOCK, + receive_mode: Union[ + ServiceBusReceiveMode, str + ] = ServiceBusReceiveMode.PEEK_LOCK, max_wait_time: Optional[float] = None, auto_lock_renewer: Optional["AutoLockRenewer"] = None, prefetch_count: int = 0, @@ -185,12 +199,18 @@ def __init__( ) else: if queue_name and topic_name: - raise ValueError("Queue/Topic name can not be specified simultaneously.") + raise ValueError( + "Queue/Topic name can not be specified simultaneously." + ) if topic_name and not subscription_name: - raise ValueError("Subscription name is missing for the topic. Please specify subscription_name.") + raise ValueError( + "Subscription name is missing for the topic. Please specify subscription_name." + ) entity_name = queue_name or topic_name if not entity_name: - raise ValueError("Queue/Topic name is missing. Please specify queue_name/topic_name.") + raise ValueError( + "Queue/Topic name is missing. Please specify queue_name/topic_name." + ) super(ServiceBusReceiver, self).__init__( fully_qualified_namespace=fully_qualified_namespace, @@ -216,13 +236,19 @@ def __init__( prefetch_count=prefetch_count, **kwargs, ) - self._session = None if self._session_id is None else ServiceBusSession(cast(str, self._session_id), self) + self._session = ( + None + if self._session_id is None + else ServiceBusSession(cast(str, self._session_id), self) + ) self._receive_context = threading.Event() self._handler: Union["pyamqp_ReceiveClientSync", "uamqp_ReceiveClientSync"] self._build_received_message = functools.partial( self._amqp_transport.build_received_message, self, ServiceBusReceivedMessage ) - self._iter_contextual_wrapper = functools.partial(self._amqp_transport.iter_contextual_wrapper, self) + self._iter_contextual_wrapper = functools.partial( + self._amqp_transport.iter_contextual_wrapper, self + ) self._iter_next = functools.partial(self._amqp_transport.iter_next, self) def __enter__(self) -> "ServiceBusReceiver": @@ -237,12 +263,16 @@ def __enter__(self) -> "ServiceBusReceiver": def __iter__(self) -> Iterator["ServiceBusReceivedMessage"]: return self._iter_contextual_wrapper() - def _inner_next(self, wait_time: Optional[float] = None) -> "ServiceBusReceivedMessage": + def _inner_next( + self, wait_time: Optional[float] = None + ) -> "ServiceBusReceivedMessage": # We do this weird wrapping such that an imperitive next() call, and a generator-based iter both trace sanely. self._check_live() while True: try: - return self._do_retryable_operation(self._iter_next, wait_time=wait_time) + return self._do_retryable_operation( + self._iter_next, wait_time=wait_time + ) except StopIteration: self._message_iter = None raise @@ -261,7 +291,11 @@ def __next__(self) -> ServiceBusReceivedMessage: next = __next__ # for python2.7 @classmethod - def _from_connection_string(cls, conn_str: str, **kwargs: Any) -> "ServiceBusReceiver": # pylint: disable=docstring-keyword-should-match-keyword-only + def _from_connection_string( + cls, conn_str: str, **kwargs: Any + ) -> ( + "ServiceBusReceiver" + ): # pylint: disable=docstring-keyword-should-match-keyword-only """Create a ServiceBusReceiver from a connection string. :param conn_str: The connection string of a Service Bus. @@ -325,10 +359,14 @@ def _from_connection_string(cls, conn_str: str, **kwargs: Any) -> "ServiceBusRec raise ValueError("Queue entity does not have subscription.") if kwargs.get("topic_name") and not kwargs.get("subscription_name"): - raise ValueError("Subscription name is missing for the topic. Please specify subscription_name.") + raise ValueError( + "Subscription name is missing for the topic. Please specify subscription_name." + ) return cls(**constructor_args) - def _create_handler(self, auth: Union["pyamqp_JWTTokenAuth", "uamqp_JWTTokenAuth"]) -> None: + def _create_handler( + self, auth: Union["pyamqp_JWTTokenAuth", "uamqp_JWTTokenAuth"] + ) -> None: self._handler = self._amqp_transport.create_receive_client( receiver=self, @@ -339,12 +377,18 @@ def _create_handler(self, auth: Union["pyamqp_JWTTokenAuth", "uamqp_JWTTokenAuth retry_policy=self._error_policy, client_name=self._name, receive_mode=self._receive_mode, - timeout=self._max_wait_time * self._amqp_transport.TIMEOUT_FACTOR if self._max_wait_time else 0, + timeout=( + self._max_wait_time * self._amqp_transport.TIMEOUT_FACTOR + if self._max_wait_time + else 0 + ), # set link_credit to at least 1 so that messages can be received link_credit=self._prefetch_count + 1, # If prefetch is "off", then keep_alive coroutine frequently listens on the connection for messages and # releases right away, since no "prefetched" messages should be in the internal buffer. - keep_alive_interval=self._config.keep_alive if self._prefetch_count != 0 else 5, + keep_alive_interval=( + self._config.keep_alive if self._prefetch_count != 0 else 5 + ), shutdown_after_timeout=False, link_properties={CONSUMER_IDENTIFIER: self._name}, ) @@ -353,9 +397,14 @@ def _create_handler(self, auth: Union["pyamqp_JWTTokenAuth", "uamqp_JWTTokenAuth # If RECEIVE_AND_DELETE mode, messages are settled and removed from the Service Bus entity immediately, # so the regular _message_received callback should be used. This will ensure that all messages are added # to the internal buffer since they cannot be re-received, even if not received during an active receive call. - if self._prefetch_count == 0 and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK: + if ( + self._prefetch_count == 0 + and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK + ): # pylint: disable=protected-access - self._handler._message_received = functools.partial(self._amqp_transport.enhanced_message_received, self) + self._handler._message_received = functools.partial( + self._amqp_transport.enhanced_message_received, self + ) def _open(self, timeout: Optional[float] = None) -> None: # pylint: disable=protected-access @@ -421,7 +470,9 @@ def _receive( receive_deadline = abs_timeout batch: Union[List["uamqp_Message"], List["pyamqp_Message"]] = [] - while not received_messages_queue.empty() and len(batch) < max_message_count: + while ( + not received_messages_queue.empty() and len(batch) < max_message_count + ): batch.append(received_messages_queue.get()) received_messages_queue.task_done() if len(batch) >= max_message_count: @@ -432,19 +483,29 @@ def _receive( expired = remaining <= 0 if not expired and max_message_count and self._prefetch_count == 0 and max_message_count >= 1: link_credit_needed = max_message_count - len(batch) - self._amqp_transport.reset_link_credit(amqp_receive_client, link_credit_needed) + self._amqp_transport.reset_link_credit( + amqp_receive_client, link_credit_needed + ) first_message_received = False receiving = True while receiving and not expired and len(batch) < max_message_count: while receiving and received_messages_queue.qsize() < max_message_count: - if abs_timeout and self._amqp_transport.get_current_time(amqp_receive_client) > abs_timeout: + if ( + abs_timeout + and self._amqp_transport.get_current_time(amqp_receive_client) + > abs_timeout + ): expired = True break before = received_messages_queue.qsize() receiving = amqp_receive_client.do_work() received = received_messages_queue.qsize() - before - if not first_message_received and received_messages_queue.qsize() > 0 and received > 0: + if ( + not first_message_received + and received_messages_queue.qsize() > 0 + and received > 0 + ): # first message(s) received, continue receiving for some time first_message_received = True abs_timeout = min( @@ -452,7 +513,10 @@ def _receive( + self._further_pull_receive_timeout, receive_deadline, ) - while not received_messages_queue.empty() and len(batch) < max_message_count: + while ( + not received_messages_queue.empty() + and len(batch) < max_message_count + ): batch.append(received_messages_queue.get()) received_messages_queue.task_done() return [self._build_received_message(message) for message in batch] @@ -469,7 +533,9 @@ def _settle_message_with_retry( # pylint: disable=protected-access self._check_live() if not isinstance(message, ServiceBusReceivedMessage): - raise TypeError("Parameter 'message' must be of type ServiceBusReceivedMessage") + raise TypeError( + "Parameter 'message' must be of type ServiceBusReceivedMessage" + ) self._check_message_alive(message, settle_operation) # The following condition check is a hot fix for settling a message received for non-session queue after @@ -525,7 +591,8 @@ def _settle_message( dead_letter_details = ( { MGMT_REQUEST_DEAD_LETTER_REASON: dead_letter_reason or "", - MGMT_REQUEST_DEAD_LETTER_ERROR_DESCRIPTION: dead_letter_error_description or "", + MGMT_REQUEST_DEAD_LETTER_ERROR_DESCRIPTION: dead_letter_error_description + or "", } if settle_operation == MESSAGE_DEAD_LETTER else None @@ -551,7 +618,9 @@ def _settle_message_via_mgmt_link( ) -> Any: message = { MGMT_REQUEST_DISPOSITION_STATUS: settlement, - MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens), + MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE( + lock_tokens + ), } self._populate_message_properties(message) @@ -572,7 +641,9 @@ def _settle_message_via_mgmt_link( def _renew_locks(self, *lock_tokens: str, **kwargs: Any) -> Any: timeout = kwargs.pop("timeout", None) - message = {MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens)} + message = { + MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens) + } return self._mgmt_request_response_with_retry( REQUEST_RESPONSE_RENEWLOCK_OPERATION, message, @@ -618,7 +689,9 @@ def close(self) -> None: super(ServiceBusReceiver, self).close() self._message_iter = None - def _get_streaming_message_iter(self, max_wait_time: Optional[float] = None) -> Iterator[ServiceBusReceivedMessage]: + def _get_streaming_message_iter( + self, max_wait_time: Optional[float] = None + ) -> Iterator[ServiceBusReceivedMessage]: """Receive messages from an iterator indefinitely, or if a max_wait_time is specified, until such a timeout occurs. @@ -765,7 +838,9 @@ def receive_deferred_messages( MGMT_REQUEST_SEQUENCE_NUMBERS: self._amqp_transport.AMQP_ARRAY_VALUE( [self._amqp_transport.AMQP_LONG_VALUE(s) for s in sequence_numbers] ), - MGMT_REQUEST_RECEIVER_SETTLE_MODE: self._amqp_transport.AMQP_UINT_VALUE(receive_mode), + MGMT_REQUEST_RECEIVER_SETTLE_MODE: self._amqp_transport.AMQP_UINT_VALUE( + receive_mode + ), } self._populate_message_properties(message) @@ -785,7 +860,10 @@ def receive_deferred_messages( ) links = get_receive_links(messages) with receive_trace_context_manager( - self, span_name=SPAN_NAME_RECEIVE_DEFERRED, links=links, start_time=start_time + self, + span_name=SPAN_NAME_RECEIVE_DEFERRED, + links=links, + start_time=start_time, ): if ( self._auto_lock_renewer @@ -796,6 +874,138 @@ def receive_deferred_messages( self._auto_lock_renewer.register(self, message) return messages + def delete_messages( + self, + message_count: int, + *, + before_enqueued_time: Optional[datetime.datetime] = None, + timeout: Optional[float] = None, + ) -> DeleteMessagesResult: + """Permanently delete up to the requested number of eligible messages. + + Large messages can cause the service to delete fewer messages than requested. Locked, deferred, + and scheduled messages are not eligible. Currently, batch delete is not supported when partitioning + is enabled. + A dispatched request is not automatically retried. After an error, cancellation, or timeout, + the deletion outcome is unknown and no count is available. + + :param int message_count: The positive 32-bit maximum number of messages to delete. The service limit + is 500 for Basic and Standard and 4,000 for Premium. + :keyword Optional[datetime.datetime] before_enqueued_time: Only messages enqueued before this UTC time + can be deleted. The operation start time is used when omitted. + :keyword Optional[float] timeout: The operation timeout in seconds. The value must be greater than 0 + if specified. + :returns: The delete result containing the number of messages actually deleted. + :rtype: ~azure.servicebus.DeleteMessagesResult + """ + self._check_live() + if isinstance(message_count, bool) or not isinstance(message_count, int): + raise TypeError("The message_count must be an integer.") + if not 1 <= message_count <= 2_147_483_647: + raise ValueError("The message_count must be between 1 and 2147483647.") + if timeout is not None and timeout <= 0: + raise ValueError("The timeout must be greater than 0.") + + if before_enqueued_time is not None and not isinstance(before_enqueued_time, datetime.datetime): + raise TypeError("The before_enqueued_time must be a datetime.") + cutoff = ( + datetime.datetime.now(datetime.timezone.utc) if before_enqueued_time is None else before_enqueued_time + ) + message = { + MGMT_REQUEST_MESSAGE_COUNT: self._amqp_transport.AMQP_INT_VALUE( + message_count + ), + MGMT_REQUEST_ENQUEUED_TIME_UTC: self._amqp_transport.AMQP_TIMESTAMP_VALUE( + datetime_to_timestamp_ms(cutoff) + ), + } + start_time = time.monotonic() + self._open_with_retry( + timeout=timeout, + suppress_next_session_timeout_message=True, + ) + remaining_timeout = ( + None if timeout is None else timeout - (time.monotonic() - start_time) + ) + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + self._open_mgmt_link_with_retry( + timeout=remaining_timeout, + suppress_next_session_timeout_message=True, + ) + remaining_timeout = ( + None if timeout is None else timeout - (time.monotonic() - start_time) + ) + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + self._populate_message_properties(message) + try: + deleted_count = self._mgmt_request_response( + REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION, + message, + functools.partial( + mgmt_handlers.batch_delete_op, max_message_count=message_count + ), + timeout=get_attempt_timeout(remaining_timeout, self._config.try_timeout), + ) + except Exception as exception: # pylint: disable=broad-except + raise self._handle_exception(exception) from exception + return DeleteMessagesResult(cast(int, deleted_count)) + + def purge_messages( + self, + *, + before_enqueued_time: Optional[datetime.datetime] = None, + max_message_count_per_batch: int = 500, + timeout: Optional[float] = None, + ) -> PurgeMessagesResult: + """Permanently delete eligible messages enqueued before the purge started or configured time. + + The enqueue-time threshold stays unchanged for every request, so newer messages remain. Large messages + can produce smaller batches, which purge continues processing. Locked, deferred, and scheduled messages + remain. Currently, purge is not supported when partitioning is enabled. + If an error, cancellation, or timeout occurs after dispatch, the purge can be partial and its exact + deletion outcome is unknown. + + :keyword Optional[datetime.datetime] before_enqueued_time: Only messages enqueued before this UTC time + can be deleted. The purge start time is used when omitted. + :keyword int max_message_count_per_batch: The maximum number of messages requested in + each batch-delete call. The default is 500. The service limit is 500 for Basic and Standard and 4,000 + for Premium. + :keyword Optional[float] timeout: The timeout in seconds for the entire purge operation. The value must + be greater than 0 if specified. + :returns: The purge result containing the total number of messages actually deleted. + :rtype: ~azure.servicebus.PurgeMessagesResult + """ + if isinstance(max_message_count_per_batch, bool) or not isinstance( + max_message_count_per_batch, int + ): + raise TypeError("The max_message_count_per_batch must be an integer.") + if not 1 <= max_message_count_per_batch <= 2_147_483_647: + raise ValueError( + "The max_message_count_per_batch must be between 1 and 2147483647." + ) + if timeout is not None and timeout <= 0: + raise ValueError("The timeout must be greater than 0.") + + cutoff = ( + datetime.datetime.now(datetime.timezone.utc) if before_enqueued_time is None else before_enqueued_time + ) + deadline = None if timeout is None else time.monotonic() + timeout + deleted_count = 0 + while True: + remaining_timeout = None if deadline is None else deadline - time.monotonic() + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + result = self.delete_messages( + max_message_count_per_batch, + before_enqueued_time=cutoff, + timeout=remaining_timeout, + ) + deleted_count += result.deleted_message_count + if result.deleted_message_count == 0: + return PurgeMessagesResult(deleted_count) + def peek_messages( self, max_message_count: int = 1, @@ -842,18 +1052,24 @@ def peek_messages( raise ValueError("max_message_count must be 1 or greater.") message = { - MGMT_REQUEST_FROM_SEQUENCE_NUMBER: self._amqp_transport.AMQP_LONG_VALUE(sequence_number), + MGMT_REQUEST_FROM_SEQUENCE_NUMBER: self._amqp_transport.AMQP_LONG_VALUE( + sequence_number + ), MGMT_REQUEST_MAX_MESSAGE_COUNT: max_message_count, } self._populate_message_properties(message) - handler = functools.partial(mgmt_handlers.peek_op, receiver=self, amqp_transport=self._amqp_transport) + handler = functools.partial( + mgmt_handlers.peek_op, receiver=self, amqp_transport=self._amqp_transport + ) start_time = time.time_ns() messages = self._mgmt_request_response_with_retry( REQUEST_RESPONSE_PEEK_OPERATION, message, handler, timeout=timeout ) links = get_receive_links(messages) - with receive_trace_context_manager(self, span_name=SPAN_NAME_PEEK, links=links, start_time=start_time): + with receive_trace_context_manager( + self, span_name=SPAN_NAME_PEEK, links=links, start_time=start_time + ): return messages def complete_message(self, message: ServiceBusReceivedMessage) -> None: @@ -1027,7 +1243,9 @@ def renew_message_lock( raise ValueError("The timeout must be greater than 0.") expiry = self._renew_locks(token, timeout=timeout) # type: ignore - message._expiry = utc_from_timestamp(expiry[MGMT_RESPONSE_MESSAGE_EXPIRATION][0] / 1000.0) + message._expiry = utc_from_timestamp( + expiry[MGMT_RESPONSE_MESSAGE_EXPIRATION][0] / 1000.0 + ) return message._expiry # type: ignore @@ -1041,4 +1259,6 @@ def client_identifier(self) -> str: return self._name def __str__(self) -> str: - return f"Receiver client id: {self.client_identifier}, entity: {self.entity_path}" + return ( + f"Receiver client id: {self.client_identifier}, entity: {self.entity_path}" + ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py index af2cb53835e3..7ca9da800b37 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- from __future__ import annotations -from typing import Union, TYPE_CHECKING, Dict, Any, Callable +from typing import Union, TYPE_CHECKING, Dict, Any, Callable, Optional from abc import ABC, abstractmethod if TYPE_CHECKING: @@ -384,6 +384,16 @@ def create_mgmt_msg(message, application_properties, config, reply_to, **kwargs) :rtype: uamqp.Message or pyamqp.Message """ + @staticmethod + @abstractmethod + def mgmt_client_setup(mgmt_client, *, node: Union[str, bytes], timeout: Optional[float]) -> None: + """Open the management link without dispatching a request. + + :param AMQPClient mgmt_client: Client used to open the management link. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + @staticmethod @abstractmethod def mgmt_client_request(mgmt_client, mgmt_msg, *, operation, operation_type, node, timeout, callback): diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py index 6bad5dfffa17..be3c9feaefe5 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py @@ -1109,6 +1109,19 @@ def mgmt_client_request( ) return callback(status, response, description, amqp_transport=PyamqpTransport) + @staticmethod + def mgmt_client_setup(mgmt_client: "AMQPClient", *, node: Union[str, bytes], timeout: Optional[float]) -> None: + """Open the pyamqp management link without dispatching a request. + + :param ~pyamqp.AMQPClient mgmt_client: Client used for management requests. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + try: + mgmt_client.open_mgmt_link(node=node, timeout=timeout or 0) + except TimeoutError as exception: + raise OperationTimeoutError(error=exception) from exception + @staticmethod def _handle_amqp_exception_with_condition( logger: "Logger", diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py index 5f988cf4158f..092659579d40 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py @@ -1084,6 +1084,16 @@ def mgmt_client_request( callback=functools.partial(callback, amqp_transport=UamqpTransport), ) + @staticmethod + def mgmt_client_setup(mgmt_client, *, node: Union[str, bytes], timeout: Optional[float]) -> None: + """uAMQP does not expose management-link setup separately from request dispatch. + + :param ~uamqp.AMQPClient mgmt_client: Client used for management requests. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + del mgmt_client, node, timeout + @staticmethod def _handle_amqp_exception_with_condition( logger: "Logger", diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_version.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_version.py index 81e4031bc24b..1dfca6143b25 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_version.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_version.py @@ -3,4 +3,4 @@ # Licensed under the MIT License. # ------------------------------------ -VERSION = "7.15.1" +VERSION = "7.16.0b1" diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_base_handler_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_base_handler_async.py index 23d39ac3e140..a22acbee35a2 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_base_handler_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_base_handler_async.py @@ -233,6 +233,9 @@ def _check_live(self): async def _do_retryable_operation(self, operation: Callable, timeout: Optional[float] = None, **kwargs: Any) -> Any: require_last_exception = kwargs.pop("require_last_exception", False) operation_requires_timeout = kwargs.pop("operation_requires_timeout", False) + suppress_next_session_timeout_message = kwargs.pop( + "suppress_next_session_timeout_message", False + ) # Opt-in per call site so long-poll operations stay bounded only by the caller. apply_try_timeout = kwargs.pop("apply_try_timeout", False) try_timeout = self._config.try_timeout if apply_try_timeout else None @@ -267,7 +270,10 @@ async def _do_retryable_operation(self, operation: Callable, timeout: Optional[f ) if isinstance(last_exception, OperationTimeoutError): description = str(last_exception) - if getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION: + if ( + getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION + and not suppress_next_session_timeout_message + ): description += ( " If trying to receive from NEXT_AVAILABLE_SESSION, use max_wait_time" " on the ServiceBusReceiver to control the timeout." @@ -281,9 +287,19 @@ async def _do_retryable_operation(self, operation: Callable, timeout: Optional[f retried_times=retried_times, last_exception=last_exception, abs_timeout_time=abs_timeout_time, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), ) - async def _backoff(self, retried_times, last_exception, abs_timeout_time=None, entity_name=None): + async def _backoff( + self, + retried_times, + last_exception, + abs_timeout_time=None, + entity_name=None, + suppress_next_session_timeout_message=False, + ): entity_name = entity_name or self._container_id backoff = _get_backoff_time( self._config.retry_mode, @@ -308,7 +324,10 @@ async def _backoff(self, retried_times, last_exception, abs_timeout_time=None, e ) if isinstance(last_exception, OperationTimeoutError): description = str(last_exception) - if getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION: + if ( + getattr(self, "_session_id", None) == NEXT_AVAILABLE_SESSION + and not suppress_next_session_timeout_message + ): description += ( " If trying to receive from NEXT_AVAILABLE_SESSION, use max_wait_time" " on the ServiceBusReceiver to control the timeout." @@ -414,8 +433,69 @@ async def _mgmt_request_response_with_retry( async def _open(self, timeout: Optional[float] = None): raise ValueError("Subclass should override the method.") - async def _open_with_retry(self): - return await self._do_retryable_operation(self._open, operation_requires_timeout=True, apply_try_timeout=True) + async def _open_with_timeout(self, timeout: float): + return await self._open(timeout=timeout) + + async def _open_with_retry( + self, + timeout: Optional[float] = None, + *, + suppress_next_session_timeout_message: bool = False, + ): + async def open_with_timeout(timeout: Optional[float] = None): + if timeout is not None and timeout <= 0: + raise OperationTimeoutError() + if timeout is None: + return await self._open() + return await self._open_with_timeout(timeout) + + return await self._do_retryable_operation( + open_with_timeout, + timeout=timeout, + operation_requires_timeout=True, + apply_try_timeout=True, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), + ) + + async def _open_mgmt_link_with_retry( + self, + timeout: Optional[float] = None, + *, + suppress_next_session_timeout_message: bool = False, + ): + async def open_mgmt_link(timeout: Optional[float] = None): + if timeout is not None and timeout <= 0: + raise OperationTimeoutError() + start_time = time.monotonic() + deadline = None if timeout is None else start_time + timeout + if timeout is None: + await self._open() + else: + await self._open_with_timeout(timeout) + timeout -= time.monotonic() - start_time + if timeout <= 0: + raise OperationTimeoutError() + return await await_with_deadline( + self._amqp_transport.mgmt_client_setup_async( + self._handler, + node=self._mgmt_target.encode(self._config.encoding), + timeout=timeout, + ), + deadline, + "Timed out waiting for the management link to open.", + ) + + return await self._do_retryable_operation( + open_mgmt_link, + timeout=timeout, + operation_requires_timeout=True, + apply_try_timeout=True, + suppress_next_session_timeout_message=( + suppress_next_session_timeout_message + ), + ) async def _close_handler(self): if self._handler: diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py index 4ba334887efb..72fff7a18eed 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py @@ -10,12 +10,21 @@ import time import warnings from enum import Enum -from typing import Any, List, Optional, AsyncIterator as AsyncIteratorType, Union, TYPE_CHECKING, cast +from typing import ( + Any, + List, + Optional, + AsyncIterator as AsyncIteratorType, + Union, + TYPE_CHECKING, + cast, +) -from ..exceptions import MessageLockLostError +from ..exceptions import MessageLockLostError, OperationTimeoutError from ._servicebus_session_async import ServiceBusSession from ._base_handler_async import BaseHandler from .._common.message import ServiceBusReceivedMessage +from .._models import DeleteMessagesResult, PurgeMessagesResult from .._common.receiver_mixins import ReceiverMixin from .._common.constants import ( CONSUMER_IDENTIFIER, @@ -25,6 +34,7 @@ REQUEST_RESPONSE_PEEK_OPERATION, REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER, REQUEST_RESPONSE_RENEWLOCK_OPERATION, + REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION, ServiceBusReceiveMode, MGMT_REQUEST_DISPOSITION_STATUS, MGMT_REQUEST_LOCK_TOKENS, @@ -32,6 +42,8 @@ MGMT_REQUEST_RECEIVER_SETTLE_MODE, MGMT_REQUEST_FROM_SEQUENCE_NUMBER, MGMT_REQUEST_MAX_MESSAGE_COUNT, + MGMT_REQUEST_MESSAGE_COUNT, + MGMT_REQUEST_ENQUEUED_TIME_UTC, MESSAGE_COMPLETE, MESSAGE_DEAD_LETTER, MESSAGE_ABANDON, @@ -44,6 +56,7 @@ ) from .._common import mgmt_handlers from .._common.utils import ( + datetime_to_timestamp_ms, utc_from_timestamp, get_attempt_timeout, get_link_ready_deadline, @@ -67,7 +80,9 @@ if TYPE_CHECKING: try: - from uamqp.async_ops.client_async import ReceiveClientAsync as uamqp_ReceiveClientAsync + from uamqp.async_ops.client_async import ( + ReceiveClientAsync as uamqp_ReceiveClientAsync, + ) from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAuthAsync from uamqp.message import Message as uamqp_Message except ImportError: @@ -75,7 +90,9 @@ from ._transport._base_async import AmqpTransportAsync from .._pyamqp.message import Message as pyamqp_Message from .._pyamqp.aio import ReceiveClientAsync as pyamqp_ReceiveClientAsync - from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync as pyamqp_JWTTokenAuthAsync + from .._pyamqp.aio._authentication_async import ( + JWTTokenAuthAsync as pyamqp_JWTTokenAuthAsync, + ) from azure.core.credentials_async import AsyncTokenCredential from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential from ._async_auto_lock_renewer import AutoLockRenewer @@ -160,19 +177,25 @@ class ServiceBusReceiver(AsyncIterator, BaseHandler, ReceiverMixin): def __init__( self, fully_qualified_namespace: str, - credential: Union["AsyncTokenCredential", "AzureSasCredential", "AzureNamedKeyCredential"], + credential: Union[ + "AsyncTokenCredential", "AzureSasCredential", "AzureNamedKeyCredential" + ], *, queue_name: Optional[str] = None, topic_name: Optional[str] = None, subscription_name: Optional[str] = None, - receive_mode: Union[ServiceBusReceiveMode, str] = ServiceBusReceiveMode.PEEK_LOCK, + receive_mode: Union[ + ServiceBusReceiveMode, str + ] = ServiceBusReceiveMode.PEEK_LOCK, max_wait_time: Optional[float] = None, auto_lock_renewer: Optional["AutoLockRenewer"] = None, prefetch_count: int = 0, **kwargs: Any, ) -> None: self._session_id = None - self._message_iter: Optional[AsyncIteratorType[Union["uamqp_Message", "pyamqp_Message"]]] = None + self._message_iter: Optional[ + AsyncIteratorType[Union["uamqp_Message", "pyamqp_Message"]] + ] = None self._amqp_transport: "AmqpTransportAsync" if kwargs.get("entity_name"): super(ServiceBusReceiver, self).__init__( @@ -189,11 +212,17 @@ def __init__( ) else: if queue_name and topic_name: - raise ValueError("Queue/Topic name can not be specified simultaneously.") + raise ValueError( + "Queue/Topic name can not be specified simultaneously." + ) if not (queue_name or topic_name): - raise ValueError("Queue/Topic name is missing. Please specify queue_name/topic_name.") + raise ValueError( + "Queue/Topic name is missing. Please specify queue_name/topic_name." + ) if topic_name and not subscription_name: - raise ValueError("Subscription name is missing for the topic. Please specify subscription_name.") + raise ValueError( + "Subscription name is missing for the topic. Please specify subscription_name." + ) entity_name = queue_name or topic_name @@ -221,14 +250,20 @@ def __init__( prefetch_count=prefetch_count, **kwargs, ) - self._session = None if self._session_id is None else ServiceBusSession(cast(str, self._session_id), self) + self._session = ( + None + if self._session_id is None + else ServiceBusSession(cast(str, self._session_id), self) + ) self._receive_context = asyncio.Event() self._handler: Union["pyamqp_ReceiveClientAsync", "uamqp_ReceiveClientAsync"] self._build_received_message = functools.partial( self._amqp_transport.build_received_message, self, ServiceBusReceivedMessage ) - self._iter_contextual_wrapper = functools.partial(self._amqp_transport.iter_contextual_wrapper_async, self) + self._iter_contextual_wrapper = functools.partial( + self._amqp_transport.iter_contextual_wrapper_async, self + ) self._iter_next = functools.partial(self._amqp_transport.iter_next_async, self) async def __aenter__(self) -> "ServiceBusReceiver": @@ -242,12 +277,16 @@ async def __aenter__(self) -> "ServiceBusReceiver": def __aiter__(self) -> AsyncIteratorType[ServiceBusReceivedMessage]: return self._iter_contextual_wrapper() - async def _inner_anext(self, wait_time: Optional[float] = None) -> ServiceBusReceivedMessage: + async def _inner_anext( + self, wait_time: Optional[float] = None + ) -> ServiceBusReceivedMessage: # We do this weird wrapping such that an imperitive next() call, and a generator-based iter both trace sanely. self._check_live() while True: try: - return await self._do_retryable_operation(self._iter_next, wait_time=wait_time) + return await self._do_retryable_operation( + self._iter_next, wait_time=wait_time + ) except StopAsyncIteration: self._message_iter = None raise @@ -263,7 +302,11 @@ async def __anext__(self) -> ServiceBusReceivedMessage: self._receive_context.clear() @classmethod - def _from_connection_string(cls, conn_str: str, **kwargs: Any) -> "ServiceBusReceiver": # pylint: disable=docstring-keyword-should-match-keyword-only + def _from_connection_string( + cls, conn_str: str, **kwargs: Any + ) -> ( + "ServiceBusReceiver" + ): # pylint: disable=docstring-keyword-should-match-keyword-only """Create a ServiceBusReceiver from a connection string. :param str conn_str: The connection string of a Service Bus. @@ -326,10 +369,14 @@ def _from_connection_string(cls, conn_str: str, **kwargs: Any) -> "ServiceBusRec raise ValueError("Queue entity does not have subscription.") if kwargs.get("topic_name") and not kwargs.get("subscription_name"): - raise ValueError("Subscription name is missing for the topic. Please specify subscription_name.") + raise ValueError( + "Subscription name is missing for the topic. Please specify subscription_name." + ) return cls(**constructor_args) - def _create_handler(self, auth: Union["pyamqp_JWTTokenAuthAsync", "uamqp_JWTTokenAuthAsync"]) -> None: + def _create_handler( + self, auth: Union["pyamqp_JWTTokenAuthAsync", "uamqp_JWTTokenAuthAsync"] + ) -> None: self._handler = self._amqp_transport.create_receive_client_async( receiver=self, @@ -340,12 +387,18 @@ def _create_handler(self, auth: Union["pyamqp_JWTTokenAuthAsync", "uamqp_JWTToke retry_policy=self._error_policy, client_name=self._name, receive_mode=self._receive_mode, - timeout=self._max_wait_time * self._amqp_transport.TIMEOUT_FACTOR if self._max_wait_time else 0, + timeout=( + self._max_wait_time * self._amqp_transport.TIMEOUT_FACTOR + if self._max_wait_time + else 0 + ), # set link_credit to at least 1 so that messages can be received link_credit=self._prefetch_count + 1, # If prefetch is 0, then keep_alive coroutine frequently listens on the connection for messages and # releases right away, since no "prefetched" messages should be in the internal buffer. - keep_alive_interval=self._config.keep_alive if self._prefetch_count != 0 else 5, + keep_alive_interval=( + self._config.keep_alive if self._prefetch_count != 0 else 5 + ), shutdown_after_timeout=False, link_properties={CONSUMER_IDENTIFIER: self._name}, ) @@ -354,7 +407,10 @@ def _create_handler(self, auth: Union["pyamqp_JWTTokenAuthAsync", "uamqp_JWTToke # If RECEIVE_AND_DELETE mode, messages are settled and removed from the Service Bus entity immediately, # so the regular _message_received callback should be used. This will ensure that all messages are added # to the internal buffer since they cannot be re-received, even if not received during an active receive call. - if self._prefetch_count == 0 and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK: + if ( + self._prefetch_count == 0 + and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK + ): self._amqp_transport.set_handler_message_received_async(self) async def _open(self, timeout: Optional[float] = None) -> None: @@ -427,7 +483,9 @@ async def _receive( batch: Union[List["uamqp_Message"], List["pyamqp_Message"]] = [] - while not received_messages_queue.empty() and len(batch) < max_message_count: + while ( + not received_messages_queue.empty() and len(batch) < max_message_count + ): batch.append(received_messages_queue.get()) received_messages_queue.task_done() if len(batch) >= max_message_count: @@ -438,19 +496,29 @@ async def _receive( expired = remaining <= 0 if not expired and max_message_count and self._prefetch_count == 0 and max_message_count >= 1: link_credit_needed = max_message_count - len(batch) - await self._amqp_transport.reset_link_credit_async(amqp_receive_client, link_credit_needed) + await self._amqp_transport.reset_link_credit_async( + amqp_receive_client, link_credit_needed + ) first_message_received = False receiving = True while receiving and not expired and len(batch) < max_message_count: while receiving and received_messages_queue.qsize() < max_message_count: - if abs_timeout and self._amqp_transport.get_current_time(amqp_receive_client) > abs_timeout: + if ( + abs_timeout + and self._amqp_transport.get_current_time(amqp_receive_client) + > abs_timeout + ): expired = True break before = received_messages_queue.qsize() receiving = await amqp_receive_client.do_work_async() received = received_messages_queue.qsize() - before - if not first_message_received and received_messages_queue.qsize() > 0 and received > 0: + if ( + not first_message_received + and received_messages_queue.qsize() > 0 + and received > 0 + ): # first message(s) received, continue receiving for some time first_message_received = True abs_timeout = min( @@ -458,7 +526,10 @@ async def _receive( + self._further_pull_receive_timeout, receive_deadline, ) - while not received_messages_queue.empty() and len(batch) < max_message_count: + while ( + not received_messages_queue.empty() + and len(batch) < max_message_count + ): batch.append(received_messages_queue.get()) received_messages_queue.task_done() return [self._build_received_message(message) for message in batch] @@ -475,7 +546,9 @@ async def _settle_message_with_retry( # pylint: disable=protected-access self._check_live() if not isinstance(message, ServiceBusReceivedMessage): - raise TypeError("Parameter 'message' must be of type ServiceBusReceivedMessage") + raise TypeError( + "Parameter 'message' must be of type ServiceBusReceivedMessage" + ) self._check_message_alive(message, settle_operation) # The following condition check is a hot fix for settling a message received for non-session queue after @@ -531,7 +604,8 @@ async def _settle_message( # type: ignore dead_letter_details = ( { MGMT_REQUEST_DEAD_LETTER_REASON: dead_letter_reason or "", - MGMT_REQUEST_DEAD_LETTER_ERROR_DESCRIPTION: dead_letter_error_description or "", + MGMT_REQUEST_DEAD_LETTER_ERROR_DESCRIPTION: dead_letter_error_description + or "", } if settle_operation == MESSAGE_DEAD_LETTER else None @@ -549,10 +623,14 @@ async def _settle_message( # type: ignore ) raise - async def _settle_message_via_mgmt_link(self, settlement, lock_tokens, dead_letter_details=None): + async def _settle_message_via_mgmt_link( + self, settlement, lock_tokens, dead_letter_details=None + ): message = { MGMT_REQUEST_DISPOSITION_STATUS: settlement, - MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens), + MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE( + lock_tokens + ), } self._populate_message_properties(message) @@ -570,8 +648,12 @@ async def _settle_message_via_mgmt_link(self, settlement, lock_tokens, dead_lett ), ) - async def _renew_locks(self, *lock_tokens: str, timeout: Optional[float] = None) -> Any: - message = {MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens)} + async def _renew_locks( + self, *lock_tokens: str, timeout: Optional[float] = None + ) -> Any: + message = { + MGMT_REQUEST_LOCK_TOKENS: self._amqp_transport.AMQP_ARRAY_VALUE(lock_tokens) + } return await self._mgmt_request_response_with_retry( REQUEST_RESPONSE_RENEWLOCK_OPERATION, message, @@ -717,7 +799,11 @@ async def receive_messages( return messages async def receive_deferred_messages( - self, sequence_numbers: Union[int, List[int]], *, timeout: Optional[float] = None, **kwargs: Any + self, + sequence_numbers: Union[int, List[int]], + *, + timeout: Optional[float] = None, + **kwargs: Any, ) -> List[ServiceBusReceivedMessage]: """Receive messages that have previously been deferred. @@ -760,7 +846,9 @@ async def receive_deferred_messages( MGMT_REQUEST_SEQUENCE_NUMBERS: self._amqp_transport.AMQP_ARRAY_VALUE( [self._amqp_transport.AMQP_LONG_VALUE(s) for s in sequence_numbers] ), - MGMT_REQUEST_RECEIVER_SETTLE_MODE: self._amqp_transport.AMQP_UINT_VALUE(receive_mode), + MGMT_REQUEST_RECEIVER_SETTLE_MODE: self._amqp_transport.AMQP_UINT_VALUE( + receive_mode + ), } self._populate_message_properties(message) @@ -780,7 +868,10 @@ async def receive_deferred_messages( ) links = get_receive_links(messages) with receive_trace_context_manager( - self, span_name=SPAN_NAME_RECEIVE_DEFERRED, links=links, start_time=start_time + self, + span_name=SPAN_NAME_RECEIVE_DEFERRED, + links=links, + start_time=start_time, ): if ( self._auto_lock_renewer @@ -791,8 +882,146 @@ async def receive_deferred_messages( self._auto_lock_renewer.register(self, message) return messages + async def delete_messages( + self, + message_count: int, + *, + before_enqueued_time: Optional[datetime.datetime] = None, + timeout: Optional[float] = None, + ) -> DeleteMessagesResult: + """Permanently delete up to the requested number of eligible messages. + + Large messages can cause the service to delete fewer messages than requested. Locked, deferred, and + scheduled messages are not eligible. Currently, batch delete is not supported when partitioning is + enabled. A dispatched request is not automatically retried; an error, cancellation, or timeout leaves + an unknown deletion outcome. + After management-link setup fails, error-path cleanup can take up to one additional second. + + :param int message_count: The positive 32-bit maximum number of messages to delete. The service limit + is 500 for Basic and Standard and 4,000 for Premium. + :keyword Optional[datetime.datetime] before_enqueued_time: Only messages enqueued before this UTC time + can be deleted. The operation start time is used when omitted. + :keyword Optional[float] timeout: The operation timeout in seconds. The value must be greater than 0 + if specified. + :returns: The delete result containing the number of messages actually deleted. + :rtype: ~azure.servicebus.DeleteMessagesResult + """ + self._check_live() + if isinstance(message_count, bool) or not isinstance(message_count, int): + raise TypeError("The message_count must be an integer.") + if not 1 <= message_count <= 2_147_483_647: + raise ValueError("The message_count must be between 1 and 2147483647.") + if timeout is not None and timeout <= 0: + raise ValueError("The timeout must be greater than 0.") + + if before_enqueued_time is not None and not isinstance(before_enqueued_time, datetime.datetime): + raise TypeError("The before_enqueued_time must be a datetime.") + cutoff = ( + datetime.datetime.now(datetime.timezone.utc) if before_enqueued_time is None else before_enqueued_time + ) + message = { + MGMT_REQUEST_MESSAGE_COUNT: self._amqp_transport.AMQP_INT_VALUE( + message_count + ), + MGMT_REQUEST_ENQUEUED_TIME_UTC: self._amqp_transport.AMQP_TIMESTAMP_VALUE( + datetime_to_timestamp_ms(cutoff) + ), + } + start_time = time.monotonic() + await self._open_with_retry( + timeout=timeout, + suppress_next_session_timeout_message=True, + ) + remaining_timeout = ( + None if timeout is None else timeout - (time.monotonic() - start_time) + ) + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + await self._open_mgmt_link_with_retry( + timeout=remaining_timeout, + suppress_next_session_timeout_message=True, + ) + remaining_timeout = ( + None if timeout is None else timeout - (time.monotonic() - start_time) + ) + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + self._populate_message_properties(message) + try: + deleted_count = await self._mgmt_request_response( + REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION, + message, + functools.partial( + mgmt_handlers.batch_delete_op, max_message_count=message_count + ), + timeout=get_attempt_timeout(remaining_timeout, self._config.try_timeout), + ) + except Exception as exception: # pylint: disable=broad-except + raise await self._handle_exception(exception) from exception + return DeleteMessagesResult(cast(int, deleted_count)) + + async def purge_messages( + self, + *, + before_enqueued_time: Optional[datetime.datetime] = None, + max_message_count_per_batch: int = 500, + timeout: Optional[float] = None, + ) -> PurgeMessagesResult: + """Permanently delete eligible messages enqueued before the purge started or configured time. + + The enqueue-time threshold stays unchanged for every request, so newer messages remain. Large messages + can produce smaller batches, which purge continues processing. Locked, deferred, and scheduled messages + remain. Currently, purge is not supported when partitioning is enabled. + If an error, cancellation, or timeout occurs after dispatch, the purge can be partial and its exact + deletion outcome is unknown. + After management-link setup fails, error-path cleanup can take up to one additional second. + + :keyword Optional[datetime.datetime] before_enqueued_time: Only messages enqueued before this UTC time + can be deleted. The purge start time is used when omitted. + :keyword int max_message_count_per_batch: The maximum number of messages requested in + each batch-delete call. The default is 500. The service limit is 500 for Basic and Standard and 4,000 + for Premium. + :keyword Optional[float] timeout: The timeout in seconds for the entire purge operation. The value must + be greater than 0 if specified. + :returns: The purge result containing the total number of messages actually deleted. + :rtype: ~azure.servicebus.PurgeMessagesResult + """ + if isinstance(max_message_count_per_batch, bool) or not isinstance( + max_message_count_per_batch, int + ): + raise TypeError("The max_message_count_per_batch must be an integer.") + if not 1 <= max_message_count_per_batch <= 2_147_483_647: + raise ValueError( + "The max_message_count_per_batch must be between 1 and 2147483647." + ) + if timeout is not None and timeout <= 0: + raise ValueError("The timeout must be greater than 0.") + + cutoff = ( + datetime.datetime.now(datetime.timezone.utc) if before_enqueued_time is None else before_enqueued_time + ) + deadline = None if timeout is None else time.monotonic() + timeout + deleted_count = 0 + while True: + remaining_timeout = None if deadline is None else deadline - time.monotonic() + if remaining_timeout is not None and remaining_timeout <= 0: + raise OperationTimeoutError() + result = await self.delete_messages( + max_message_count_per_batch, + before_enqueued_time=cutoff, + timeout=remaining_timeout, + ) + deleted_count += result.deleted_message_count + if result.deleted_message_count == 0: + return PurgeMessagesResult(deleted_count) + async def peek_messages( - self, max_message_count: int = 1, *, sequence_number: int = 0, timeout: Optional[float] = None, **kwargs: Any + self, + max_message_count: int = 1, + *, + sequence_number: int = 0, + timeout: Optional[float] = None, + **kwargs: Any, ) -> List[ServiceBusReceivedMessage]: """Browse messages currently pending in the queue. @@ -832,18 +1061,24 @@ async def peek_messages( message = { - MGMT_REQUEST_FROM_SEQUENCE_NUMBER: self._amqp_transport.AMQP_LONG_VALUE(sequence_number), + MGMT_REQUEST_FROM_SEQUENCE_NUMBER: self._amqp_transport.AMQP_LONG_VALUE( + sequence_number + ), MGMT_REQUEST_MAX_MESSAGE_COUNT: max_message_count, } self._populate_message_properties(message) - handler = functools.partial(mgmt_handlers.peek_op, receiver=self, amqp_transport=self._amqp_transport) + handler = functools.partial( + mgmt_handlers.peek_op, receiver=self, amqp_transport=self._amqp_transport + ) start_time = time.time_ns() messages = await self._mgmt_request_response_with_retry( REQUEST_RESPONSE_PEEK_OPERATION, message, handler, timeout=timeout ) links = get_receive_links(messages) - with receive_trace_context_manager(self, span_name=SPAN_NAME_PEEK, links=links, start_time=start_time): + with receive_trace_context_manager( + self, span_name=SPAN_NAME_PEEK, links=links, start_time=start_time + ): return messages async def complete_message(self, message: ServiceBusReceivedMessage) -> None: @@ -920,7 +1155,10 @@ async def defer_message(self, message: ServiceBusReceivedMessage) -> None: await self._settle_message_with_retry(message, MESSAGE_DEFER) async def dead_letter_message( - self, message: ServiceBusReceivedMessage, reason: Optional[str] = None, error_description: Optional[str] = None + self, + message: ServiceBusReceivedMessage, + reason: Optional[str] = None, + error_description: Optional[str] = None, ) -> None: """Move the message to the Dead Letter queue. @@ -955,7 +1193,11 @@ async def dead_letter_message( ) async def renew_message_lock( - self, message: ServiceBusReceivedMessage, *, timeout: Optional[float] = None, **kwargs: Any + self, + message: ServiceBusReceivedMessage, + *, + timeout: Optional[float] = None, + **kwargs: Any, ) -> datetime.datetime: # pylint: disable=protected-access """Renew the message lock. @@ -1024,4 +1266,6 @@ def client_identifier(self) -> str: return self._name def __str__(self) -> str: - return f"Receiver client id: {self.client_identifier}, entity: {self.entity_path}" + return ( + f"Receiver client id: {self.client_identifier}, entity: {self.entity_path}" + ) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py index 8d4ca77dacf2..a1941544343d 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py @@ -4,7 +4,7 @@ # -------------------------------------------------------------------------------------------- from __future__ import annotations from abc import ABC, abstractmethod -from typing import Tuple, Union, TYPE_CHECKING, Any, Dict, Callable +from typing import Tuple, Union, TYPE_CHECKING, Any, Dict, Callable, Optional from typing_extensions import Literal if TYPE_CHECKING: @@ -39,6 +39,8 @@ class AmqpTransportAsync(ABC): # pylint: disable=too-many-public-methods AMQP_LONG_VALUE: Callable AMQP_ARRAY_VALUE: Callable AMQP_UINT_VALUE: Callable + AMQP_INT_VALUE: Callable + AMQP_TIMESTAMP_VALUE: Callable @staticmethod @abstractmethod @@ -320,6 +322,16 @@ async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwa :rtype: ~pyamqp.aio._authentication_async.JWTTokenAuth or ~uamqp.authentication.JWTTokenAuth """ + @staticmethod + @abstractmethod + async def mgmt_client_setup_async(mgmt_client, *, node: Union[str, bytes], timeout: Optional[float]) -> None: + """Open the management link without dispatching a request. + + :param AMQPClient mgmt_client: Client used to open the management link. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + @staticmethod @abstractmethod async def mgmt_client_request_async(mgmt_client, mgmt_msg, *, operation, operation_type, node, timeout, callback): diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py index 45b975ef9cfd..1e22021d30f7 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py @@ -555,3 +555,18 @@ async def mgmt_client_request_async( timeout=timeout, # TODO: check if this should be seconds * 1000 if timeout else None, ) return callback(status, response, description, amqp_transport=PyamqpTransportAsync) + + @staticmethod + async def mgmt_client_setup_async( + mgmt_client: "AMQPClientAsync", *, node: Union[str, bytes], timeout: Optional[float] + ) -> None: + """Open the pyamqp management link without dispatching a request. + + :param ~pyamqp.aio.AMQPClientAsync mgmt_client: Client used for management requests. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + try: + await mgmt_client.open_mgmt_link_async(node=node, timeout=timeout or 0) + except TimeoutError as exception: + raise OperationTimeoutError(error=exception) from exception diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py index bd3b42aaeed6..cdb961b51d2d 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py @@ -379,5 +379,15 @@ async def mgmt_client_request_async( callback=functools.partial(callback, amqp_transport=UamqpTransportAsync), ) + @staticmethod + async def mgmt_client_setup_async(mgmt_client, *, node: Union[str, bytes], timeout: Optional[float]) -> None: + """uAMQP does not expose management-link setup separately from request dispatch. + + :param ~uamqp.AMQPClientAsync mgmt_client: Client used for management requests. + :keyword str or bytes node: Management target. + :keyword float or None timeout: Timeout in seconds, or None for no timeout. + """ + del mgmt_client, node, timeout + except ImportError: pass diff --git a/sdk/servicebus/azure-servicebus/pyproject.toml b/sdk/servicebus/azure-servicebus/pyproject.toml index 4d1d9557b729..96b975ea21f0 100644 --- a/sdk/servicebus/azure-servicebus/pyproject.toml +++ b/sdk/servicebus/azure-servicebus/pyproject.toml @@ -12,7 +12,7 @@ keywords = ["azure", "azure sdk"] requires-python = ">=3.9" license = "MIT" classifiers = [ - "Development Status :: 5 - Production/Stable", + "Development Status :: 4 - Beta", "Programming Language :: Python", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3", diff --git a/sdk/servicebus/azure-servicebus/samples/async_samples/sample_code_servicebus_async.py b/sdk/servicebus/azure-servicebus/samples/async_samples/sample_code_servicebus_async.py index fcae29cc7907..f525f9e20044 100644 --- a/sdk/servicebus/azure-servicebus/samples/async_samples/sample_code_servicebus_async.py +++ b/sdk/servicebus/azure-servicebus/samples/async_samples/sample_code_servicebus_async.py @@ -35,7 +35,9 @@ def example_create_servicebus_client_async(): from azure.servicebus.aio import ServiceBusClient servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) # [END create_sb_client_from_conn_str_async] # [START create_sb_client_async] @@ -61,7 +63,9 @@ async def example_create_servicebus_sender_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - queue_sender = ServiceBusSender._from_connection_string(conn_str=servicebus_connection_str, queue_name=queue_name) + queue_sender = ServiceBusSender._from_connection_string( + conn_str=servicebus_connection_str, queue_name=queue_name + ) # [END create_servicebus_sender_from_conn_str_async] # [START create_servicebus_sender_from_sb_client_async] @@ -70,7 +74,9 @@ async def example_create_servicebus_sender_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: queue_sender = servicebus_client.get_queue_sender(queue_name=queue_name) # [END create_servicebus_sender_from_sb_client_async] @@ -81,7 +87,9 @@ async def example_create_servicebus_sender_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: topic_sender = servicebus_client.get_topic_sender(topic_name=topic_name) # [END create_topic_sender_from_sb_client_async] @@ -111,7 +119,9 @@ async def example_create_servicebus_receiver_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: queue_receiver = servicebus_client.get_queue_receiver( queue_name=queue_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER @@ -124,7 +134,9 @@ async def example_create_servicebus_receiver_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: queue_receiver = servicebus_client.get_queue_receiver(queue_name=queue_name) # [END create_servicebus_receiver_from_sb_client_async] @@ -137,10 +149,14 @@ async def example_create_servicebus_receiver_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] subscription_name = os.environ["SERVICEBUS_SUBSCRIPTION_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: subscription_receiver = servicebus_client.get_subscription_receiver( - topic_name=topic_name, subscription_name=subscription_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER + topic_name=topic_name, + subscription_name=subscription_name, + sub_queue=ServiceBusSubQueue.DEAD_LETTER, ) # [END create_subscription_deadletter_receiver_from_sb_client_async] @@ -151,7 +167,9 @@ async def example_create_servicebus_receiver_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] subscription_name = os.environ["SERVICEBUS_SUBSCRIPTION_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) async with servicebus_client: subscription_receiver = servicebus_client.get_subscription_receiver( topic_name=topic_name, @@ -244,7 +262,9 @@ async def example_send_and_receive_async(): lock_renewal = AutoLockRenewer() async with servicebus_receiver: async for message in servicebus_receiver: - lock_renewal.register(servicebus_receiver, message, max_lock_renewal_duration=60) + lock_renewal.register( + servicebus_receiver, message, max_lock_renewal_duration=60 + ) await process_message(message) await servicebus_receiver.complete_message(message) # [END auto_lock_renew_message_async] @@ -281,21 +301,29 @@ async def example_receive_deadletter_async(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - async with ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) as servicebus_client: + async with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: async with servicebus_client.get_queue_sender(queue_name) as servicebus_sender: await servicebus_sender.send_messages(ServiceBusMessage("Hello World")) # [START receive_deadletter_async] - async with servicebus_client.get_queue_receiver(queue_name) as servicebus_receiver: + async with servicebus_client.get_queue_receiver( + queue_name + ) as servicebus_receiver: messages = await servicebus_receiver.receive_messages(max_wait_time=5) for message in messages: await servicebus_receiver.dead_letter_message( - message, reason="reason for dead lettering", error_description="description for dead lettering" + message, + reason="reason for dead lettering", + error_description="description for dead lettering", ) async with servicebus_client.get_queue_receiver( queue_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER ) as servicebus_deadletter_receiver: - messages = await servicebus_deadletter_receiver.receive_messages(max_wait_time=5) + messages = await servicebus_deadletter_receiver.receive_messages( + max_wait_time=5 + ) for message in messages: await servicebus_deadletter_receiver.complete_message(message) # [END receive_deadletter_async] @@ -306,30 +334,40 @@ async def example_session_ops_async(): queue_name = os.environ["SERVICEBUS_SESSION_QUEUE_NAME"] session_id = "" - async with ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) as servicebus_client: + async with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: async with servicebus_client.get_queue_sender(queue_name=queue_name) as sender: await sender.send_messages(ServiceBusMessage("msg", session_id=session_id)) # [START get_session_async] - async with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session # [END get_session_async] # [START get_session_state_async] - async with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session session_state = await session.get_state() # [END get_session_state_async] # [START set_session_state_async] - async with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session await session.set_state("START") # [END set_session_state_async] # [START session_renew_lock_async] - async with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session await session.renew_lock() # [END session_renew_lock_async] @@ -338,7 +376,9 @@ async def example_session_ops_async(): from azure.servicebus.aio import AutoLockRenewer lock_renewal = AutoLockRenewer() - async with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session # Auto renew session lock for 2 minutes lock_renewal.register(receiver, session, max_lock_renewal_duration=120) @@ -349,13 +389,82 @@ async def example_session_ops_async(): break +async def example_delete_and_purge_messages_async(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] + + # [START delete_and_purge_messages_async] + async with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + async with servicebus_client.get_queue_receiver(queue_name) as receiver: + requested_count = 100 + delete_result = await receiver.delete_messages(requested_count) + # Any request can return fewer deletions than requested, especially when messages are large. + print( + f"Requested {requested_count}; the service deleted {delete_result.deleted_message_count}." + ) + + # The default purge uses 500-message batches. + purge_result = await receiver.purge_messages() + print( + f"The service purged {purge_result.deleted_message_count} remaining messages." + ) + # [END delete_and_purge_messages_async] + + +async def example_purge_messages_advanced_async(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] + + # [START purge_messages_advanced_async] + async with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + async with servicebus_client.get_queue_receiver(queue_name) as receiver: + enqueue_time_threshold = datetime.datetime.now(datetime.timezone.utc) + # Premium supports up to 4,000 messages per request. + result = await receiver.purge_messages( + before_enqueued_time=enqueue_time_threshold, + max_message_count_per_batch=4000, + ) + print( + f"Purged {result.deleted_message_count} messages enqueued before {enqueue_time_threshold}." + ) + + # If a destructive call raises after dispatch, its exact outcome can be unknown. + # Inspect your application state before deciding whether another purge is appropriate. + # [END purge_messages_advanced_async] + + +async def example_purge_messages_from_session_async(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_SESSION_QUEUE_NAME"] + session_id = os.environ["SERVICEBUS_SESSION_ID"] + + # [START purge_messages_from_session_async] + async with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + async with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as session_receiver: + result = await session_receiver.purge_messages() + print( + f"Removed {result.deleted_message_count} messages from session {session_id}." + ) + # [END purge_messages_from_session_async] + + async def example_schedule_ops_async(): servicebus_sender = await example_create_servicebus_sender_async() # [START scheduling_messages_async] async with servicebus_sender: scheduled_time_utc = datetime.datetime.utcnow() + datetime.timedelta(seconds=30) scheduled_messages = [ServiceBusMessage("Scheduled message") for _ in range(10)] - sequence_nums = await servicebus_sender.schedule_messages(scheduled_messages, scheduled_time_utc) + sequence_nums = await servicebus_sender.schedule_messages( + scheduled_messages, scheduled_time_utc + ) # [END scheduling_messages_async] servicebus_sender = await example_create_servicebus_sender_async() @@ -371,3 +480,5 @@ async def example_schedule_ops_async(): asyncio.run(example_schedule_ops_async()) asyncio.run(example_receive_deadletter_async()) asyncio.run(example_session_ops_async()) + # Delete and purge examples are intentionally not run automatically. + # Call them explicitly only against disposable test entities. diff --git a/sdk/servicebus/azure-servicebus/samples/sync_samples/sample_code_servicebus.py b/sdk/servicebus/azure-servicebus/samples/sync_samples/sample_code_servicebus.py index 1c8e9fe0dae5..91610cdb1dcc 100644 --- a/sdk/servicebus/azure-servicebus/samples/sync_samples/sample_code_servicebus.py +++ b/sdk/servicebus/azure-servicebus/samples/sync_samples/sample_code_servicebus.py @@ -32,7 +32,9 @@ def example_create_servicebus_client_sync(): from azure.servicebus import ServiceBusClient servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) # [END create_sb_client_from_conn_str_sync] # [START create_sb_client_sync] @@ -68,7 +70,9 @@ def example_create_servicebus_sender_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - queue_sender = ServiceBusSender._from_connection_string(conn_str=servicebus_connection_str, queue_name=queue_name) + queue_sender = ServiceBusSender._from_connection_string( + conn_str=servicebus_connection_str, queue_name=queue_name + ) # [END create_servicebus_sender_from_conn_str_sync] # [START create_servicebus_sender_from_sb_client_sync] @@ -77,7 +81,9 @@ def example_create_servicebus_sender_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: queue_sender = servicebus_client.get_queue_sender(queue_name=queue_name) # [END create_servicebus_sender_from_sb_client_sync] @@ -88,7 +94,9 @@ def example_create_servicebus_sender_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: topic_sender = servicebus_client.get_topic_sender(topic_name=topic_name) # [END create_topic_sender_from_sb_client_sync] @@ -117,7 +125,9 @@ def example_create_servicebus_receiver_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: queue_dlq_receiver = servicebus_client.get_queue_receiver( queue_name=queue_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER @@ -130,7 +140,9 @@ def example_create_servicebus_receiver_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: queue_receiver = servicebus_client.get_queue_receiver(queue_name=queue_name) # [END create_servicebus_receiver_from_sb_client_sync] @@ -142,7 +154,9 @@ def example_create_servicebus_receiver_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] subscription_name = os.environ["SERVICEBUS_SUBSCRIPTION_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: subscription_receiver = servicebus_client.get_subscription_receiver( topic_name=topic_name, @@ -157,10 +171,14 @@ def example_create_servicebus_receiver_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] topic_name = os.environ["SERVICEBUS_TOPIC_NAME"] subscription_name = os.environ["SERVICEBUS_SUBSCRIPTION_NAME"] - servicebus_client = ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) + servicebus_client = ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) with servicebus_client: subscription_dlq_receiver = servicebus_client.get_subscription_receiver( - topic_name=topic_name, subscription_name=subscription_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER + topic_name=topic_name, + subscription_name=subscription_name, + sub_queue=ServiceBusSubQueue.DEAD_LETTER, ) # [END create_subscription_deadletter_receiver_from_sb_client_sync] @@ -212,7 +230,9 @@ def example_send_and_receive_sync(): with servicebus_receiver: for message in servicebus_receiver: # Auto renew message for 1 minute. - lock_renewal.register(servicebus_receiver, message, max_lock_renewal_duration=60) + lock_renewal.register( + servicebus_receiver, message, max_lock_renewal_duration=60 + ) process_message(message) servicebus_receiver.complete_message(message) # [END auto_lock_renew_message_sync] @@ -231,12 +251,16 @@ def example_send_and_receive_sync(): from typing import List from azure.servicebus import ServiceBusReceivedMessage - messages_complex: List[ServiceBusReceivedMessage] = servicebus_receiver.receive_messages(max_wait_time=5) + messages_complex: List[ServiceBusReceivedMessage] = ( + servicebus_receiver.receive_messages(max_wait_time=5) + ) for message in messages_complex: print("Receiving: {}".format(message)) print("Time to live: {}".format(message.time_to_live)) print("Sequence number: {}".format(message.sequence_number)) - print("Enqueued Sequence number: {}".format(message.enqueued_sequence_number)) + print( + "Enqueued Sequence number: {}".format(message.enqueued_sequence_number) + ) print("Partition Key: {}".format(message.partition_key)) print("Application Properties: {}".format(message.application_properties)) print("Delivery count: {}".format(message.delivery_count)) @@ -315,7 +339,9 @@ def example_receive_deadletter_sync(): servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] - with ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) as servicebus_client: + with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: with servicebus_client.get_queue_sender(queue_name) as servicebus_sender: servicebus_sender.send_messages(ServiceBusMessage("Hello World")) # [START receive_deadletter_sync] @@ -323,13 +349,17 @@ def example_receive_deadletter_sync(): messages = servicebus_receiver.receive_messages(max_wait_time=5) for message in messages: servicebus_receiver.dead_letter_message( - message, reason="reason for dead lettering", error_description="description for dead lettering" + message, + reason="reason for dead lettering", + error_description="description for dead lettering", ) with servicebus_client.get_queue_receiver( queue_name, sub_queue=ServiceBusSubQueue.DEAD_LETTER ) as servicebus_deadletter_receiver: - messages_deadletter = servicebus_deadletter_receiver.receive_messages(max_wait_time=5) + messages_deadletter = servicebus_deadletter_receiver.receive_messages( + max_wait_time=5 + ) for message in messages_deadletter: servicebus_deadletter_receiver.complete_message(message) # [END receive_deadletter_sync] @@ -340,30 +370,40 @@ def example_session_ops_sync(): queue_name = os.environ["SERVICEBUS_SESSION_QUEUE_NAME"] session_id = os.environ["SERVICEBUS_SESSION_ID"] - with ServiceBusClient.from_connection_string(conn_str=servicebus_connection_str) as servicebus_client: + with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: with servicebus_client.get_queue_sender(queue_name=queue_name) as sender: sender.send_messages(ServiceBusMessage("msg", session_id=session_id)) # [START get_session_sync] - with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session # [END get_session_sync] # [START get_session_state_sync] - with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session session_state = session.get_state() # [END get_session_state_sync] # [START set_session_state_sync] - with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session session.set_state("START") # [END set_session_state_sync] # [START session_renew_lock_sync] - with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session session.renew_lock() # [END session_renew_lock_sync] @@ -372,7 +412,9 @@ def example_session_ops_sync(): from azure.servicebus import AutoLockRenewer lock_renewal = AutoLockRenewer(max_workers=4) - with servicebus_client.get_queue_receiver(queue_name=queue_name, session_id=session_id) as receiver: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as receiver: session = receiver.session # Auto renew session lock for 2 minutes lock_renewal.register(receiver, session, max_lock_renewal_duration=120) @@ -383,13 +425,82 @@ def example_session_ops_sync(): break +def example_delete_and_purge_messages_sync(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] + + # [START delete_and_purge_messages_sync] + with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + with servicebus_client.get_queue_receiver(queue_name) as receiver: + requested_count = 100 + delete_result = receiver.delete_messages(requested_count) + # Any request can return fewer deletions than requested, especially when messages are large. + print( + f"Requested {requested_count}; the service deleted {delete_result.deleted_message_count}." + ) + + # The default purge uses 500-message batches. + purge_result = receiver.purge_messages() + print( + f"The service purged {purge_result.deleted_message_count} remaining messages." + ) + # [END delete_and_purge_messages_sync] + + +def example_purge_messages_advanced_sync(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_QUEUE_NAME"] + + # [START purge_messages_advanced_sync] + with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + with servicebus_client.get_queue_receiver(queue_name) as receiver: + enqueue_time_threshold = datetime.datetime.now(datetime.timezone.utc) + # Premium supports up to 4,000 messages per request. + result = receiver.purge_messages( + before_enqueued_time=enqueue_time_threshold, + max_message_count_per_batch=4000, + ) + print( + f"Purged {result.deleted_message_count} messages enqueued before {enqueue_time_threshold}." + ) + + # If a destructive call raises after dispatch, its exact outcome can be unknown. + # Inspect your application state before deciding whether another purge is appropriate. + # [END purge_messages_advanced_sync] + + +def example_purge_messages_from_session_sync(): + servicebus_connection_str = os.environ["SERVICEBUS_CONNECTION_STR"] + queue_name = os.environ["SERVICEBUS_SESSION_QUEUE_NAME"] + session_id = os.environ["SERVICEBUS_SESSION_ID"] + + # [START purge_messages_from_session_sync] + with ServiceBusClient.from_connection_string( + conn_str=servicebus_connection_str + ) as servicebus_client: + with servicebus_client.get_queue_receiver( + queue_name=queue_name, session_id=session_id + ) as session_receiver: + result = session_receiver.purge_messages() + print( + f"Removed {result.deleted_message_count} messages from session {session_id}." + ) + # [END purge_messages_from_session_sync] + + def example_schedule_ops_sync(): servicebus_sender = example_create_servicebus_sender_sync() # [START scheduling_messages] with servicebus_sender: scheduled_time_utc = datetime.datetime.utcnow() + datetime.timedelta(seconds=30) scheduled_messages = [ServiceBusMessage("Scheduled message") for _ in range(10)] - sequence_nums = servicebus_sender.schedule_messages(scheduled_messages, scheduled_time_utc) + sequence_nums = servicebus_sender.schedule_messages( + scheduled_messages, scheduled_time_utc + ) # [END scheduling_messages] servicebus_sender = example_create_servicebus_sender_sync() @@ -405,3 +516,5 @@ def example_schedule_ops_sync(): example_schedule_ops_sync() example_receive_deadletter_sync() example_session_ops_sync() +# Delete and purge examples are intentionally not run automatically. +# Call them explicitly only against disposable test entities. diff --git a/sdk/servicebus/azure-servicebus/tests/async_tests/test_queues_async.py b/sdk/servicebus/azure-servicebus/tests/async_tests/test_queues_async.py index feb3ec9f7e4f..d32be627d56f 100644 --- a/sdk/servicebus/azure-servicebus/tests/async_tests/test_queues_async.py +++ b/sdk/servicebus/azure-servicebus/tests/async_tests/test_queues_async.py @@ -82,6 +82,32 @@ class TestServiceBusQueueAsync(AzureMgmtRecordedTestCase): + @pytest.mark.asyncio + @pytest.mark.liveTest + @pytest.mark.live_test_only + @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest") + @CachedServiceBusNamespacePreparer(name_prefix="servicebustest") + @ServiceBusQueuePreparer(name_prefix="batchdelete") + @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids) + @ArgPasserAsync() + async def test_async_queue_batch_delete_and_purge_wire_contract( + self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs + ): + fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}" + async with ServiceBusClient( + fully_qualified_namespace, get_credential(is_async=True), uamqp_transport=uamqp_transport + ) as sb_client: + async with sb_client.get_queue_sender(servicebus_queue.name) as sender: + await sender.send_messages([ServiceBusMessage(str(i)) for i in range(3)], timeout=30) + async with sb_client.get_queue_receiver(servicebus_queue.name) as receiver: + result = await receiver.delete_messages(1, timeout=30) + assert result.deleted_message_count == 1 + result = await receiver.purge_messages(max_message_count_per_batch=1, timeout=30) + assert result.deleted_message_count == 2 + assert (await receiver.delete_messages(1, timeout=30)).deleted_message_count == 0 + assert (await receiver.purge_messages(timeout=30)).deleted_message_count == 0 + assert await receiver.peek_messages(max_message_count=1, timeout=30) == [] + @pytest.mark.asyncio @pytest.mark.liveTest @pytest.mark.live_test_only diff --git a/sdk/servicebus/azure-servicebus/tests/test_queues.py b/sdk/servicebus/azure-servicebus/tests/test_queues.py index fee5bd018844..663748222dcd 100644 --- a/sdk/servicebus/azure-servicebus/tests/test_queues.py +++ b/sdk/servicebus/azure-servicebus/tests/test_queues.py @@ -87,6 +87,31 @@ # are ported to offline-compatible code. class TestServiceBusQueue(AzureMgmtRecordedTestCase): + @pytest.mark.liveTest + @pytest.mark.live_test_only + @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest") + @CachedServiceBusNamespacePreparer(name_prefix="servicebustest") + @ServiceBusQueuePreparer(name_prefix="batchdelete") + @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids) + @ArgPasser() + def test_queue_batch_delete_and_purge_wire_contract( + self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs + ): + fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}" + with ServiceBusClient( + fully_qualified_namespace, get_credential(), uamqp_transport=uamqp_transport + ) as sb_client: + with sb_client.get_queue_sender(servicebus_queue.name) as sender: + sender.send_messages([ServiceBusMessage(str(i)) for i in range(3)], timeout=30) + with sb_client.get_queue_receiver(servicebus_queue.name) as receiver: + result = receiver.delete_messages(1, timeout=30) + assert result.deleted_message_count == 1 + result = receiver.purge_messages(max_message_count_per_batch=1, timeout=30) + assert result.deleted_message_count == 2 + assert receiver.delete_messages(1, timeout=30).deleted_message_count == 0 + assert receiver.purge_messages(timeout=30).deleted_message_count == 0 + assert receiver.peek_messages(max_message_count=1, timeout=30) == [] + @pytest.mark.liveTest @pytest.mark.live_test_only @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest") diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_batch_delete.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_batch_delete.py new file mode 100644 index 000000000000..38014183dfb2 --- /dev/null +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_batch_delete.py @@ -0,0 +1,1247 @@ +# ------------------------------------------------------------------------ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- +import ast +import asyncio +import threading +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from azure.servicebus import DeleteMessagesResult, ServiceBusClient, ServiceBusReceiver, ServiceBusSubQueue +from azure.servicebus.aio import ServiceBusClient as AsyncServiceBusClient, ServiceBusReceiver as AsyncServiceBusReceiver +from azure.servicebus._common.constants import ( + ERROR_CODE_MESSAGE_NOT_FOUND, + MGMT_REQUEST_ENQUEUED_TIME_UTC, + MGMT_REQUEST_MESSAGE_COUNT, + MGMT_REQUEST_SESSION_ID, + MGMT_RESPONSE_MESSAGE_ERROR_CONDITION, + REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION, + NEXT_AVAILABLE_SESSION, + ServiceBusReceiveMode, +) +from azure.servicebus._common import mgmt_handlers +from azure.servicebus._base_handler import BaseHandler +from azure.servicebus.aio._base_handler_async import BaseHandler as AsyncBaseHandler +from azure.servicebus._pyamqp._encode import encode_payload +from azure.servicebus._pyamqp.client import AMQPClient +from azure.servicebus._pyamqp.error import AMQPConnectionError +from azure.servicebus._pyamqp.aio._client_async import AMQPClientAsync +from azure.servicebus._transport._pyamqp_transport import PyamqpTransport +from azure.servicebus.aio._transport._pyamqp_transport_async import PyamqpTransportAsync +from azure.servicebus.exceptions import OperationTimeoutError, ServiceBusConnectionError + + +@pytest.mark.asyncio +@pytest.mark.parametrize("handler_type", [BaseHandler, AsyncBaseHandler, ServiceBusReceiver, AsyncServiceBusReceiver]) +async def test_default_open_with_timeout_forwards_budget(handler_type): + is_async = handler_type in (AsyncBaseHandler, AsyncServiceBusReceiver) + handler = object.__new__(handler_type) + handler._open = AsyncMock() if is_async else MagicMock() + if is_async: + await handler._open_with_timeout(1.5) + handler._open.assert_awaited_once_with(timeout=1.5) + else: + handler._open_with_timeout(1.5) + handler._open.assert_called_once_with(timeout=1.5) + + +@pytest.mark.parametrize("sample", [ + ("sync_samples", "sample_code_servicebus.py"), + ("async_samples", "sample_code_servicebus_async.py"), +]) +def test_destructive_sample_examples_are_standalone(sample): + source = Path(__file__).resolve().parents[2].joinpath("samples", *sample).read_text(encoding="utf-8") + tree = ast.parse(source) + startup = ast.Module( + body=[node for node in tree.body if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))], + type_ignores=[], + ) + assert not any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id.startswith(("example_delete", "example_purge")) + for node in ast.walk(startup) + ) + + +@pytest.mark.parametrize("uamqp_transport", [False, True], ids=["pyamqp", "uamqp"]) +def test_batch_delete_request_keys_encode_as_amqp_strings(uamqp_transport): + transport = PyamqpTransport + if uamqp_transport: + pytest.importorskip("uamqp") + from azure.servicebus._transport._uamqp_transport import UamqpTransport + transport = UamqpTransport + message = transport.create_mgmt_msg( + message={ + MGMT_REQUEST_MESSAGE_COUNT: transport.AMQP_INT_VALUE(1), + MGMT_REQUEST_ENQUEUED_TIME_UTC: transport.AMQP_TIMESTAMP_VALUE(2), + }, + application_properties={}, + config=MagicMock(encoding="UTF-8"), + reply_to="queue/$management", + ) + if uamqp_transport: + encoded = message.encode_message() + else: + encoded = bytearray() + encode_payload(encoded, message) + + assert b"\xa1\x0dmessage-count" in encoded + assert b"\xa1\x11enqueued-time-utc" in encoded + assert b"\xa0\x0dmessage-count" not in encoded + assert b"\xa0\x11enqueued-time-utc" not in encoded + + +@pytest.mark.parametrize("deleted_count", [0, 2]) +def test_batch_delete_uamqp_response_round_trip(deleted_count): + uamqp = pytest.importorskip("uamqp") + from azure.servicebus._transport._uamqp_transport import UamqpTransport + response = uamqp.Message(body={b"message-count": deleted_count}, application_properties={b"status-code": 200}) + decoded = uamqp.Message.decode_from_bytes(response.encode_message()) + assert mgmt_handlers.batch_delete_op(200, decoded, None, UamqpTransport, 10) == deleted_count + + +def test_batch_delete_handler_returns_actual_count_and_rejects_204(): + message = MagicMock() + message.value = {b"message-count": 2} + message.application_properties = {} + transport = MagicMock() + + assert mgmt_handlers.batch_delete_op(200, message, None, transport, 10) == 2 + + mgmt_handlers.batch_delete_op(204, message, "missing count", transport, 10) + transport.handle_amqp_mgmt_error.assert_called_once() + transport.reset_mock() + mgmt_handlers.batch_delete_op(202, message, "unexpected", transport, 10) + transport.handle_amqp_mgmt_error.assert_called_once() + + +@pytest.mark.parametrize("deleted_count", [-1, 1.5, 11, True, None]) +def test_batch_delete_handler_rejects_malformed_count(deleted_count): + message = MagicMock() + message.value = {b"message-count": deleted_count} + message.application_properties = {} + + with pytest.raises(ValueError, match="valid message-count"): + mgmt_handlers.batch_delete_op(200, message, None, MagicMock(), 10) + + +def test_batch_delete_handler_maps_message_not_found_to_aggregate(): + message = MagicMock() + message.value = {b"message-count": 2} + message.application_properties = { + MGMT_RESPONSE_MESSAGE_ERROR_CONDITION: ERROR_CODE_MESSAGE_NOT_FOUND + } + + assert ( + mgmt_handlers.batch_delete_op(404, message, None, MagicMock(), 10) == 2 + ) + + +@pytest.mark.parametrize("value", [None, {}, {b"message-count": "2"}]) +def test_batch_delete_handler_rejects_message_not_found_without_valid_count(value): + message = MagicMock() + message.value = value + message.application_properties = { + MGMT_RESPONSE_MESSAGE_ERROR_CONDITION: ERROR_CODE_MESSAGE_NOT_FOUND + } + + with pytest.raises(ValueError, match="valid message-count"): + mgmt_handlers.batch_delete_op(404, message, None, MagicMock(), 10) + + +def _receiver(receiver_type, session_id=None): + receiver = object.__new__(receiver_type) + receiver._check_live = lambda: None + receiver._session_id = session_id + receiver._session = object() if session_id else None + receiver._handler = MagicMock() + receiver._config = MagicMock(try_timeout=None, custom_endpoint_address=None) + receiver._handle_exception = ( + AsyncMock(side_effect=lambda error: error) + if receiver_type is AsyncServiceBusReceiver else MagicMock(side_effect=lambda error: error) + ) + receiver._amqp_transport = MagicMock() + receiver._amqp_transport.AMQP_INT_VALUE = lambda value: value + receiver._amqp_transport.AMQP_TIMESTAMP_VALUE = lambda value: value + receiver._open_with_retry = ( + AsyncMock() if receiver_type is AsyncServiceBusReceiver else MagicMock() + ) + receiver._open_mgmt_link_with_retry = ( + AsyncMock() if receiver_type is AsyncServiceBusReceiver else MagicMock() + ) + receiver._mgmt_request_response_with_retry = MagicMock( + side_effect=AssertionError("destructive requests must not be retried") + ) + return receiver + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("operation", ["delete_messages", "purge_messages"]) +@pytest.mark.parametrize("cutoff", [False, 0, 0.0, "", [], {}, "2026-01-02"]) +async def test_batch_delete_rejects_invalid_cutoff_before_open(is_async, operation, cutoff): + receiver = _receiver(AsyncServiceBusReceiver if is_async else ServiceBusReceiver) + receiver._mgmt_request_response = AsyncMock(return_value=0) if is_async else MagicMock(return_value=0) + args = (1,) if operation == "delete_messages" else () + with pytest.raises(TypeError, match="before_enqueued_time"): + result = getattr(receiver, operation)(*args, before_enqueued_time=cutoff) + if is_async: + await result + receiver._open_with_retry.assert_not_called() + receiver._open_mgmt_link_with_retry.assert_not_called() + receiver._mgmt_request_response.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("timeout", [None, 300, 2]) +async def test_batch_delete_dispatch_honors_try_timeout(is_async, timeout): + receiver = _receiver(AsyncServiceBusReceiver if is_async else ServiceBusReceiver) + receiver._config.try_timeout = 5 + dispatch = AsyncMock(return_value=1) if is_async else MagicMock(return_value=1) + receiver._mgmt_request_response = dispatch + if is_async: + await receiver.delete_messages(1, timeout=timeout) + else: + receiver.delete_messages(1, timeout=timeout) + budget = dispatch.call_args.kwargs["timeout"] + assert 0 < budget <= (2 if timeout == 2 else 5) + if timeout != 2: + assert budget == 5 + dispatch.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("session", [False, True]) +@pytest.mark.parametrize("uamqp_transport", [False, True]) +async def test_batch_delete_dispatch_translates_connection_failure(is_async, session, uamqp_transport): + receiver = _receiver(AsyncServiceBusReceiver if is_async else ServiceBusReceiver) + handler = AsyncBaseHandler if is_async else BaseHandler + receiver._handle_exception = handler._handle_exception.__get__(receiver) + error = AMQPConnectionError(condition=b"amqp:connection:forced") + transport = PyamqpTransportAsync if is_async else PyamqpTransport + if uamqp_transport: + uamqp = pytest.importorskip("uamqp") + from azure.servicebus._transport._uamqp_transport import UamqpTransport + error = uamqp.errors.AMQPConnectionError("connection lost") + transport = UamqpTransport + receiver._amqp_transport.create_servicebus_exception = transport.create_servicebus_exception + receiver._session = MagicMock(_lock_lost=False) if session else None + receiver._running = True + receiver._close_handler = AsyncMock() if is_async else MagicMock() + dispatch = AsyncMock(side_effect=error) if is_async else MagicMock(side_effect=error) + receiver._mgmt_request_response = dispatch + with pytest.raises(ServiceBusConnectionError): + if is_async: + await receiver.delete_messages(1) + else: + receiver.delete_messages(1) + dispatch.assert_called_once() + receiver._close_handler.assert_called_once() + if session: + assert receiver._session._lock_lost is True + receiver._mgmt_request_response_with_retry.assert_not_called() + + +@pytest.mark.parametrize("cleanup_error", [None, RuntimeError("close failed")]) +def test_failed_management_link_open_is_evicted_before_retry(cleanup_error): + client = object.__new__(AMQPClient) + client._mgmt_links = {} + client._mgmt_link_lock = threading.Lock() + client._session = MagicMock() + client.auth_complete = MagicMock(return_value=True) + failed_link = MagicMock() + failed_link.open.side_effect = RuntimeError("attach failed") + failed_link.close.side_effect = cleanup_error + ready_link = MagicMock() + ready_link.ready.return_value = True + + with patch( + "azure.servicebus._pyamqp.client.ManagementOperation", + side_effect=[failed_link, ready_link], + ) as operation_type: + with pytest.raises(RuntimeError, match="attach failed"): + client.open_mgmt_link() + assert client._mgmt_links == {} + failed_link.close.assert_called_once() + + assert client.open_mgmt_link() is ready_link + + assert operation_type.call_count == 2 + + +def test_management_link_setup_converts_timeout_error(): + client = MagicMock() + client.open_mgmt_link.side_effect = TimeoutError("setup timed out") + + with pytest.raises(OperationTimeoutError): + PyamqpTransport.mgmt_client_setup(client, node="$management", timeout=1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + "setup_error", [TimeoutError("setup timed out"), RuntimeError("setup failed"), asyncio.CancelledError()] +) +async def test_cached_management_link_is_not_closed_by_failed_borrower(is_async, setup_error): + client = object.__new__(AMQPClientAsync if is_async else AMQPClient) + link = MagicMock() + client._mgmt_links = {"$management": link} + if is_async: + client._mgmt_link_lock_async = asyncio.Lock() + client.auth_complete_async = AsyncMock(return_value=True) + link.ready = AsyncMock(side_effect=setup_error) + link.close = AsyncMock() + else: + client._mgmt_link_lock = threading.Lock() + client.auth_complete = MagicMock(return_value=True) + link.ready.side_effect = setup_error + with pytest.raises(type(setup_error)) as caught: + if is_async: + await client.open_mgmt_link_async() + else: + client.open_mgmt_link() + assert caught.value is setup_error + assert client._mgmt_links == {"$management": link} + link.close.assert_not_called() + link.ready.side_effect = None + link.ready.return_value = True + if is_async: + assert await client.open_mgmt_link_async() is link + else: + assert client.open_mgmt_link() is link + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setup_error", [RuntimeError("attach failed"), asyncio.CancelledError()]) +@pytest.mark.parametrize("cleanup_error", [None, RuntimeError("close failed"), asyncio.CancelledError()]) +async def test_failed_async_management_link_open_is_evicted_before_retry(setup_error, cleanup_error): + client = object.__new__(AMQPClientAsync) + client._mgmt_links = {} + client._mgmt_link_lock_async = asyncio.Lock() + client._session = MagicMock() + client.auth_complete_async = AsyncMock(return_value=True) + failed_link = MagicMock() + failed_link.open = AsyncMock(side_effect=setup_error) + failed_link.close = AsyncMock(side_effect=cleanup_error) + ready_link = MagicMock() + ready_link.open = AsyncMock() + ready_link.ready = AsyncMock(return_value=True) + + with patch( + "azure.servicebus._pyamqp.aio._client_async.ManagementOperation", + side_effect=[failed_link, ready_link], + ) as operation_type: + with pytest.raises(type(setup_error)) as caught: + await client.open_mgmt_link_async() + assert caught.value is setup_error + assert client._mgmt_links == {} + failed_link.close.assert_awaited_once() + + assert await client.open_mgmt_link_async() is ready_link + + assert operation_type.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout", [None, 0]) +@pytest.mark.parametrize("cancel_setup", [False, True]) +async def test_unbounded_async_setup_does_not_stall_on_cleanup(timeout, cancel_setup): + client = object.__new__(AMQPClientAsync) + client._mgmt_links = {} + client._mgmt_link_lock_async = asyncio.Lock() + client._session = MagicMock() + opening = asyncio.Event() + setup_error = RuntimeError("attach failed") + + async def open_link(): + opening.set() + if not cancel_setup: + raise setup_error + await asyncio.Event().wait() + + link = MagicMock(open=AsyncMock(side_effect=open_link), close=AsyncMock(side_effect=asyncio.Event().wait)) + with patch("azure.servicebus._pyamqp.aio._client_async.ManagementOperation", return_value=link): + task = asyncio.create_task(client.open_mgmt_link_async(timeout=timeout)) + try: + await opening.wait() + if cancel_setup: + task.cancel() + done, _ = await asyncio.wait({task}, timeout=2) + assert task in done, "Management cleanup must finish without a second cancellation." + with pytest.raises(asyncio.CancelledError if cancel_setup else RuntimeError) as caught: + await task + if not cancel_setup: + assert caught.value is setup_error + assert client._mgmt_links == {} + link.close.assert_awaited_once() + finally: + if not task.done(): + task.cancel() + with pytest.raises((RuntimeError, asyncio.CancelledError)): + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slow_close", [False, True]) +@pytest.mark.parametrize("setup_error", [RuntimeError("attach failed after deadline"), asyncio.CancelledError()]) +async def test_batch_delete_expired_setup_attempts_bounded_link_cleanup(slow_close, setup_error): + client = object.__new__(AMQPClientAsync) + client._mgmt_links = {} + client._mgmt_link_lock_async = asyncio.Lock() + client._session = MagicMock() + cleanup_finished = asyncio.Event() + + async def close_link(): + try: + if slow_close: + await asyncio.Event().wait() + finally: + cleanup_finished.set() + + link = MagicMock(open=AsyncMock(side_effect=setup_error), close=AsyncMock(side_effect=close_link)) + with patch("azure.servicebus._pyamqp.aio._client_async.ManagementOperation", return_value=link), patch( + "azure.servicebus._pyamqp.aio._client_async.time", + MagicMock(monotonic=MagicMock(side_effect=[0, 2])), + ), patch("azure.servicebus._pyamqp.aio._client_async.asyncio.wait_for", wraps=asyncio.wait_for) as wait_for: + task = asyncio.create_task(client.open_mgmt_link_async(timeout=1)) + done, _ = await asyncio.wait({task}, timeout=2) + try: + assert task in done, "Cleanup must not outlive its separate one-second grace." + with pytest.raises(type(setup_error)) as caught: + await task + assert caught.value is setup_error + assert client._mgmt_links == {} + link.close.assert_awaited_once() + assert cleanup_finished.is_set() + assert wait_for.call_args.kwargs["timeout"] == 1 + finally: + if not task.done(): + task.cancel() + with pytest.raises((RuntimeError, asyncio.CancelledError)): + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_type", [ServiceBusClient, AsyncServiceBusClient]) +@pytest.mark.parametrize("subscription", [False, True]) +@pytest.mark.parametrize("operation", ["delete_messages", "purge_messages"]) +@pytest.mark.parametrize("sub_queue,suffix", [ + (ServiceBusSubQueue.DEAD_LETTER, "/$DeadLetterQueue"), + (ServiceBusSubQueue.TRANSFER_DEAD_LETTER, "/$Transfer/$DeadLetterQueue"), +]) +async def test_client_created_subqueue_batch_delete_routing(client_type, subscription, operation, sub_queue, suffix): + client = client_type("example.servicebus.windows.net", MagicMock()) + is_async = client_type is AsyncServiceBusClient + try: + if subscription: + receiver = client.get_subscription_receiver("topic", "subscription", sub_queue=sub_queue) + entity_path = "topic/Subscriptions/subscription" + else: + receiver = client.get_queue_receiver("queue", sub_queue=sub_queue) + entity_path = "queue" + receiver._open = AsyncMock() if is_async else MagicMock() + expected_node = (entity_path + suffix + "/$management").encode("UTF-8") + setup_method = "mgmt_client_setup_async" if is_async else "mgmt_client_setup" + request_method = "mgmt_client_request_async" if is_async else "mgmt_client_request" + with patch.object(receiver._amqp_transport, setup_method) as setup, patch.object( + receiver._amqp_transport, request_method, return_value=0 + ) as dispatch: + args = (1,) if operation == "delete_messages" else () + if is_async: + await getattr(receiver, operation)(*args) + else: + getattr(receiver, operation)(*args) + assert setup.call_args.kwargs["node"] == expected_node + assert dispatch.call_args.kwargs["node"] == expected_node + assert dispatch.call_args.kwargs["operation"] == REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION + finally: + if is_async: + await client.close() + else: + client.close() + + +@pytest.mark.asyncio +async def test_async_management_link_setup_converts_timeout_error(): + client = MagicMock() + client.open_mgmt_link_async = AsyncMock(side_effect=TimeoutError("setup timed out")) + + with pytest.raises(OperationTimeoutError): + await PyamqpTransportAsync.mgmt_client_setup_async( + client, node="$management", timeout=1 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["lock", "open", "auth", "ready", "open_cleanup", "auth_cleanup_lock"]) +async def test_async_management_setup_deadline_prevents_dispatch(phase): + client = object.__new__(AMQPClientAsync) + client._mgmt_links = {} + client._mgmt_link_lock_async = asyncio.Lock() + client._session = MagicMock() + client.auth_complete_async = AsyncMock(return_value=True) + link = MagicMock(open=AsyncMock(), ready=AsyncMock(return_value=True), close=AsyncMock()) + blocked = AsyncMock(side_effect=asyncio.Event().wait) + if phase == "lock": + await client._mgmt_link_lock_async.acquire() + client._mgmt_links["$management"] = link + elif phase == "auth": + client.auth_complete_async = blocked + elif phase == "open_cleanup": + link.open = blocked + link.close = AsyncMock(side_effect=asyncio.Event().wait) + elif phase == "auth_cleanup_lock": + async def blocked_auth(): + await client._mgmt_link_lock_async.acquire() + await blocked() + + client.auth_complete_async = AsyncMock(side_effect=blocked_auth) + else: + setattr(link, phase, blocked) + receiver = _receiver(AsyncServiceBusReceiver) + receiver._config = MagicMock(retry_total=0, try_timeout=None, encoding="UTF-8") + receiver._container_id = "receiver" + receiver._mgmt_target = "$management" + receiver._handler = client + receiver._amqp_transport = PyamqpTransportAsync + receiver._open_with_timeout = AsyncMock() + receiver._handle_exception = AsyncMock(side_effect=lambda error: error) + receiver._mgmt_request_response = AsyncMock() + receiver._open_mgmt_link_with_retry = AsyncServiceBusReceiver._open_mgmt_link_with_retry.__get__(receiver) + try: + with patch("azure.servicebus._pyamqp.aio._client_async.ManagementOperation", return_value=link): + with pytest.raises(OperationTimeoutError): + await asyncio.wait_for(receiver.delete_messages(1, timeout=0.05), timeout=2) + receiver._mgmt_request_response.assert_not_awaited() + if phase == "lock": + assert client._mgmt_links == {"$management": link} + link.close.assert_not_called() + else: + blocked.assert_awaited_once() + assert client._mgmt_links == {} + link.close.assert_called_once() + finally: + if phase in ("lock", "auth_cleanup_lock"): + client._mgmt_link_lock_async.release() + + +class TestDeleteMessages: + @pytest.mark.parametrize( + "session_id,suppress_session_advice,expects_session_advice", + [ + (None, False, False), + (NEXT_AVAILABLE_SESSION, False, True), + (NEXT_AVAILABLE_SESSION, True, False), + ], + ) + def test_setup_timeout_advice_only_applies_to_next_available_session( + self, session_id, suppress_session_advice, expects_session_advice + ): + receiver = _receiver(ServiceBusReceiver, session_id=session_id) + receiver._config = MagicMock(retry_total=0) + receiver._container_id = "receiver" + receiver._handle_exception = lambda error: error + + def time_out(): + raise OperationTimeoutError() + + with pytest.raises(OperationTimeoutError) as exc_info: + receiver._do_retryable_operation( + time_out, + suppress_next_session_timeout_message=suppress_session_advice, + ) + + assert ("NEXT_AVAILABLE_SESSION" in str(exc_info.value)) is expects_session_advice + + def test_returns_actual_count_and_uses_one_shot_dispatch(self): + receiver = _receiver(ServiceBusReceiver) + cutoff = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) + calls = [] + + def dispatch(operation, message, callback, **kwargs): + calls.append((operation, message, callback, kwargs)) + return 7 + + receiver._mgmt_request_response = dispatch + + result = receiver.delete_messages(10, before_enqueued_time=cutoff, timeout=12) + + assert result.deleted_message_count == 7 + assert len(calls) == 1 + assert calls[0][0] == REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION + assert calls[0][2].func is mgmt_handlers.batch_delete_op + assert calls[0][2].keywords["max_message_count"] == 10 + assert calls[0][1][MGMT_REQUEST_MESSAGE_COUNT] == 10 + assert calls[0][1][MGMT_REQUEST_ENQUEUED_TIME_UTC] == int( + cutoff.timestamp() * 1000 + ) + assert 0 < calls[0][3]["timeout"] <= 12 + receiver._open_with_retry.assert_called_once_with( + timeout=12, suppress_next_session_timeout_message=True + ) + receiver._open_mgmt_link_with_retry.assert_called_once() + setup_kwargs = receiver._open_mgmt_link_with_retry.call_args.kwargs + assert setup_kwargs["suppress_next_session_timeout_message"] is True + assert 0 < setup_kwargs["timeout"] <= 12 + receiver._mgmt_request_response_with_retry.assert_not_called() + + def test_setup_failure_does_not_dispatch(self): + receiver = _receiver(ServiceBusReceiver) + receiver._open_with_retry.side_effect = RuntimeError("open failed") + receiver._mgmt_request_response = MagicMock() + + with pytest.raises(RuntimeError, match="open failed"): + receiver.delete_messages(1) + + receiver._mgmt_request_response.assert_not_called() + + def test_management_setup_failure_does_not_dispatch(self): + receiver = _receiver(ServiceBusReceiver) + receiver._open_mgmt_link_with_retry.side_effect = RuntimeError("management open failed") + receiver._mgmt_request_response = MagicMock() + + with pytest.raises(RuntimeError, match="management open failed"): + receiver.delete_messages(1) + + receiver._mgmt_request_response.assert_not_called() + + def test_dispatch_failure_is_not_retried(self): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock( + side_effect=RuntimeError("dispatch failed") + ) + + with pytest.raises(RuntimeError, match="dispatch failed"): + receiver.delete_messages(1) + + receiver._open_with_retry.assert_called_once_with( + timeout=None, suppress_next_session_timeout_message=True + ) + receiver._mgmt_request_response.assert_called_once() + receiver._mgmt_request_response_with_retry.assert_not_called() + + def test_dispatch_uses_remaining_timeout_and_stops_when_setup_exhausts_it(self): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock(return_value=1) + + with patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[10.0, 11.0, 11.5], + ): + receiver.delete_messages(1, timeout=2) + + assert receiver._mgmt_request_response.call_args.kwargs[ + "timeout" + ] == pytest.approx(0.5) + assert receiver._open_mgmt_link_with_retry.call_args.kwargs[ + "timeout" + ] == pytest.approx(1.0) + + receiver._mgmt_request_response.reset_mock() + with patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[20.0, 21.0], + ): + with pytest.raises(OperationTimeoutError, match="Operation timed out"): + receiver.delete_messages(1, timeout=0.5) + + receiver._mgmt_request_response.assert_not_called() + + def test_management_setup_shares_timeout_with_receiver_reopen(self): + receiver = _receiver(ServiceBusReceiver) + receiver._config = MagicMock( + retry_total=0, encoding="UTF-8", try_timeout=None + ) + receiver._container_id = "receiver" + receiver._mgmt_target = "$management" + receiver._open_with_timeout = MagicMock() + receiver._amqp_transport.mgmt_client_setup = MagicMock() + + with patch( + "azure.servicebus._base_handler.time.monotonic", + side_effect=[10.0, 10.0, 10.0, 10.5], + ): + ServiceBusReceiver._open_mgmt_link_with_retry(receiver, timeout=2) + + receiver._open_with_timeout.assert_called_once() + receiver_timeout = receiver._open_with_timeout.call_args.args[0] + assert 0 < receiver_timeout <= 2 + receiver._amqp_transport.mgmt_client_setup.assert_called_once() + setup_args = receiver._amqp_transport.mgmt_client_setup.call_args + assert setup_args.args == (receiver._handler,) + assert setup_args.kwargs["node"] == receiver._mgmt_target.encode("UTF-8") + assert setup_args.kwargs["timeout"] == pytest.approx(receiver_timeout - 0.5) + + def test_forwards_session_id(self): + receiver = _receiver(ServiceBusReceiver, session_id="session-a") + captured = {} + + def dispatch(operation, message, callback, **kwargs): + captured.update(message) + return 1 + + receiver._mgmt_request_response = dispatch + + receiver.delete_messages(1) + + assert captured[MGMT_REQUEST_SESSION_ID] == "session-a" + + def test_populates_session_id_after_setup_resolves_next_session(self): + receiver = _receiver(ServiceBusReceiver, session_id="") + receiver._open_with_retry.side_effect = lambda **_: setattr( + receiver, "_session_id", "session-a" + ) + receiver._mgmt_request_response = MagicMock(return_value=1) + + receiver.delete_messages(1) + + message = receiver._mgmt_request_response.call_args.args[1] + assert message[MGMT_REQUEST_SESSION_ID] == "session-a" + + def test_open_readiness_respects_timeout(self): + receiver = _receiver(ServiceBusReceiver) + receiver._running = False + receiver._connection = None + receiver._auto_lock_renewer = None + receiver._receive_mode = ServiceBusReceiveMode.PEEK_LOCK + receiver._handler._shutdown = False + receiver._handler.client_ready.return_value = False + receiver._create_handler = MagicMock() + + with patch( + "azure.servicebus._servicebus_receiver.create_authentication", + return_value=None, + ), patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[10.0, 11.0], + ), patch( + "azure.servicebus._servicebus_receiver.time.sleep" + ): + with pytest.raises(OperationTimeoutError): + receiver._open(timeout=0.5) + + def test_open_authentication_consumes_timeout(self): + receiver = _receiver(ServiceBusReceiver) + receiver._running = False + receiver._connection = None + receiver._handler = None + receiver._create_handler = MagicMock() + + with patch( + "azure.servicebus._servicebus_receiver.create_authentication", + return_value=None, + ), patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[10.0, 11.0], + ): + with pytest.raises(OperationTimeoutError): + receiver._open(timeout=0.5) + + receiver._create_handler.assert_not_called() + + def test_supports_premium_count(self): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock(return_value=4000) + + result = receiver.delete_messages(4000) + + assert result.deleted_message_count == 4000 + assert ( + receiver._mgmt_request_response.call_args.args[1][ + MGMT_REQUEST_MESSAGE_COUNT + ] + == 4000 + ) + + @pytest.mark.parametrize("message_count", [0, -1, 2_147_483_648]) + def test_rejects_counts_outside_service_limit(self, message_count): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock() + + with pytest.raises(ValueError): + receiver.delete_messages(message_count) + + receiver._mgmt_request_response.assert_not_called() + + @pytest.mark.parametrize("message_count", [True, 1.5, "1"]) + def test_rejects_non_integer_counts(self, message_count): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock() + + with pytest.raises(TypeError): + receiver.delete_messages(message_count) + + receiver._mgmt_request_response.assert_not_called() + + +class TestPurgeMessages: + def test_keeps_one_cutoff_and_stops_only_on_zero(self): + receiver = _receiver(ServiceBusReceiver) + counts = iter([500, 2, 0]) + calls = [] + + def dispatch(operation, message, callback, **kwargs): + calls.append((message.copy(), kwargs)) + return next(counts) + + receiver._mgmt_request_response = dispatch + + cutoff = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) + result = receiver.purge_messages(before_enqueued_time=cutoff, timeout=18) + + assert result.deleted_message_count == 502 + assert len(calls) == 3 + assert [call[0][MGMT_REQUEST_MESSAGE_COUNT] for call in calls] == [ + 500, + 500, + 500, + ] + cutoffs = [call[0][MGMT_REQUEST_ENQUEUED_TIME_UTC] for call in calls] + assert cutoffs[0] == cutoffs[1] == cutoffs[2] + assert cutoffs[0] == int(cutoff.timestamp() * 1000) + assert all(0 < call[1]["timeout"] <= 18 for call in calls) + + def test_supports_premium_batch_size(self): + receiver = _receiver(ServiceBusReceiver) + counts = iter([4000, 2, 0]) + calls = [] + + def dispatch(operation, message, callback, **kwargs): + calls.append(message.copy()) + return next(counts) + + receiver._mgmt_request_response = dispatch + + result = receiver.purge_messages(max_message_count_per_batch=4000) + + assert result.deleted_message_count == 4002 + assert [call[MGMT_REQUEST_MESSAGE_COUNT] for call in calls] == [ + 4000, + 4000, + 4000, + ] + cutoffs = [call[MGMT_REQUEST_ENQUEUED_TIME_UTC] for call in calls] + assert cutoffs[0] == cutoffs[1] == cutoffs[2] + + def test_allows_service_to_enforce_batch_size(self): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock(return_value=0) + + receiver.purge_messages(max_message_count_per_batch=4001) + + assert receiver._mgmt_request_response.call_args.args[1][MGMT_REQUEST_MESSAGE_COUNT] == 4001 + + @pytest.mark.parametrize("message_count", [True, 0, 2_147_483_648]) + def test_rejects_invalid_batch_size(self, message_count): + receiver = _receiver(ServiceBusReceiver) + receiver._mgmt_request_response = MagicMock() + + with pytest.raises((TypeError, ValueError)): + receiver.purge_messages(max_message_count_per_batch=message_count) + + receiver._mgmt_request_response.assert_not_called() + + def test_uses_one_operation_deadline(self): + receiver = _receiver(ServiceBusReceiver) + receiver.delete_messages = MagicMock( + side_effect=[DeleteMessagesResult(1), DeleteMessagesResult(0)] + ) + + with patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[10.0, 11.0, 12.0], + ): + result = receiver.purge_messages(timeout=5) + + assert result.deleted_message_count == 1 + assert [call.kwargs["timeout"] for call in receiver.delete_messages.call_args_list] == [4.0, 3.0] + + def test_operation_deadline_stops_before_another_dispatch(self): + receiver = _receiver(ServiceBusReceiver) + receiver.delete_messages = MagicMock(return_value=DeleteMessagesResult(1)) + + with patch( + "azure.servicebus._servicebus_receiver.time.monotonic", + side_effect=[20.0, 21.0, 25.1], + ): + with pytest.raises(OperationTimeoutError): + receiver.purge_messages(timeout=5) + + receiver.delete_messages.assert_called_once() + + +class TestDeleteMessagesAsync: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "session_id,suppress_session_advice,expects_session_advice", + [ + (None, False, False), + (NEXT_AVAILABLE_SESSION, False, True), + (NEXT_AVAILABLE_SESSION, True, False), + ], + ) + async def test_setup_timeout_advice_only_applies_to_next_available_session( + self, session_id, suppress_session_advice, expects_session_advice + ): + receiver = _receiver(AsyncServiceBusReceiver, session_id=session_id) + receiver._config = MagicMock(retry_total=0) + receiver._container_id = "receiver" + + async def handle_exception(error): + return error + + async def time_out(): + raise OperationTimeoutError() + + receiver._handle_exception = handle_exception + with pytest.raises(OperationTimeoutError) as exc_info: + await receiver._do_retryable_operation( + time_out, + suppress_next_session_timeout_message=suppress_session_advice, + ) + + assert ("NEXT_AVAILABLE_SESSION" in str(exc_info.value)) is expects_session_advice + + @pytest.mark.asyncio + async def test_returns_actual_count_and_uses_one_shot_dispatch(self): + receiver = _receiver(AsyncServiceBusReceiver) + cutoff = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) + calls = [] + + async def dispatch(operation, message, callback, **kwargs): + calls.append((operation, message, callback, kwargs)) + return 7 + + receiver._mgmt_request_response = dispatch + + result = await receiver.delete_messages( + 10, before_enqueued_time=cutoff, timeout=12 + ) + + assert result.deleted_message_count == 7 + assert len(calls) == 1 + assert calls[0][0] == REQUEST_RESPONSE_BATCH_DELETE_MESSAGES_OPERATION + assert calls[0][2].func is mgmt_handlers.batch_delete_op + assert calls[0][2].keywords["max_message_count"] == 10 + assert calls[0][1][MGMT_REQUEST_MESSAGE_COUNT] == 10 + assert calls[0][1][MGMT_REQUEST_ENQUEUED_TIME_UTC] == int( + cutoff.timestamp() * 1000 + ) + assert 0 < calls[0][3]["timeout"] <= 12 + receiver._open_with_retry.assert_awaited_once_with( + timeout=12, suppress_next_session_timeout_message=True + ) + receiver._open_mgmt_link_with_retry.assert_awaited_once() + setup_kwargs = receiver._open_mgmt_link_with_retry.await_args.kwargs + assert setup_kwargs["suppress_next_session_timeout_message"] is True + assert 0 < setup_kwargs["timeout"] <= 12 + receiver._mgmt_request_response_with_retry.assert_not_called() + + @pytest.mark.asyncio + async def test_setup_failure_does_not_dispatch(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._open_with_retry.side_effect = RuntimeError("open failed") + receiver._mgmt_request_response = AsyncMock() + + with pytest.raises(RuntimeError, match="open failed"): + await receiver.delete_messages(1) + + receiver._mgmt_request_response.assert_not_awaited() + + @pytest.mark.asyncio + async def test_management_setup_failure_does_not_dispatch(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._open_mgmt_link_with_retry.side_effect = RuntimeError("management open failed") + receiver._mgmt_request_response = AsyncMock() + + with pytest.raises(RuntimeError, match="management open failed"): + await receiver.delete_messages(1) + + receiver._mgmt_request_response.assert_not_awaited() + + @pytest.mark.asyncio + async def test_populates_session_id_after_setup_resolves_next_session(self): + receiver = _receiver( + AsyncServiceBusReceiver, session_id="" + ) + + async def resolve_session(**_): + receiver._session_id = "session-a" + + receiver._open_with_retry.side_effect = resolve_session + receiver._mgmt_request_response = AsyncMock(return_value=1) + + await receiver.delete_messages(1) + + message = receiver._mgmt_request_response.call_args.args[1] + assert message[MGMT_REQUEST_SESSION_ID] == "session-a" + + @pytest.mark.asyncio + async def test_open_readiness_respects_timeout(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._running = False + receiver._connection = None + receiver._auto_lock_renewer = None + receiver._receive_mode = ServiceBusReceiveMode.PEEK_LOCK + receiver._handler._shutdown = False + receiver._handler.close_async = AsyncMock() + receiver._handler.open_async = AsyncMock() + receiver._handler.client_ready_async = AsyncMock(return_value=False) + receiver._amqp_transport.drain_and_release_messages_async = AsyncMock() + receiver._create_handler = MagicMock() + + async def await_operation(awaitable, _deadline, _message): + return await awaitable + + async def discard_cleanup(close_coro, _deadline): + close_coro.close() + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.create_authentication", + new=AsyncMock(return_value=None), + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.get_link_ready_deadline", + return_value=10.5, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.check_link_ready_deadline", + side_effect=[None, None, None, OperationTimeoutError()], + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.close_handler_with_deadline", + new=AsyncMock(), + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.open_handler_with_deadline", + new=AsyncMock(), + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.close_handler_for_cleanup", + side_effect=discard_cleanup, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.await_with_deadline", + side_effect=await_operation, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.asyncio.sleep", + new=AsyncMock(), + ): + with pytest.raises(OperationTimeoutError): + await receiver._open(timeout=0.5) + + @pytest.mark.asyncio + async def test_open_authentication_consumes_timeout(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._running = False + receiver._connection = None + receiver._handler = None + receiver._create_handler = MagicMock() + + async def authentication_timeout(awaitable, _deadline, _message): + awaitable.close() + raise OperationTimeoutError() + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.create_authentication", + new=AsyncMock(return_value=None), + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.get_link_ready_deadline", + return_value=10.5, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.check_link_ready_deadline", + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.await_with_deadline", + side_effect=authentication_timeout, + ): + with pytest.raises(OperationTimeoutError): + await receiver._open(timeout=0.5) + + receiver._create_handler.assert_not_called() + + @pytest.mark.asyncio + async def test_open_does_not_create_awaitable_after_timeout(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._running = False + receiver._connection = None + receiver._handler = None + authentication = AsyncMock(return_value=None) + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.create_authentication", + new=authentication, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.get_link_ready_deadline", + return_value=10.5, + ), patch( + "azure.servicebus.aio._servicebus_receiver_async.check_link_ready_deadline", + side_effect=OperationTimeoutError(), + ): + with pytest.raises(OperationTimeoutError): + await receiver._open(timeout=0.5) + + authentication.assert_not_called() + + @pytest.mark.asyncio + async def test_dispatch_failure_is_not_retried(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._mgmt_request_response = AsyncMock( + side_effect=RuntimeError("dispatch failed") + ) + + with pytest.raises(RuntimeError, match="dispatch failed"): + await receiver.delete_messages(1) + + receiver._open_with_retry.assert_awaited_once_with( + timeout=None, suppress_next_session_timeout_message=True + ) + receiver._mgmt_request_response.assert_awaited_once() + receiver._mgmt_request_response_with_retry.assert_not_called() + + @pytest.mark.asyncio + async def test_dispatch_uses_remaining_timeout_and_stops_when_setup_exhausts_it( + self, + ): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._mgmt_request_response = AsyncMock(return_value=1) + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.time.monotonic", + side_effect=[10.0, 11.0, 11.5], + ): + await receiver.delete_messages(1, timeout=2) + + assert receiver._mgmt_request_response.call_args.kwargs[ + "timeout" + ] == pytest.approx(0.5) + assert receiver._open_mgmt_link_with_retry.call_args.kwargs[ + "timeout" + ] == pytest.approx(1.0) + + receiver._mgmt_request_response.reset_mock() + with patch( + "azure.servicebus.aio._servicebus_receiver_async.time.monotonic", + side_effect=[20.0, 21.0], + ): + with pytest.raises(OperationTimeoutError, match="Operation timed out"): + await receiver.delete_messages(1, timeout=0.5) + + receiver._mgmt_request_response.assert_not_awaited() + + @pytest.mark.asyncio + async def test_management_setup_shares_timeout_with_receiver_reopen(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._config = MagicMock( + retry_total=0, encoding="UTF-8", try_timeout=None + ) + receiver._container_id = "receiver" + receiver._mgmt_target = "$management" + clock = MagicMock(return_value=10.0) + + async def open_receiver(_timeout): + clock.return_value = 10.5 + + receiver._open_with_timeout = AsyncMock(side_effect=open_receiver) + receiver._amqp_transport.mgmt_client_setup_async = AsyncMock() + + with patch( + "azure.servicebus.aio._base_handler_async.time.monotonic", + clock, + ): + await AsyncServiceBusReceiver._open_mgmt_link_with_retry( + receiver, timeout=2 + ) + + receiver._open_with_timeout.assert_awaited_once() + receiver_timeout = receiver._open_with_timeout.await_args.args[0] + assert 0 < receiver_timeout <= 2 + receiver._amqp_transport.mgmt_client_setup_async.assert_awaited_once() + setup_args = receiver._amqp_transport.mgmt_client_setup_async.await_args + assert setup_args.args == (receiver._handler,) + assert setup_args.kwargs["node"] == receiver._mgmt_target.encode("UTF-8") + assert setup_args.kwargs["timeout"] == pytest.approx(receiver_timeout - 0.5) + + @pytest.mark.asyncio + async def test_keeps_one_cutoff_and_stops_only_on_zero(self): + receiver = _receiver(AsyncServiceBusReceiver) + counts = iter([500, 2, 0]) + calls = [] + + async def dispatch(operation, message, callback, **kwargs): + calls.append((message.copy(), kwargs)) + return next(counts) + + receiver._mgmt_request_response = dispatch + + cutoff = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) + result = await receiver.purge_messages(before_enqueued_time=cutoff, timeout=18) + + assert result.deleted_message_count == 502 + assert len(calls) == 3 + assert [call[0][MGMT_REQUEST_MESSAGE_COUNT] for call in calls] == [ + 500, + 500, + 500, + ] + cutoffs = [call[0][MGMT_REQUEST_ENQUEUED_TIME_UTC] for call in calls] + assert cutoffs[0] == cutoffs[1] == cutoffs[2] + assert cutoffs[0] == int(cutoff.timestamp() * 1000) + assert all(0 < call[1]["timeout"] <= 18 for call in calls) + + @pytest.mark.asyncio + async def test_purge_supports_premium_batch_size(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._mgmt_request_response = AsyncMock(side_effect=[4000, 2, 0]) + + result = await receiver.purge_messages(max_message_count_per_batch=4000) + + assert result.deleted_message_count == 4002 + calls = receiver._mgmt_request_response.await_args_list + assert [call.args[1][MGMT_REQUEST_MESSAGE_COUNT] for call in calls] == [ + 4000, + 4000, + 4000, + ] + cutoffs = [call.args[1][MGMT_REQUEST_ENQUEUED_TIME_UTC] for call in calls] + assert cutoffs[0] == cutoffs[1] == cutoffs[2] + + @pytest.mark.asyncio + async def test_purge_allows_service_to_enforce_batch_size(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver._mgmt_request_response = AsyncMock(return_value=0) + + await receiver.purge_messages(max_message_count_per_batch=4001) + + assert receiver._mgmt_request_response.call_args.args[1][MGMT_REQUEST_MESSAGE_COUNT] == 4001 + + @pytest.mark.asyncio + async def test_purge_uses_one_operation_deadline(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver.delete_messages = AsyncMock( + side_effect=[DeleteMessagesResult(1), DeleteMessagesResult(0)] + ) + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.time.monotonic", + side_effect=[10.0, 11.0, 12.0], + ): + result = await receiver.purge_messages(timeout=5) + + assert result.deleted_message_count == 1 + assert [call.kwargs["timeout"] for call in receiver.delete_messages.await_args_list] == [4.0, 3.0] + + @pytest.mark.asyncio + async def test_purge_deadline_stops_before_another_dispatch(self): + receiver = _receiver(AsyncServiceBusReceiver) + receiver.delete_messages = AsyncMock(return_value=DeleteMessagesResult(1)) + + with patch( + "azure.servicebus.aio._servicebus_receiver_async.time.monotonic", + side_effect=[20.0, 21.0, 25.1], + ): + with pytest.raises(OperationTimeoutError): + await receiver.purge_messages(timeout=5) + + receiver.delete_messages.assert_awaited_once()