def _build_handler(self):
     auth = None if self.connection else authentication.SASTokenAsync.from_shared_access_key(**self.auth_config)
     self._handler = ReceiveClientAsync(
         self.endpoint,
         auth=auth,
         debug=self.debug,
         properties=self.properties,
         error_policy=self.error_policy,
         client_name=self.name,
         auto_complete=False,
         encoding=self.encoding,
         loop=self.loop,
         **self.handler_kwargs)
예제 #2
0
 def _create_handler(self, auth):
     self._handler = ReceiveClientAsync(
         self._get_source(),
         auth=auth,
         debug=self._config.logging_enable,
         properties=self._properties,
         error_policy=self._error_policy,
         client_name=self._name,
         on_attach=self._on_attach,
         auto_complete=False,
         encoding=self._config.encoding,
         receive_settle_mode=ServiceBusToAMQPReceiveModeMap[self._receive_mode],
         send_settle_mode=SenderSettleMode.Settled
         if self._receive_mode == ServiceBusReceiveMode.RECEIVE_AND_DELETE
         else None,
         timeout=self._max_wait_time * 1000 if self._max_wait_time else 0,
         prefetch=self._prefetch_count,
         keep_alive_interval=self._config.keep_alive,
         shutdown_after_timeout=False,
     )
예제 #3
0
 def _create_handler(self):
     alt_creds = {
         "username": self.client._auth_config.get("iot_username"),
         "password": self.client._auth_config.get("iot_password")
     }
     source = Source(self.source)
     if self.offset is not None:
         source.set_filter(self.offset._selector())
     self._handler = ReceiveClientAsync(
         source,
         auth=self.client.get_auth(**alt_creds),
         debug=self.client.config.network_tracing,
         prefetch=self.prefetch,
         link_properties=self._link_properties,
         timeout=self.timeout,
         error_policy=self.retry_policy,
         keep_alive_interval=self.keep_alive,
         client_name=self.name,
         properties=self.client._create_properties(
             self.client.config.user_agent),  # pylint: disable=protected-access
         loop=self.loop)
     self.messages_iter = None
    def _create_handler(self):
        source = Source(self._source)
        if self._offset is not None:
            source.set_filter(self._offset._selector())  # pylint:disable=protected-access

        if StrictVersion(uamqp.__version__) < StrictVersion(
                "1.2.3"):  # backward compatible until uamqp 1.2.3 released
            desired_capabilities = {}
        elif self._track_last_enqueued_event_properties:
            symbol_array = [
                types.AMQPSymbol(self._receiver_runtime_metric_symbol)
            ]
            desired_capabilities = {
                "desired_capabilities":
                utils.data_factory(types.AMQPArray(symbol_array))
            }
        else:
            desired_capabilities = {"desired_capabilities": None}

        self._handler = ReceiveClientAsync(
            source,
            auth=self._client._create_auth(),  # pylint:disable=protected-access
            debug=self._client._config.network_tracing,  # pylint:disable=protected-access
            prefetch=self._prefetch,
            link_properties=self._link_properties,
            timeout=self._timeout,
            error_policy=self._retry_policy,
            keep_alive_interval=self._keep_alive,
            client_name=self._name,
            receive_settle_mode=uamqp.constants.ReceiverSettleMode.
            ReceiveAndDelete,
            auto_complete=False,
            properties=self._client._create_properties(  # pylint:disable=protected-access
                self._client._config.user_agent),  # pylint:disable=protected-access
            **desired_capabilities,  # pylint:disable=protected-access
            loop=self._loop)
        self._messages_iter = None
    def _create_handler(self, auth: "JWTTokenAsync") -> None:
        source = Source(self._source)
        if self._offset is not None:
            source.set_filter(
                event_position_selector(self._offset, self._offset_inclusive)
            )
        desired_capabilities = None
        if self._track_last_enqueued_event_properties:
            symbol_array = [types.AMQPSymbol(RECEIVER_RUNTIME_METRIC_SYMBOL)]
            desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array))

        properties = create_properties(
            self._client._config.user_agent  # pylint:disable=protected-access
        )
        self._handler = ReceiveClientAsync(
            source,
            auth=auth,
            debug=self._client._config.network_tracing,  # pylint:disable=protected-access
            prefetch=self._prefetch,
            link_properties=self._link_properties,
            timeout=self._timeout,
            idle_timeout=self._idle_timeout,
            error_policy=self._retry_policy,
            keep_alive_interval=self._keep_alive,
            client_name=self._name,
            receive_settle_mode=uamqp.constants.ReceiverSettleMode.ReceiveAndDelete,
            auto_complete=False,
            properties=properties,
            desired_capabilities=desired_capabilities,
            loop=self._loop,
        )

        self._handler._streaming_receive = True  # pylint:disable=protected-access
        self._handler._message_received_callback = (  # pylint:disable=protected-access
            self._message_received
        )
예제 #6
0
class ServiceBusReceiver(collections.abc.AsyncIterator, BaseHandlerAsync,
                         ReceiverMixin):
    """The ServiceBusReceiver class defines a high level interface for
    receiving messages from the Azure Service Bus Queue or Topic Subscription.

    :ivar fully_qualified_namespace: The fully qualified host name for the Service Bus namespace.
     The namespace format is: `<yournamespace>.servicebus.windows.net`.
    :vartype fully_qualified_namespace: str
    :ivar entity_path: The path of the entity that the client connects to.
    :vartype entity_path: str

    :param str fully_qualified_namespace: The fully qualified host name for the Service Bus namespace.
     The namespace format is: `<yournamespace>.servicebus.windows.net`.
    :param ~azure.core.credentials.TokenCredential credential: The credential object used for authentication which
     implements a particular interface for getting tokens. It accepts
     :class:`ServiceBusSharedKeyCredential<azure.servicebus.ServiceBusSharedKeyCredential>`, or credential objects
     generated by the azure-identity library and objects that implement the `get_token(self, *scopes)` method.
    :keyword str queue_name: The path of specific Service Bus Queue the client connects to.
    :keyword str topic_name: The path of specific Service Bus Topic which contains the Subscription
     the client connects to.
    :keyword str subscription_name: The path of specific Service Bus Subscription under the
     specified Topic the client connects to.
    :keyword mode: The mode with which messages will be retrieved from the entity. The two options
     are PeekLock and ReceiveAndDelete. Messages received with PeekLock must be settled within a given
     lock period before they will be removed from the queue. Messages received with ReceiveAndDelete
     will be immediately removed from the queue, and cannot be subsequently rejected or re-received if
     the client fails to process the message. The default mode is PeekLock.
    :paramtype mode: ~azure.servicebus.ReceiveSettleMode
    :keyword int prefetch: The maximum number of messages to cache with each request to the service.
     The default value is 0 meaning messages will be received from the service and processed
     one at a time. Increasing this value will improve message throughput performance but increase
     the change that messages will expire while they are cached if they're not processed fast enough.
    :keyword float idle_timeout: The timeout in seconds between received messages after which the receiver will
     automatically shutdown. The default value is 0, meaning no timeout.
    :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
    :keyword int retry_total: The total number of attempts to redo a failed operation when an error occurs.
     Default value is 3.
    :keyword transport_type: The type of transport protocol that will be used for communicating with
     the Service Bus service. Default is `TransportType.Amqp`.
    :paramtype transport_type: ~azure.servicebus.TransportType
    :keyword dict http_proxy: HTTP proxy settings. This must be a dictionary with the following
     keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value).
     Additionally the following keys may also be present: `'username', 'password'`.

    .. admonition:: Example:

        .. literalinclude:: ../samples/async_samples/sample_code_servicebus_async.py
            :start-after: [START create_servicebus_receiver_async]
            :end-before: [END create_servicebus_receiver_async]
            :language: python
            :dedent: 4
            :caption: Create a new instance of the ServiceBusReceiver.

    """
    def __init__(self, fully_qualified_namespace: str,
                 credential: "TokenCredential", **kwargs: Any):
        if kwargs.get("entity_name"):
            super(ServiceBusReceiver, self).__init__(
                fully_qualified_namespace=fully_qualified_namespace,
                credential=credential,
                **kwargs)
        else:
            queue_name = kwargs.get("queue_name")
            topic_name = kwargs.get("topic_name")
            subscription_name = kwargs.get("subscription_name")
            if queue_name and topic_name:
                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."
                )
            if topic_name and not subscription_name:
                raise ValueError(
                    "Subscription name is missing for the topic. Please specify subscription_name."
                )

            entity_name = queue_name or topic_name

            super(ServiceBusReceiver, self).__init__(
                fully_qualified_namespace=fully_qualified_namespace,
                credential=credential,
                entity_name=str(entity_name),
                **kwargs)
        self._message_iter = None
        self._create_attribute(**kwargs)
        self._connection = kwargs.get("connection")

    async def __anext__(self):
        self._check_live()
        while True:
            try:
                return await self._do_retryable_operation(self._iter_next)
            except StopAsyncIteration:
                await self.close()
                raise

    async def _iter_next(self):
        await self._open()
        uamqp_message = await self._message_iter.__anext__()
        message = self._build_message(uamqp_message, ReceivedMessage)
        return message

    def _create_handler(self, auth):
        self._handler = ReceiveClientAsync(
            self._get_source(),
            auth=auth,
            debug=self._config.logging_enable,
            properties=self._properties,
            error_policy=self._error_policy,
            client_name=self._name,
            on_attach=self._on_attach,
            auto_complete=False,
            encoding=self._config.encoding,
            receive_settle_mode=self._mode.value,
            send_settle_mode=SenderSettleMode.Settled
            if self._mode == ReceiveSettleMode.ReceiveAndDelete else None,
            timeout=self._config.idle_timeout *
            1000 if self._config.idle_timeout else 0,
            prefetch=self._config.prefetch)

    async def _open(self):
        if self._running:
            return
        if self._handler:
            await self._handler.close_async()
        auth = None if self._connection else (await
                                              create_authentication(self))
        self._create_handler(auth)
        try:
            await self._handler.open_async(connection=self._connection)
            self._message_iter = self._handler.receive_messages_iter_async()
            while not await self._handler.client_ready_async():
                await asyncio.sleep(0.05)
            self._running = True
        except:
            await self.close()
            raise

    async def _receive(self, max_batch_size=None, timeout=None):
        await self._open()
        max_batch_size = max_batch_size or self._handler._prefetch  # pylint: disable=protected-access

        timeout_ms = 1000 * (timeout or self._config.idle_timeout) if (
            timeout or self._config.idle_timeout) else 0
        batch = await self._handler.receive_message_batch_async(
            max_batch_size=max_batch_size, timeout=timeout_ms)

        return [
            self._build_message(message, ReceivedMessage) for message in batch
        ]

    async def _settle_message(self,
                              settlement,
                              lock_tokens,
                              dead_letter_details=None):
        message = {
            MGMT_REQUEST_DISPOSITION_STATUS: settlement,
            MGMT_REQUEST_LOCK_TOKENS: types.AMQPArray(lock_tokens)
        }

        self._populate_message_properties(message)
        if dead_letter_details:
            message.update(dead_letter_details)

        return await self._mgmt_request_response_with_retry(
            REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION, message,
            mgmt_handlers.default)

    async def _renew_locks(self, *lock_tokens):
        message = {MGMT_REQUEST_LOCK_TOKENS: types.AMQPArray(lock_tokens)}
        return await self._mgmt_request_response_with_retry(
            REQUEST_RESPONSE_RENEWLOCK_OPERATION, message,
            mgmt_handlers.lock_renew_op)

    @classmethod
    def from_connection_string(cls, conn_str: str,
                               **kwargs: Any) -> "ServiceBusReceiver":
        """Create a ServiceBusReceiver from a connection string.

        :param conn_str: The connection string of a Service Bus.
        :keyword str queue_name: The path of specific Service Bus Queue the client connects to.
        :keyword str topic_name: The path of specific Service Bus Topic which contains the Subscription
         the client connects to.
        :keyword str subscription_name: The path of specific Service Bus Subscription under the
         specified Topic the client connects to.
        :keyword mode: The mode with which messages will be retrieved from the entity. The two options
         are PeekLock and ReceiveAndDelete. Messages received with PeekLock must be settled within a given
         lock period before they will be removed from the queue. Messages received with ReceiveAndDelete
         will be immediately removed from the queue, and cannot be subsequently rejected or re-received if
         the client fails to process the message. The default mode is PeekLock.
        :paramtype mode: ~azure.servicebus.ReceiveSettleMode
        :keyword int prefetch: The maximum number of messages to cache with each request to the service.
         The default value is 0, meaning messages will be received from the service and processed
         one at a time. Increasing this value will improve message throughput performance but increase
         the change that messages will expire while they are cached if they're not processed fast enough.
        :keyword float idle_timeout: The timeout in seconds between received messages after which the receiver will
         automatically shutdown. The default value is 0, meaning no timeout.
        :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
        :keyword int retry_total: The total number of attempts to redo a failed operation when an error occurs.
         Default value is 3.
        :keyword transport_type: The type of transport protocol that will be used for communicating with
         the Service Bus service. Default is `TransportType.Amqp`.
        :paramtype transport_type: ~azure.servicebus.TransportType
        :keyword dict http_proxy: HTTP proxy settings. This must be a dictionary with the following
         keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value).
         Additionally the following keys may also be present: `'username', 'password'`.
        :rtype: ~azure.servicebus.aio.ServiceBusReceiver

        .. admonition:: Example:

            .. literalinclude:: ../samples/async_samples/sample_code_servicebus_async.py
                :start-after: [START create_servicebus_receiver_from_conn_str_async]
                :end-before: [END create_servicebus_receiver_from_conn_str_async]
                :language: python
                :dedent: 4
                :caption: Create a new instance of the ServiceBusReceiver from connection string.

        """
        constructor_args = cls._from_connection_string(conn_str, **kwargs)
        if kwargs.get("queue_name") and kwargs.get("subscription_name"):
            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."
            )
        return cls(**constructor_args)

    async def receive(self, max_batch_size=None, max_wait_time=None):
        # type: (int, float) -> List[ReceivedMessage]
        """Receive a batch of messages at once.

        This approach it optimal if you wish to process multiple messages simultaneously. Note that the
        number of messages retrieved in a single batch will be dependent on
        whether `prefetch` was set for the receiver. This call will prioritize returning
        quickly over meeting a specified batch size, and so will return as soon as at least
        one message is received and there is a gap in incoming messages regardless
        of the specified batch size.

        :param int max_batch_size: Maximum number of messages in the batch. Actual number
         returned will depend on prefetch size and incoming stream rate.
        :param float max_wait_time: Maximum time to wait in seconds for the first message to arrive.
         If no messages arrive, and no timeout is specified, this call will not return
         until the connection is closed. If specified, an no messages arrive within the
         timeout period, an empty list will be returned.
        :rtype: list[~azure.servicebus.aio.ReceivedMessage]

        .. admonition:: Example:

            .. literalinclude:: ../samples/async_samples/sample_code_servicebus_async.py
                :start-after: [START receive_async]
                :end-before: [END receive_async]
                :language: python
                :dedent: 4
                :caption: Receive messages from ServiceBus.

        """
        self._check_live()
        if max_batch_size and self._config.prefetch < max_batch_size:
            raise ValueError(
                "max_batch_size should be less than or equal to prefetch of ServiceBusReceiver, or you "
                "could set a larger prefetch value when you're constructing the ServiceBusReceiver."
            )
        return await self._do_retryable_operation(
            self._receive,
            max_batch_size=max_batch_size,
            timeout=max_wait_time,
            require_timeout=True)

    async def receive_deferred_messages(self, sequence_numbers):
        # type: (List[int]) -> List[ReceivedMessage]
        """Receive messages that have previously been deferred.

        When receiving deferred messages from a partitioned entity, all of the supplied
        sequence numbers must be messages from the same partition.

        :param list[int] sequence_numbers: A list of the sequence numbers of messages that have been
         deferred.
        :rtype: list[~azure.servicebus.aio.ReceivedMessage]

        .. admonition:: Example:

            .. literalinclude:: ../samples/async_samples/sample_code_servicebus_async.py
                :start-after: [START receive_defer_async]
                :end-before: [END receive_defer_async]
                :language: python
                :dedent: 4
                :caption: Receive deferred messages from ServiceBus.

        """
        self._check_live()
        if not sequence_numbers:
            raise ValueError("At least one sequence number must be specified.")
        await self._open()
        try:
            receive_mode = self._mode.value.value
        except AttributeError:
            receive_mode = int(self._mode)
        message = {
            MGMT_REQUEST_SEQUENCE_NUMBERS:
            types.AMQPArray([types.AMQPLong(s) for s in sequence_numbers]),
            MGMT_REQUEST_RECEIVER_SETTLE_MODE:
            types.AMQPuInt(receive_mode)
        }

        self._populate_message_properties(message)

        handler = functools.partial(mgmt_handlers.deferred_message_op,
                                    mode=self._mode,
                                    message_type=ReceivedMessage)
        messages = await self._mgmt_request_response_with_retry(
            REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER, message, handler)
        for m in messages:
            m._receiver = self  # pylint: disable=protected-access
        return messages

    async def peek(self, message_count=1, sequence_number=0):
        """Browse messages currently pending in the queue.

        Peeked messages are not removed from queue, nor are they locked. They cannot be completed,
        deferred or dead-lettered.

        :param int message_count: The maximum number of messages to try and peek. The default
         value is 1.
        :param int sequence_number: A message sequence number from which to start browsing messages.
        :rtype: list[~azure.servicebus.PeekMessage]

        .. admonition:: Example:

            .. literalinclude:: ../samples/async_samples/sample_code_servicebus_async.py
                :start-after: [START peek_messages_async]
                :end-before: [END peek_messages_async]
                :language: python
                :dedent: 4
                :caption: Peek messages in the queue.
        """
        self._check_live()
        if not sequence_number:
            sequence_number = self._last_received_sequenced_number or 1
        if int(message_count) < 1:
            raise ValueError("count must be 1 or greater.")
        if int(sequence_number) < 1:
            raise ValueError("start_from must be 1 or greater.")

        await self._open()

        message = {
            MGMT_REQUEST_FROM_SEQUENCE_NUMBER: types.AMQPLong(sequence_number),
            MGMT_REQUEST_MESSAGE_COUNT: message_count
        }

        self._populate_message_properties(message)

        return await self._mgmt_request_response_with_retry(
            REQUEST_RESPONSE_PEEK_OPERATION, message, mgmt_handlers.peek_op)
예제 #7
0
 async def _reconnect_async(self):  # pylint: disable=too-many-statements
     # pylint: disable=protected-access
     alt_creds = {
         "username": self.client._auth_config.get("iot_username"),
         "password": self.client._auth_config.get("iot_password")
     }
     await self._handler.close_async()
     source = Source(self.source)
     if self.offset is not None:
         source.set_filter(self.offset.selector())
     self._handler = ReceiveClientAsync(
         source,
         auth=self.client.get_auth(**alt_creds),
         debug=self.client.debug,
         prefetch=self.prefetch,
         link_properties=self.properties,
         timeout=self.timeout,
         error_policy=self.retry_policy,
         keep_alive_interval=self.keep_alive,
         client_name=self.name,
         properties=self.client.create_properties(),
         loop=self.loop)
     try:
         await self._handler.open_async()
         while not await self._handler.client_ready_async():
             await asyncio.sleep(0.05)
         return True
     except errors.TokenExpired as shutdown:
         log.info(
             "AsyncReceiver disconnected due to token expiry. Shutting down."
         )
         error = EventHubError(str(shutdown), shutdown)
         await self.close_async(exception=error)
         raise error
     except (errors.LinkDetach, errors.ConnectionClose) as shutdown:
         if shutdown.action.retry and self.auto_reconnect:
             log.info("AsyncReceiver detached. Attempting reconnect.")
             return False
         log.info("AsyncReceiver detached. Shutting down.")
         error = EventHubError(str(shutdown), shutdown)
         await self.close_async(exception=error)
         raise error
     except errors.MessageHandlerError as shutdown:
         if self.auto_reconnect:
             log.info("AsyncReceiver detached. Attempting reconnect.")
             return False
         log.info("AsyncReceiver detached. Shutting down.")
         error = EventHubError(str(shutdown), shutdown)
         await self.close_async(exception=error)
         raise error
     except errors.AMQPConnectionError as shutdown:
         if str(shutdown).startswith("Unable to open authentication session"
                                     ) and self.auto_reconnect:
             log.info(
                 "AsyncReceiver couldn't authenticate. Attempting reconnect."
             )
             return False
         log.info("AsyncReceiver connection error (%r). Shutting down.",
                  shutdown)
         error = EventHubError(str(shutdown))
         await self.close_async(exception=error)
         raise error
     except Exception as e:
         log.info("Unexpected error occurred (%r). Shutting down.", e)
         error = EventHubError("Receiver reconnect failed: {}".format(e))
         await self.close_async(exception=error)
         raise error
class Receiver(collections.abc.AsyncIterator, BaseHandler):  # pylint: disable=too-many-instance-attributes
    """A message receiver.

    This receive handler acts as an iterable message stream for retrieving
    messages for a Service Bus entity. It operates a single connection that must be opened and
    closed on completion. The service connection will remain open for the entirety of the iterator.
    If you find yourself only partially iterating the message stream, you should run the receiver
    in a `with` statement to ensure the connection is closed.
    The Receiver should not be instantiated directly, and should be accessed from a `QueueClient` or
    `SubscriptionClient` using the `get_receiver()` method.

    .. note:: This object is not thread-safe.

    :param handler_id: The ID used as the connection name for the Receiver.
    :type handler_id: str
    :param source: The endpoint from which to receive messages.
    :type source: ~uamqp.Source
    :param auth_config: The SASL auth credentials.
    :type auth_config: dict[str, str]
    :param loop: An async event loop
    :type loop: ~asyncio.EventLoop
    :param connection: A shared connection [not yet supported].
    :type connection: ~uamqp.Connection
    :param mode: The receive connection mode. Value must be either PeekLock or ReceiveAndDelete.
    :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
    :param encoding: The encoding used for string properties. Default is 'UTF-8'.
    :type encoding: str
    :param debug: Whether to enable network trace debug logs.
    :type debug: bool

    Example:
        .. literalinclude:: ../examples/async_examples/test_examples_async.py
            :start-after: [START open_close_receiver_context]
            :end-before: [END open_close_receiver_context]
            :language: python
            :dedent: 4
            :caption: Running a queue receiver within a context manager.

    """

    def __init__(
            self, handler_id, source, auth_config, *, loop=None, connection=None,
            mode=ReceiveSettleMode.PeekLock, encoding='UTF-8', debug=False, **kwargs):
        self._used = asyncio.Event()
        self.name = "SBReceiver-{}".format(handler_id)
        self.last_received = None
        self.mode = mode
        self.message_iter = None
        super(Receiver, self).__init__(
            source, auth_config, loop=loop, connection=connection, encoding=encoding, debug=debug, **kwargs)

    async def __anext__(self):
        await self._can_run()
        while True:
            if self.receiver_shutdown:
                await self.close()
                raise StopAsyncIteration
            try:
                received = await self.message_iter.__anext__()
                wrapped = self._build_message(received)
                return wrapped
            except StopAsyncIteration:
                await self.close()
                raise
            except Exception as e:  # pylint: disable=broad-except
                await self._handle_exception(e)

    def _build_handler(self):
        auth = None if self.connection else authentication.SASTokenAsync.from_shared_access_key(**self.auth_config)
        self._handler = ReceiveClientAsync(
            self.endpoint,
            auth=auth,
            debug=self.debug,
            properties=self.properties,
            error_policy=self.error_policy,
            client_name=self.name,
            auto_complete=False,
            encoding=self.encoding,
            loop=self.loop,
            **self.handler_kwargs)

    async def _build_receiver(self):
        """This is a temporary patch pending a fix in uAMQP."""
        # pylint: disable=protected-access
        self._handler.message_handler = self._handler.receiver_type(
            self._handler._session,
            self._handler._remote_address,
            self._handler._name,
            on_message_received=self._handler._message_received,
            name='receiver-link-{}'.format(uuid.uuid4()),
            debug=self._handler._debug_trace,
            prefetch=self._handler._prefetch,
            max_message_size=self._handler._max_message_size,
            properties=self._handler._link_properties,
            error_policy=self._handler._error_policy,
            encoding=self._handler._encoding,
            loop=self._handler.loop)
        if self.mode != ReceiveSettleMode.PeekLock:
            self._handler.message_handler.send_settle_mode = constants.SenderSettleMode.Settled
            self._handler.message_handler.receive_settle_mode = constants.ReceiverSettleMode.ReceiveAndDelete
            self._handler.message_handler._settle_mode = constants.ReceiverSettleMode.ReceiveAndDelete
        await self._handler.message_handler.open_async()

    def _build_message(self, received):
        message = Message(None, message=received)
        message._receiver = self  # pylint: disable=protected-access
        self.last_received = message.sequence_number
        return message

    async def _can_run(self):
        if self._used.is_set():
            raise InvalidHandlerState("Receiver has already closed.")
        if self.receiver_shutdown:
            await self.close()
            raise InvalidHandlerState("Receiver has already closed.")
        if not self.running:
            await self.open()

    async def _renew_locks(self, *lock_tokens):
        message = {'lock-tokens': types.AMQPArray(lock_tokens)}
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_RENEWLOCK_OPERATION,
            message,
            mgmt_handlers.lock_renew_op)

    async def _settle_deferred(self, settlement, lock_tokens, dead_letter_details=None):
        message = {
            'disposition-status': settlement,
            'lock-tokens': types.AMQPArray(lock_tokens)}
        if dead_letter_details:
            message.update(dead_letter_details)
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION,
            message,
            mgmt_handlers.default)

    @property
    def receiver_shutdown(self):
        """Whether the receiver connection has been marked for shutdown.

        If this value is `True` - it does not indicate that the connection
        has yet been closed.
        This property is used internally and should not be relied upon to asses
        the status of the connection.

        :rtype: bool
        """
        if self._handler:
            return self._handler._shutdown  # pylint: disable=protected-access
        return True

    @receiver_shutdown.setter
    def receiver_shutdown(self, value):
        """Mark the connection as ready for shutdown.

        This property is used internally and should not be set in normal usage.

        :param bool value: Whether to shutdown the connection.
        """
        if self._handler:
            self._handler._shutdown = value  # pylint: disable=protected-access
        else:
            raise ValueError("Receiver has no AMQP handler")

    @property
    def queue_size(self):
        """The current size of the unprocessed message queue.

        :rtype: int
        """
        # pylint: disable=protected-access
        if self._handler._received_messages:
            return self._handler._received_messages.qsize()
        return 0

    async def open(self):
        """Open receiver connection and authenticate session.

        If the receiver is already open, this operation will do nothing.
        This method will be called automatically when one starts to iterate
        messages in the receiver, so there should be no need to call it directly.
        A receiver opened with this method must be explicitly closed.
        It is recommended to open a handler within a context manager as
        opposed to calling the method directly.

        .. note:: This operation is not thread-safe.

        """
        if self.running:
            return
        self.running = True
        try:
            await self._handler.open_async(connection=self.connection)
            self.message_iter = self._handler.receive_messages_iter_async()
            while not await self._handler.auth_complete_async():
                await asyncio.sleep(0.05)
            await self._build_receiver()
            while not await self._handler.client_ready_async():
                await asyncio.sleep(0.05)
        except Exception as e:  # pylint: disable=broad-except
            try:
                await self._handle_exception(e)
            except:
                self.running = False
                raise

    async def close(self, exception=None):
        """Close down the receiver connection.

        If the receiver has already closed, this operation will do nothing. An optional
        exception can be passed in to indicate that the handler was shutdown due to error.
        It is recommended to open a handler within a context manager as
        opposed to calling the method directly.
        The receiver will be implicitly closed on completion of the message iterator,
        however this method will need to be called explicitly if the message iterator is not run
        to completion.

        .. note:: This operation is not thread-safe.

        :param exception: An optional exception if the handler is closing
         due to an error.
        :type exception: Exception

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START open_close_receiver_directly]
                :end-before: [END open_close_receiver_directly]
                :language: python
                :dedent: 4
                :caption: Iterate then explicitly close a Receiver.

        """
        if not self.running:
            return
        self.running = False
        self.receiver_shutdown = True
        self._used.set()
        await super(Receiver, self).close(exception=exception)

    async def peek(self, count=1, start_from=0):
        """Browse messages currently pending in the queue.

        Peeked messages are not removed from queue, nor are they locked. They cannot be completed,
        deferred or dead-lettered.

        :param count: The maximum number of messages to try and peek. The default
         value is 1.
        :type count: int
        :param start_from: A message sequence number from which to start browsing messages.
        :type start_from: int
        :rtype: list[~azure.servicebus.common.message.PeekMessage]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_peek_messages]
                :end-before: [END receiver_peek_messages]
                :language: python
                :dedent: 4
                :caption: Peek messages in the queue.

        """
        await self._can_run()
        if not start_from:
            start_from = self.last_received or 1
        if int(count) < 1:
            raise ValueError("count must be 1 or greater.")
        if int(start_from) < 1:
            raise ValueError("start_from must be 1 or greater.")

        message = {
            'from-sequence-number': types.AMQPLong(start_from),
            'message-count': count
        }
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_PEEK_OPERATION,
            message,
            mgmt_handlers.peek_op)

    async def receive_deferred_messages(self, sequence_numbers, mode=ReceiveSettleMode.PeekLock):
        """Receive messages that have previously been deferred.

        When receiving deferred messages from a partitioned entity, all of the supplied
        sequence numbers must be messages from the same partition.

        :param sequence_numbers: A list of the sequence numbers of messages that have been
         deferred.
        :type sequence_numbers: list[int]
        :param mode: The receive mode, default value is PeekLock.
        :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
        :rtype: list[~azure.servicebus.aio.async_message.DeferredMessage]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_defer_messages]
                :end-before: [END receiver_defer_messages]
                :language: python
                :dedent: 8
                :caption: Defer messages, then retrieve them by sequence number.

        """
        if not sequence_numbers:
            raise ValueError("At least one sequence number must be specified.")
        await self._can_run()
        try:
            receive_mode = mode.value.value
        except AttributeError:
            receive_mode = int(mode)
        message = {
            'sequence-numbers': types.AMQPArray([types.AMQPLong(s) for s in sequence_numbers]),
            'receiver-settle-mode': types.AMQPuInt(receive_mode)
        }
        handler = functools.partial(mgmt_handlers.deferred_message_op, mode=receive_mode, message_type=DeferredMessage)
        messages = await self._mgmt_request_response(
            REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER,
            message,
            handler)
        for m in messages:
            m._receiver = self  # pylint: disable=protected-access
        return messages

    async def fetch_next(self, max_batch_size=None, timeout=None):
        """Receive a batch of messages at once.

        This approach it optimal if you wish to process multiple messages simultaneously.
        Note that the number of messages retrieved in a single batch will be dependent on
        whether `prefetch` was set for the receiver. This call will prioritize returning
        quickly over meeting a specified batch size, and so will return as soon as at least
        one message is received and there is a gap in incoming messages regardless
        of the specified batch size.

        :param max_batch_size: Maximum number of messages in the batch. Actual number
         returned will depend on prefetch size and incoming stream rate.
        :type max_batch_size: int
        :param timeout: The time to wait in seconds for the first message to arrive.
         If no messages arrive, and no timeout is specified, this call will not return
         until the connection is closed. If specified, an no messages arrive within the
         timeout period, an empty list will be returned.
        :rtype: list[~azure.servicebus.aio.async_message.Message]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_fetch_batch]
                :end-before: [END receiver_fetch_batch]
                :language: python
                :dedent: 4
                :caption: Fetch a batch of messages.

        """
        await self._can_run()
        wrapped_batch = []
        max_batch_size = max_batch_size or self._handler._prefetch  # pylint: disable=protected-access
        try:
            timeout_ms = 1000 * timeout if timeout else 0
            batch = await self._handler.receive_message_batch_async(
                max_batch_size=max_batch_size,
                timeout=timeout_ms)
            for received in batch:
                message = self._build_message(received)
                wrapped_batch.append(message)
        except Exception as e:  # pylint: disable=broad-except
            await self._handle_exception(e)
        return wrapped_batch
class Receiver(collections.abc.AsyncIterator, BaseHandler):  # pylint: disable=too-many-instance-attributes
    """A message receiver.

    This receive handler acts as an iterable message stream for retrieving
    messages for a Service Bus entity. It operates a single connection that must be opened and
    closed on completion. The service connection will remain open for the entirety of the iterator.
    If you find yourself only partially iterating the message stream, you should run the receiver
    in a `with` statement to ensure the connection is closed.
    The Receiver should not be instantiated directly, and should be accessed from a `QueueClient` or
    `SubscriptionClient` using the `get_receiver()` method.

    .. note:: This object is not thread-safe.

    :param handler_id: The ID used as the connection name for the Receiver.
    :type handler_id: str
    :param source: The endpoint from which to receive messages.
    :type source: ~uamqp.Source
    :param auth_config: The SASL auth credentials.
    :type auth_config: dict[str, str]
    :param loop: An async event loop
    :type loop: ~asyncio.EventLoop
    :param connection: A shared connection [not yet supported].
    :type connection: ~uamqp.Connection
    :param mode: The receive connection mode. Value must be either PeekLock or ReceiveAndDelete.
    :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
    :param encoding: The encoding used for string properties. Default is 'UTF-8'.
    :type encoding: str
    :param debug: Whether to enable network trace debug logs.
    :type debug: bool

    Example:
        .. literalinclude:: ../examples/async_examples/test_examples_async.py
            :start-after: [START open_close_receiver_context]
            :end-before: [END open_close_receiver_context]
            :language: python
            :dedent: 4
            :caption: Running a queue receiver within a context manager.

    """
    def __init__(self,
                 handler_id,
                 source,
                 auth_config,
                 *,
                 loop=None,
                 connection=None,
                 mode=ReceiveSettleMode.PeekLock,
                 encoding='UTF-8',
                 debug=False,
                 **kwargs):
        self._used = asyncio.Event()
        self.name = "SBReceiver-{}".format(handler_id)
        self.last_received = None
        self.mode = mode
        self.message_iter = None
        super(Receiver, self).__init__(source,
                                       auth_config,
                                       loop=loop,
                                       connection=connection,
                                       encoding=encoding,
                                       debug=debug,
                                       **kwargs)

    async def __anext__(self):
        await self._can_run()
        while True:
            if self.receiver_shutdown:
                await self.close()
                raise StopAsyncIteration
            try:
                received = await self.message_iter.__anext__()
                wrapped = self._build_message(received)
                return wrapped
            except StopAsyncIteration:
                await self.close()
                raise
            except Exception as e:  # pylint: disable=broad-except
                await self._handle_exception(e)

    def _build_handler(self):
        auth = None if self.connection else authentication.SASTokenAsync.from_shared_access_key(
            **self.auth_config)
        self._handler = ReceiveClientAsync(self.endpoint,
                                           auth=auth,
                                           debug=self.debug,
                                           properties=self.properties,
                                           error_policy=self.error_policy,
                                           client_name=self.name,
                                           auto_complete=False,
                                           encoding=self.encoding,
                                           loop=self.loop,
                                           **self.handler_kwargs)

    async def _build_receiver(self):
        """This is a temporary patch pending a fix in uAMQP."""
        # pylint: disable=protected-access
        self._handler.message_handler = self._handler.receiver_type(
            self._handler._session,
            self._handler._remote_address,
            self._handler._name,
            on_message_received=self._handler._message_received,
            name='receiver-link-{}'.format(uuid.uuid4()),
            debug=self._handler._debug_trace,
            prefetch=self._handler._prefetch,
            max_message_size=self._handler._max_message_size,
            properties=self._handler._link_properties,
            error_policy=self._handler._error_policy,
            encoding=self._handler._encoding,
            loop=self._handler.loop)
        if self.mode != ReceiveSettleMode.PeekLock:
            self._handler.message_handler.send_settle_mode = constants.SenderSettleMode.Settled
            self._handler.message_handler.receive_settle_mode = constants.ReceiverSettleMode.ReceiveAndDelete
            self._handler.message_handler._settle_mode = constants.ReceiverSettleMode.ReceiveAndDelete
        await self._handler.message_handler.open_async()

    def _build_message(self, received):
        message = Message(None, message=received)
        message._receiver = self  # pylint: disable=protected-access
        self.last_received = message.sequence_number
        return message

    async def _can_run(self):
        if self._used.is_set():
            raise InvalidHandlerState("Receiver has already closed.")
        if self.receiver_shutdown:
            await self.close()
            raise InvalidHandlerState("Receiver has already closed.")
        if not self.running:
            await self.open()

    async def _renew_locks(self, *lock_tokens):
        message = {'lock-tokens': types.AMQPArray(lock_tokens)}
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_RENEWLOCK_OPERATION, message,
            mgmt_handlers.lock_renew_op)

    async def _settle_deferred(self,
                               settlement,
                               lock_tokens,
                               dead_letter_details=None):
        message = {
            'disposition-status': settlement,
            'lock-tokens': types.AMQPArray(lock_tokens)
        }
        if dead_letter_details:
            message.update(dead_letter_details)
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION, message,
            mgmt_handlers.default)

    @property
    def receiver_shutdown(self):
        """Whether the receiver connection has been marked for shutdown.

        If this value is `True` - it does not indicate that the connection
        has yet been closed.
        This property is used internally and should not be relied upon to asses
        the status of the connection.

        :rtype: bool
        """
        if self._handler:
            return self._handler._shutdown  # pylint: disable=protected-access
        return True

    @receiver_shutdown.setter
    def receiver_shutdown(self, value):
        """Mark the connection as ready for shutdown.

        This property is used internally and should not be set in normal usage.

        :param bool value: Whether to shutdown the connection.
        """
        if self._handler:
            self._handler._shutdown = value  # pylint: disable=protected-access
        else:
            raise ValueError("Receiver has no AMQP handler")

    @property
    def queue_size(self):
        """The current size of the unprocessed message queue.

        :rtype: int
        """
        # pylint: disable=protected-access
        if self._handler._received_messages:
            return self._handler._received_messages.qsize()
        return 0

    async def open(self):
        """Open receiver connection and authenticate session.

        If the receiver is already open, this operation will do nothing.
        This method will be called automatically when one starts to iterate
        messages in the receiver, so there should be no need to call it directly.
        A receiver opened with this method must be explicitly closed.
        It is recommended to open a handler within a context manager as
        opposed to calling the method directly.

        .. note:: This operation is not thread-safe.

        """
        if self.running:
            return
        self.running = True
        try:
            await self._handler.open_async(connection=self.connection)
            self.message_iter = self._handler.receive_messages_iter_async()
            while not await self._handler.auth_complete_async():
                await asyncio.sleep(0.05)
            await self._build_receiver()
            while not await self._handler.client_ready_async():
                await asyncio.sleep(0.05)
        except Exception as e:  # pylint: disable=broad-except
            try:
                await self._handle_exception(e)
            except:
                self.running = False
                raise

    async def close(self, exception=None):
        """Close down the receiver connection.

        If the receiver has already closed, this operation will do nothing. An optional
        exception can be passed in to indicate that the handler was shutdown due to error.
        It is recommended to open a handler within a context manager as
        opposed to calling the method directly.
        The receiver will be implicitly closed on completion of the message iterator,
        however this method will need to be called explicitly if the message iterator is not run
        to completion.

        .. note:: This operation is not thread-safe.

        :param exception: An optional exception if the handler is closing
         due to an error.
        :type exception: Exception

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START open_close_receiver_directly]
                :end-before: [END open_close_receiver_directly]
                :language: python
                :dedent: 4
                :caption: Iterate then explicitly close a Receiver.

        """
        if not self.running:
            return
        self.running = False
        self.receiver_shutdown = True
        self._used.set()
        await super(Receiver, self).close(exception=exception)

    async def peek(self, count=1, start_from=0):
        """Browse messages currently pending in the queue.

        Peeked messages are not removed from queue, nor are they locked. They cannot be completed,
        deferred or dead-lettered.

        :param count: The maximum number of messages to try and peek. The default
         value is 1.
        :type count: int
        :param start_from: A message sequence number from which to start browsing messages.
        :type start_from: int
        :rtype: list[~azure.servicebus.common.message.PeekMessage]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_peek_messages]
                :end-before: [END receiver_peek_messages]
                :language: python
                :dedent: 4
                :caption: Peek messages in the queue.

        """
        await self._can_run()
        if not start_from:
            start_from = self.last_received or 1
        if int(count) < 1:
            raise ValueError("count must be 1 or greater.")
        if int(start_from) < 1:
            raise ValueError("start_from must be 1 or greater.")

        message = {
            'from-sequence-number': types.AMQPLong(start_from),
            'message-count': count
        }
        return await self._mgmt_request_response(
            REQUEST_RESPONSE_PEEK_OPERATION, message, mgmt_handlers.peek_op)

    async def receive_deferred_messages(self,
                                        sequence_numbers,
                                        mode=ReceiveSettleMode.PeekLock):
        """Receive messages that have previously been deferred.

        When receiving deferred messages from a partitioned entity, all of the supplied
        sequence numbers must be messages from the same partition.

        :param sequence_numbers: A list of the sequence numbers of messages that have been
         deferred.
        :type sequence_numbers: list[int]
        :param mode: The receive mode, default value is PeekLock.
        :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
        :rtype: list[~azure.servicebus.aio.async_message.DeferredMessage]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_defer_messages]
                :end-before: [END receiver_defer_messages]
                :language: python
                :dedent: 8
                :caption: Defer messages, then retrieve them by sequence number.

        """
        if not sequence_numbers:
            raise ValueError("At least one sequence number must be specified.")
        await self._can_run()
        try:
            receive_mode = mode.value.value
        except AttributeError:
            receive_mode = int(mode)
        message = {
            'sequence-numbers':
            types.AMQPArray([types.AMQPLong(s) for s in sequence_numbers]),
            'receiver-settle-mode':
            types.AMQPuInt(receive_mode)
        }
        handler = functools.partial(mgmt_handlers.deferred_message_op,
                                    mode=receive_mode,
                                    message_type=DeferredMessage)
        messages = await self._mgmt_request_response(
            REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER, message, handler)
        for m in messages:
            m._receiver = self  # pylint: disable=protected-access
        return messages

    async def fetch_next(self, max_batch_size=None, timeout=None):
        """Receive a batch of messages at once.

        This approach it optimal if you wish to process multiple messages simultaneously.
        Note that the number of messages retrieved in a single batch will be dependent on
        whether `prefetch` was set for the receiver. This call will prioritize returning
        quickly over meeting a specified batch size, and so will return as soon as at least
        one message is received and there is a gap in incoming messages regardless
        of the specified batch size.

        :param max_batch_size: Maximum number of messages in the batch. Actual number
         returned will depend on prefetch size and incoming stream rate.
        :type max_batch_size: int
        :param timeout: The time to wait in seconds for the first message to arrive.
         If no messages arrive, and no timeout is specified, this call will not return
         until the connection is closed. If specified, an no messages arrive within the
         timeout period, an empty list will be returned.
        :rtype: list[~azure.servicebus.aio.async_message.Message]

        Example:
            .. literalinclude:: ../examples/async_examples/test_examples_async.py
                :start-after: [START receiver_fetch_batch]
                :end-before: [END receiver_fetch_batch]
                :language: python
                :dedent: 4
                :caption: Fetch a batch of messages.

        """
        await self._can_run()
        wrapped_batch = []
        max_batch_size = max_batch_size or self._handler._prefetch  # pylint: disable=protected-access
        try:
            timeout_ms = 1000 * timeout if timeout else 0
            batch = await self._handler.receive_message_batch_async(
                max_batch_size=max_batch_size, timeout=timeout_ms)
            for received in batch:
                message = self._build_message(received)
                wrapped_batch.append(message)
        except Exception as e:  # pylint: disable=broad-except
            await self._handle_exception(e)
        return wrapped_batch