diff --git a/eng/pipelines/templates/steps/build-artifacts.yml b/eng/pipelines/templates/steps/build-artifacts.yml index fd7b2dc23847..2b8d5b625acb 100644 --- a/eng/pipelines/templates/steps/build-artifacts.yml +++ b/eng/pipelines/templates/steps/build-artifacts.yml @@ -90,9 +90,9 @@ steps: Write-Host "##vso[task.setvariable variable=PIP_INDEX_URL]https://pypi.python.org/simple" displayName: Reset PIP Index For APIStubGen -# - template: /eng/pipelines/templates/steps/run_apistub.yml -# parameters: -# ServiceDirectory: ${{ parameters.ServiceDirectory }} + #- template: /eng/pipelines/templates/steps/run_apistub.yml + # parameters: + # ServiceDirectory: ${{ parameters.ServiceDirectory }} - ${{ parameters.BeforePublishSteps }} diff --git a/sdk/eventhub/azure-eventhub/CHANGELOG.md b/sdk/eventhub/azure-eventhub/CHANGELOG.md index 7a5b2230c7f8..f51053596b3c 100644 --- a/sdk/eventhub/azure-eventhub/CHANGELOG.md +++ b/sdk/eventhub/azure-eventhub/CHANGELOG.md @@ -1,6 +1,6 @@ # Release History -## 5.8.0b6 (Unreleased) +## 5.8.0b1 (Unreleased) ### Features Added @@ -10,6 +10,8 @@ ### Other Changes +- Added the `uamqp_transport` optional parameter to the clients, to allow switching to the `uamqp` library as the transport. + ## 5.8.0a5 (2022-07-19) ### Bugs Fixed @@ -21,6 +23,26 @@ - Added logging added in to track proper token refreshes & fetches, output exception reason for producer init failure. +## 5.10.0 (2022-06-08) + +### Features Added + +- Includes the following features related to buffered sending of events: + - A new method `send_event` to `EventHubProducerClient` which allows sending single `EventData` or `AmqpAnnotatedMessage`. + - Buffered mode sending to `EventHubProducerClient` which is intended to allow for efficient publishing of events + without having to explicitly manage batches in the application. + - The constructor of `EventHubProducerClient` and `from_connection_string` method takes the following new keyword arguments + for configuration: + - `buffered_mode`: The flag to enable/disable buffered mode sending. + - `on_success`: The callback to be called once events have been successfully published. + - `on_error`: The callback to be called once events have failed to be published. + - `max_buffer_length`: The total number of events per partition that can be buffered before a flush will be triggered. + - `max_wait_time`: The amount of time to wait for a batch to be built with events in the buffer before publishing. + - A new method `EventHubProducerClient.flush` which flushes events in the buffer to be sent immediately. + - A new method `EventHubProducerClient.get_buffered_event_count` which returns the number of events that are buffered and waiting to be published for a given partition. + - A new property `EventHubProducerClient.total_buffered_event_count` which returns the total number of events that are currently buffered and waiting to be published, across all partitions. + - A new boolean keyword argument `flush` to `EventHubProducerClient.close` which indicates whether to flush the buffer or not while closing. + ## 5.8.0a4 (2022-06-07) ### Features Added @@ -28,6 +50,39 @@ - Added support for connection using websocket and http proxy. - Added support for custom endpoint connection over websocket. +## 5.9.0 (2022-05-10) + +### Features Added + +- The classmethod `from_message_content` has been added to `EventData` for interoperability with the Schema Registry Avro Encoder library, and takes `content` and `content_type` as positional parameters. + +### Other Changes + +- Features related to buffered sending of events are still in beta and will not be included in this release. + +## 5.9.0b3 (2022-04-20) + +### Features Added + +- Introduced new method `send_event` to `EventHubProducerClient` which allows sending single `EventData` or `AmqpAnnotatedMessage`. +- Introduced buffered mode sending to `EventHubProducerClient` which is intended to allow for efficient publishing of events + without having to explicitly manage batches in the application. + - The constructor of `EventHubProducerClient` and `from_connection_string` method now takes the following new keyword arguments + for configuration: + - `buffered_mode`: The flag to enable/disable buffered mode sending. + - `on_success`: The callback to be called once events have been successfully published. + - `on_error`: The callback to be called once events have failed to be published. + - `max_buffer_length`: The total number of events per partition that can be buffered before a flush will be triggered. + - `max_wait_time`: The amount of time to wait for a batch to be built with events in the buffer before publishing. + - Introduced new method `EventHubProducerClient.flush` which flushes events in the buffer to be sent immediately. + - Introduced new method `EventHubProducerClient.get_buffered_event_count` which returns the number of events that are buffered and waiting to be published for a given partition. + - Introduced new property `EventHubProducerClient.total_buffered_event_count` which returns the total number of events that are currently buffered and waiting to be published, across all partitions. + - Introduced new boolean keyword argument `flush` to `EventHubProducerClient.close` which indicates whether to flush the buffer or not while closing. + +### Other Changes + +- Updated `EventData` internals for interoperability with the Schema Registry Avro Encoder library. + ## 5.8.0a3 (2022-03-08) ### Other Changes @@ -40,7 +95,7 @@ - Added support for async `EventHubProducerClient` and `EventHubConsumerClient`. -### Breaking changes +## 5.9.0b1 (2022-02-09) - The following features have been temporarily pulled out of async `EventHubProducerClient` and `EventHubConsumerClient` which will be added back in future previews as we work towards a stable release: - Passing the following keyword arguments to the constructors and `from_connection_string` methods of the `EventHubProducerClient` and `EventHubConsumerClient` is not supported: `transport_type`, `http_proxy`, `custom_endpoint_address`, and `connection_verify`. diff --git a/sdk/eventhub/azure-eventhub/README.md b/sdk/eventhub/azure-eventhub/README.md index 3923e456f503..0a702d3bb91a 100644 --- a/sdk/eventhub/azure-eventhub/README.md +++ b/sdk/eventhub/azure-eventhub/README.md @@ -26,7 +26,7 @@ Please refer to the changelog for more details._ ### Prerequisites -- Python 3.6 or later. +- Python 3.7 or later. - **Microsoft Azure Subscription:** To use Azure services, including Azure Event Hubs, you'll need a subscription. If you do not have an existing Azure account, you may sign up for a free trial or use your MSDN subscriber benefits when you [create an account](https://account.windowsazure.com/Home/Index). @@ -281,22 +281,7 @@ partition_ids = client.get_partition_ids() ## Troubleshooting -### General - -The Event Hubs APIs generate the following exceptions in azure.eventhub.exceptions - -- **AuthenticationError:** Failed to authenticate because of wrong address, SAS policy/key pair, SAS token or azure identity. -- **ConnectError:** Failed to connect to the EventHubs. The AuthenticationError is a type of ConnectError. -- **ConnectionLostError:** Lose connection after a connection has been built. -- **EventDataError:** The EventData to be sent fails data validation. For instance, this error is raised if you try to send an EventData that is already sent. -- **EventDataSendError:** The Eventhubs service responds with an error when an EventData is sent. -- **OperationTimeoutError:** EventHubConsumer.send() times out. -- **EventHubError:** All other Eventhubs related errors. It is also the root error class of all the errors described above. - -### Logging - -- Enable `azure.eventhub` logger to collect traces from the library. -- Enable AMQP frame level trace by setting `logging_enable=True` when creating the client. +See the `azure-eventhubs` [troubleshooting guide](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/eventhub/azure-eventhub/TROUBLESHOOTING.md) for details on how to diagnose various failure scenarios. ## Next steps diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py index c88edfa66292..6645bf9ea577 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py @@ -7,6 +7,7 @@ __version__ = VERSION +from ._constants import TransportType from ._producer_client import EventHubProducerClient from ._consumer_client import EventHubConsumerClient from ._client_base import EventHubSharedKeyCredential @@ -15,9 +16,8 @@ from ._eventprocessor.partition_context import PartitionContext from ._connection_string_parser import ( parse_connection_string, - EventHubConnectionStringProperties + EventHubConnectionStringProperties, ) -from ._constants import TransportType __all__ = [ "EventData", diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/__init__.py new file mode 100644 index 000000000000..bfee862537a1 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/__init__.py @@ -0,0 +1,13 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from ._buffered_producer import BufferedProducer +from ._partition_resolver import PartitionResolver +from ._buffered_producer_dispatcher import BufferedProducerDispatcher + +__all__ = [ + "BufferedProducer", + "PartitionResolver", + "BufferedProducerDispatcher", +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py new file mode 100644 index 000000000000..f8d327760a2d --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py @@ -0,0 +1,218 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +import time +import queue +import logging +from threading import RLock +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Callable, TYPE_CHECKING + +from .._producer import EventHubProducer +from .._common import EventDataBatch +from ..exceptions import OperationTimeoutError + +if TYPE_CHECKING: + from .._transport._base import AmqpTransport + from .._producer_client import SendEventTypes + +_LOGGER = logging.getLogger(__name__) + + +class BufferedProducer: + # pylint: disable=too-many-instance-attributes + def __init__( + self, + producer: EventHubProducer, + partition_id: str, + on_success: Callable[["SendEventTypes", Optional[str]], None], + on_error: Callable[["SendEventTypes", Optional[str], Exception], None], + max_message_size_on_link: int, + executor: ThreadPoolExecutor, + *, + amqp_transport: AmqpTransport, + max_buffer_length: int, + max_wait_time: float = 1 + ): + self._buffered_queue: queue.Queue = queue.Queue() + self._max_buffer_len = max_buffer_length + self._cur_buffered_len = 0 + self._executor: ThreadPoolExecutor = executor + self._producer: EventHubProducer = producer + self._lock = RLock() + self._max_wait_time = max_wait_time + self._on_success = self.failsafe_callback(on_success) + self._on_error = self.failsafe_callback(on_error) + self._last_send_time = None + self._running = False + self._cur_batch: Optional[EventDataBatch] = None + self._max_message_size_on_link = max_message_size_on_link + self._check_max_wait_time_future = None + self.partition_id = partition_id + self._amqp_transport = amqp_transport + + def start(self): + with self._lock: + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + self._running = True + if self._max_wait_time: + self._last_send_time = time.time() + self._check_max_wait_time_future = self._executor.submit( + self.check_max_wait_time_worker + ) + + def stop(self, flush=True, timeout_time=None, raise_error=False): + self._running = False + if flush: + with self._lock: + self.flush(timeout_time=timeout_time, raise_error=raise_error) + else: + if self._cur_buffered_len: + _LOGGER.warning( + "Shutting down Partition %r. There are still %r events in the buffer which will be lost", + self.partition_id, + self._cur_buffered_len, + ) + if self._check_max_wait_time_future: + remain_timeout = timeout_time - time.time() if timeout_time else None + try: + self._check_max_wait_time_future.result(remain_timeout) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.warning( + "Partition %r stopped with error %r", self.partition_id, exc + ) + self._producer.close() + + def put_events(self, events, timeout_time=None): + # Put single event or EventDataBatch into the queue. + # This method would raise OperationTimeout if the queue does not have enough space for the input and + # flush cannot finish in timeout. + try: + new_events_len = len(events) + except TypeError: + new_events_len = 1 + if self._max_buffer_len - self._cur_buffered_len < new_events_len: + _LOGGER.info( + "The buffer for partition %r is full. Attempting to flush before adding %r events.", + self.partition_id, + new_events_len, + ) + # flush the buffer + self.flush(timeout_time=timeout_time) + if timeout_time and time.time() > timeout_time: + raise OperationTimeoutError( + "Failed to enqueue events into buffer due to timeout." + ) + try: + # add single event into current batch + self._cur_batch.add(events) + except AttributeError: # if the input events is a EventDataBatch, put the whole into the buffer + # if there are events in cur_batch, enqueue cur_batch to the buffer + with self._lock: + if self._cur_batch: + self._buffered_queue.put(self._cur_batch) + self._buffered_queue.put(events) + # create a new batch for incoming events + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + except ValueError: + # add single event exceeds the cur batch size, create new batch + with self._lock: + self._buffered_queue.put(self._cur_batch) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + self._cur_batch.add(events) + with self._lock: + self._cur_buffered_len += new_events_len + + def failsafe_callback(self, callback): + def wrapper_callback(*args, **kwargs): + try: + callback(*args, **kwargs) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.warning( + "On partition %r, callback %r encountered exception %r", + callback.__name__, + exc, + self.partition_id, + ) + + return wrapper_callback + + def flush(self, timeout_time=None, raise_error=True): + # pylint: disable=protected-access + # try flushing all the buffered batch within given time + with self._lock: + _LOGGER.info("Partition: %r started flushing.", self.partition_id) + if self._cur_batch: # if there is batch, enqueue it to the buffer first + self._buffered_queue.put(self._cur_batch) + while self._buffered_queue.qsize() > 0: + remaining_time = timeout_time - time.time() if timeout_time else None + if (remaining_time and remaining_time > 0) or remaining_time is None: + try: + batch = self._buffered_queue.get(block=False) + except queue.Empty: + break + self._buffered_queue.task_done() + try: + _LOGGER.info("Partition %r is sending.", self.partition_id) + self._producer.send( + batch, + timeout=timeout_time - time.time() + if timeout_time + else None, + ) + _LOGGER.info( + "Partition %r sending %r events succeeded.", + self.partition_id, + len(batch), + ) + self._on_success(batch._internal_events, self.partition_id) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info( + "Partition %r sending %r events failed due to exception: %r ", + self.partition_id, + len(batch), + exc, + ) + self._on_error(batch._internal_events, self.partition_id, exc) + finally: + self._cur_buffered_len -= len(batch) + else: + _LOGGER.info( + "Partition %r fails to flush due to timeout.", self.partition_id + ) + if raise_error: + raise OperationTimeoutError( + "Failed to flush {!r} within {}".format( + self.partition_id, timeout_time + ) + ) + break + # after finishing flushing, reset cur batch and put it into the buffer + self._last_send_time = time.time() + #reset buffered count + self._cur_buffered_len = 0 + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + _LOGGER.info("Partition %r finished flushing.", self.partition_id) + + def check_max_wait_time_worker(self): + while self._running: + if self._cur_buffered_len > 0: + now_time = time.time() + _LOGGER.info( + "Partition %r worker is checking max_wait_time.", self.partition_id + ) + # flush the partition if the producer is running beyond the waiting time + # or the buffer is at max capacity + if (now_time - self._last_send_time > self._max_wait_time) or ( + self._cur_buffered_len >= self._max_buffer_len + ): + # in the worker, not raising error for flush, users can not handle this + with self._lock: + self.flush(raise_error=False) + time.sleep(min(self._max_wait_time, 5)) + + @property + def buffered_event_count(self): + return self._cur_buffered_len diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py new file mode 100644 index 000000000000..7bad79bead64 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py @@ -0,0 +1,187 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +import logging +from threading import Lock +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, Optional, List, Callable, Union, TYPE_CHECKING + +from ._partition_resolver import PartitionResolver +from ._buffered_producer import BufferedProducer +from .._producer import EventHubProducer +from ..exceptions import EventDataSendError, ConnectError, EventHubError + +if TYPE_CHECKING: + from .._transport._base import AmqpTransport + from .._producer_client import SendEventTypes + +_LOGGER = logging.getLogger(__name__) + + +class BufferedProducerDispatcher: + # pylint: disable=too-many-instance-attributes + def __init__( + self, + partitions: List[str], + on_success: Callable[["SendEventTypes", Optional[str]], None], + on_error: Callable[["SendEventTypes", Optional[str], Exception], None], + create_producer: Callable[..., EventHubProducer], + eventhub_name: str, + max_message_size_on_link: int, + *, + amqp_transport: AmqpTransport, + max_buffer_length: int = 1500, + max_wait_time: float = 1, + executor: Optional[Union[ThreadPoolExecutor, int]] = None + ): + self._buffered_producers: Dict[str, BufferedProducer] = {} + self._partition_ids: List[str] = partitions + self._lock = Lock() + self._on_success = on_success + self._on_error = on_error + self._create_producer = create_producer + self._eventhub_name = eventhub_name + self._max_message_size_on_link = max_message_size_on_link + self._partition_resolver = PartitionResolver(self._partition_ids) + self._max_wait_time = max_wait_time + self._max_buffer_length = max_buffer_length + self._existing_executor = False + self._amqp_transport = amqp_transport + + if not executor: + self._executor = ThreadPoolExecutor() + elif isinstance(executor, ThreadPoolExecutor): + self._existing_executor = True + self._executor = executor + elif isinstance(executor, int): + self._executor = ThreadPoolExecutor(executor) + + def _get_partition_id(self, partition_id, partition_key): + if partition_id: + if partition_id not in self._partition_ids: + raise ConnectError( + "Invalid partition {} for the event hub {}".format( + partition_id, self._eventhub_name + ) + ) + return partition_id + if isinstance(partition_key, str): + return self._partition_resolver.get_partition_id_by_partition_key( + partition_key + ) + return self._partition_resolver.get_next_partition_id() + + def enqueue_events( + self, events, *, partition_id=None, partition_key=None, timeout_time=None + ): + pid = self._get_partition_id(partition_id, partition_key) + with self._lock: + try: + self._buffered_producers[pid].put_events(events, timeout_time) + except KeyError: + buffered_producer = BufferedProducer( + self._create_producer(pid), + pid, + self._on_success, + self._on_error, + self._max_message_size_on_link, + executor=self._executor, + max_wait_time=self._max_wait_time, + max_buffer_length=self._max_buffer_length, + amqp_transport = self._amqp_transport, + ) + buffered_producer.start() + self._buffered_producers[pid] = buffered_producer + buffered_producer.put_events(events, timeout_time) + + def flush(self, timeout_time=None): + # flush all the buffered producer, the method will block until finishes or times out + with self._lock: + futures = [] + for pid, producer in self._buffered_producers.items(): + # call each producer's flush method + futures.append( + ( + pid, + self._executor.submit( + producer.flush, timeout_time=timeout_time + ), + ) + ) + + # gather results + exc_results = {} + for pid, future in futures: + try: + future.result() + except Exception as exc: # pylint: disable=broad-except + exc_results[pid] = exc + + if not exc_results: + _LOGGER.info("Flushing all partitions succeeded") + return + + _LOGGER.warning( + "Flushing all partitions partially failed with result %r.", exc_results + ) + raise EventDataSendError( + message="Flushing all partitions partially failed, failed partitions are {!r}" + " Exception details are {!r}".format(exc_results.keys(), exc_results) + ) + + def close(self, *, flush=True, timeout_time=None, raise_error=False): + + with self._lock: + + futures = [] + # stop all buffered producers + for pid, producer in self._buffered_producers.items(): + futures.append( + ( + pid, + self._executor.submit( + producer.stop, + flush=flush, + timeout_time=timeout_time, + raise_error=raise_error, + ), + ) + ) + + exc_results = {} + # gather results + for pid, future in futures: + try: + future.result() + except Exception as exc: # pylint: disable=broad-except + exc_results[pid] = exc + + if exc_results: + _LOGGER.warning( + "Stopping all partitions partially failed with result %r.", + exc_results, + ) + if raise_error: + raise EventHubError( + message="Stopping all partitions partially failed, failed partitions are {!r}" + " Exception details are {!r}".format( + exc_results.keys(), exc_results + ) + ) + + if not self._existing_executor: + self._executor.shutdown() + + def get_buffered_event_count(self, pid): + try: + return self._buffered_producers[pid].buffered_event_count + except KeyError: + return 0 + + @property + def total_buffered_event_count(self): + return sum( + [self.get_buffered_event_count(pid) for pid in self._buffered_producers] + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_partition_resolver.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_partition_resolver.py new file mode 100644 index 000000000000..7d647a8553b6 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_partition_resolver.py @@ -0,0 +1,287 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +""" +jenkins-hash lookup3 algorithm implementation +""" + +from threading import Lock +import struct + +c_signed_short = struct.Struct(">h") + + +def rot(x, k): + return (x << k) | (x >> (32 - k)) + + +def mix(a, b, c): + a &= 0xFFFFFFFF + b &= 0xFFFFFFFF + c &= 0xFFFFFFFF + a -= c + a &= 0xFFFFFFFF + a ^= rot(c, 4) + a &= 0xFFFFFFFF + c += b + c &= 0xFFFFFFFF + b -= a + b &= 0xFFFFFFFF + b ^= rot(a, 6) + b &= 0xFFFFFFFF + a += c + a &= 0xFFFFFFFF + c -= b + c &= 0xFFFFFFFF + c ^= rot(b, 8) + c &= 0xFFFFFFFF + b += a + b &= 0xFFFFFFFF + a -= c + a &= 0xFFFFFFFF + a ^= rot(c, 16) + a &= 0xFFFFFFFF + c += b + c &= 0xFFFFFFFF + b -= a + b &= 0xFFFFFFFF + b ^= rot(a, 19) + b &= 0xFFFFFFFF + a += c + a &= 0xFFFFFFFF + c -= b + c &= 0xFFFFFFFF + c ^= rot(b, 4) + c &= 0xFFFFFFFF + b += a + b &= 0xFFFFFFFF + return a, b, c + + +def final(a, b, c): + a &= 0xFFFFFFFF + b &= 0xFFFFFFFF + c &= 0xFFFFFFFF + c ^= b + c &= 0xFFFFFFFF + c -= rot(b, 14) + c &= 0xFFFFFFFF + a ^= c + a &= 0xFFFFFFFF + a -= rot(c, 11) + a &= 0xFFFFFFFF + b ^= a + b &= 0xFFFFFFFF + b -= rot(a, 25) + b &= 0xFFFFFFFF + c ^= b + c &= 0xFFFFFFFF + c -= rot(b, 16) + c &= 0xFFFFFFFF + a ^= c + a &= 0xFFFFFFFF + a -= rot(c, 4) + a &= 0xFFFFFFFF + b ^= a + b &= 0xFFFFFFFF + b -= rot(a, 14) + b &= 0xFFFFFFFF + c ^= b + c &= 0xFFFFFFFF + c -= rot(b, 24) + c &= 0xFFFFFFFF + return a, b, c + + +def compute_hash(data, init_val=0, init_val2=0): + # pylint: disable=too-many-statements + """ + implementation by: + https://stackoverflow.com/questions/3279615/python-implementation-of-jenkins-hash + """ + length = lenpos = len(data) + + a = b = c = 0xDEADBEEF + length + init_val + + c += init_val2 + c &= 0xFFFFFFFF + + p = 0 # string offset + while lenpos > 12: + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + a &= 0xFFFFFFFF + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + b &= 0xFFFFFFFF + c += ( + ord(data[p + 8]) + + (ord(data[p + 9]) << 8) + + (ord(data[p + 10]) << 16) + + (ord(data[p + 11]) << 24) + ) + c &= 0xFFFFFFFF + a, b, c = mix(a, b, c) + p += 12 + lenpos -= 12 + + if lenpos == 12: + c += ( + ord(data[p + 8]) + + (ord(data[p + 9]) << 8) + + (ord(data[p + 10]) << 16) + + (ord(data[p + 11]) << 24) + ) + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 11: + c += ord(data[p + 8]) + (ord(data[p + 9]) << 8) + (ord(data[p + 10]) << 16) + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 10: + c += ord(data[p + 8]) + (ord(data[p + 9]) << 8) + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 9: + c += ord(data[p + 8]) + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 8: + b += ( + ord(data[p + 4]) + + (ord(data[p + 5]) << 8) + + (ord(data[p + 6]) << 16) + + (ord(data[p + 7]) << 24) + ) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 7: + b += ord(data[p + 4]) + (ord(data[p + 5]) << 8) + (ord(data[p + 6]) << 16) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 6: + b += (ord(data[p + 5]) << 8) + ord(data[p + 4]) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 5: + b += ord(data[p + 4]) + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 4: + a += ( + ord(data[p + 0]) + + (ord(data[p + 1]) << 8) + + (ord(data[p + 2]) << 16) + + (ord(data[p + 3]) << 24) + ) + if lenpos == 3: + a += ord(data[p + 0]) + (ord(data[p + 1]) << 8) + (ord(data[p + 2]) << 16) + if lenpos == 2: + a += ord(data[p + 0]) + (ord(data[p + 1]) << 8) + if lenpos == 1: + a += ord(data[p + 0]) + + a &= 0xFFFFFFFF + b &= 0xFFFFFFFF + c &= 0xFFFFFFFF + if lenpos == 0: + return c, b + + a, b, c = final(a, b, c) + + return c, b + + +def generate_hash_code(partition_key): + if not partition_key: + return 0 + + hash_tuple = compute_hash(partition_key, 0, 0) + hash_value = (hash_tuple[0] ^ hash_tuple[1]) & 0xFFFF + return c_signed_short.unpack(struct.pack(">H", hash_value))[0] + + +class PartitionResolver: + def __init__(self, partitions): + self._idx = -1 + self._partitions = partitions + self._partitions_cnt = len(self._partitions) + self._lock = Lock() + + def get_next_partition_id(self): + """ + round-robin partition assignment + """ + with self._lock: + self._idx += 1 + self._idx %= self._partitions_cnt + return self._partitions[self._idx] + + def get_partition_id_by_partition_key(self, partition_key): + hash_code = generate_hash_code(partition_key) + return self._partitions[abs(hash_code % self._partitions_cnt)] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py index ede4c672d4a3..1a39889b52f1 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations import logging import uuid @@ -10,9 +10,12 @@ import functools import collections from typing import Any, Dict, Tuple, List, Optional, TYPE_CHECKING, cast, Union -from datetime import timedelta, datetime +try: + from typing import TypeAlias +except ImportError: + from typing_extensions import TypeAlias +from datetime import timedelta from urllib.parse import urlparse -import six from azure.core.credentials import ( AccessToken, @@ -23,26 +26,36 @@ from azure.core.pipeline.policies import RetryMode -from ._pyamqp.client import AMQPClient -from ._pyamqp.message import Message -from ._pyamqp import constants, error as errors, utils as pyamqp_utils -from ._pyamqp.authentication import JWTTokenAuth -from .exceptions import _handle_exception, ClientClosedError +try: + from ._transport._uamqp_transport import UamqpTransport +except ImportError: + UamqpTransport = None +from ._transport._pyamqp_transport import PyamqpTransport +from .exceptions import ClientClosedError from ._configuration import Configuration from ._utils import utc_from_timestamp, parse_sas_credential +from ._pyamqp.utils import generate_sas_token from ._connection_manager import get_connection_manager from ._constants import ( CONTAINER_PREFIX, JWT_TOKEN_SCOPE, - MGMT_OPERATION, - MGMT_PARTITION_OPERATION, + READ_OPERATION, MGMT_STATUS_CODE, MGMT_STATUS_DESC, - READ_OPERATION + MGMT_OPERATION, + MGMT_PARTITION_OPERATION, ) if TYPE_CHECKING: from azure.core.credentials import TokenCredential + try: + from uamqp import Message as uamqp_Message + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_Message = None + uamqp_JWTTokenAuth = None + from ._pyamqp.message import Message + from ._pyamqp.authentication import JWTTokenAuth _LOGGER = logging.getLogger(__name__) _Address = collections.namedtuple("_Address", "hostname path") @@ -135,7 +148,7 @@ def _parse_conn_str(conn_str, **kwargs): def _generate_sas_token(uri, policy, key, expiry=None): # type: (str, str, str, Optional[timedelta]) -> AccessToken - """Create a shared access signiture token as a string literal. + """Create a shared access signature token as a string literal. :returns: SAS token as string literal. :rtype: str """ @@ -144,7 +157,7 @@ def _generate_sas_token(uri, policy, key, expiry=None): abs_expiry = int(time.time()) + expiry.seconds - token = pyamqp_utils.generate_sas_token(uri, policy, key, abs_expiry).encode() + token = generate_sas_token(uri, policy, key, abs_expiry).encode() return AccessToken(token=token, expires_on=abs_expiry) @@ -184,7 +197,6 @@ def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument # type: (str, Any) -> AccessToken if not scopes: raise ValueError("No token scope provided.") - return _generate_sas_token(scopes[0], self.policy, self.key) @@ -261,9 +273,31 @@ def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument return AccessToken(signature, expiry) +# separate TYPE_CHECKING block here for EventHubSharedKeyCredential, o/w mypy raised error even with forward referencing +if TYPE_CHECKING: + from azure.core.credentials import TokenCredential + + CredentialTypes: TypeAlias = Union[ + AzureSasCredential, + AzureNamedKeyCredential, + EventHubSharedKeyCredential, + TokenCredential, + ] + + class ClientBase(object): # pylint:disable=too-many-instance-attributes - def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwargs): - # type: (str, str, Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential], Any) -> None + def __init__( + self, + fully_qualified_namespace: str, + eventhub_name: str, + credential: CredentialTypes, + **kwargs: Any, + ) -> None: + uamqp_transport = kwargs.pop("uamqp_transport", False) + if uamqp_transport and not UamqpTransport: + raise ValueError("To use the uAMQP transport, please install `uamqp>=1.6.0,<2.0.0`.") + self._amqp_transport = kwargs.pop("amqp_transport", UamqpTransport if uamqp_transport else PyamqpTransport) + self.eventhub_name = eventhub_name if not eventhub_name: raise ValueError("The eventhub name can not be None or empty.") @@ -278,13 +312,17 @@ def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwarg self._credential = credential # type: ignore self._keep_alive = kwargs.get("keep_alive", 30) self._auto_reconnect = kwargs.get("auto_reconnect", True) - self._mgmt_target = "amqps://{}/{}".format( - self._address.hostname, self.eventhub_name + self._auth_uri = f"sb://{self._address.hostname}{self._address.path}" + self._config = Configuration( + uamqp_transport=uamqp_transport, + hostname=self._address.hostname, + **kwargs, ) - self._auth_uri = "sb://{}{}".format(self._address.hostname, self._address.path) - self._config = Configuration(**kwargs) self._debug = self._config.network_tracing - self._conn_manager = get_connection_manager(**kwargs) + kwargs["custom_endpoint_address"] = self._config.custom_endpoint_address + self._conn_manager = get_connection_manager( + amqp_transport=self._amqp_transport, + **kwargs) self._idle_timeout = kwargs.get("idle_timeout", None) @staticmethod @@ -301,11 +339,10 @@ def _from_connection_string(conn_str, **kwargs): kwargs["credential"] = EventHubSharedKeyCredential(policy, key) return kwargs - def _create_auth(self): - # type: () -> JWTTokenAuth + def _create_auth(self) -> Union[uamqp_JWTTokenAuth, JWTTokenAuth]: """ - Create an ~uamqp.authentication.SASTokenAuth instance to authenticate - the session. + Create an ~uamqp.authentication.SASTokenAuth instance + to authenticate the session. """ try: # ignore mypy's warning because token_type is Optional @@ -313,20 +350,19 @@ def _create_auth(self): except AttributeError: token_type = b"jwt" if token_type == b"servicebus.windows.net:sastoken": - return JWTTokenAuth( - self._auth_uri, + return self._amqp_transport.create_token_auth( self._auth_uri, - functools.partial(self._credential.get_token, self._auth_uri) + functools.partial(self._credential.get_token, self._auth_uri), + token_type=token_type, + config=self._config, + update_token=True, ) - return JWTTokenAuth( - self._auth_uri, + return self._amqp_transport.create_token_auth( self._auth_uri, functools.partial(self._credential.get_token, JWT_TOKEN_SCOPE), token_type=token_type, - timeout=self._config.auth_timeout, - custom_endpoint_hostname=self._config.custom_endpoint_hostname, - port=self._config.connection_port, - verify=self._config.connection_verify, + config=self._config, + update_token=False, ) def _close_connection(self): @@ -361,81 +397,45 @@ def _backoff( ) raise last_exception - def _management_request(self, mgmt_msg, op_type): - # type: (Message, bytes) -> Any + def _management_request( + self, mgmt_msg: Union[uamqp_Message, Message], op_type: bytes + ) -> Any: + # pylint:disable=assignment-from-none retried_times = 0 last_exception = None while retried_times <= self._config.max_retries: mgmt_auth = self._create_auth() - hostname = self._address.hostname - custom_endpoint_address = self._config.custom_endpoint_address - if self._config.transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - mgmt_client = AMQPClient( - hostname, - auth=mgmt_auth, - network_trace=self._config.network_tracing, - transport_type=self._config.transport_type, - http_proxy=self._config.http_proxy, - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._config.connection_verify + mgmt_client = self._amqp_transport.create_mgmt_client( + self._address, mgmt_auth=mgmt_auth, config=self._config ) try: - mgmt_client.open() + conn = self._conn_manager.get_connection( # pylint:disable=assignment-from-none + host=self._address.hostname, auth=mgmt_auth + ) + mgmt_client.open(connection=conn) while not mgmt_client.client_ready(): time.sleep(0.05) - access_token = mgmt_auth.get_token() - - if not access_token: - _LOGGER.debug("Management client received an empty access token object") - - elif not access_token.token: - _LOGGER.debug("Management client received an empty token") - - else: - _LOGGER.debug(f"Management client token expires on: {datetime.fromtimestamp(access_token.expires_on)}") - - mgmt_msg.application_properties["security_token"] = access_token.token - - status_code, description, response = mgmt_client.mgmt_request( + mgmt_msg.application_properties[ + "security_token" + ] = self._amqp_transport.get_updated_token(mgmt_auth) + status_code, description, response = self._amqp_transport.mgmt_client_request( + mgmt_client, mgmt_msg, - operation=READ_OPERATION.decode(), - operation_type=op_type.decode(), + operation=READ_OPERATION, + operation_type=op_type, status_code_field=MGMT_STATUS_CODE, description_fields=MGMT_STATUS_DESC, ) status_code = int(status_code) - if description and isinstance(description, six.binary_type): + if description and isinstance(description, bytes): description = description.decode("utf-8") if status_code < 400: return response - if status_code in [401]: - raise errors.AuthenticationException( - errors.ErrorCondition.UnauthorizedAccess, - description="Management authentication failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) - if status_code in [404]: - raise errors.AMQPConnectionError( - errors.ErrorCondition.NotFound, - description="Management connection failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) - raise errors.AMQPConnectionError( - errors.ErrorCondition.UnknownError, - description="Management operation failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) + raise self._amqp_transport.get_error(status_code, description) except Exception as exception: # pylint: disable=broad-except - last_exception = _handle_exception(exception, self) + last_exception = self._amqp_transport._handle_exception( # pylint: disable=protected-access + exception, self + ) self._backoff( retried_times=retried_times, last_exception=last_exception ) @@ -454,12 +454,13 @@ def _add_span_request_attributes(self, span): span.add_attribute("message_bus.destination", self._address.path) span.add_attribute("peer.address", self._address.hostname) - def _get_eventhub_properties(self): - # type:() -> Dict[str, Any] - mgmt_msg = Message(application_properties={"name": self.eventhub_name}) + def _get_eventhub_properties(self) -> Dict[str, Any]: + mgmt_msg = self._amqp_transport.build_message( + application_properties={"name": self.eventhub_name} + ) response = self._management_request(mgmt_msg, op_type=MGMT_OPERATION) output = {} - eh_info = response.value # type: Dict[bytes, Any] + eh_info: Dict[bytes, Any] = response.value if eh_info: output["eventhub_name"] = eh_info[b"name"].decode("utf-8") output["created_at"] = utc_from_timestamp( @@ -468,7 +469,7 @@ def _get_eventhub_properties(self): output["partition_ids"] = [ p.decode("utf-8") for p in eh_info[b"partition_ids"] ] - return output + return output def _get_partition_ids(self): # type:() -> List[str] @@ -476,7 +477,7 @@ def _get_partition_ids(self): def _get_partition_properties(self, partition_id): # type:(str) -> Dict[str, Any] - mgmt_msg = Message( + mgmt_msg = self._amqp_transport.build_message( application_properties={ "name": self.eventhub_name, "partition": partition_id, @@ -521,9 +522,7 @@ def _create_handler(self, auth): def _check_closed(self): if self.closed: raise ClientClosedError( - "{} has been closed. Please create a new one to handle event data.".format( - self._name - ) + f"{self._name} has been closed. Please create a new one to handle event data." ) def _open(self): @@ -534,13 +533,16 @@ def _open(self): self._handler.close() auth = self._client._create_auth() self._create_handler(auth) - self._handler.open() + conn = self._client._conn_manager.get_connection( # pylint: disable=protected-access + host=self._client._address.hostname, auth=auth + ) + self._handler.open(connection=conn) while not self._handler.client_ready(): time.sleep(0.05) self._max_message_size_on_link = ( - self._handler._link.remote_max_message_size - or constants.MAX_FRAME_SIZE_BYTES - ) # pylint: disable=protected-access + self._amqp_transport.get_remote_max_message_size(self._handler) + or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES + ) self.running = True def _close_handler(self): @@ -553,12 +555,10 @@ def _close_connection(self): self._client._conn_manager.reset_connection_if_broken() # pylint: disable=protected-access def _handle_exception(self, exception): - if not self.running and isinstance(exception, TimeoutError): - exception = errors.AuthenticationException( - errors.ErrorCondition.InternalError, - description="Authorization timeout." - ) - return _handle_exception(exception, self) + exception = self._amqp_transport.check_timeout_exception(self, exception) + return self._amqp_transport._handle_exception( # pylint: disable=protected-access + exception, self + ) def _do_retryable_operation(self, operation, timeout=None, **kwargs): # pylint:disable=protected-access @@ -576,7 +576,7 @@ def _do_retryable_operation(self, operation, timeout=None, **kwargs): return operation( timeout_time=timeout_time, last_exception=last_exception, - **kwargs + **kwargs, ) return operation() except Exception as exception: # pylint:disable=broad-except diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py index 3ce98d6efda4..447049f8f840 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_common.py @@ -2,9 +2,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations import json +import warnings +import datetime import logging import uuid from typing import ( @@ -18,10 +20,9 @@ TYPE_CHECKING, cast, ) -import six +from typing_extensions import TypedDict from ._utils import ( - set_message_partition_key, trace_message, utc_from_timestamp, transform_outbound_single_message, @@ -52,23 +53,31 @@ AmqpMessageHeader, AmqpMessageProperties, ) - -from ._pyamqp import constants, utils as pyutils -from ._pyamqp.message import BatchMessage, Message +from ._pyamqp._message_backcompat import LegacyMessage, LegacyBatchMessage +from ._pyamqp.message import Message +from ._transport._pyamqp_transport import PyamqpTransport if TYPE_CHECKING: - import datetime - -PrimitiveTypes = Optional[Union[ - int, - float, - bytes, - bool, - str, - Dict, - List, - uuid.UUID, -]] + try: + from uamqp import Message as uamqp_Message, BatchMessage + except ImportError: + uamqp_Message = None + BatchMessage = None + from ._transport._base import AmqpTransport + +MessageContent = TypedDict("MessageContent", {"content": bytes, "content_type": str}) +PrimitiveTypes = Optional[ + Union[ + int, + float, + bytes, + bool, + str, + Dict, + List, + uuid.UUID, + ] +] _LOGGER = logging.getLogger(__name__) @@ -122,69 +131,94 @@ def __init__( self._raw_amqp_message = AmqpAnnotatedMessage( # type: ignore data_body=body, annotations={}, application_properties={} ) - self.message = (self._raw_amqp_message._message) # pylint:disable=protected-access + self._uamqp_message = None + self._message = None self._raw_amqp_message.header = AmqpMessageHeader() self._raw_amqp_message.properties = AmqpMessageProperties() self.message_id = None self.content_type = None self.correlation_id = None - def __repr__(self): - # type: () -> str + def __repr__(self) -> str: # pylint: disable=bare-except try: - # TODO: below call won't work b/c pyamqp.message.message doesn't have body_type body_str = self.body_as_str() except: body_str = "" - event_repr = "body='{}'".format(body_str) + event_repr = f"body='{body_str}'" try: - event_repr += ", properties={}".format(self.properties) + event_repr += f", properties={self.properties}" except: event_repr += ", properties=" try: - event_repr += ", offset={}".format(self.offset) + event_repr += f", offset={self.offset}" except: event_repr += ", offset=" try: - event_repr += ", sequence_number={}".format(self.sequence_number) + event_repr += f", sequence_number={self.sequence_number}" except: event_repr += ", sequence_number=" try: - event_repr += ", partition_key={!r}".format(self.partition_key) + event_repr += f", partition_key={self.partition_key!r}" except: event_repr += ", partition_key=" try: - event_repr += ", enqueued_time={!r}".format(self.enqueued_time) + event_repr += f", enqueued_time={self.enqueued_time!r}" except: event_repr += ", enqueued_time=" - return "EventData({})".format(event_repr) + return f"EventData({event_repr})" - def __str__(self): - # type: () -> str + def __str__(self) -> str: try: body_str = self.body_as_str() except: # pylint: disable=bare-except body_str = "" - event_str = "{{ body: '{}'".format(body_str) + event_str = f"{{ body: '{body_str}'" try: - event_str += ", properties: {}".format(self.properties) + event_str += f", properties: {self.properties}" if self.offset: - event_str += ", offset: {}".format(self.offset) + event_str += f", offset: {self.offset}" if self.sequence_number: - event_str += ", sequence_number: {}".format(self.sequence_number) + event_str += f", sequence_number: {self.sequence_number}" if self.partition_key: - event_str += ", partition_key={!r}".format(self.partition_key) + event_str += f", partition_key={self.partition_key!r}" if self.enqueued_time: - event_str += ", enqueued_time={!r}".format(self.enqueued_time) + event_str += f", enqueued_time={self.enqueued_time!r}" except: # pylint: disable=bare-except pass event_str += " }" return event_str + def __message_content__(self) -> MessageContent: + if self.body_type != AmqpMessageBodyType.DATA: + raise TypeError("`body_type` must be `AmqpMessageBodyType.DATA`.") + content = bytearray() + for c in self.body: # type: ignore + content += c # type: ignore + content_type = cast(str, self.content_type) + return {"content": bytes(content), "content_type": content_type} + + @classmethod + def from_message_content( # pylint: disable=unused-argument + cls, content: bytes, content_type: str, **kwargs: Any + ) -> "EventData": + """ + Creates an EventData object given content type and a content value to be set as body. + + :param bytes content: The content value to be set as the body of the message. + :param str content_type: The content type to be set on the message. + :rtype: ~azure.eventhub.EventData + """ + event_data = cls(content) + event_data.content_type = content_type + return event_data + @classmethod - def _from_message(cls, message, raw_amqp_message=None): - # type: (Message, Optional[AmqpAnnotatedMessage]) -> EventData + def _from_message( + cls, + message: Union[uamqp_Message, Message], + raw_amqp_message: Optional[AmqpAnnotatedMessage] = None, + ) -> EventData: # pylint:disable=protected-access """Internal use only. @@ -195,18 +229,12 @@ def _from_message(cls, message, raw_amqp_message=None): :rtype: ~azure.eventhub.EventData """ event_data = cls(body="") - event_data.message = message # pylint: disable=protected-access + event_data._message = message event_data._raw_amqp_message = raw_amqp_message if raw_amqp_message else AmqpAnnotatedMessage(message=message) return event_data - def _encode_message(self): - # type: () -> bytes - # pylint: disable=protected-access - return self._raw_amqp_message._message.encode_message() - - def _decode_non_data_body_as_str(self, encoding="UTF-8"): - # type: (str) -> str + def _decode_non_data_body_as_str(self, encoding: str = "UTF-8") -> str: # pylint: disable=protected-access body = self.raw_amqp_message.body if self.body_type == AmqpMessageBodyType.VALUE: @@ -217,20 +245,37 @@ def _decode_non_data_body_as_str(self, encoding="UTF-8"): seq_list = [d for seq_section in body for d in seq_section] return str(decode_with_recurse(seq_list, encoding)) - def _to_outgoing_message(self): - # type: () -> EventData - self.message = (self._raw_amqp_message._to_outgoing_amqp_message()) # pylint:disable=protected-access - return self + @property + def message(self) -> LegacyMessage: + """ + DEPRECATED: Get the underlying LegacyMessage. + This is deprecated and will be removed in a later release. + :rtype: LegacyMessage + """ + warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + if not self._uamqp_message: + self._uamqp_message = LegacyMessage( + self._raw_amqp_message, + to_outgoing_amqp_message=PyamqpTransport().to_outgoing_amqp_message, + ) + return self._uamqp_message + + @message.setter + def message(self, value: "uamqp_Message") -> None: + """ + DEPRECATED: Set the underlying Message. + This is deprecated and will be removed in a later release. + """ + warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + self._uamqp_message = value @property - def raw_amqp_message(self): - # type: () -> AmqpAnnotatedMessage + def raw_amqp_message(self) -> AmqpAnnotatedMessage: """Advanced usage only. The internal AMQP message payload that is sent or received.""" return self._raw_amqp_message @property - def sequence_number(self): - # type: () -> Optional[int] + def sequence_number(self) -> Optional[int]: """The sequence number of the event. :rtype: int @@ -238,8 +283,7 @@ def sequence_number(self): return self._raw_amqp_message.annotations.get(PROP_SEQ_NUMBER, None) @property - def offset(self): - # type: () -> Optional[str] + def offset(self) -> Optional[str]: """The offset of the event. :rtype: str @@ -250,8 +294,7 @@ def offset(self): return None @property - def enqueued_time(self): - # type: () -> Optional[datetime.datetime] + def enqueued_time(self) -> Optional[datetime.datetime]: """The enqueued timestamp of the event. :rtype: datetime.datetime @@ -262,20 +305,15 @@ def enqueued_time(self): return None @property - def partition_key(self): - # type: () -> Optional[bytes] + def partition_key(self) -> Optional[bytes]: """The partition key of the event. :rtype: bytes """ - try: - return self._raw_amqp_message.annotations[PROP_PARTITION_KEY] - except KeyError: - return self._raw_amqp_message.annotations.get(PROP_PARTITION_KEY, None) + return self._raw_amqp_message.annotations.get(PROP_PARTITION_KEY, None) @property - def properties(self): - # type: () -> Dict[Union[str, bytes], Any] + def properties(self) -> Dict[Union[str, bytes], Any]: """Application-defined properties on the event. :rtype: dict @@ -283,8 +321,7 @@ def properties(self): return self._raw_amqp_message.application_properties @properties.setter - def properties(self, value): - # type: (Dict[Union[str, bytes], Any]) -> None + def properties(self, value: Dict[Union[str, bytes], Any]): """Application-defined properties on the event. :param dict value: The application properties for the EventData. @@ -293,8 +330,7 @@ def properties(self, value): self._raw_amqp_message.application_properties = properties @property - def system_properties(self): - # type: () -> Dict[bytes, Any] + def system_properties(self) -> Dict[bytes, Any]: """Metadata set by the Event Hubs Service associated with the event. An EventData could have some or all of the following meta data depending on the source @@ -332,8 +368,7 @@ def system_properties(self): return self._sys_properties @property - def body(self): - # type: () -> PrimitiveTypes + def body(self) -> PrimitiveTypes: """The body of the Message. The format may vary depending on the body type: For :class:`azure.eventhub.amqp.AmqpMessageBodyType.DATA`, the body could be bytes or Iterable[bytes]. @@ -350,16 +385,14 @@ def body(self): raise ValueError("Event content empty.") @property - def body_type(self): - # type: () -> AmqpMessageBodyType + def body_type(self) -> AmqpMessageBodyType: """The body type of the underlying AMQP message. :rtype: ~azure.eventhub.amqp.AmqpMessageBodyType """ return self._raw_amqp_message.body_type - def body_as_str(self, encoding="UTF-8"): - # type: (str) -> str + def body_as_str(self, encoding: str = "UTF-8") -> str: """The content of the event as a string, if the data is of a compatible type. :param encoding: The encoding to use for decoding event data. @@ -372,18 +405,15 @@ def body_as_str(self, encoding="UTF-8"): return self._decode_non_data_body_as_str(encoding=encoding) return "".join(b.decode(encoding) for b in cast(Iterable[bytes], data)) except TypeError: - return six.text_type(data) + return str(data) except: # pylint: disable=bare-except pass try: return cast(bytes, data).decode(encoding) except Exception as e: - raise TypeError( - "Message data is not compatible with string type: {}".format(e) - ) + raise TypeError(f"Message data is not compatible with string type: {e}") - def body_as_json(self, encoding="UTF-8"): - # type: (str) -> Dict[str, Any] + def body_as_json(self, encoding: str = "UTF-8") -> Dict[str, Any]: """The content of the event loaded as a JSON object, if the data is compatible. :param encoding: The encoding to use for decoding event data. @@ -394,11 +424,10 @@ def body_as_json(self, encoding="UTF-8"): try: return json.loads(data_str) except Exception as e: - raise TypeError("Event data is not compatible with JSON type: {}".format(e)) + raise TypeError(f"Event data is not compatible with JSON type: {e}") @property - def content_type(self): - # type: () -> Optional[str] + def content_type(self) -> Optional[str]: """The content type descriptor. Optionally describes the payload of the message, with a descriptor following the format of RFC2045, Section 5, for example "application/json". @@ -412,15 +441,13 @@ def content_type(self): return self._raw_amqp_message.properties.content_type @content_type.setter - def content_type(self, value): - # type: (str) -> None + def content_type(self, value: str) -> None: if not self._raw_amqp_message.properties: self._raw_amqp_message.properties = AmqpMessageProperties() self._raw_amqp_message.properties.content_type = value @property - def correlation_id(self): - # type: () -> Optional[str] + def correlation_id(self) -> Optional[str]: """The correlation identifier. Allows an application to specify a context for the message for the purposes of correlation, for example reflecting the MessageId of a message that is being replied to. @@ -434,15 +461,13 @@ def correlation_id(self): return self._raw_amqp_message.properties.correlation_id @correlation_id.setter - def correlation_id(self, value): - # type: (str) -> None + def correlation_id(self, value: str) -> None: if not self._raw_amqp_message.properties: self._raw_amqp_message.properties = AmqpMessageProperties() self._raw_amqp_message.properties.correlation_id = value @property - def message_id(self): - # type: () -> Optional[str] + def message_id(self) -> Optional[str]: """The id to identify the message. The message identifier is an application-defined value that uniquely identifies the message and its payload. The identifier is a free-form string and can reflect a GUID or an identifier derived from the @@ -458,7 +483,7 @@ def message_id(self): return self._raw_amqp_message.properties.message_id @message_id.setter - def message_id(self, value): + def message_id(self, value: str) -> None: if not self._raw_amqp_message.properties: self._raw_amqp_message.properties = AmqpMessageProperties() self._raw_amqp_message.properties.message_id = value @@ -489,12 +514,16 @@ class EventDataBatch(object): Event Hub decided by the service. """ - def __init__(self, max_size_in_bytes=None, partition_id=None, partition_key=None): - # type: (Optional[int], Optional[str], Optional[Union[str, bytes]]) -> None + def __init__( + self, + max_size_in_bytes: Optional[int] = None, + partition_id: Optional[str] = None, + partition_key: Optional[Union[str, bytes]] = None, + **kwargs, + ) -> None: + self._amqp_transport = kwargs.pop("amqp_transport", PyamqpTransport) - if partition_key and not isinstance( - partition_key, (six.text_type, six.binary_type) - ): + if partition_key and not isinstance(partition_key, (str, bytes)): _LOGGER.info( "WARNING: Setting partition_key of non-string value on the events to be sent is discouraged " "as the partition_key will be ignored by the Event Hub service and events will be assigned " @@ -502,33 +531,52 @@ def __init__(self, max_size_in_bytes=None, partition_id=None, partition_key=None "partition_key to only be string type, they might fail to parse the non-string value." ) - self.max_size_in_bytes = max_size_in_bytes or constants.MAX_FRAME_SIZE_BYTES - self.message = BatchMessage(data=[]) self._partition_id = partition_id self._partition_key = partition_key - self.message = set_message_partition_key(self.message, self._partition_key) - self._size = pyutils.get_message_encoded_size(self.message) + + self._message = self._amqp_transport.build_batch_message(data=[]) + self._message = self._amqp_transport.set_message_partition_key(self._message, self._partition_key) + self._size = self._amqp_transport.get_batch_message_encoded_size(self._message) + self.max_size_in_bytes = max_size_in_bytes or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES + self._count = 0 + self._internal_events: List[Union[EventData, AmqpAnnotatedMessage]] = [] + self._uamqp_message = ( + None if PyamqpTransport.TIMEOUT_FACTOR == self._amqp_transport.TIMEOUT_FACTOR else self._message + ) - def __repr__(self): - # type: () -> str - batch_repr = "max_size_in_bytes={}, partition_id={}, partition_key={!r}, event_count={}".format( - self.max_size_in_bytes, self._partition_id, self._partition_key, self._count + def __repr__(self) -> str: + batch_repr = ( + f"max_size_in_bytes={self.max_size_in_bytes}, partition_id={self._partition_id}, " + f"partition_key={self._partition_key!r}, event_count={self._count}" ) - return "EventDataBatch({})".format(batch_repr) + return f"EventDataBatch({batch_repr})" - def __len__(self): + def __len__(self) -> int: return self._count @classmethod - def _from_batch(cls, batch_data, partition_key=None): - # type: (Iterable[EventData], Optional[AnyStr]) -> EventDataBatch - outgoing_batch_data = [transform_outbound_single_message(m, EventData) for m in batch_data] - batch_data_instance = cls(partition_key=partition_key) - + def _from_batch( + cls, + batch_data: Iterable[EventData], + amqp_transport: AmqpTransport, + partition_key: Optional[AnyStr] = None, + *, + max_size_in_bytes: Optional[int] = None, + partition_id: Optional[str] = None, + ) -> EventDataBatch: + outgoing_batch_data = [ + transform_outbound_single_message(m, EventData, amqp_transport.to_outgoing_amqp_message) for m in batch_data + ] + batch_data_instance = cls( + partition_key=partition_key, + amqp_transport=amqp_transport, + max_size_in_bytes=max_size_in_bytes, + partition_id=partition_id, + ) + for event_data in outgoing_batch_data: batch_data_instance.add(event_data) - return batch_data_instance def _load_events(self, events): @@ -543,16 +591,38 @@ def _load_events(self, events): ) @property - def size_in_bytes(self): - # type: () -> int + def message(self) -> Union["BatchMessage", LegacyBatchMessage]: + """ + DEPRECATED: Get the underlying uamqp.BatchMessage or LegacyBatchMessage. + This is deprecated and will be removed in a later release. + :rtype: uamqp.BatchMessage or LegacyBatchMessage + """ + warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + if not self._uamqp_message: + message = AmqpAnnotatedMessage(message=Message(*self._message)) + self._uamqp_message = LegacyBatchMessage( + message, to_outgoing_amqp_message=PyamqpTransport().to_outgoing_amqp_message + ) + return self._uamqp_message + + @message.setter + def message(self, value: "BatchMessage") -> None: + """ + DEPRECATED: Set the underlying BatchMessage. + This is deprecated and will be removed in a later release. + """ + warnings.warn("The `message` property is deprecated and will be removed in future versions.", DeprecationWarning) + self._uamqp_message = value + + @property + def size_in_bytes(self) -> int: """The combined size of the events in the batch, in bytes. :rtype: int """ return self._size - def add(self, event_data): - # type: (Union[EventData, AmqpAnnotatedMessage]) -> None + def add(self, event_data: Union[EventData, AmqpAnnotatedMessage]) -> None: """Try to add an EventData to the batch. The total size of an added event is the sum of its body, properties, etc. @@ -565,38 +635,32 @@ def add(self, event_data): :raise: :class:`ValueError`, when exceeding the size limit. """ - outgoing_event_data = transform_outbound_single_message(event_data, EventData) + outgoing_event_data = transform_outbound_single_message( + event_data, EventData, self._amqp_transport.to_outgoing_amqp_message + ) if self._partition_key: - if ( - outgoing_event_data.partition_key - and outgoing_event_data.partition_key != self._partition_key - ): - raise ValueError( - "The partition key of event_data does not match the partition key of this batch." - ) + if outgoing_event_data.partition_key and outgoing_event_data.partition_key != self._partition_key: + raise ValueError("The partition key of event_data does not match the partition key of this batch.") if not outgoing_event_data.partition_key: - outgoing_event_data.message = set_message_partition_key( - outgoing_event_data.message, self._partition_key + outgoing_event_data._message = self._amqp_transport.set_message_partition_key( # pylint: disable=protected-access + outgoing_event_data._message, # pylint: disable=protected-access + self._partition_key, ) trace_message(outgoing_event_data) - event_data_size = pyutils.get_message_encoded_size(outgoing_event_data.message) + event_data_size = self._amqp_transport.get_message_encoded_size( + outgoing_event_data._message # pylint: disable=protected-access + ) # For a BatchMessage, if the encoded_message_size of event_data is < 256, then the overhead cost to encode that # message into the BatchMessage would be 5 bytes, if >= 256, it would be 8 bytes. size_after_add = ( - self._size - + event_data_size - + _BATCH_MESSAGE_OVERHEAD_COST[0 if (event_data_size < 256) else 1] + self._size + event_data_size + _BATCH_MESSAGE_OVERHEAD_COST[0 if (event_data_size < 256) else 1] ) if size_after_add > self.max_size_in_bytes: - raise ValueError( - "EventDataBatch has reached its size limit: {}".format( - self.max_size_in_bytes - ) - ) + raise ValueError(f"EventDataBatch has reached its size limit: {self.max_size_in_bytes}") - pyutils.add_batch(self.message, outgoing_event_data.message) + self._amqp_transport.add_batch(self, outgoing_event_data, event_data) self._size = size_after_add self._count += 1 diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py index e9aaeb17e1a9..fe87d65b60cd 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py @@ -6,10 +6,10 @@ from urllib.parse import urlparse from azure.core.pipeline.policies import RetryMode - from ._constants import TransportType, DEFAULT_AMQPS_PORT, DEFAULT_AMQP_WSS_PORT + class Configuration(object): # pylint:disable=too-many-instance-attributes def __init__(self, **kwargs): self.user_agent = kwargs.get("user_agent") # type: Optional[str] @@ -34,10 +34,14 @@ def __init__(self, **kwargs): self.connection_verify = kwargs.get("connection_verify") # type: Optional[str] self.connection_port = DEFAULT_AMQPS_PORT self.custom_endpoint_hostname = None + self.hostname = kwargs.pop("hostname") + uamqp_transport = kwargs.pop("uamqp_transport") - if self.http_proxy or self.transport_type == TransportType.AmqpOverWebsocket: + if self.http_proxy or self.transport_type.value == TransportType.AmqpOverWebsocket.value: self.transport_type = TransportType.AmqpOverWebsocket self.connection_port = DEFAULT_AMQP_WSS_PORT + if not uamqp_transport: + self.hostname += "/$servicebus/websocket" # custom end point if self.custom_endpoint_address: @@ -48,5 +52,7 @@ def __init__(self, **kwargs): endpoint = urlparse(self.custom_endpoint_address) self.transport_type = TransportType.AmqpOverWebsocket self.custom_endpoint_hostname = endpoint.hostname + if not uamqp_transport: + self.custom_endpoint_address += "/$servicebus/websocket" # in case proxy and custom endpoint are both provided, we default port to 443 if it's not provided self.connection_port = endpoint.port or DEFAULT_AMQP_WSS_PORT diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py index 66dace638f3e..a00c489568b3 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py @@ -3,15 +3,18 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from typing import TYPE_CHECKING +from __future__ import annotations +from typing import TYPE_CHECKING, Optional, Union from threading import Lock from enum import Enum -from ._pyamqp._connection import Connection, _CLOSING_STATES from ._constants import TransportType if TYPE_CHECKING: - from uamqp.authentication import JWTTokenAuth + from ._pyamqp.authentication import JWTTokenAuth + from ._pyamqp._connection import Connection + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + from uamqp import Connection as uamqp_Connection try: from typing_extensions import Protocol @@ -19,8 +22,13 @@ Protocol = object # type: ignore class ConnectionManager(Protocol): - def get_connection(self, host, auth): - # type: (str, 'JWTTokenAuth') -> Connection + def get_connection( + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, + ) -> Union[Connection, uamqp_Connection]: pass def close_connection(self): @@ -38,9 +46,10 @@ class _ConnectionMode(Enum): class _SharedConnectionManager(object): # pylint:disable=too-many-instance-attributes def __init__(self, **kwargs): self._lock = Lock() - self._conn = None # type: Connection + self._conn: Union[Connection, uamqp_Connection] = None self._container_id = kwargs.get("container_id") + self._custom_endpoint_address = kwargs.get("custom_endpoint_address") self._debug = kwargs.get("debug") self._error_policy = kwargs.get("error_policy") self._properties = kwargs.get("properties") @@ -50,16 +59,23 @@ def __init__(self, **kwargs): self._max_frame_size = kwargs.get("max_frame_size") self._channel_max = kwargs.get("channel_max") self._idle_timeout = kwargs.get("idle_timeout") - self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get( - "remote_idle_timeout_empty_frame_send_ratio" - ) - - def get_connection(self, endpoint): - # type: (str, JWTTokenAuth) -> Connection + self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get("remote_idle_timeout_empty_frame_send_ratio") + self._amqp_transport = kwargs.get("amqp_transport") + + def get_connection( + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, + ) -> Union[Connection, uamqp_Connection]: with self._lock: if self._conn is None: - self._conn = Connection( - endpoint, + self._conn = self._amqp_transport.create_connection( + host=host, + auth=auth, + endpoint=endpoint, + custom_endpoint_address=self._custom_endpoint_address, container_id=self._container_id, max_frame_size=self._max_frame_size, channel_max=self._channel_max, @@ -76,13 +92,14 @@ def close_connection(self): # type: () -> None with self._lock: if self._conn: - self._conn.close() + self._amqp_transport.close_connection(self._conn) self._conn = None def reset_connection_if_broken(self): # type: () -> None with self._lock: - if self._conn and self._conn.state in _CLOSING_STATES: + conn_state = self._amqp_transport.get_connection_state(self._conn) + if self._conn and conn_state in self._amqp_transport.CONNECTION_CLOSING_STATES: self._conn = None @@ -90,8 +107,13 @@ class _SeparateConnectionManager(object): def __init__(self, **kwargs): pass - def get_connection(self, endpoint): # pylint:disable=unused-argument, no-self-use - # type: (str) -> None + def get_connection( # pylint:disable=unused-argument, no-self-use + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuth, uamqp_JWTTokenAuth]] = None, + endpoint: Optional[str] = None, + ) -> None: return None def close_connection(self): @@ -105,7 +127,7 @@ def reset_connection_if_broken(self): def get_connection_manager(**kwargs): # type: (...) -> 'ConnectionManager' - connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) + connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore if connection_mode == _ConnectionMode.ShareConnection: return _SharedConnectionManager(**kwargs) return _SeparateConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_string_parser.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_string_parser.py index 66027b8ccbaf..b0a0743ae4f4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_string_parser.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_connection_string_parser.py @@ -66,7 +66,9 @@ def parse_connection_string(conn_str): :type conn_str: str :rtype: ~azure.eventhub.EventHubConnectionStringProperties """ - fully_qualified_namespace, policy, key, entity, signature = _parse_conn_str(conn_str, check_case=True)[:-1] + fully_qualified_namespace, policy, key, entity, signature = _parse_conn_str( + conn_str, check_case=True + )[:-1] endpoint = "sb://" + fully_qualified_namespace + "/" props = { "fully_qualified_namespace": fully_qualified_namespace, diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py index 5806480b002b..de5659411a84 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_constants.py @@ -41,8 +41,8 @@ JWT_TOKEN_SCOPE = "https://eventhubs.azure.net//.default" MGMT_OPERATION = b"com.microsoft:eventhub" MGMT_PARTITION_OPERATION = b"com.microsoft:partition" -MGMT_STATUS_CODE = b'status-code' -MGMT_STATUS_DESC = b'status-description' +MGMT_STATUS_CODE = b"status-code" +MGMT_STATUS_DESC = b"status-description" USER_AGENT_PREFIX = "azsdk-python-eventhubs" NO_RETRY_ERRORS = [ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py index 8790a22d2a69..d3e88d21e592 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py @@ -2,24 +2,13 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations import time import uuid import logging from collections import deque -from typing import TYPE_CHECKING, Callable, Dict, Optional, Any, Deque -from urllib.parse import urlparse - -from ._pyamqp import ( - ReceiveClient, - types, - utils as pyamqp_utils, - error, - constants as pyamqp_constants -) -from ._pyamqp.endpoints import Source, ApacheFilters -from ._pyamqp.message import Message +from typing import TYPE_CHECKING, Callable, Dict, Optional, Any, Deque, Union from ._common import EventData from ._client_base import ConsumerProducerMixin @@ -28,11 +17,21 @@ EPOCH_SYMBOL, TIMEOUT_SYMBOL, RECEIVER_RUNTIME_METRIC_SYMBOL, - NO_RETRY_ERRORS, - CUSTOM_CONDITION_BACKOFF, ) if TYPE_CHECKING: + try: + from uamqp import ReceiveClient as uamqp_ReceiveClient, Message as uamqp_Message + from uamqp.types import AMQPType as uamqp_AMQPType + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_ReceiveClient = None + uamqp_Message = None + uamqp_AMQPType = None + uamqp_JWTTokenAuth = None + from ._pyamqp import types + from ._pyamqp.client import ReceiveClient + from ._pyamqp.message import Message from ._pyamqp.authentication import JWTTokenAuth from ._consumer_client import EventHubConsumerClient @@ -60,7 +59,7 @@ class EventHubConsumer( :param client: The parent EventHubConsumerClient. :type client: ~azure.eventhub.EventHubConsumerClient :param source: The source EventHub from which to receive events. - :type source: ~azure.eventhub._pyamqp.endpoints.Source + :type source: ~azure.eventhub._pyamqp.endpoints.Source or ~uamqp.address.Source :keyword event_position: The position from which to start receiving. :paramtype event_position: int, str, datetime.datetime :keyword int prefetch: The number of events to prefetch from the service @@ -76,8 +75,7 @@ class EventHubConsumer( It is set to `False` by default. """ - def __init__(self, client, source, **kwargs): - # type: (EventHubConsumerClient, str, Any) -> None + def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs: Any) -> None: event_position = kwargs.get("event_position", None) prefetch = kwargs.get("prefetch", 300) owner_level = kwargs.get("owner_level", None) @@ -93,9 +91,10 @@ def __init__(self, client, source, **kwargs): self.stop = False # used by event processor self.handler_ready = False - self._on_event_received = kwargs[ + self._amqp_transport = kwargs.pop("amqp_transport") + self._on_event_received: Callable[[EventData], None] = kwargs[ "on_event_received" - ] # type: Callable[[EventData], None] + ] self._client = client self._source = source self._offset = event_position @@ -104,87 +103,61 @@ def __init__(self, client, source, **kwargs): self._owner_level = owner_level self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect - self._retry_policy = error.RetryPolicy( - retry_total=self._client._config.max_retries, # pylint:disable=protected-access - retry_backoff_factor=self._client._config.backoff_factor, # pylint:disable=protected-access - retry_backoff_max=self._client._config.backoff_max, # pylint:disable=protected-access - retry_mode=self._client._config.retry_mode, # pylint:disable=protected-access - no_retry_condition=NO_RETRY_ERRORS, - custom_condition_backoff=CUSTOM_CONDITION_BACKOFF, - ) + self._retry_policy = self._amqp_transport.create_retry_policy(self._client._config) self._reconnect_backoff = 1 - self._link_properties = {} # type: Dict[types.AMQPType, types.AMQPType] + link_properties: Union[Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes]] = {} self._error = None self._timeout = 0 - self._idle_timeout = idle_timeout if idle_timeout else None - partition = self._source.split("/")[-1] - self._partition = partition - self._name = "EHConsumer-{}-partition{}".format(uuid.uuid4(), partition) + self._idle_timeout = (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) if idle_timeout else None + self._partition = self._source.split("/")[-1] + self._name = f"EHConsumer-{uuid.uuid4()}-partition{self._partition}" if owner_level is not None: - self._link_properties[EPOCH_SYMBOL] = pyamqp_utils.amqp_long_value(int(owner_level)) + link_properties[EPOCH_SYMBOL] = int(owner_level) link_property_timeout_ms = ( self._client._config.receive_timeout or self._timeout # pylint:disable=protected-access - ) * 1000 - self._link_properties[TIMEOUT_SYMBOL] = pyamqp_utils.amqp_long_value(int(link_property_timeout_ms)) - self._handler = None # type: Optional[ReceiveClient] + ) * self._amqp_transport.TIMEOUT_FACTOR + link_properties[TIMEOUT_SYMBOL] = int(link_property_timeout_ms) + self._link_properties = self._amqp_transport.create_link_properties(link_properties) + self._handler: Optional[Union[uamqp_ReceiveClient, ReceiveClient]] = None self._track_last_enqueued_event_properties = ( track_last_enqueued_event_properties ) - self._message_buffer = deque() # type: Deque[Message] - self._last_received_event = None # type: Optional[EventData] - self._receive_start_time = None # type: Optional[float] - - def _create_handler(self, auth): - # type: (JWTTokenAuth) -> None - source = Source(address=self._source, filters={}) - if self._offset is not None: - filter_key = ApacheFilters.selector_filter - source.filters[filter_key] = ( - filter_key, - pyamqp_utils.amqp_string_value( - event_position_selector( - self._offset, - self._offset_inclusive - ) - ) - ) + self._message_buffer: Deque[uamqp_Message] = deque() + self._last_received_event: Optional[EventData] = None + self._receive_start_time: Optional[float]= None + + def _create_handler(self, auth: Union[uamqp_JWTTokenAuth, JWTTokenAuth]) -> None: + source = self._amqp_transport.create_source( + self._source, + self._offset, + event_position_selector(self._offset, self._offset_inclusive) + ) desired_capabilities = [RECEIVER_RUNTIME_METRIC_SYMBOL] if self._track_last_enqueued_event_properties else None - custom_endpoint_address = self._client._config.custom_endpoint_address # pylint:disable=protected-access - transport_type = self._client._config.transport_type # pylint:disable=protected-access - hostname = urlparse(source.address).hostname - if transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - - self._handler = ReceiveClient( - hostname, - source, + self._handler = self._amqp_transport.create_receive_client( + config=self._client._config, # pylint:disable=protected-access + source=source, auth=auth, - idle_timeout=self._idle_timeout, network_trace=self._client._config.network_tracing, # pylint:disable=protected-access - transport_type=transport_type, - http_proxy=self._client._config.http_proxy, # pylint:disable=protected-access link_credit=self._prefetch, link_properties=self._link_properties, + timeout=self._timeout, + idle_timeout=self._idle_timeout, retry_policy=self._retry_policy, + keep_alive_interval=self._keep_alive, client_name=self._name, - receive_settle_mode=pyamqp_constants.ReceiverSettleMode.First, - properties=create_properties(self._client._config.user_agent), # pylint:disable=protected-access + properties=create_properties( + self._client._config.user_agent, amqp_transport=self._amqp_transport # pylint:disable=protected-access + ), desired_capabilities=desired_capabilities, streaming_receive=True, message_received_callback=self._message_received, - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._client._config.connection_verify, ) - def _open_with_retry(self): - # type: () -> None + def _open_with_retry(self) -> None: self._do_retryable_operation(self._open, operation_need_param=False) - def _message_received(self, message): - # type: (Message) -> None + def _message_received(self, message: Union[uamqp_Message, Message]) -> None: # pylint:disable=protected-access self._message_buffer.append(message) @@ -195,10 +168,8 @@ def _next_message_in_buffer(self): self._last_received_event = event_data return event_data - def _open(self): - # type: () -> bool + def _open(self) -> bool: """Open the EventHubConsumer/EventHubProducer using the supplied connection. - """ # pylint: disable=protected-access if not self.running: @@ -206,7 +177,10 @@ def _open(self): self._handler.close() auth = self._client._create_auth() self._create_handler(auth) - self._handler.open() + conn = self._client._conn_manager.get_connection( # pylint: disable=protected-access + host=self._client._address.hostname, auth=auth + ) + self._handler.open(connection=conn) while not self._handler.client_ready(): time.sleep(0.05) self.handler_ready = True @@ -220,7 +194,9 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): self._client._config.max_retries # pylint:disable=protected-access ) self._receive_start_time = self._receive_start_time or time.time() - deadline = self._receive_start_time + (max_wait_time or 0) # max_wait_time can be None + deadline = self._receive_start_time + ( + max_wait_time or 0 + ) if len(self._message_buffer) < max_batch_size: # TODO: the retry here is a bit tricky as we are using low-level api from the amqp client. # Currently we create a new client with the latest received event's offset per retry. @@ -231,11 +207,7 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): self._handler.do_work(batch=self._prefetch) # type: ignore break except Exception as exception: # pylint: disable=broad-except - if ( - isinstance(exception, error.AMQPLinkError) - and exception.condition == error.ErrorCondition.LinkStolen # pylint: disable=no-member - ): - raise self._handle_exception(exception) + self._amqp_transport.check_link_stolen(self, exception) if not self.running: # exit by close return if self._last_received_event: @@ -249,9 +221,11 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): last_exception, ) raise last_exception - if len(self._message_buffer) >= max_batch_size \ - or (self._message_buffer and not max_wait_time) \ - or (deadline <= time.time() and max_wait_time): + if ( + len(self._message_buffer) >= max_batch_size + or (self._message_buffer and not max_wait_time) + or (deadline <= time.time() and max_wait_time) + ): if batch: events_for_callback = [] for _ in range(min(max_batch_size, len(self._message_buffer))): @@ -260,5 +234,7 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None): ) self._on_event_received(events_for_callback) else: - self._on_event_received(self._next_message_in_buffer() if self._message_buffer else None) + self._on_event_received( + self._next_message_in_buffer() if self._message_buffer else None + ) self._receive_start_time = None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py index fa00d9aa5dd5..07309007369d 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py @@ -31,7 +31,9 @@ _LOGGER = logging.getLogger(__name__) -class EventHubConsumerClient(ClientBase): # pylint: disable=client-accepts-api-version-keyword +class EventHubConsumerClient( + ClientBase +): # pylint: disable=client-accepts-api-version-keyword """The EventHubConsumerClient class defines a high level interface for receiving events from the Azure Event Hubs service. @@ -85,6 +87,14 @@ class EventHubConsumerClient(ClientBase): # pylint: disable=client-accepts-api :keyword float idle_timeout: Timeout, in seconds, after which this client will close the underlying connection if there is no further activity. By default the value is None, meaning that the client will not shutdown due to inactivity unless initiated by the service. + :keyword transport_type: The type of transport protocol that will be used for communicating with + the Event Hubs service. Default is `TransportType.Amqp` in which case port 5671 is used. + If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could + be used instead which uses port 443 for communication. + :paramtype transport_type: ~azure.eventhub.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'`. :keyword checkpoint_store: A manager that stores the partition load-balancing and checkpoint data when receiving events. The checkpoint store will be used in both cases of receiving from all partitions or a single partition. In the latter case load-balancing does not apply. @@ -116,6 +126,8 @@ class EventHubConsumerClient(ClientBase): # pylint: disable=client-accepts-api :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -136,6 +148,7 @@ def __init__( **kwargs # type: Any ): # type: (...) -> None + self._checkpoint_store = kwargs.pop("checkpoint_store", None) self._load_balancing_interval = kwargs.pop("load_balancing_interval", None) if self._load_balancing_interval is None: @@ -200,6 +213,7 @@ def _create_consumer( prefetch=prefetch, idle_timeout=self._idle_timeout, track_last_enqueued_event_properties=track_last_enqueued_event_properties, + amqp_transport=self._amqp_transport, ) return handler diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py index ebd1f90bedfb..01a8dda668e4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/_eventprocessor_mixin.py @@ -99,7 +99,9 @@ def _context(self, links=None): if span_impl_type is None: yield else: - child = span_impl_type(name="Azure.EventHubs.process", kind=SpanKind.CONSUMER, links=links) + child = span_impl_type( + name="Azure.EventHubs.process", kind=SpanKind.CONSUMER, links=links + ) self._eventhub_client._add_span_request_attributes(child) # type: ignore # pylint: disable=protected-access with child: yield diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/checkpoint_store.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/checkpoint_store.py index 5560caf72206..72f61ce1fee5 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/checkpoint_store.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/checkpoint_store.py @@ -14,7 +14,9 @@ class CheckpointStore(object): """ @abstractmethod - def list_ownership(self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs): + def list_ownership( + self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs + ): # type: (str, str, str, Any) -> Iterable[Dict[str, Any]] """Retrieves a complete ownership list from the chosen storage service. diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py index 8490f331268c..d0bb447d08da 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/event_processor.py @@ -76,7 +76,9 @@ def __init__( self._partition_close_handler = kwargs.get( "on_partition_close", None ) # type: Optional[Callable[[PartitionContext, CloseReason], None]] - checkpoint_store = kwargs.get("checkpoint_store") # type: Optional[CheckpointStore] + checkpoint_store = kwargs.get( + "checkpoint_store" + ) # type: Optional[CheckpointStore] self._checkpoint_store = checkpoint_store or InMemoryCheckpointStore() self._initial_event_position = kwargs.get( "initial_event_position", "@latest" @@ -88,9 +90,12 @@ def __init__( self._load_balancing_interval = kwargs.get( "load_balancing_interval", 10.0 ) # type: float - self._load_balancing_strategy = kwargs.get("load_balancing_strategy") or LoadBalancingStrategy.GREEDY + self._load_balancing_strategy = ( + kwargs.get("load_balancing_strategy") or LoadBalancingStrategy.GREEDY + ) self._ownership_timeout = kwargs.get( - "partition_ownership_expiration_interval", self._load_balancing_interval * 6) + "partition_ownership_expiration_interval", self._load_balancing_interval * 6 + ) self._partition_contexts = {} # type: Dict[str, PartitionContext] @@ -142,7 +147,7 @@ def _cancel_tasks_for_partitions(self, to_cancel_partitions): _LOGGER.debug( "EventProcessor %r tries to cancel partitions %r", self._id, - to_cancel_partitions + to_cancel_partitions, ) for partition_id in to_cancel_partitions: if partition_id in self._consumers: @@ -150,13 +155,15 @@ def _cancel_tasks_for_partitions(self, to_cancel_partitions): _LOGGER.info( "EventProcessor %r has cancelled partition %r", self._id, - partition_id + partition_id, ) def _initialize_partition_consumer(self, partition_id): if self._partition_initialize_handler: try: - self._partition_initialize_handler(self._partition_contexts[partition_id]) + self._partition_initialize_handler( + self._partition_contexts[partition_id] + ) except Exception as err: # pylint:disable=broad-except _LOGGER.warning( "EventProcessor instance %r of eventhub %r partition %r consumer group %r. " @@ -169,9 +176,7 @@ def _initialize_partition_consumer(self, partition_id): ) self._process_error(self._partition_contexts[partition_id], err) _LOGGER.info( - "EventProcessor %r has claimed partition %r", - self._id, - partition_id + "EventProcessor %r has claimed partition %r", self._id, partition_id ) def _create_tasks_for_claimed_ownership(self, claimed_partitions, checkpoints=None): @@ -180,13 +185,15 @@ def _create_tasks_for_claimed_ownership(self, claimed_partitions, checkpoints=No _LOGGER.debug( "EventProcessor %r tries to claim partition %r", self._id, - claimed_partitions + claimed_partitions, ) for partition_id in claimed_partitions: if partition_id not in self._consumers: if partition_id in self._partition_contexts: partition_context = self._partition_contexts[partition_id] - partition_context._last_received_event = None # pylint:disable=protected-access + partition_context._last_received_event = ( # pylint:disable=protected-access + None + ) else: partition_context = PartitionContext( self._namespace, @@ -284,7 +291,7 @@ def _load_balancing(self): self._eventhub_name, self._consumer_group, err, - load_balancing_interval + load_balancing_interval, ) self._process_error(None, err) # type: ignore @@ -308,7 +315,9 @@ def _close_consumer(self, partition_id, consumer, reason): if self._partition_close_handler: try: - self._partition_close_handler(self._partition_contexts[partition_id], reason) + self._partition_close_handler( + self._partition_contexts[partition_id], reason + ) except Exception as err: # pylint:disable=broad-except _LOGGER.warning( "EventProcessor instance %r of eventhub %r partition %r consumer group %r. " @@ -325,8 +334,7 @@ def _close_consumer(self, partition_id, consumer, reason): def _do_receive(self, partition_id, consumer): # type: (str, EventHubConsumer) -> None - """Call the consumer.receive() and handle exceptions if any after it exhausts retries. - """ + """Call the consumer.receive() and handle exceptions if any after it exhausts retries.""" try: consumer.receive(self._batch, self._max_batch_size, self._max_wait_time) except Exception as error: # pylint:disable=broad-except @@ -340,9 +348,7 @@ def _do_receive(self, partition_id, consumer): error, ) self._process_error(self._partition_contexts[partition_id], error) - self._close_consumer( - partition_id, consumer, CloseReason.OWNERSHIP_LOST - ) + self._close_consumer(partition_id, consumer, CloseReason.OWNERSHIP_LOST) def start(self): # type: () -> None diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/in_memory_checkpoint_store.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/in_memory_checkpoint_store.py index cb6a65aee2a6..dffac7041471 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/in_memory_checkpoint_store.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/in_memory_checkpoint_store.py @@ -41,7 +41,7 @@ def _lookup_trie(root, path, path_index): if path_index == len(path) - 1: return root.get(path[path_index]) if path[path_index] in root: - return _lookup_trie(root.children[path[path_index]], path, path_index+1) + return _lookup_trie(root.children[path[path_index]], path, path_index + 1) return None @@ -84,21 +84,33 @@ def set_ele(self, ele): class InMemoryCheckpointStore(CheckpointStore): - def __init__( - self - ): + def __init__(self): self._ownerships_trie = _DictTrie( "ownerships_trie", - keys_path=("fully_qualified_namespace", "eventhub_name", "consumer_group", "partition_id") + keys_path=( + "fully_qualified_namespace", + "eventhub_name", + "consumer_group", + "partition_id", + ), ) self._checkpoints_trie = _DictTrie( "checkpoints_trie", - keys_path=("fully_qualified_namespace", "eventhub_name", "consumer_group", "partition_id") + keys_path=( + "fully_qualified_namespace", + "eventhub_name", + "consumer_group", + "partition_id", + ), ) - def list_ownership(self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs): + def list_ownership( + self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs + ): # type: (str, str, str, Any) -> Iterable[Dict[str, Any]] - consumer_group_node = self._ownerships_trie.lookup((fully_qualified_namespace, eventhub_name, consumer_group)) + consumer_group_node = self._ownerships_trie.lookup( + (fully_qualified_namespace, eventhub_name, consumer_group) + ) return self._ownerships_trie.list_leaves(consumer_group_node) def claim_ownership(self, ownership_list, **kwargs): @@ -119,15 +131,15 @@ def claim_ownership(self, ownership_list, **kwargs): ownership["etag"] = str(uuid.uuid4()) ownership["last_modified_time"] = time.time() old_ownership["etag"] = ownership["etag"] - old_ownership["last_modified_time"] = ownership["last_modified_time"] + old_ownership["last_modified_time"] = ownership[ + "last_modified_time" + ] old_ownership["owner_id"] = ownership["owner_id"] result.append(old_ownership) else: ownership["etag"] = str(uuid.uuid4()) ownership["last_modified_time"] = time.time() - self._ownerships_trie.set_ele( - ownership - ) + self._ownerships_trie.set_ele(ownership) result.append(ownership) return result @@ -136,7 +148,9 @@ def update_checkpoint(self, checkpoint, **kwargs): return self._checkpoints_trie.set_ele(checkpoint) def list_checkpoints( - self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs + self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs ): - consumer_group_node = self._checkpoints_trie.lookup((fully_qualified_namespace, eventhub_name, consumer_group)) + consumer_group_node = self._checkpoints_trie.lookup( + (fully_qualified_namespace, eventhub_name, consumer_group) + ) return self._checkpoints_trie.list_leaves(consumer_group_node) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/ownership_manager.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/ownership_manager.py index 7a566ddfd1a2..9d6913d4fc53 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/ownership_manager.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_eventprocessor/ownership_manager.py @@ -9,6 +9,7 @@ from typing import List, Iterable, Optional, Dict, Any, Union, TYPE_CHECKING from .common import LoadBalancingStrategy + if TYPE_CHECKING: from .checkpoint_store import CheckpointStore from .._consumer_client import EventHubConsumerClient @@ -52,8 +53,7 @@ def __init__( def claim_ownership(self): # type: () -> List[str] - """Claims ownership for this EventProcessor - """ + """Claims ownership for this EventProcessor""" if not self.cached_parition_ids: self._retrieve_partition_ids() @@ -102,16 +102,14 @@ def release_ownership(self, partition_id): def _retrieve_partition_ids(self): # type: () -> None - """List all partition ids of the event hub that the EventProcessor is working on. - """ + """List all partition ids of the event hub that the EventProcessor is working on.""" self.cached_parition_ids = self.eventhub_client.get_partition_ids() def _balance_ownership( # pylint:disable=too-many-locals self, ownership_list, all_partition_ids ): # type: (Iterable[Dict[str, Any]], List[str]) -> List[Dict[str, Any]] - """Balances and claims ownership of partitions for this EventProcessor. - """ + """Balances and claims ownership of partitions for this EventProcessor.""" now = time.time() ownership_dict = { x["partition_id"]: x for x in ownership_list @@ -148,7 +146,7 @@ def _balance_ownership( # pylint:disable=too-many-locals expected_count_per_owner = all_partition_count // owners_count # Py2 math.ceil() returns float, a/b return int if not divisable. # Py3 math.ceil() returns int, a/b return float if not divisable. - max_count_per_owner = int(math.ceil(all_partition_count*1.0 / owners_count)) + max_count_per_owner = int(math.ceil(all_partition_count * 1.0 / owners_count)) # end of calculating expected count per owner to_claim = active_ownership_self @@ -157,9 +155,11 @@ def _balance_ownership( # pylint:disable=too-many-locals if self.load_balancing_strategy is LoadBalancingStrategy.GREEDY: # Greedily claim more partitions if there are claimable partitions to_greedy_claim_ids = random.sample( - claimable_partition_ids, k=min( - max_count_per_owner - len(active_ownership_self), len(claimable_partition_ids) - ) + claimable_partition_ids, + k=min( + max_count_per_owner - len(active_ownership_self), + len(claimable_partition_ids), + ), ) if to_greedy_claim_ids: for pid in to_greedy_claim_ids: @@ -195,13 +195,9 @@ def _balance_ownership( # pylint:disable=too-many-locals active_ownership_count_group_by_owner = Counter( dict((x, len(y)) for x, y in active_ownership_by_owner.items()) ) - most_frequent_owner_id = active_ownership_count_group_by_owner.most_common( - 1 - )[ - 0 - ][ - 0 - ] + most_frequent_owner_id = ( + active_ownership_count_group_by_owner.most_common(1)[0][0] + ) # randomly choose a partition to steal from the most_frequent_owner to_steal_partition = random.choice( active_ownership_by_owner[most_frequent_owner_id] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py index 0ebc1f62a548..990e2d46b622 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer.py @@ -2,11 +2,10 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations import uuid import logging -import time import threading from typing import ( Iterable, @@ -18,45 +17,51 @@ TYPE_CHECKING, ) # pylint: disable=unused-import -from azure.core.tracing import AbstractSpan - -from .exceptions import OperationTimeoutError from ._common import EventData, EventDataBatch from ._client_base import ConsumerProducerMixin from ._utils import ( create_properties, - set_message_partition_key, trace_message, send_context_manager, transform_outbound_single_message, ) -from ._constants import ( - TIMEOUT_SYMBOL, - NO_RETRY_ERRORS, - CUSTOM_CONDITION_BACKOFF -) -from ._pyamqp import ( - error, - utils as pyamqp_utils, - SendClient -) +from ._constants import TIMEOUT_SYMBOL +from .amqp import AmqpAnnotatedMessage _LOGGER = logging.getLogger(__name__) if TYPE_CHECKING: - from uamqp.authentication import JWTTokenAuth # pylint: disable=ungrouped-imports + from azure.core.tracing import AbstractSpan + + try: + from uamqp import SendClient as uamqp_SendClient + from uamqp.constants import MessageSendResult as uamqp_MessageSendResult + from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth + except ImportError: + uamqp_MessageSendResult = None + uamqp_SendClient = None + uamqp_JWTTokenAuth = None + from ._pyamqp.client import SendClient + from ._pyamqp.authentication import JWTTokenAuth + from ._transport._base import AmqpTransport from ._producer_client import EventHubProducerClient +_LOGGER = logging.getLogger(__name__) + -def _set_partition_key(event_datas, partition_key): - # type: (Iterable[EventData], AnyStr) -> Iterable[EventData] +def _set_partition_key( + event_datas: Iterable[EventData], + partition_key: AnyStr, + amqp_transport: AmqpTransport, +) -> Iterable[EventData]: for ed in iter(event_datas): - set_message_partition_key(ed.message, partition_key) + amqp_transport.set_message_partition_key(ed._message, partition_key) # pylint: disable=protected-access yield ed -def _set_trace_message(event_datas, parent_span=None): - # type: (Iterable[EventData], Optional[AbstractSpan]) -> Iterable[EventData] +def _set_trace_message( + event_datas: Iterable[EventData], parent_span: Optional["AbstractSpan"] = None +) -> Iterable[EventData]: for ed in iter(event_datas): trace_message(ed, parent_span) yield ed @@ -87,8 +92,11 @@ class EventHubProducer( Default value is `True`. """ - def __init__(self, client, target, **kwargs): - # type: (EventHubProducerClient, str, Any) -> None + def __init__( + self, client: "EventHubProducerClient", target: str, **kwargs: Any + ) -> None: + + self._amqp_transport = kwargs.pop("amqp_transport") partition = kwargs.get("partition", None) send_timeout = kwargs.get("send_timeout", 60) keep_alive = kwargs.get("keep_alive", None) @@ -103,115 +111,151 @@ def __init__(self, client, target, **kwargs): self._target = target self._partition = partition self._timeout = send_timeout - self._idle_timeout = idle_timeout if idle_timeout else None + self._idle_timeout = ( + (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) + if idle_timeout + else None + ) self._error = None self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect - self._retry_policy = error.RetryPolicy( - retry_total=self._client._config.max_retries, # pylint: disable=protected-access - no_retry_condition=NO_RETRY_ERRORS, - custom_condition_backoff=CUSTOM_CONDITION_BACKOFF + self._retry_policy = self._amqp_transport.create_retry_policy( + config=self._client._config ) self._reconnect_backoff = 1 - self._name = "EHProducer-{}".format(uuid.uuid4()) - self._unsent_events = [] # type: List[Any] + self._name = f"EHProducer-{uuid.uuid4()}" + self._unsent_events: List[Any] = [] if partition: self._target += "/Partitions/" + partition - self._name += "-partition{}".format(partition) - self._handler = None # type: Optional[SendClient] - self._condition = None # type: Optional[Exception] + self._name += f"-partition{partition}" + self._handler: Optional[Union[uamqp_SendClient, SendClient]] = None + self._outcome: Optional[uamqp_MessageSendResult] = None + self._condition: Optional[Exception] = None self._lock = threading.Lock() - self._link_properties = {TIMEOUT_SYMBOL: pyamqp_utils.amqp_long_value(int(self._timeout * 1000))} - - def _create_handler(self, auth): - # type: (JWTTokenAuth) -> None - transport_type = self._client._config.transport_type # pylint:disable=protected-access - custom_endpoint_address = self._client._config.custom_endpoint_address # pylint: disable=protected-access - hostname = self._client._address.hostname # pylint: disable=protected-access - if transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - self._handler = SendClient( - hostname, - self._target, + self._link_properties = self._amqp_transport.create_link_properties( + {TIMEOUT_SYMBOL: int(self._timeout * self._amqp_transport.TIMEOUT_FACTOR)} + ) + + def _create_handler( + self, auth: Union[uamqp_JWTTokenAuth, JWTTokenAuth] + ) -> None: + self._handler = self._amqp_transport.create_send_client( + config=self._client._config, # pylint:disable=protected-access + target=self._target, auth=auth, + network_trace=self._client._config.network_tracing, # pylint:disable=protected-access idle_timeout=self._idle_timeout, - network_trace=self._client._config.network_tracing, # pylint:disable=protected-access - transport_type=transport_type, - http_proxy=self._client._config.http_proxy, # pylint:disable=protected-access retry_policy=self._retry_policy, keep_alive_interval=self._keep_alive, client_name=self._name, link_properties=self._link_properties, - properties=create_properties(self._client._config.user_agent), # pylint: disable=protected-access - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._client._config.connection_verify + properties=create_properties( + self._client._config.user_agent, # pylint: disable=protected-access + amqp_transport=self._amqp_transport, + ), + msg_timeout=self._timeout * self._amqp_transport.TIMEOUT_FACTOR, ) - def _open_with_retry(self): - # type: () -> None + def _open_with_retry(self) -> None: return self._do_retryable_operation(self._open, operation_need_param=False) - def _send_event_data(self, timeout_time=None): - # type: (Optional[float]) -> None + def _on_outcome( + self, + outcome: uamqp_MessageSendResult, + condition: Optional[Exception], + ) -> None: + """ + ONLY USED FOR uamqp_transport=True. Called when the outcome is received for a delivery. + + :param outcome: The outcome of the message delivery - success or failure. + :type outcome: ~uamqp.constants.MessageSendResult + :param condition: Detail information of the outcome. + + """ + self._outcome = outcome + self._condition = condition + + def _send_event_data( + self, + timeout_time: Optional[float] = None, + last_exception: Optional[Exception] = None, + ) -> None: if self._unsent_events: - self._open() - timeout = timeout_time - time.time() if timeout_time else 0 - self._handler.send_message(self._unsent_events[0], timeout=timeout) - self._unsent_events = None + self._amqp_transport.send_messages( + self, timeout_time, last_exception, _LOGGER + ) - def _send_event_data_with_retry(self, timeout=None): - # type: (Optional[float]) -> None + def _send_event_data_with_retry(self, timeout: Optional[float] = None) -> None: return self._do_retryable_operation(self._send_event_data, timeout=timeout) - @staticmethod def _wrap_eventdata( - event_data, # type: Union[EventData, EventDataBatch, Iterable[EventData]] - span, # type: Optional[AbstractSpan] - partition_key, # type: Optional[AnyStr] - ): - # type: (...) -> Union[EventData, EventDataBatch] - if isinstance(event_data, EventData): - outgoing_event_data = transform_outbound_single_message(event_data, EventData) + self, + event_data: Union[EventData, EventDataBatch, Iterable[EventData], AmqpAnnotatedMessage], + span: Optional["AbstractSpan"], + partition_key: Optional[AnyStr], + ) -> Union[EventData, EventDataBatch]: + if isinstance(event_data, (EventData, AmqpAnnotatedMessage)): + outgoing_event_data = transform_outbound_single_message( + event_data, EventData, self._amqp_transport.to_outgoing_amqp_message + ) if partition_key: - set_message_partition_key(outgoing_event_data.message, partition_key) + self._amqp_transport.set_message_partition_key( + outgoing_event_data._message, partition_key # pylint: disable=protected-access + ) wrapper_event_data = outgoing_event_data trace_message(wrapper_event_data, span) else: if isinstance( event_data, EventDataBatch ): # The partition_key in the param will be omitted. + if not event_data: + return event_data + # If AmqpTransports are not the same, create batch with correct BatchMessage. + if self._amqp_transport.TIMEOUT_FACTOR != event_data._amqp_transport.TIMEOUT_FACTOR: # pylint: disable=protected-access + # pylint: disable=protected-access + event_data = EventDataBatch._from_batch( + event_data._internal_events, + amqp_transport=self._amqp_transport, + partition_key=event_data._partition_key, + partition_id=event_data._partition_id, + max_size_in_bytes=event_data.max_size_in_bytes + ) if ( - partition_key and partition_key != event_data._partition_key # pylint: disable=protected-access + partition_key + and partition_key + != event_data._partition_key # pylint: disable=protected-access ): raise ValueError( "The partition_key does not match the one of the EventDataBatch" ) - - for event in event_data.message.data: # pylint: disable=protected-access + for ( + event + ) in event_data._message.data: # pylint: disable=protected-access trace_message(event, span) wrapper_event_data = event_data # type:ignore else: if partition_key: - event_data = _set_partition_key(event_data, partition_key) + event_data = _set_partition_key( + event_data, partition_key, self._amqp_transport + ) event_data = _set_trace_message(event_data, span) - wrapper_event_data = EventDataBatch._from_batch(event_data, partition_key) # type: ignore # pylint: disable=protected-access + wrapper_event_data = EventDataBatch._from_batch( # type: ignore # pylint: disable=protected-access + event_data, self._amqp_transport, partition_key=partition_key + ) return wrapper_event_data def send( self, - event_data, # type: Union[EventData, EventDataBatch, Iterable[EventData]] - partition_key=None, # type: Optional[AnyStr] - timeout=None, # type: Optional[float] - ): - # type:(...) -> None + event_data: Union[EventData, EventDataBatch, Iterable[EventData], AmqpAnnotatedMessage], + partition_key: Optional[AnyStr] = None, + timeout: Optional[float] = None, + ) -> None: """ Sends an event data and blocks until acknowledgement is received or operation times out. :param event_data: The event to be sent. It can be an EventData object, or iterable of EventData objects - :type event_data: ~azure.eventhub.common.EventData, Iterator, Generator, list + :type event_data: ~azure.eventhub.common.EventData, Iterator, Generator, list or AmqpAnnotatedMessage :param partition_key: With the given partition_key, event data will land to a particular partition of the Event Hub decided by the service. partition_key could be omitted if event_data is of type ~azure.eventhub.EventDataBatch. @@ -233,22 +277,21 @@ def send( with self._lock: with send_context_manager() as child: self._check_closed() - wrapper_event_data = self._wrap_eventdata(event_data, child, partition_key) + wrapper_event_data = self._wrap_eventdata( + event_data, child, partition_key + ) + + if not wrapper_event_data: + return + + self._unsent_events = [wrapper_event_data._message] # pylint: disable=protected-access if child: self._client._add_span_request_attributes( # pylint: disable=protected-access child ) + self._send_event_data_with_retry(timeout=timeout) - try: - self._open() - self._handler.send_message(wrapper_event_data.message, timeout=timeout) - except TimeoutError as exception: - raise OperationTimeoutError(message=str(exception), details=exception) - except Exception as exception: # pylint:disable=broad-except - raise self._handle_exception(exception) - - def close(self): - # type:() -> None + def close(self) -> None: """ Close down the handler. If the handler has already closed, this will be a no op. diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py index 577462b6d0db..2f4a461f3877 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py @@ -2,31 +2,44 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from concurrent.futures import ThreadPoolExecutor import logging import threading +import time +from typing import ( + Any, + Union, + TYPE_CHECKING, + Dict, + List, + Optional, + Callable, + cast, + overload, +) +from typing_extensions import Literal -from typing import Any, Union, TYPE_CHECKING, Dict, List, Optional, cast - -from .exceptions import ConnectError, EventHubError -from .amqp import AmqpAnnotatedMessage from ._client_base import ClientBase from ._producer import EventHubProducer -from ._constants import ALL_PARTITIONS, MAX_MESSAGE_LENGTH_BYTES +from ._constants import ALL_PARTITIONS from ._common import EventDataBatch, EventData +from ._buffered_producer import BufferedProducerDispatcher +from ._utils import set_event_partition_key +from .amqp import AmqpAnnotatedMessage +from .exceptions import ConnectError, EventHubError if TYPE_CHECKING: - from azure.core.credentials import ( - TokenCredential, - AzureSasCredential, - AzureNamedKeyCredential, - ) + from ._client_base import CredentialTypes SendEventTypes = List[Union[EventData, AmqpAnnotatedMessage]] _LOGGER = logging.getLogger(__name__) -class EventHubProducerClient(ClientBase): +class EventHubProducerClient( + ClientBase +): # pylint: disable=client-accepts-api-version-keyword + # pylint: disable=too-many-instance-attributes """The EventHubProducerClient class defines a high level interface for sending events to the Azure Event Hubs service. @@ -39,6 +52,40 @@ class EventHubProducerClient(ClientBase): by the azure-identity library and objects that implement the `get_token(self, *scopes)` method. :type credential: ~azure.core.credentials.TokenCredential or ~azure.core.credentials.AzureSasCredential or ~azure.core.credentials.AzureNamedKeyCredential + :keyword bool buffered_mode: If True, the producer client will collect events in a buffer, efficiently batch, + then publish. Default is False. + :keyword buffer_concurrency: The ThreadPoolExecutor to be used for publishing events + or the number of workers for the ThreadPoolExecutor. Default is None and a ThreadPoolExecutor + with the default number of workers will be created per + https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor + :paramtype buffer_concurrency: ThreadPoolExecutor or int or None + :keyword on_success: The callback to be called once a batch has been successfully published. + The callback takes two parameters: + - `events`: The list of events that have been successfully published + - `partition_id`: The partition id that the events in the list have been published to. + The callback function should be defined like: `on_success(events, partition_id)`. + It is required when `buffered_mode` is True while optional if `buffered_mode` is False. + :paramtype on_success: Optional[Callable[[SendEventTypes, Optional[str]], None]] + :keyword on_error: The callback to be called once a batch has failed to be published. + The callback function should be defined like: `on_error(events, partition_id, error)`, where: + - `events`: The list of events that failed to be published, + - `partition_id`: The partition id that the events in the list have been tried to be published to and + - `error`: The exception related to the sending failure. + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + :paramtype on_error: Optional[Callable[[SendEventTypes, Optional[str], Exception], None]] + :keyword int max_buffer_length: Buffered mode only. + The total number of events per partition that can be buffered before a flush will be triggered. + The default value is 1500 in buffered mode. + :keyword Optional[float] max_wait_time: Buffered mode only. + The amount of time to wait for a batch to be built with events in the buffer before publishing. + The default value is 1 in buffered mode. :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`. :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service. The default value is 60 seconds. If set to 0, no timeout will be enforced from the client. @@ -74,6 +121,8 @@ class EventHubProducerClient(ClientBase): :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -86,14 +135,50 @@ class EventHubProducerClient(ClientBase): """ + @overload def __init__( self, - fully_qualified_namespace, # type: str - eventhub_name, # type: str - credential, # type: Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential] - **kwargs # type: Any - ): - # type:(...) -> None + fully_qualified_namespace: str, + eventhub_name: str, + credential: "CredentialTypes", + *, + buffered_mode: Literal[False] = False, + **kwargs: Any + ) -> None: + ... + + @overload + def __init__( + self, + fully_qualified_namespace: str, + eventhub_name: str, + credential: "CredentialTypes", + *, + buffered_mode: Literal[True], + buffer_concurrency: Optional[Union[ThreadPoolExecutor, int]] = None, + on_error: Callable[[SendEventTypes, Optional[str], Exception], None], + on_success: Callable[[SendEventTypes, Optional[str]], None], + max_buffer_length: int = 1500, + max_wait_time: float = 1, + **kwargs: Any + ) -> None: + ... + + def __init__( + self, + fully_qualified_namespace: str, + eventhub_name: str, + credential: "CredentialTypes", + *, + buffered_mode: bool = False, + on_error: Optional[ + Callable[[SendEventTypes, Optional[str], Exception], None] + ] = None, + on_success: Optional[Callable[[SendEventTypes, Optional[str]], None]] = None, + max_buffer_length: Optional[int] = None, + max_wait_time: Optional[float] = None, + **kwargs: Any + ) -> None: super(EventHubProducerClient, self).__init__( fully_qualified_namespace=fully_qualified_namespace, eventhub_name=eventhub_name, @@ -107,6 +192,41 @@ def __init__( self._max_message_size_on_link = 0 self._partition_ids = None # Optional[List[str]] self._lock = threading.Lock() + self._buffered_mode = buffered_mode + self._on_success = on_success + self._on_error = on_error + self._buffered_producer_dispatcher = None + self._max_wait_time = max_wait_time + self._max_buffer_length = max_buffer_length + self._executor = kwargs.get("buffer_concurrency") + + if self._buffered_mode: + setattr(self, "send_batch", self._buffered_send_batch) + setattr(self, "send_event", self._buffered_send_event) + if not self._on_error: + raise TypeError( + "EventHubProducerClient in buffered mode missing 1 required keyword argument: 'on_error'" + ) + if not self._on_success: + raise TypeError( + "EventHubProducerClient in buffered mode missing 1 required keyword argument: 'on_success'" + ) + if self._max_wait_time is None: + self._max_wait_time = 1 + if self._max_wait_time <= 0: + raise ValueError( + "'max_wait_time' must be a float greater than 0 in buffered mode" + ) + if self._max_buffer_length is None: + self._max_buffer_length = 1500 + if self._max_buffer_length <= 0: + raise ValueError( + "'max_buffer_length' must be an integer greater than 0 in buffered mode" + ) + if isinstance(self._executor, int) and self._executor <= 0: + raise ValueError( + "'buffer_concurrency' must be an integer greater than 0 in buffered mode" + ) def __enter__(self): return self @@ -114,6 +234,79 @@ def __enter__(self): def __exit__(self, *args): self.close() + def _buffered_send(self, events, **kwargs): + try: + self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) + except AttributeError: + self._get_partitions() + self._get_max_message_size() + self._buffered_producer_dispatcher = BufferedProducerDispatcher( + self._partition_ids, + self._on_success, + self._on_error, + self._create_producer, + self.eventhub_name, + self._max_message_size_on_link, + max_wait_time=self._max_wait_time, + max_buffer_length=self._max_buffer_length, + executor=self._executor, + amqp_transport=self._amqp_transport + ) + self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) + + def _batch_preparer(self, event_data_batch, **kwargs): + partition_id = kwargs.pop("partition_id", None) + partition_key = kwargs.pop("partition_key", None) + + if isinstance(event_data_batch, EventDataBatch): + if partition_id or partition_key: + raise TypeError( + "partition_id and partition_key should be None when sending an EventDataBatch " + "because type EventDataBatch itself may have partition_id or partition_key" + ) + to_send_batch = event_data_batch + else: + to_send_batch = self.create_batch( + partition_id=partition_id, partition_key=partition_key + ) + to_send_batch._load_events( # pylint:disable=protected-access + event_data_batch + ) + + return ( + to_send_batch, + to_send_batch._partition_id, # pylint:disable=protected-access + partition_key, + ) + + def _buffered_send_batch(self, event_data_batch, **kwargs): + batch, pid, pkey = self._batch_preparer(event_data_batch, **kwargs) + + if len(batch) == 0: + return + + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + + self._buffered_send( + event_data_batch, + partition_id=pid, + partition_key=pkey, + timeout_time=timeout_time, + ) + + def _buffered_send_event(self, event, **kwargs): + partition_key = kwargs.get("partition_key") + set_event_partition_key(event, partition_key, self._amqp_transport) + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + self._buffered_send( + event, + partition_id=kwargs.get("partition_id"), + partition_key=partition_key, + timeout_time=timeout_time, + ) + def _get_partitions(self): # type: () -> None if not self._partition_ids: @@ -131,10 +324,10 @@ def _get_max_message_size(self): EventHubProducer, self._producers[ALL_PARTITIONS] )._open_with_retry() self._max_message_size_on_link = ( - self._producers[ # type: ignore - ALL_PARTITIONS - ]._handler._link.remote_max_message_size - or MAX_MESSAGE_LENGTH_BYTES + self._amqp_transport.get_remote_max_message_size( + self._producers[ALL_PARTITIONS]._handler # type: ignore + ) + or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES ) def _start_producer(self, partition_id, send_timeout): @@ -175,17 +368,97 @@ def _create_producer(self, partition_id=None, send_timeout=None): partition=partition_id, send_timeout=send_timeout, idle_timeout=self._idle_timeout, + amqp_transport=self._amqp_transport, ) return handler @classmethod - def from_connection_string(cls, conn_str, **kwargs): - # type: (str, Any) -> EventHubProducerClient + @overload + def from_connection_string( + cls, + conn_str: str, + *, + eventhub_name: Optional[str] = None, + buffered_mode: Literal[False] = False, + **kwargs: Any + ) -> "EventHubProducerClient": + ... + + @classmethod + @overload + def from_connection_string( + cls, + conn_str: str, + *, + eventhub_name: Optional[str] = None, + buffered_mode: Literal[True], + buffer_concurrency: Optional[Union[ThreadPoolExecutor, int]] = None, + on_error: Callable[[SendEventTypes, Optional[str], Exception], None], + on_success: Callable[[SendEventTypes, Optional[str]], None], + max_buffer_length: int = 1500, + max_wait_time: float = 1, + **kwargs: Any + ) -> "EventHubProducerClient": + ... + + @classmethod + def from_connection_string( + cls, + conn_str: str, + *, + eventhub_name: Optional[str] = None, + buffered_mode: bool = False, + on_error: Optional[ + Callable[[SendEventTypes, Optional[str], Exception], None] + ] = None, + on_success: Optional[Callable[[SendEventTypes, Optional[str]], None]] = None, + max_buffer_length: Optional[int] = None, + max_wait_time: Optional[float] = None, + **kwargs: Any + ) -> "EventHubProducerClient": """Create an EventHubProducerClient from a connection string. :param str conn_str: The connection string of an Event Hub. :keyword str eventhub_name: The path of the specific Event Hub to connect the client to. + :keyword bool buffered_mode: If True, the producer client will collect events in a buffer, efficiently batch, + then publish. Default is False. + :keyword buffer_concurrency: The ThreadPoolExecutor to be used for publishing events + or the number of workers for the ThreadPoolExecutor. Default is None and a ThreadPoolExecutor + with the default number of workers will be created per + https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor + :paramtype buffer_concurrency: ThreadPoolExecutor or int or None + :keyword on_success: The callback to be called once a batch has been successfully published. + The callback takes two parameters: + - `events`: The list of events that have been successfully published + - `partition_id`: The partition id that the events in the list have been published to. + The callback function should be defined like: `on_success(events, partition_id)`. + Required when `buffered_mode` is True while optional if `buffered_mode` is False. + :paramtype on_success: Optional[Callable[[SendEventTypes, Optional[str]], None]] + :keyword on_error: The callback to be called once a batch has failed to be published. + Required when in `buffered_mode` is True while optional if `buffered_mode` is False. + The callback function should be defined like: `on_error(events, partition_id, error)`, where: + - `events`: The list of events that failed to be published, + - `partition_id`: The partition id that the events in the list have been tried to be published to and + - `error`: The exception related to the sending failure. + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + :paramtype on_error: Optional[Callable[[SendEventTypes, Optional[str], Exception], None]] + :keyword int max_buffer_length: Buffered mode only. + The total number of events per partition that can be buffered before a flush will be triggered. + The default value is 1500 in buffered mode. + :keyword Optional[float] max_wait_time: Buffered mode only. + The amount of time to wait for a batch to be built with events in the buffer before publishing. + The default value is 1 in buffered mode. :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`. + :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'`. :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service. The default value is 60 seconds. If set to 0, no timeout will be enforced from the client. :keyword str user_agent: If specified, this will be added in front of the user agent string. @@ -231,12 +504,110 @@ def from_connection_string(cls, conn_str, **kwargs): :dedent: 4 :caption: Create a new instance of the EventHubProducerClient from connection string. """ - constructor_args = cls._from_connection_string(conn_str, **kwargs) + constructor_args = cls._from_connection_string( + conn_str, + eventhub_name=eventhub_name, + buffered_mode=buffered_mode, + on_success=on_success, + on_error=on_error, + max_buffer_length=max_buffer_length, + max_wait_time=max_wait_time, + **kwargs + ) return cls(**constructor_args) + def send_event(self, event_data, **kwargs): + # type: (Union[EventData, AmqpAnnotatedMessage], Any) -> None + """ + Sends an event data. + By default, the method will block until acknowledgement is received or operation times out. + If the `EventHubProducerClient` is configured to run in buffered mode, the method will try enqueuing + the events into buffer within the given time if specified and return. + The producer will do automatic sending in the background in buffered mode. + + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + + :param event_data: The `EventData` object to be sent. + :type event_data: Union[~azure.eventhub.EventData, ~azure.eventhub.amqp.AmqpAnnotatedMessage] + :keyword float timeout: The maximum wait time to send the event data in non-buffered mode or the + maximum wait time to enqueue the event data into the buffer in buffered mode. + In non-buffered mode, the default wait time specified when the producer + was created will be used. In buffered mode, the default wait time is None. + :keyword str partition_id: The specific partition ID to send to. Default is None, in which case the service + will assign to all partitions using round-robin. + A `TypeError` will be raised if partition_id is specified and event_data_batch is an `EventDataBatch` because + `EventDataBatch` itself has partition_id. + :keyword str partition_key: With the given partition_key, event data will be sent to + a particular partition of the Event Hub decided by the service. + A `TypeError` will be raised if partition_key is specified and event_data_batch is an `EventDataBatch` because + `EventDataBatch` itself has partition_key. + If both partition_id and partition_key are provided, the partition_id will take precedence. + **WARNING: Setting partition_key of non-string value on the events to be sent is discouraged + as the partition_key will be ignored by the Event Hub service and events will be assigned + to all partitions using round-robin. Furthermore, there are SDKs for consuming events which expect + partition_key to only be string type, they might fail to parse the non-string value.** + :rtype: None + :raises: :class:`AuthenticationError` + :class:`ConnectError` + :class:`ConnectionLostError` + :class:`EventDataError` + :class:`EventDataSendError` + :class:`EventHubError` + :raises OperationTimeoutError: If the value specified by the timeout parameter elapses before the event can be + sent in non-buffered mode or the events can be enqueued into the buffered in buffered mode. + """ + input_pid = kwargs.get("partition_id") + pid = input_pid or ALL_PARTITIONS + partition_key = kwargs.get("partition_key") + send_timeout = kwargs.get("timeout") + try: + try: + cast(EventHubProducer, self._producers[pid]).send( + event_data, partition_key=partition_key, timeout=send_timeout + ) + except (KeyError, AttributeError, EventHubError): + self._start_producer(pid, send_timeout) + cast(EventHubProducer, self._producers[pid]).send( + event_data, partition_key=partition_key, timeout=send_timeout + ) + if self._on_success: + self._on_success([event_data], input_pid) + except Exception as exc: # pylint: disable=broad-except + if self._on_error: + self._on_error([event_data], input_pid, exc) + else: + raise + def send_batch(self, event_data_batch, **kwargs): # type: (Union[EventDataBatch, SendEventTypes], Any) -> None - """Sends event data and blocks until acknowledgement is received or operation times out. + # pylint: disable=protected-access + """ + Sends a batch of event data. + By default, the method will block until acknowledgement is received or operation times out. + If the `EventHubProducerClient` is configured to run in buffered mode, the method will try enqueuing + the events into buffer within the given time if specified and return. + The producer will do automatic sending in the background in buffered mode. + + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + + In buffered mode, sending a batch will remain intact and sent as a single unit. + The batch will not be rearranged. This may result in inefficiency of sending events. If you're sending a finite list of `EventData` or `AmqpAnnotatedMessage` and you know it's within the event hub frame size limit, you can send them with a `send_batch` call. Otherwise, use :meth:`create_batch` @@ -247,8 +618,10 @@ def send_batch(self, event_data_batch, **kwargs): All `EventData` or `AmqpAnnotatedMessage` in the list or `EventDataBatch` will land on the same partition. :type event_data_batch: Union[~azure.eventhub.EventDataBatch, List[Union[~azure.eventhub.EventData, ~azure.eventhub.amqp.AmqpAnnotatedMessage]] - :keyword float timeout: The maximum wait time to send the event data. - If not specified, the default wait time specified when the producer was created will be used. + :keyword float timeout: The maximum wait time to send the event data in non-buffered mode or the + maximum wait time to enqueue the event data into the buffer in buffered mode. + In non-buffered mode, the default wait time specified when the producer + was created will be used. In buffered mode, the default wait time is None. :keyword str partition_id: The specific partition ID to send to. Default is None, in which case the service will assign to all partitions using round-robin. A `TypeError` will be raised if partition_id is specified and event_data_batch is an `EventDataBatch` because @@ -271,6 +644,8 @@ def send_batch(self, event_data_batch, **kwargs): :class:`EventHubError` :class:`ValueError` :class:`TypeError` + :raises OperationTimeoutError: If the value specified by the timeout parameter elapses before the event can be + sent in non-buffered mode or the events can not be enqueued into the buffered in buffered mode. .. admonition:: Example: @@ -282,43 +657,35 @@ def send_batch(self, event_data_batch, **kwargs): :caption: Sends event data """ - partition_id = kwargs.get("partition_id") - partition_key = kwargs.get("partition_key") - - if isinstance(event_data_batch, EventDataBatch): - if partition_id or partition_key: - raise TypeError( - "partition_id and partition_key should be None when sending an EventDataBatch " - "because type EventDataBatch itself may have partition_id or partition_key" - ) - to_send_batch = event_data_batch - else: - to_send_batch = self.create_batch( - partition_id=partition_id, partition_key=partition_key - ) - to_send_batch._load_events( # pylint:disable=protected-access - event_data_batch - ) - partition_id = ( - to_send_batch._partition_id # pylint:disable=protected-access - or ALL_PARTITIONS - ) + batch, pid, pkey = self._batch_preparer(event_data_batch, **kwargs) - if len(to_send_batch) == 0: + if len(batch) == 0: return + partition_id = pid or ALL_PARTITIONS send_timeout = kwargs.pop("timeout", None) + try: - cast(EventHubProducer, self._producers[partition_id]).send( - to_send_batch, timeout=send_timeout - ) - except (KeyError, AttributeError, EventHubError) as e: - _LOGGER.debug( + try: + cast(EventHubProducer, self._producers[partition_id]).send( + batch, partition_key=pkey, timeout=send_timeout + ) + if self._on_success: + self._on_success(batch._internal_events, pid) + except (KeyError, AttributeError, EventHubError) as e: + _LOGGER.debug( "Producer for partition ID '{}' not available: {}. Rebuilding new producer.".format(partition_id, e)) - self._start_producer(partition_id, send_timeout) - cast(EventHubProducer, self._producers[partition_id]).send( - to_send_batch, timeout=send_timeout - ) + self._start_producer(partition_id, send_timeout) + cast(EventHubProducer, self._producers[partition_id]).send( + batch, partition_key=pkey, timeout=send_timeout + ) + if self._on_success: + self._on_success(batch._internal_events, pid) + except Exception as exc: # pylint: disable=broad-except + if self._on_error: + self._on_error(batch._internal_events, pid, exc) + else: + raise def create_batch(self, **kwargs): # type:(Any) -> EventDataBatch @@ -367,6 +734,7 @@ def create_batch(self, **kwargs): max_size_in_bytes=(max_size_in_bytes or self._max_message_size_on_link), partition_id=partition_id, partition_key=partition_key, + amqp_transport=self._amqp_transport ) return event_data_batch @@ -418,11 +786,32 @@ def get_partition_properties(self, partition_id): partition_id ) - def close(self): - # type: () -> None + def flush(self, **kwargs: Any) -> None: + """ + Buffered mode only. + Flush events in the buffer to be sent immediately if the client is working in buffered mode. + + :keyword Optional[float] timeout: Timeout to flush the buffered events, default is None which means no timeout. + :rtype: None + :raises EventDataSendError: If the producer fails to flush the buffer within the given timeout + in buffered mode. + """ + with self._lock: + if self._buffered_mode and self._buffered_producer_dispatcher: + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + self._buffered_producer_dispatcher.flush(timeout_time=timeout_time) + + def close(self, *, flush: bool = True, **kwargs: Any) -> None: """Close the Producer client underlying AMQP connection and links. + :keyword bool flush: Buffered mode only. If set to True, events in the buffer will be sent + immediately. Default is True. + :keyword Optional[float] timeout: Buffered mode only. Timeout to close the producer. + Default is None which means no timeout. :rtype: None + :raises EventHubError: If an error occurred when flushing the buffer if `flush` is set to True or closing the + underlying AMQP connections in buffered mode. .. admonition:: Example: @@ -434,10 +823,61 @@ def close(self): :caption: Close down the client. """ - _LOGGER.info("Closing ProducerClient") with self._lock: + if self._buffered_mode and self._buffered_producer_dispatcher: + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + self._buffered_producer_dispatcher.close( + flush=flush, timeout_time=timeout_time, raise_error=True + ) + self._buffered_producer_dispatcher = None + for pid in self._producers: if self._producers[pid]: self._producers[pid].close() # type: ignore self._producers[pid] = None super(EventHubProducerClient, self)._close() + + def get_buffered_event_count(self, partition_id: str) -> Optional[int]: + """ + The number of events that are buffered and waiting to be published for a given partition. + Returns None in non-buffered mode. **NOTE: The event buffer is processed in a background + thread, therefore the number of events in the buffer reported by this API should be + considered only an approximation and is only recommend for use in debugging. For a + partition ID that has no events buffered, 0 will be returned regardless of whether that + partition ID actually exists within the Event Hub.** + + + :param str partition_id: The target partition ID. + :rtype: int or None + """ + if not self._buffered_mode: + return None + + try: + return cast( + BufferedProducerDispatcher, self._buffered_producer_dispatcher + ).get_buffered_event_count(partition_id) + except AttributeError: + return 0 + + @property + def total_buffered_event_count(self) -> Optional[int]: + """ + The total number of events that are currently buffered and waiting to be published, + across all partitions. Returns None in non-buffered mode. **NOTE: The event buffer + is processed in a background thread, therefore the number of events in the buffer + reported by this API should be considered only an approximation and is only recommend + for use in debugging.** + + :rtype: int or None + """ + if not self._buffered_mode: + return None + + try: + return cast( + BufferedProducerDispatcher, self._buffered_producer_dispatcher + ).total_buffered_event_count + except AttributeError: + return 0 diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py index d4160e1a96da..fc9544449266 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/__init__.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- __version__ = "2.0.0a1" @@ -11,3 +11,11 @@ from ._transport import SSLTransport from .client import AMQPClient, ReceiveClient, SendClient + +__all__ = [ + "Connection", + "SSLTransport", + "AMQPClient", + "ReceiveClient", + "SendClient", +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py index 942381be00d3..7ba726e8d817 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_connection.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging @@ -24,14 +24,10 @@ HEADER_FRAME, ConnectionState, EMPTY_FRAME, - TransportType + TransportType, ) -from .error import ( - ErrorCondition, - AMQPConnectionError, - AMQPError -) +from .error import ErrorCondition, AMQPConnectionError, AMQPError _LOGGER = logging.getLogger(__name__) _CLOSING_STATES = ( @@ -39,7 +35,7 @@ ConnectionState.CLOSE_PIPE, ConnectionState.DISCARDING, ConnectionState.CLOSE_SENT, - ConnectionState.END + ConnectionState.END, ) @@ -57,7 +53,7 @@ def get_local_timeout(now, idle_timeout, last_frame_received_time): return False -class Connection(object): # pylint:disable=too-many-instance-attributes +class Connection(object): # pylint:disable=too-many-instance-attributes """An AMQP Connection. :ivar str state: The connection state. @@ -87,14 +83,14 @@ class Connection(object): # pylint:disable=too-many-instance-attributes Additionally the following keys may also be present: `'username', 'password'`. """ - def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements + def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements # type(str, Any) -> None parsed_url = urlparse(endpoint) self._hostname = parsed_url.hostname endpoint = self._hostname if parsed_url.port: self._port = parsed_url.port - elif parsed_url.scheme == 'amqps': + elif parsed_url.scheme == "amqps": self._port = SECURE_PORT else: self._port = PORT @@ -108,48 +104,41 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements custom_port = custom_parsed_url.port or WEBSOCKET_PORT custom_endpoint = "{}:{}{}".format(custom_parsed_url.hostname, custom_port, custom_parsed_url.path) - transport = kwargs.get('transport') - self._transport_type = kwargs.pop('transport_type', TransportType.Amqp) + transport = kwargs.get("transport") + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) if transport: self._transport = transport - elif 'sasl_credential' in kwargs: + elif "sasl_credential" in kwargs: sasl_transport = SASLTransport - if self._transport_type.name == 'AmqpOverWebsocket' or kwargs.get("http_proxy"): + if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get("http_proxy"): sasl_transport = SASLWithWebSocket endpoint = parsed_url.hostname + parsed_url.path self._transport = sasl_transport( - host=endpoint, - credential=kwargs['sasl_credential'], - custom_endpoint=custom_endpoint, - **kwargs + host=endpoint, credential=kwargs["sasl_credential"], custom_endpoint=custom_endpoint, **kwargs ) else: self._transport = Transport(parsed_url.netloc, transport_type=self._transport_type, **kwargs) - self._container_id = kwargs.pop('container_id', None) or str(uuid.uuid4()) # type: str - self._max_frame_size = kwargs.pop('max_frame_size', MAX_FRAME_SIZE_BYTES) # type: int + self._container_id = kwargs.pop("container_id", None) or str(uuid.uuid4()) # type: str + self._max_frame_size = kwargs.pop("max_frame_size", MAX_FRAME_SIZE_BYTES) # type: int self._remote_max_frame_size = None # type: Optional[int] - self._channel_max = kwargs.pop('channel_max', MAX_CHANNELS) # type: int - self._idle_timeout = kwargs.pop('idle_timeout', None) # type: Optional[int] - self._outgoing_locales = kwargs.pop('outgoing_locales', None) # type: Optional[List[str]] - self._incoming_locales = kwargs.pop('incoming_locales', None) # type: Optional[List[str]] + self._channel_max = kwargs.pop("channel_max", MAX_CHANNELS) # type: int + self._idle_timeout = kwargs.pop("idle_timeout", None) # type: Optional[int] + self._outgoing_locales = kwargs.pop("outgoing_locales", None) # type: Optional[List[str]] + self._incoming_locales = kwargs.pop("incoming_locales", None) # type: Optional[List[str]] self._offered_capabilities = None # type: Optional[str] - self._desired_capabilities = kwargs.pop('desired_capabilities', None) # type: Optional[str] - self._properties = kwargs.pop('properties', None) # type: Optional[Dict[str, str]] + self._desired_capabilities = kwargs.pop("desired_capabilities", None) # type: Optional[str] + self._properties = kwargs.pop("properties", None) # type: Optional[Dict[str, str]] - self._allow_pipelined_open = kwargs.pop('allow_pipelined_open', True) # type: bool + self._allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) # type: bool self._remote_idle_timeout = None # type: Optional[int] self._remote_idle_timeout_send_frame = None # type: Optional[int] - self._idle_timeout_empty_frame_send_ratio = kwargs.get('idle_timeout_empty_frame_send_ratio', 0.5) + self._idle_timeout_empty_frame_send_ratio = kwargs.get("idle_timeout_empty_frame_send_ratio", 0.5) self._last_frame_received_time = None # type: Optional[float] self._last_frame_sent_time = None # type: Optional[float] - self._idle_wait_time = kwargs.get('idle_wait_time', 0.1) # type: float - self._network_trace = kwargs.get('network_trace', False) - self._network_trace_params = { - 'connection': self._container_id, - 'session': None, - 'link': None - } + self._idle_wait_time = kwargs.get("idle_wait_time", 0.1) # type: float + self._network_trace = kwargs.get("network_trace", False) + self._network_trace_params = {"connection": self._container_id, "session": None, "link": None} self._error = None self._outgoing_endpoints = {} # type: Dict[int, Session] self._incoming_endpoints = {} # type: Dict[int, Session] @@ -200,10 +189,10 @@ def _connect(self): raise AMQPConnectionError( ErrorCondition.SocketError, description="Failed to initiate the connection due to exception: " + str(exc), - error=exc + error=exc, ) - except Exception: # pylint:disable=try-except-raise - raise + except Exception: # pylint:disable=try-except-raise + raise def _disconnect(self): # type: () -> None @@ -231,9 +220,9 @@ def _read_frame(self, wait=True, **kwargs): descriptor and field values. """ if self._can_read(): - if wait is False: # pylint:disable=no-else-return + if wait is False: return self._transport.receive_frame(**kwargs) - elif wait is True: + if wait is True: with self._transport.block(): return self._transport.receive_frame(**kwargs) else: @@ -272,9 +261,9 @@ def _send_frame(self, channel, frame, timeout=None, **kwargs): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send frame out due to exception: " + str(exc), - error=exc + error=exc, ) - except Exception: # pylint:disable=try-except-raise + except Exception: # pylint:disable=try-except-raise raise else: _LOGGER.warning("Cannot write frame in current state: %r", self.state) @@ -309,9 +298,9 @@ def _outgoing_empty(self): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send empty frame due to exception: " + str(exc), - error=exc + error=exc, ) - except Exception: # pylint:disable=try-except-raise + except Exception: # pylint:disable=try-except-raise raise def _outgoing_header(self): @@ -382,8 +371,7 @@ def _incoming_open(self, channel, frame): _LOGGER.error("OPEN frame received on a channel that is not 0.") self.close( error=AMQPError( - condition=ErrorCondition.NotAllowed, - description="OPEN frame received on a channel that is not 0." + condition=ErrorCondition.NotAllowed, description="OPEN frame received on a channel that is not 0." ) ) self._set_state(ConnectionState.END) @@ -391,16 +379,19 @@ def _incoming_open(self, channel, frame): _LOGGER.error("OPEN frame received in the OPENED state.") self.close() if frame[4]: - self._remote_idle_timeout = frame[4]/1000 # Convert to seconds + self._remote_idle_timeout = frame[4] / 1000 # Convert to seconds self._remote_idle_timeout_send_frame = self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout if frame[2] < 512: # Ensure minimum max frame size. - #Close with error - #Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] - #Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] - self.close(error=AMQPConnectionError( - condition=ErrorCondition.InvalidField, - description="connection_endpoint_frame_received::failed parsing OPEN frame")) + # Close with error + # Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] + # Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] + self.close( + error=AMQPConnectionError( + condition=ErrorCondition.InvalidField, + description="connection_endpoint_frame_received::failed parsing OPEN frame", + ) + ) _LOGGER.error("connection_endpoint_frame_received::failed parsing OPEN frame") else: self._remote_max_frame_size = frame[2] @@ -440,7 +431,7 @@ def _incoming_close(self, channel, frame): ConnectionState.HDR_EXCH, ConnectionState.OPEN_RCVD, ConnectionState.CLOSE_SENT, - ConnectionState.DISCARDING + ConnectionState.DISCARDING, ] if self.state in disconnect_states: self._disconnect() @@ -458,12 +449,8 @@ def _incoming_close(self, channel, frame): self._set_state(ConnectionState.END) if frame[0]: - self._error = AMQPConnectionError( - condition=frame[0][0], - description=frame[0][1], - info=frame[0][2] - ) - _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation + self._error = AMQPConnectionError(condition=frame[0][0], description=frame[0][1], info=frame[0][2]) + _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation def _incoming_begin(self, channel, frame): # type: (int, Tuple[Any, ...]) -> None @@ -516,7 +503,7 @@ def _incoming_end(self, channel, frame): _LOGGER.error(f"Invalid channel {channel} ") self.close(error=end_error) - def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements + def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements # type: (int, Optional[Union[bytes, Tuple[int, Tuple[Any, ...]]]]) -> bool """Process an incoming frame, either directly or by passing to the necessary Session. @@ -565,13 +552,15 @@ def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-ret if performative == 0: self._incoming_header(channel, fields) return True - if performative == 1: # pylint:disable=no-else-return + if performative == 1: # pylint:disable=no-else-return return False # TODO: incoming EMPTY else: - _LOGGER.error("Unrecognized incoming frame: {}".format(frame)) # pylint:disable=logging-format-interpolation + _LOGGER.error( + "Unrecognized incoming frame: {}".format(frame) + ) # pylint:disable=logging-format-interpolation return True except KeyError: - return True #TODO: channel error + return True # TODO: channel error def _process_outgoing_frame(self, channel, frame): # type: (int, NamedTuple) -> None @@ -589,9 +578,9 @@ def _process_outgoing_frame(self, channel, frame): # TODO: check error condition error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, - description="No frame received for the idle timeout." + description="No frame received for the idle timeout.", ), - wait=False + wait=False, ) return self._send_frame(channel, frame) @@ -659,21 +648,24 @@ def listen(self, wait=False, batch=1, **kwargs): try: if self.state not in _CLOSING_STATES: now = time.time() - if get_local_timeout(now, self._idle_timeout, self._last_frame_received_time) or self._get_remote_timeout(now): # pylint:disable=line-too-long + if get_local_timeout( + now, self._idle_timeout, self._last_frame_received_time + ) or self._get_remote_timeout( + now + ): # pylint:disable=line-too-long # TODO: check error condition self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, - description="No frame received for the idle timeout." + description="No frame received for the idle timeout.", ), - wait=False + wait=False, ) return if self.state == ConnectionState.END: # TODO: check error condition self._error = AMQPConnectionError( - condition=ErrorCondition.ConnectionCloseForced, - description="Connection was already closed." + condition=ErrorCondition.ConnectionCloseForced, description="Connection was already closed." ) return for _ in range(batch): @@ -684,9 +676,9 @@ def listen(self, wait=False, batch=1, **kwargs): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send frame out due to exception: " + str(exc), - error=exc + error=exc, ) - except Exception: # pylint:disable=try-except-raise + except Exception: # pylint:disable=try-except-raise raise def create_session(self, **kwargs): @@ -711,14 +703,15 @@ def create_session(self, **kwargs): will be logged at the logging.INFO level. Default value is that configured for the connection. """ assigned_channel = self._get_next_outgoing_channel() - kwargs['allow_pipelined_open'] = self._allow_pipelined_open - kwargs['idle_wait_time'] = self._idle_wait_time + kwargs["allow_pipelined_open"] = self._allow_pipelined_open + kwargs["idle_wait_time"] = self._idle_wait_time session = Session( self, assigned_channel, - network_trace=kwargs.pop('network_trace', self._network_trace), + network_trace=kwargs.pop("network_trace", self._network_trace), network_trace_params=dict(self._network_trace_params), - **kwargs) + **kwargs + ) self._outgoing_endpoints[assigned_channel] = session return session @@ -759,9 +752,7 @@ def close(self, error=None, wait=False): self._outgoing_close(error=error) if error: self._error = AMQPConnectionError( - condition=error.condition, - description=error.description, - info=error.info + condition=error.condition, description=error.description, info=error.info ) if self.state == ConnectionState.OPEN_PIPE: self._set_state(ConnectionState.OC_PIPE) @@ -772,7 +763,7 @@ def close(self, error=None, wait=False): else: self._set_state(ConnectionState.CLOSE_SENT) self._wait_for_response(wait, ConnectionState.END) - except Exception as exc: # pylint:disable=broad-except + except Exception as exc: # pylint:disable=broad-except # If error happened during closing, ignore the error and set state to END _LOGGER.info("An error occurred when closing the connection: %r", exc) self._set_state(ConnectionState.END) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py index 140736c790c5..8ecca3adb2cb 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_encode.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import calendar import struct @@ -13,20 +13,16 @@ import six from .types import TYPE, VALUE, AMQPTypes, FieldDefinition, ObjDefinition, ConstructorBytes -from .message import Header, Properties, Message +from .message import Message from . import performatives -from . import outcomes -from . import endpoints -from . import error - _FRAME_OFFSET = b"\x02" -_FRAME_TYPE = b'\x00' +_FRAME_TYPE = b"\x00" def _construct(byte, construct): # type: (bytes, bool) -> bytes - return byte if construct else b'' + return byte if construct else b"" def encode_null(output, *args, **kwargs): # pylint: disable=unused-argument @@ -48,7 +44,7 @@ def encode_boolean(output, value, with_constructor=True, **kwargs): # pylint: d value = bool(value) if with_constructor: output.extend(_construct(ConstructorBytes.bool, with_constructor)) - output.extend(b'\x01' if value else b'\x00') + output.extend(b"\x01" if value else b"\x00") return output.extend(ConstructorBytes.bool_true if value else ConstructorBytes.bool_false) @@ -65,7 +61,7 @@ def encode_ubyte(output, value, with_constructor=True, **kwargs): # pylint: dis value = ord(value) try: output.extend(_construct(ConstructorBytes.ubyte, with_constructor)) - output.extend(struct.pack('>B', abs(value))) + output.extend(struct.pack(">B", abs(value))) except struct.error: raise ValueError("Unsigned byte value must be 0-255") @@ -78,7 +74,7 @@ def encode_ushort(output, value, with_constructor=True, **kwargs): # pylint: di value = int(value) try: output.extend(_construct(ConstructorBytes.ushort, with_constructor)) - output.extend(struct.pack('>H', abs(value))) + output.extend(struct.pack(">H", abs(value))) except struct.error: raise ValueError("Unsigned byte value must be 0-65535") @@ -98,10 +94,10 @@ def encode_uint(output, value, with_constructor=True, use_smallest=True): try: if use_smallest and value <= 255: output.extend(_construct(ConstructorBytes.uint_small, with_constructor)) - output.extend(struct.pack('>B', abs(value))) + output.extend(struct.pack(">B", abs(value))) return output.extend(_construct(ConstructorBytes.uint_large, with_constructor)) - output.extend(struct.pack('>I', abs(value))) + output.extend(struct.pack(">I", abs(value))) except struct.error: raise ValueError("Value supplied for unsigned int invalid: {}".format(value)) @@ -124,10 +120,10 @@ def encode_ulong(output, value, with_constructor=True, use_smallest=True): try: if use_smallest and value <= 255: output.extend(_construct(ConstructorBytes.ulong_small, with_constructor)) - output.extend(struct.pack('>B', abs(value))) + output.extend(struct.pack(">B", abs(value))) return output.extend(_construct(ConstructorBytes.ulong_large, with_constructor)) - output.extend(struct.pack('>Q', abs(value))) + output.extend(struct.pack(">Q", abs(value))) except struct.error: raise ValueError("Value supplied for unsigned long invalid: {}".format(value)) @@ -140,7 +136,7 @@ def encode_byte(output, value, with_constructor=True, **kwargs): # pylint: disa value = int(value) try: output.extend(_construct(ConstructorBytes.byte, with_constructor)) - output.extend(struct.pack('>b', value)) + output.extend(struct.pack(">b", value)) except struct.error: raise ValueError("Byte value must be -128-127") @@ -153,7 +149,7 @@ def encode_short(output, value, with_constructor=True, **kwargs): # pylint: dis value = int(value) try: output.extend(_construct(ConstructorBytes.short, with_constructor)) - output.extend(struct.pack('>h', value)) + output.extend(struct.pack(">h", value)) except struct.error: raise ValueError("Short value must be -32768-32767") @@ -168,10 +164,10 @@ def encode_int(output, value, with_constructor=True, use_smallest=True): try: if use_smallest and (-128 <= value <= 127): output.extend(_construct(ConstructorBytes.int_small, with_constructor)) - output.extend(struct.pack('>b', value)) + output.extend(struct.pack(">b", value)) return output.extend(_construct(ConstructorBytes.int_large, with_constructor)) - output.extend(struct.pack('>i', value)) + output.extend(struct.pack(">i", value)) except struct.error: raise ValueError("Value supplied for int invalid: {}".format(value)) @@ -183,7 +179,7 @@ def encode_long(output, value, with_constructor=True, use_smallest=True): """ if isinstance(value, datetime): - value = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond/1000) + value = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond / 1000) try: value = long(value) except NameError: @@ -191,13 +187,14 @@ def encode_long(output, value, with_constructor=True, use_smallest=True): try: if use_smallest and (-128 <= value <= 127): output.extend(_construct(ConstructorBytes.long_small, with_constructor)) - output.extend(struct.pack('>b', value)) + output.extend(struct.pack(">b", value)) return output.extend(_construct(ConstructorBytes.long_large, with_constructor)) - output.extend(struct.pack('>q', value)) + output.extend(struct.pack(">q", value)) except struct.error: raise ValueError("Value supplied for long invalid: {}".format(value)) + def encode_float(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument # type: (bytearray, float, bool, Any) -> None """ @@ -205,7 +202,7 @@ def encode_float(output, value, with_constructor=True, **kwargs): # pylint: dis """ value = float(value) output.extend(_construct(ConstructorBytes.float, with_constructor)) - output.extend(struct.pack('>f', value)) + output.extend(struct.pack(">f", value)) def encode_double(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument @@ -215,7 +212,7 @@ def encode_double(output, value, with_constructor=True, **kwargs): # pylint: di """ value = float(value) output.extend(_construct(ConstructorBytes.double, with_constructor)) - output.extend(struct.pack('>d', value)) + output.extend(struct.pack(">d", value)) def encode_timestamp(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument @@ -225,10 +222,10 @@ def encode_timestamp(output, value, with_constructor=True, **kwargs): # pylint: label="64-bit two's-complement integer representing milliseconds since the unix epoch"/> """ if isinstance(value, datetime): - value = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond/1000) + value = (calendar.timegm(value.utctimetuple()) * 1000) + (value.microsecond / 1000) value = int(value) output.extend(_construct(ConstructorBytes.timestamp, with_constructor)) - output.extend(struct.pack('>q', value)) + output.extend(struct.pack(">q", value)) def encode_uuid(output, value, with_constructor=True, **kwargs): # pylint: disable=unused-argument @@ -257,12 +254,12 @@ def encode_binary(output, value, with_constructor=True, use_smallest=True): length = len(value) if use_smallest and length <= 255: output.extend(_construct(ConstructorBytes.binary_small, with_constructor)) - output.extend(struct.pack('>B', length)) + output.extend(struct.pack(">B", length)) output.extend(value) return try: output.extend(_construct(ConstructorBytes.binary_large, with_constructor)) - output.extend(struct.pack('>L', length)) + output.extend(struct.pack(">L", length)) output.extend(value) except struct.error: raise ValueError("Binary data to long to encode") @@ -277,16 +274,16 @@ def encode_string(output, value, with_constructor=True, use_smallest=True): label="up to 2^32 - 1 octets worth of UTF-8 Unicode (with no byte order mark)"/> """ if isinstance(value, six.text_type): - value = value.encode('utf-8') + value = value.encode("utf-8") length = len(value) if use_smallest and length <= 255: output.extend(_construct(ConstructorBytes.string_small, with_constructor)) - output.extend(struct.pack('>B', length)) + output.extend(struct.pack(">B", length)) output.extend(value) return try: output.extend(_construct(ConstructorBytes.string_large, with_constructor)) - output.extend(struct.pack('>L', length)) + output.extend(struct.pack(">L", length)) output.extend(value) except struct.error: raise ValueError("String value too long to encode.") @@ -301,16 +298,16 @@ def encode_symbol(output, value, with_constructor=True, use_smallest=True): label="up to 2^32 - 1 seven bit ASCII characters representing a symbolic value"/> """ if isinstance(value, six.text_type): - value = value.encode('utf-8') + value = value.encode("utf-8") length = len(value) if use_smallest and length <= 255: output.extend(_construct(ConstructorBytes.symbol_small, with_constructor)) - output.extend(struct.pack('>B', length)) + output.extend(struct.pack(">B", length)) output.extend(value) return try: output.extend(_construct(ConstructorBytes.symbol_large, with_constructor)) - output.extend(struct.pack('>L', length)) + output.extend(struct.pack(">L", length)) output.extend(value) except struct.error: raise ValueError("Symbol value too long to encode.") @@ -337,18 +334,17 @@ def encode_list(output, value, with_constructor=True, use_smallest=True): encoded_size += len(encoded_values) if use_smallest and count <= 255 and encoded_size < 255: output.extend(_construct(ConstructorBytes.list_small, with_constructor)) - output.extend(struct.pack('>B', encoded_size + 1)) - output.extend(struct.pack('>B', count)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) else: try: output.extend(_construct(ConstructorBytes.list_large, with_constructor)) - output.extend(struct.pack('>L', encoded_size + 4)) - output.extend(struct.pack('>L', count)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) except struct.error: raise ValueError("List is too large or too long to be encoded.") output.extend(encoded_values) - def encode_map(output, value, with_constructor=True, use_smallest=True): # type: (bytearray, Union[Dict[Any, Any], Iterable[Tuple[Any, Any]]], bool, bool) -> None """ @@ -370,13 +366,13 @@ def encode_map(output, value, with_constructor=True, use_smallest=True): encoded_size = len(encoded_values) if use_smallest and count <= 255 and encoded_size < 255: output.extend(_construct(ConstructorBytes.map_small, with_constructor)) - output.extend(struct.pack('>B', encoded_size + 1)) - output.extend(struct.pack('>B', count)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) else: try: output.extend(_construct(ConstructorBytes.map_large, with_constructor)) - output.extend(struct.pack('>L', encoded_size + 4)) - output.extend(struct.pack('>L', count)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) except struct.error: raise ValueError("Map is too large or too long to be encoded.") output.extend(encoded_values) @@ -386,11 +382,11 @@ def encode_map(output, value, with_constructor=True, use_smallest=True): def _check_element_type(item, element_type): if not element_type: try: - return item['TYPE'] + return item["TYPE"] except (KeyError, TypeError): return type(item) try: - if item['TYPE'] != element_type: + if item["TYPE"] != element_type: raise TypeError("All elements in an array must be the same type.") except (KeyError, TypeError): if not isinstance(item, element_type): @@ -421,13 +417,13 @@ def encode_array(output, value, with_constructor=True, use_smallest=True): encoded_size += len(encoded_values) if use_smallest and count <= 255 and encoded_size < 255: output.extend(_construct(ConstructorBytes.array_small, with_constructor)) - output.extend(struct.pack('>B', encoded_size + 1)) - output.extend(struct.pack('>B', count)) + output.extend(struct.pack(">B", encoded_size + 1)) + output.extend(struct.pack(">B", count)) else: try: output.extend(_construct(ConstructorBytes.array_large, with_constructor)) - output.extend(struct.pack('>L', encoded_size + 4)) - output.extend(struct.pack('>L', count)) + output.extend(struct.pack(">L", encoded_size + 4)) + output.extend(struct.pack(">L", count)) except struct.error: raise ValueError("Array is too large or too long to be encoded.") output.extend(encoded_values) @@ -452,10 +448,10 @@ def encode_fields(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE:[]} + fields = {TYPE: AMQPTypes.map, VALUE: []} for key, data in value.items(): if isinstance(key, six.text_type): - key = key.encode('utf-8') + key = key.encode("utf-8") fields[VALUE].append(({TYPE: AMQPTypes.symbol, VALUE: key}, data)) return fields @@ -473,7 +469,7 @@ def encode_annotations(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE:[]} + fields = {TYPE: AMQPTypes.map, VALUE: []} for key, data in value.items(): if isinstance(key, int): field_key = {TYPE: AMQPTypes.ulong, VALUE: key} @@ -500,7 +496,7 @@ def encode_application_properties(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE:[]} + fields = {TYPE: AMQPTypes.map, VALUE: []} for key, data in value.items(): fields[VALUE].append(({TYPE: AMQPTypes.string, VALUE: key}, data)) return fields @@ -516,11 +512,11 @@ def encode_message_id(value): """ if isinstance(value, int): return {TYPE: AMQPTypes.ulong, VALUE: value} - elif isinstance(value, uuid.UUID): + if isinstance(value, uuid.UUID): return {TYPE: AMQPTypes.uuid, VALUE: value} - elif isinstance(value, six.binary_type): + if isinstance(value, six.binary_type): return {TYPE: AMQPTypes.binary, VALUE: value} - elif isinstance(value, six.text_type): + if isinstance(value, six.text_type): return {TYPE: AMQPTypes.string, VALUE: value} raise TypeError("Unsupported Message ID type.") @@ -530,16 +526,16 @@ def encode_node_properties(value): """Properties of a node. - + A symbol-keyed map containing properties of a node used when requesting creation or reporting the creation of a dynamic node. The following common properties are defined:: - + - `lifetime-policy`: The lifetime of a dynamically generated node. Definitionally, the lifetime will never be less than the lifetime of the link which caused its creation, however it is possible to extend the lifetime of dynamically created node using a lifetime policy. The value of this entry MUST be of a type which provides the lifetime-policy archetype. The following standard lifetime-policies are defined below: delete-on-close, delete-on-no-links, delete-on-no-messages or delete-on-no-links-or-messages. - + - `supported-dist-modes`: The distribution modes that the node supports. The value of this entry MUST be one or more symbols which are valid distribution-modes. That is, the value MUST be of the same type as would be valid in a field defined with the following attributes: @@ -548,7 +544,7 @@ def encode_node_properties(value): if not value: return {TYPE: AMQPTypes.null, VALUE: None} # TODO - fields = {TYPE: AMQPTypes.map, VALUE:[]} + fields = {TYPE: AMQPTypes.map, VALUE: []} # fields[{TYPE: AMQPTypes.symbol, VALUE: b'lifetime-policy'}] = { # TYPE: AMQPTypes.described, # VALUE: ( @@ -577,21 +573,18 @@ def encode_filter_set(value): """ if not value: return {TYPE: AMQPTypes.null, VALUE: None} - fields = {TYPE: AMQPTypes.map, VALUE:[]} + fields = {TYPE: AMQPTypes.map, VALUE: []} for name, data in value.items(): if data is None: described_filter = {TYPE: AMQPTypes.null, VALUE: None} else: if isinstance(name, six.text_type): - name = name.encode('utf-8') + name = name.encode("utf-8") try: descriptor, filter_value = data described_filter = { TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.symbol, VALUE: descriptor}, - filter_value - ) + VALUE: ({TYPE: AMQPTypes.symbol, VALUE: descriptor}, filter_value), } except ValueError: described_filter = data @@ -678,7 +671,7 @@ def describe_performative(performative): # type: (Performative) -> Tuple(bytes, bytes) body = [] for index, value in enumerate(performative): - field = performative._definition[index] + field = performative._definition[index] # pylint: disable=protected-access if value is None: body.append({TYPE: AMQPTypes.null, VALUE: None}) elif field is None: @@ -698,10 +691,7 @@ def describe_performative(performative): return { TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: performative._code}, - {TYPE: AMQPTypes.list, VALUE: body} - ) + VALUE: ({TYPE: AMQPTypes.ulong, VALUE: performative._code}, {TYPE: AMQPTypes.list, VALUE: body}), # pylint: disable=protected-access } @@ -724,13 +714,16 @@ def encode_payload(output, payload): encode_value(output, describe_performative(payload[0])) if payload[2]: # message annotations - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000072}, - encode_annotations(payload[2]), - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000072}, + encode_annotations(payload[2]), + ), + }, + ) if payload[3]: # properties # TODO: Header and Properties encoding can be optimized to @@ -739,51 +732,54 @@ def encode_payload(output, payload): encode_value(output, describe_performative(payload[3])) if payload[4]: # application properties - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000074}, - encode_application_properties(payload[4]) - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000074}, encode_application_properties(payload[4])), + }, + ) if payload[5]: # data for item_value in payload[5]: - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000075}, - {TYPE: AMQPTypes.binary, VALUE: item_value} - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000075}, {TYPE: AMQPTypes.binary, VALUE: item_value}), + }, + ) if payload[6]: # sequence for item_value in payload[6]: - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000076}, - {TYPE: None, VALUE: item_value} - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000076}, {TYPE: None, VALUE: item_value}), + }, + ) if payload[7]: # value - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000077}, - {TYPE: None, VALUE: payload[7]} - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ({TYPE: AMQPTypes.ulong, VALUE: 0x00000077}, {TYPE: None, VALUE: payload[7]}), + }, + ) if payload[8]: # footer - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000078}, - encode_annotations(payload[8]), - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000078}, + encode_annotations(payload[8]), + ), + }, + ) # TODO: # currently the delivery annotations must be finally encoded instead of being encoded at the 2nd position @@ -791,13 +787,16 @@ def encode_payload(output, payload): # -- received message doesn't have it populated # check with service team? if payload[1]: # delivery annotations - encode_value(output, { - TYPE: AMQPTypes.described, - VALUE: ( - {TYPE: AMQPTypes.ulong, VALUE: 0x00000071}, - encode_annotations(payload[1]), - ) - }) + encode_value( + output, + { + TYPE: AMQPTypes.described, + VALUE: ( + {TYPE: AMQPTypes.ulong, VALUE: 0x00000071}, + encode_annotations(payload[1]), + ), + }, + ) return output @@ -807,7 +806,7 @@ def encode_frame(frame, frame_type=_FRAME_TYPE): # TODO: allow passing type specific bytes manually, e.g. Empty Frame needs padding if frame is None: size = 8 - header = size.to_bytes(4, 'big') + _FRAME_OFFSET + frame_type + header = size.to_bytes(4, "big") + _FRAME_OFFSET + frame_type return header, None frame_description = describe_performative(frame) @@ -817,5 +816,5 @@ def encode_frame(frame, frame_type=_FRAME_TYPE): frame_data += frame.payload size = len(frame_data) + 8 - header = size.to_bytes(4, 'big') + _FRAME_OFFSET + frame_type + header = size.to_bytes(4, "big") + _FRAME_OFFSET + frame_type return header, frame_data diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py index 8aefedd9123a..20f3c595d455 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_message_backcompat.py @@ -5,13 +5,13 @@ # -------------------------------------------------------------------------- # pylint: disable=too-many-lines +from typing import Callable from enum import Enum from ._encode import encode_payload from .utils import get_message_encoded_size from .error import AMQPError -from .message import Message, Header, Properties, BatchMessage -#from uamqp import constants, errors +from .message import Header, Properties class MessageState(Enum): @@ -49,16 +49,17 @@ def __init__(self, message, **kwargs): self.delivery_no = kwargs.get('delivery_no') self.delivery_tag = kwargs.get('delivery_tag') or None self.on_send_complete = None - self.properties = LegacyMessageProperties(self._message.properties) - self.application_properties = self._message.application_properties - self.annotations = self._message.annotations - self.header = LegacyMessageHeader(self._message.header) + self.properties = LegacyMessageProperties(self._message.properties) if self._message.properties else None + self.application_properties = self._message.application_properties if any(self._message.application_properties) else None + self.annotations = self._message.annotations if any(self._message.annotations) else None + self.header = LegacyMessageHeader(self._message.header) if self._message.header else None self.footer = self._message.footer self.delivery_annotations = self._message.delivery_annotations if self._settler: self.state = MessageState.ReceivedUnsettled elif self.delivery_no: self.state = MessageState.ReceivedSettled + self._to_outgoing_amqp_message: Callable = kwargs.get('to_outgoing_amqp_message') def __str__(self): return str(self._message) @@ -77,11 +78,11 @@ def settled(self): return True def get_message_encoded_size(self): - return get_message_encoded_size(self._message._to_outgoing_amqp_message()) + return get_message_encoded_size(self._to_outgoing_amqp_message(self._message)) def encode_message(self): output = bytearray() - encode_payload(output, self._message._to_outgoing_amqp_message()) + encode_payload(output, self._to_outgoing_amqp_message(self._message)) return bytes(output) def get_data(self): @@ -97,7 +98,7 @@ def gather(self): return [self] def get_message(self): - return self._message._to_outgoing_amqp_message() + return self._to_outgoing_amqp_message(self._message) def accept(self): if self._can_settle_message(): @@ -211,7 +212,7 @@ def get_properties_obj(self): class LegacyMessageHeader(object): def __init__(self, header): - self.delivery_count = header.delivery_count # or 0 + self.delivery_count = header.delivery_count or 0 self.time_to_live = header.time_to_live self.first_acquirer = header.first_acquirer self.durable = header.durable diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py index a775fe4ca335..d65a5cfdd9b3 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py @@ -1,4 +1,4 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # This is a fork of the transport.py which was originally written by Barry Pederson and # maintained by the Celery project: https://github.com/celery/py-amqp. # @@ -30,7 +30,7 @@ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF # THE POSSIBILITY OF SUCH DAMAGE. -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- from __future__ import absolute_import, unicode_literals @@ -48,7 +48,7 @@ import certifi -from ._platform import KNOWN_TCP_OPTS, SOL_TCP, pack, unpack +from ._platform import KNOWN_TCP_OPTS, SOL_TCP from ._encode import encode_frame from ._decode import decode_frame, decode_empty_frame from .constants import TLS_HEADER_FRAME, WEBSOCKET_PORT, TransportType, AMQP_WS_SUBPROTOCOL @@ -57,7 +57,7 @@ try: import fcntl except ImportError: # pragma: no cover - fcntl = None # noqa + fcntl = None # noqa try: from os import set_cloexec # Python 3.4? except ImportError: # pragma: no cover @@ -70,7 +70,7 @@ def set_cloexec(fd, cloexec): # noqa FD_CLOEXEC = fcntl.FD_CLOEXEC except AttributeError: raise NotImplementedError( - 'close-on-exec flag not supported on this platform', + "close-on-exec flag not supported on this platform", ) flags = fcntl.fcntl(fd, fcntl.F_GETFD) if cloexec: @@ -79,25 +79,28 @@ def set_cloexec(fd, cloexec): # noqa flags &= ~FD_CLOEXEC return fcntl.fcntl(fd, fcntl.F_SETFD, flags) + _LOGGER = logging.getLogger(__name__) _UNAVAIL = {errno.EAGAIN, errno.EINTR, errno.ENOENT, errno.EWOULDBLOCK} AMQP_PORT = 5672 AMQPS_PORT = 5671 -AMQP_FRAME = memoryview(b'AMQP') +AMQP_FRAME = memoryview(b"AMQP") EMPTY_BUFFER = bytes() SIGNED_INT_MAX = 0x7FFFFFFF TIMEOUT_INTERVAL = 1 +WS_TIMEOUT_INTERVAL = 1 +READ_TIMEOUT_INTERVAL = 0.2 # Match things like: [fe80::1]:5432, from RFC 2732 -IPV6_LITERAL = re.compile(r'\[([\.0-9a-f:]+)\](?::(\d+))?') +IPV6_LITERAL = re.compile(r"\[([\.0-9a-f:]+)\](?::(\d+))?") DEFAULT_SOCKET_SETTINGS = { - 'TCP_NODELAY': 1, - 'TCP_USER_TIMEOUT': 1000, - 'TCP_KEEPIDLE': 60, - 'TCP_KEEPINTVL': 10, - 'TCP_KEEPCNT': 9, + "TCP_NODELAY": 1, + "TCP_USER_TIMEOUT": 1000, + "TCP_KEEPIDLE": 60, + "TCP_KEEPINTVL": 10, + "TCP_KEEPCNT": 9, } @@ -128,8 +131,8 @@ def to_host_port(host, port=AMQP_PORT): if m.group(2): port = int(m.group(2)) else: - if ':' in host: - host, port = host.rsplit(':', 1) + if ":" in host: + host, port = host.rsplit(":", 1) port = int(port) return host, port @@ -141,17 +144,25 @@ class UnexpectedFrame(Exception): class _AbstractTransport(object): """Common superclass for TCP and SSL transports.""" - def __init__(self, host, port=AMQP_PORT, connect_timeout=None, - read_timeout=None, write_timeout=None, - socket_settings=None, raise_on_initial_eintr=True, **kwargs): + def __init__( + self, + host, + port=AMQP_PORT, + connect_timeout=None, + read_timeout=None, + write_timeout=None, + socket_settings=None, + raise_on_initial_eintr=True, + **kwargs + ): self.connected = False self.sock = None self.raise_on_initial_eintr = raise_on_initial_eintr self._read_buffer = BytesIO() self.host, self.port = to_host_port(host, port) - + self.connect_timeout = connect_timeout or TIMEOUT_INTERVAL - self.read_timeout = read_timeout + self.read_timeout = read_timeout or READ_TIMEOUT_INTERVAL self.write_timeout = write_timeout self.socket_settings = socket_settings self.socket_lock = Lock() @@ -163,7 +174,9 @@ def connect(self): return self._connect(self.host, self.port, self.connect_timeout) self._init_socket( - self.socket_settings, self.read_timeout, self.write_timeout, + self.socket_settings, + self.read_timeout, + self.write_timeout, ) # we've sent the banner; signal connect # EINTR, EAGAIN, EWOULDBLOCK would signal that the banner @@ -188,10 +201,10 @@ def block_with_timeout(self, timeout): try: yield self.sock except SSLError as exc: - if 'timed out' in str(exc): + if "timed out" in str(exc): # http://bugs.python.org/issue10272 raise socket.timeout() - elif 'The operation did not complete' in str(exc): + if "The operation did not complete" in str(exc): # Non-blocking SSL sockets can throw SSLError raise socket.timeout() raise @@ -213,10 +226,10 @@ def block(self): try: yield self.sock except SSLError as exc: - if 'timed out' in str(exc): + if "timed out" in str(exc): # http://bugs.python.org/issue10272 raise socket.timeout() - elif 'The operation did not complete' in str(exc): + if "The operation did not complete" in str(exc): # Non-blocking SSL sockets can throw SSLError raise socket.timeout() raise @@ -238,10 +251,10 @@ def non_blocking(self): try: yield self.sock except SSLError as exc: - if 'timed out' in str(exc): + if "timed out" in str(exc): # http://bugs.python.org/issue10272 raise socket.timeout() - elif 'The operation did not complete' in str(exc): + if "The operation did not complete" in str(exc): # Non-blocking SSL sockets can throw SSLError raise socket.timeout() raise @@ -270,8 +283,7 @@ def _connect(self, host, port, timeout): for n, family in enumerate(addr_types): # first, resolve the address for a single address family try: - entries = socket.getaddrinfo( - host, port, family, socket.SOCK_STREAM, SOL_TCP) + entries = socket.getaddrinfo(host, port, family, socket.SOCK_STREAM, SOL_TCP) entries_num = len(entries) except socket.gaierror: # we may have depleted all our options @@ -279,10 +291,7 @@ def _connect(self, host, port, timeout): # if getaddrinfo succeeded before for another address # family, reraise the previous socket.error since it's more # relevant to users - raise (e - if e is not None - else socket.error( - "failed to resolve broker hostname")) + raise e if e is not None else socket.error("failed to resolve broker hostname") continue # pragma: no cover # now that we have address(es) for the hostname, connect to broker @@ -326,14 +335,14 @@ def _init_socket(self, socket_settings, read_timeout, write_timeout): self._setup_transport() # TODO: a greater timeout value is needed in long distance communication # we should either figure out a reasonable value error/dynamically adjust the timeout - # 1 second is enough for perf analysis - self.sock.settimeout(1) # set socket back to non-blocking mode + # 0.2 second is enough for perf analysis + self.sock.settimeout(read_timeout) # set socket back to non-blocking mode def _get_tcp_socket_defaults(self, sock): tcp_opts = {} for opt in KNOWN_TCP_OPTS: enum = None - if opt == 'TCP_USER_TIMEOUT': + if opt == "TCP_USER_TIMEOUT": try: from socket import TCP_USER_TIMEOUT as enum except ImportError: @@ -346,8 +355,7 @@ def _get_tcp_socket_defaults(self, sock): if opt in DEFAULT_SOCKET_SETTINGS: tcp_opts[enum] = DEFAULT_SOCKET_SETTINGS[opt] elif hasattr(socket, opt): - tcp_opts[enum] = sock.getsockopt( - SOL_TCP, getattr(socket, opt)) + tcp_opts[enum] = sock.getsockopt(SOL_TCP, getattr(socket, opt)) return tcp_opts def _set_socket_options(self, socket_settings): @@ -359,19 +367,17 @@ def _set_socket_options(self, socket_settings): def _read(self, n, initial=False): """Read exactly n bytes from the peer.""" - raise NotImplementedError('Must be overriden in subclass') + raise NotImplementedError("Must be overriden in subclass") def _setup_transport(self): """Do any additional initialization of the class.""" - pass def _shutdown_transport(self): """Do any preliminary work in shutting down the connection.""" - pass def _write(self, s): """Completely write a string to the peer.""" - raise NotImplementedError('Must be overriden in subclass') + raise NotImplementedError("Must be overriden in subclass") def close(self): if self.sock is not None: @@ -389,18 +395,18 @@ def close(self): self.sock = None self.connected = False - def read(self, verify_frame_type=0, **kwargs): + def read(self, verify_frame_type=0, **kwargs): read = self._read read_frame_buffer = BytesIO() try: frame_header = memoryview(bytearray(8)) read_frame_buffer.write(read(8, buffer=frame_header, initial=True)) - channel = struct.unpack('>H', frame_header[6:])[0] + channel = struct.unpack(">H", frame_header[6:])[0] size = frame_header[0:4] if size == AMQP_FRAME: # Empty frame or AMQP header negotiation TODO return frame_header, channel, None - size = struct.unpack('>I', size)[0] + size = struct.unpack(">I", size)[0] offset = frame_header[4] frame_type = frame_header[5] @@ -421,7 +427,7 @@ def read(self, verify_frame_type=0, **kwargs): except (OSError, IOError, SSLError, socket.error) as exc: # Don't disconnect for ssl read time outs # http://bugs.python.org/issue10272 - if isinstance(exc, SSLError) and 'timed out' in str(exc): + if isinstance(exc, SSLError) and "timed out" in str(exc): raise socket.timeout() if get_errno(exc) not in _UNAVAIL: self.connected = False @@ -455,7 +461,7 @@ def send_frame(self, channel, frame, **kwargs): if performative is None: data = header else: - encoded_channel = struct.pack('>H', channel) + encoded_channel = struct.pack(">H", channel) data = header + encoded_channel + performative self.write(data) @@ -469,17 +475,12 @@ class SSLTransport(_AbstractTransport): def __init__(self, host, port=AMQPS_PORT, connect_timeout=None, ssl=None, **kwargs): self.sslopts = ssl if isinstance(ssl, dict) else {} self._read_buffer = BytesIO() - super(SSLTransport, self).__init__( - host, - port=port, - connect_timeout=connect_timeout, - **kwargs - ) + super(SSLTransport, self).__init__(host, port=port, connect_timeout=connect_timeout, **kwargs) def _setup_transport(self): """Wrap the socket in an SSL object.""" self.sock = self._wrap_socket(self.sock, **self.sslopts) - a = self.sock.do_handshake() + self.sock.do_handshake() self._quick_recv = self.sock.recv def _wrap_socket(self, sock, context=None, **sslopts): @@ -494,11 +495,20 @@ def _wrap_context(self, sock, sslopts, check_hostname=None, **ctx_options): ctx.check_hostname = check_hostname return ctx.wrap_socket(sock, **sslopts) - def _wrap_socket_sni(self, sock, keyfile=None, certfile=None, - server_side=False, cert_reqs=ssl.CERT_REQUIRED, - ca_certs=None, do_handshake_on_connect=False, - suppress_ragged_eofs=True, server_hostname=None, - ciphers=None, ssl_version=None): + def _wrap_socket_sni( + self, + sock, + keyfile=None, + certfile=None, + server_side=False, + cert_reqs=ssl.CERT_REQUIRED, + ca_certs=None, + do_handshake_on_connect=False, + suppress_ragged_eofs=True, + server_hostname=None, + ciphers=None, + ssl_version=None, + ): """Socket wrap with SNI headers. Default `ssl.wrap_socket` method augmented with support for @@ -507,27 +517,25 @@ def _wrap_socket_sni(self, sock, keyfile=None, certfile=None, # Setup the right SSL version; default to optimal versions across # ssl implementations if ssl_version is None: - ssl_version = ssl.PROTOCOL_TLS + ssl_version = ssl.PROTOCOL_TLS opts = { - 'sock': sock, - 'keyfile': keyfile, - 'certfile': certfile, - 'server_side': server_side, - 'cert_reqs': cert_reqs, - 'ca_certs': ca_certs, - 'do_handshake_on_connect': do_handshake_on_connect, - 'suppress_ragged_eofs': suppress_ragged_eofs, - 'ciphers': ciphers, + "sock": sock, + "keyfile": keyfile, + "certfile": certfile, + "server_side": server_side, + "cert_reqs": cert_reqs, + "ca_certs": ca_certs, + "do_handshake_on_connect": do_handshake_on_connect, + "suppress_ragged_eofs": suppress_ragged_eofs, + "ciphers": ciphers, #'ssl_version': ssl_version } sock = ssl.wrap_socket(**opts) # Set SNI headers if supported - if (server_hostname is not None) and ( - hasattr(ssl, 'HAS_SNI') and ssl.HAS_SNI) and ( - hasattr(ssl, 'SSLContext')): - context = ssl.SSLContext(opts['ssl_version']) + if (server_hostname is not None) and (hasattr(ssl, "HAS_SNI") and ssl.HAS_SNI) and (hasattr(ssl, "SSLContext")): + context = ssl.SSLContext(opts["ssl_version"]) context.verify_mode = cert_reqs if cert_reqs != ssl.CERT_NONE: context.check_hostname = True @@ -544,8 +552,7 @@ def _shutdown_transport(self): except OSError: pass - def _read(self, toread, initial=False, buffer=None, - _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)): + def _read(self, toread, initial=False, buffer=None, _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)): # According to SSL_read(3), it can at most return 16kb of data. # Thus, we use an internal read buffer like TCPTransport._read # to get the exact number of bytes wanted. @@ -561,7 +568,7 @@ def _read(self, toread, initial=False, buffer=None, except socket.error as exc: # ssl.sock.read may cause a SSLerror without errno # http://bugs.python.org/issue10272 - if isinstance(exc, SSLError) and 'timed out' in str(exc): + if isinstance(exc, SSLError) and "timed out" in str(exc): raise socket.timeout() # ssl.sock.read may cause ENOENT if the # operation couldn't be performed (Issue celery#1414). @@ -571,7 +578,7 @@ def _read(self, toread, initial=False, buffer=None, continue raise if not nbytes: - raise IOError('Server unexpectedly closed connection') + raise IOError("Server unexpectedly closed connection") length += nbytes toread -= nbytes @@ -593,7 +600,7 @@ def _write(self, s): # None. n = 0 if not n: - raise IOError('Socket closed') + raise IOError("Socket closed") s = s[n:] def negotiate(self): @@ -601,8 +608,11 @@ def negotiate(self): self.write(TLS_HEADER_FRAME) channel, returned_header = self.receive_frame(verify_frame_type=None) if returned_header[1] == TLS_HEADER_FRAME: - raise ValueError("Mismatching TLS header protocol. Excpected: {}, received: {}".format( - TLS_HEADER_FRAME, returned_header[1])) + raise ValueError( + "Mismatching TLS header protocol. Excpected: {}, received: {}".format( + TLS_HEADER_FRAME, returned_header[1] + ) + ) class TCPTransport(_AbstractTransport): @@ -630,7 +640,7 @@ def _read(self, n, initial=False, _errnos=(errno.EAGAIN, errno.EINTR)): continue raise if not s: - raise IOError('Server unexpectedly closed connection') + raise IOError("Server unexpectedly closed connection") rbuf += s except: # noqa self._read_buffer = rbuf @@ -639,6 +649,7 @@ def _read(self, n, initial=False, _errnos=(errno.EAGAIN, errno.EINTR)): result, self._read_buffer = rbuf[:n], rbuf[n:] return result + def Transport(host, transport_type, connect_timeout=None, ssl=False, **kwargs): """Create transport. @@ -651,29 +662,29 @@ def Transport(host, transport_type, connect_timeout=None, ssl=False, **kwargs): transport = SSLTransport if ssl else TCPTransport return transport(host, connect_timeout=connect_timeout, ssl=ssl, **kwargs) + class WebSocketTransport(_AbstractTransport): def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs): self.sslopts = ssl if isinstance(ssl, dict) else {} - self._connect_timeout = connect_timeout or TIMEOUT_INTERVAL + self._connect_timeout = connect_timeout or WS_TIMEOUT_INTERVAL self._host = host self._custom_endpoint = kwargs.get("custom_endpoint") - super().__init__( - host, port, connect_timeout, **kwargs - ) + super().__init__(host, port, connect_timeout, **kwargs) self.ws = None - self._http_proxy = kwargs.get('http_proxy', None) + self._http_proxy = kwargs.get("http_proxy", None) def connect(self): http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None if self._http_proxy: - http_proxy_host = self._http_proxy['proxy_hostname'] - http_proxy_port = self._http_proxy['proxy_port'] - username = self._http_proxy.get('username', None) - password = self._http_proxy.get('password', None) + http_proxy_host = self._http_proxy["proxy_hostname"] + http_proxy_port = self._http_proxy["proxy_port"] + username = self._http_proxy.get("username", None) + password = self._http_proxy.get("password", None) if username or password: http_proxy_auth = (username, password) try: from websocket import create_connection + self.ws = create_connection( url="wss://{}".format(self._custom_endpoint or self._host), subprotocols=[AMQP_WS_SUBPROTOCOL], @@ -682,9 +693,8 @@ def connect(self): sslopt=self.sslopts, http_proxy_host=http_proxy_host, http_proxy_port=http_proxy_port, - http_proxy_auth=http_proxy_auth + http_proxy_auth=http_proxy_auth, ) - except ImportError: raise ValueError("Please install websocket-client library to use websocket transport.") @@ -702,10 +712,10 @@ def _read(self, n, initial=False, buffer=None, **kwargs): # pylint: disable=unu data = self.ws.recv() if len(data) <= n: - view[length: length + len(data)] = data + view[length : length + len(data)] = data n -= len(data) else: - view[length: length + n] = data[0:n] + view[length : length + n] = data[0:n] self._read_buffer = BytesIO(data[n:]) n = 0 return view diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py index c513f35b9e32..bcf047fdb428 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/__init__.py @@ -1,11 +1,12 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- from ._connection_async import Connection, ConnectionState -from ._link_async import Link, LinkDeliverySettleReason, LinkState +from ._link_async import Link, LinkState +from ..constants import LinkDeliverySettleReason from ._receiver_async import ReceiverLink from ._sasl_async import SASLPlainCredential, SASLTransport from ._sender_async import SenderLink @@ -13,3 +14,22 @@ from ._transport_async import AsyncTransport from ._client_async import AMQPClientAsync, ReceiveClientAsync, SendClientAsync from ._authentication_async import SASTokenAuthAsync + +__all__ = [ + "Connection", + "ConnectionState", + "Link", + "LinkDeliverySettleReason", + "LinkState", + "ReceiverLink", + "SASLPlainCredential", + "SASLTransport", + "SenderLink", + "Session", + "SessionState", + "AsyncTransport", + "AMQPClientAsync", + "ReceiveClientAsync", + "SendClientAsync", + "SASTokenAuthAsync", +] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py index 938fbe0a8ee3..6348008cb38f 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_authentication_async.py @@ -12,12 +12,6 @@ ) from ..constants import AUTH_DEFAULT_EXPIRATION_SECONDS -try: - from urlparse import urlparse - from urllib import quote_plus # type: ignore -except ImportError: - from urllib.parse import urlparse, quote_plus - async def _generate_sas_token_async(auth_uri, sas_name, sas_key, expiry_in=AUTH_DEFAULT_EXPIRATION_SECONDS): return _generate_sas_access_token(auth_uri, sas_name, sas_key, expiry_in=expiry_in) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py index 3164780cc474..ed4742fff3a4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_cbs_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- import logging from datetime import datetime @@ -11,12 +11,7 @@ from ..utils import utc_now, utc_from_timestamp from ._management_link_async import ManagementLink from ..message import Message, Properties -from ..error import ( - AuthenticationException, - ErrorCondition, - TokenAuthFailure, - TokenExpired -) +from ..error import AuthenticationException, ErrorCondition, TokenAuthFailure, TokenExpired from ..constants import ( CbsState, CbsAuthState, @@ -27,34 +22,26 @@ CBS_OPERATION, ManagementExecuteOperationResult, ManagementOpenResult, - DEFAULT_AUTH_TIMEOUT -) -from ..cbs import ( - check_put_timeout_status, - check_expiration_and_refresh_status + DEFAULT_AUTH_TIMEOUT, ) +from ..cbs import check_put_timeout_status, check_expiration_and_refresh_status _LOGGER = logging.getLogger(__name__) -class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes - def __init__( - self, - session, - auth, - **kwargs - ): +class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes + def __init__(self, session, auth, **kwargs): self._session = session self._connection = self._session._connection self._mgmt_link = self._session.create_request_response_link_pair( - endpoint='$cbs', + endpoint="$cbs", on_amqp_management_open_complete=self._on_amqp_management_open_complete, on_amqp_management_error=self._on_amqp_management_error, - status_code_field=b'status-code', - status_description_field=b'status-description' + status_code_field=b"status-code", + status_description_field=b"status-description", ) # type: ManagementLink - #if not auth.get_token or not asyncio.iscoroutinefunction(auth.get_token): + # if not auth.get_token or not asyncio.iscoroutinefunction(auth.get_token): # raise ValueError("get_token must be a coroutine object.") self._auth = auth @@ -80,15 +67,15 @@ async def _put_token(self, token, token_type, audience, expires_on=None): CBS_NAME: audience, CBS_OPERATION: CBS_PUT_TOKEN, CBS_TYPE: token_type, - CBS_EXPIRATION: expires_on - } + CBS_EXPIRATION: expires_on, + }, ) await self._mgmt_link.execute_operation( message, self._on_execute_operation_complete, timeout=self._auth_timeout, operation=CBS_PUT_TOKEN, - type=token_type + type=token_type, ) self._mgmt_link.next_message_id += 1 @@ -99,12 +86,15 @@ async def _on_amqp_management_open_complete(self, management_open_result): self.state = CbsState.ERROR _LOGGER.info( "Unexpected AMQP management open complete in OPEN, CBS error occurred on connection %r.", - self._connection._container_id # pylint:disable=protected-access + self._connection._container_id, # pylint:disable=protected-access ) elif self.state == CbsState.OPENING: self.state = CbsState.OPEN if management_open_result == ManagementOpenResult.OK else CbsState.CLOSED - _LOGGER.info("CBS for connection %r completed opening with status: %r", - self._connection._container_id, management_open_result) # pylint:disable=protected-access + _LOGGER.info( + "CBS for connection %r completed opening with status: %r", + self._connection._container_id, # pylint: disable=protected-access + management_open_result, + ) # pylint:disable=protected-access async def _on_amqp_management_error(self): if self.state == CbsState.CLOSED: @@ -112,22 +102,26 @@ async def _on_amqp_management_error(self): elif self.state == CbsState.OPENING: self.state = CbsState.ERROR await self._mgmt_link.close() - _LOGGER.info("CBS for connection %r failed to open with status: %r", - self._connection._container_id, ManagementOpenResult.ERROR) # pylint:disable=protected-access + _LOGGER.info( + "CBS for connection %r failed to open with status: %r", + self._connection._container_id, + ManagementOpenResult.ERROR, + ) # pylint:disable=protected-access elif self.state == CbsState.OPEN: self.state = CbsState.ERROR - _LOGGER.info("CBS error occurred on connection %r.", self._connection._container_id) # pylint:disable=protected-access + _LOGGER.info( + "CBS error occurred on connection %r.", self._connection._container_id + ) # pylint:disable=protected-access async def _on_execute_operation_complete( - self, + self, execute_operation_result, status_code, status_description, message, error_condition=None + ): # TODO: message and error_condition never used + _LOGGER.info( + "CBS Put token result (%r), status code: %s, status_description: %s.", execute_operation_result, status_code, status_description, - message, - error_condition=None - ): # TODO: message and error_condition never used - _LOGGER.info("CBS Put token result (%r), status code: %s, status_description: %s.", - execute_operation_result, status_code, status_description) + ) self._token_status_code = status_code self._token_status_description = status_description @@ -143,15 +137,17 @@ async def _on_execute_operation_complete( async def _update_status(self): if self.auth_state == CbsAuthState.OK or self.auth_state == CbsAuthState.REFRESH_REQUIRED: - _LOGGER.debug('update_status In refresh required or OK.') - is_expired, is_refresh_required = check_expiration_and_refresh_status(self._expires_on, self._refresh_window) # pylint:disable=line-too-long - _LOGGER.debug('is expired == %r, is refresh required == %r', is_expired, is_refresh_required) + _LOGGER.debug("update_status In refresh required or OK.") + is_expired, is_refresh_required = check_expiration_and_refresh_status( + self._expires_on, self._refresh_window + ) # pylint:disable=line-too-long + _LOGGER.debug("is expired == %r, is refresh required == %r", is_expired, is_refresh_required) if is_expired: self.auth_state = CbsAuthState.EXPIRED elif is_refresh_required: self.auth_state = CbsAuthState.REFRESH_REQUIRED elif self.auth_state == CbsAuthState.IN_PROGRESS: - _LOGGER.debug('In update status, in progress. token put time: %r', self._token_put_time) + _LOGGER.debug("In update status, in progress. token put time: %r", self._token_put_time) put_timeout = check_put_timeout_status(self._auth_timeout, self._token_put_time) if put_timeout: self.auth_state = CbsAuthState.TIMEOUT @@ -166,7 +162,7 @@ async def _cbs_link_ready(self): # Think how upper layer handle this exception + condition code raise AuthenticationException( condition=ErrorCondition.ClientError, - description="CBS authentication link is in broken status, please recreate the cbs link." + description="CBS authentication link is in broken status, please recreate the cbs link.", ) async def open(self): @@ -190,39 +186,38 @@ async def update_token(self): except AttributeError: self._token = access_token.token self._token_put_time = int(utc_now().timestamp()) - await self._put_token(self._token, self._auth.token_type, self._auth.audience, utc_from_timestamp(self._expires_on)) + await self._put_token( + self._token, self._auth.token_type, self._auth.audience, utc_from_timestamp(self._expires_on) + ) async def handle_token(self): - if not (await self._cbs_link_ready()): + if not await self._cbs_link_ready(): return False await self._update_status() if self.auth_state == CbsAuthState.IDLE: await self.update_token() return False - elif self.auth_state == CbsAuthState.IN_PROGRESS: + if self.auth_state == CbsAuthState.IN_PROGRESS: return False - elif self.auth_state == CbsAuthState.OK: + if self.auth_state == CbsAuthState.OK: return True - elif self.auth_state == CbsAuthState.REFRESH_REQUIRED: - _LOGGER.info("Token on connection %r will expire soon - attempting to refresh.", - self._connection._container_id) # pylint:disable=protected-access + if self.auth_state == CbsAuthState.REFRESH_REQUIRED: + _LOGGER.info( + "Token on connection %r will expire soon - attempting to refresh.", self._connection._container_id + ) # pylint:disable=protected-access await self.update_token() return False - elif self.auth_state == CbsAuthState.FAILURE: + if self.auth_state == CbsAuthState.FAILURE: raise AuthenticationException( - condition=ErrorCondition.InternalError, - description="Failed to open CBS authentication link." + condition=ErrorCondition.InternalError, description="Failed to open CBS authentication link." ) - elif self.auth_state == CbsAuthState.ERROR: + if self.auth_state == CbsAuthState.ERROR: raise TokenAuthFailure( self._token_status_code, self._token_status_description, - encoding=self._encoding # TODO: drop off all the encodings + encoding=self._encoding, # TODO: drop off all the encodings ) - elif self.auth_state == CbsAuthState.TIMEOUT: + if self.auth_state == CbsAuthState.TIMEOUT: raise TimeoutError("Authentication attempt timed-out.") - elif self.auth_state == CbsAuthState.EXPIRED: - raise TokenExpired( - condition=ErrorCondition.InternalError, - description="CBS Authentication Expired." - ) + if self.auth_state == CbsAuthState.EXPIRED: + raise TokenExpired(condition=ErrorCondition.InternalError, description="CBS Authentication Expired.") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py index a68ae7d79b74..143555c41714 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_client_async.py @@ -5,31 +5,23 @@ #-------------------------------------------------------------------------- # TODO: Check types of kwargs (issue exists for this) import asyncio -import collections.abc import logging -from typing import Any, Dict, Optional, Tuple, Union, overload -from typing_extensions import Literal -import uuid import time import queue -import certifi from functools import partial +from typing import Any, Dict, Optional, Tuple, Union, overload +from typing_extensions import Literal +import certifi from ..outcomes import Accepted, Modified, Received, Rejected, Released from ._connection_async import Connection from ._management_operation_async import ManagementOperation -from ._receiver_async import ReceiverLink -from ._sender_async import SenderLink -from ._session_async import Session from ._cbs_async import CBSAuthenticator from ..client import AMQPClient as AMQPClientSync from ..client import ReceiveClient as ReceiveClientSync from ..client import SendClient as SendClientSync from ..message import _MessageDelivery -from ..endpoints import Source, Target from ..constants import ( - SenderSettleMode, - ReceiverSettleMode, MessageDeliveryState, SEND_DISPOSITION_ACCEPT, SEND_DISPOSITION_REJECT, @@ -39,7 +31,6 @@ ) from ..error import ( AMQPError, - ErrorResponse, ErrorCondition, AMQPException, MessageException @@ -164,7 +155,7 @@ async def _client_run_async(self, **kwargs): await self._connection.listen(wait=self._socket_timeout) async def _close_link_async(self, **kwargs): - if self._link and not self._link._is_closed: + if self._link and not self._link._is_closed: # pylint: disable=protected-access await self._link.detach(close=True) self._link = None @@ -273,7 +264,7 @@ async def auth_complete_async(self): :rtype: bool """ - if self._cbs_authenticator and not (await self._cbs_authenticator.handle_token()): + if self._cbs_authenticator and not await self._cbs_authenticator.handle_token(): await self._connection.listen(wait=self._socket_timeout) return False return True diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py index 2c3072143591..5a03b1b38475 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging @@ -26,19 +26,15 @@ HEADER_FRAME, ConnectionState, EMPTY_FRAME, - TransportType + TransportType, ) -from ..error import ( - ErrorCondition, - AMQPConnectionError, - AMQPError -) +from ..error import ErrorCondition, AMQPConnectionError, AMQPError _LOGGER = logging.getLogger(__name__) -class Connection(object): # pylint:disable=too-many-instance-attributes +class Connection(object): # pylint:disable=too-many-instance-attributes """An AMQP Connection. :ivar str state: The connection state. @@ -68,14 +64,14 @@ class Connection(object): # pylint:disable=too-many-instance-attributes Additionally the following keys may also be present: `'username', 'password'`. """ - def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements + def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements # type(str, Any) -> None parsed_url = urlparse(endpoint) self._hostname = parsed_url.hostname endpoint = self._hostname if parsed_url.port: self._port = parsed_url.port - elif parsed_url.scheme == 'amqps': + elif parsed_url.scheme == "amqps": self._port = SECURE_PORT else: self._port = PORT @@ -89,48 +85,41 @@ def __init__(self, endpoint, **kwargs): # pylint:disable=too-many-statements custom_port = custom_parsed_url.port or WEBSOCKET_PORT custom_endpoint = "{}:{}{}".format(custom_parsed_url.hostname, custom_port, custom_parsed_url.path) - transport = kwargs.get('transport') - self._transport_type = kwargs.pop('transport_type', TransportType.Amqp) + transport = kwargs.get("transport") + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) if transport: self._transport = transport - elif 'sasl_credential' in kwargs: + elif "sasl_credential" in kwargs: sasl_transport = SASLTransport - if self._transport_type.name == 'AmqpOverWebsocket' or kwargs.get("http_proxy"): + if self._transport_type.name == "AmqpOverWebsocket" or kwargs.get("http_proxy"): sasl_transport = SASLWithWebSocket endpoint = parsed_url.hostname + parsed_url.path self._transport = sasl_transport( - host=endpoint, - credential=kwargs['sasl_credential'], - custom_endpoint=custom_endpoint, - **kwargs + host=endpoint, credential=kwargs["sasl_credential"], custom_endpoint=custom_endpoint, **kwargs ) else: - self._transport = AsyncTransport(parsed_url.netloc, transport_type=self._transport_type, **kwargs) + self._transport = AsyncTransport(parsed_url.netloc, transport_type=self._transport_type, **kwargs) - self._container_id = kwargs.pop('container_id', None) or str(uuid.uuid4()) # type: str - self._max_frame_size = kwargs.pop('max_frame_size', MAX_FRAME_SIZE_BYTES) # type: int + self._container_id = kwargs.pop("container_id", None) or str(uuid.uuid4()) # type: str + self._max_frame_size = kwargs.pop("max_frame_size", MAX_FRAME_SIZE_BYTES) # type: int self._remote_max_frame_size = None # type: Optional[int] - self._channel_max = kwargs.pop('channel_max', MAX_CHANNELS) # type: int - self._idle_timeout = kwargs.pop('idle_timeout', None) # type: Optional[int] - self._outgoing_locales = kwargs.pop('outgoing_locales', None) # type: Optional[List[str]] - self._incoming_locales = kwargs.pop('incoming_locales', None) # type: Optional[List[str]] + self._channel_max = kwargs.pop("channel_max", MAX_CHANNELS) # type: int + self._idle_timeout = kwargs.pop("idle_timeout", None) # type: Optional[int] + self._outgoing_locales = kwargs.pop("outgoing_locales", None) # type: Optional[List[str]] + self._incoming_locales = kwargs.pop("incoming_locales", None) # type: Optional[List[str]] self._offered_capabilities = None # type: Optional[str] - self._desired_capabilities = kwargs.pop('desired_capabilities', None) # type: Optional[str] - self._properties = kwargs.pop('properties', None) # type: Optional[Dict[str, str]] + self._desired_capabilities = kwargs.pop("desired_capabilities", None) # type: Optional[str] + self._properties = kwargs.pop("properties", None) # type: Optional[Dict[str, str]] - self._allow_pipelined_open = kwargs.pop('allow_pipelined_open', True) # type: bool + self._allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) # type: bool self._remote_idle_timeout = None # type: Optional[int] self._remote_idle_timeout_send_frame = None # type: Optional[int] - self._idle_timeout_empty_frame_send_ratio = kwargs.get('idle_timeout_empty_frame_send_ratio', 0.5) + self._idle_timeout_empty_frame_send_ratio = kwargs.get("idle_timeout_empty_frame_send_ratio", 0.5) self._last_frame_received_time = None # type: Optional[float] self._last_frame_sent_time = None # type: Optional[float] - self._idle_wait_time = kwargs.get('idle_wait_time', 0.1) # type: float - self._network_trace = kwargs.get('network_trace', False) - self._network_trace_params = { - 'connection': self._container_id, - 'session': None, - 'link': None - } + self._idle_wait_time = kwargs.get("idle_wait_time", 0.1) # type: float + self._network_trace = kwargs.get("network_trace", False) + self._network_trace_params = {"connection": self._container_id, "session": None, "link": None} self._error = None self._outgoing_endpoints = {} # type: Dict[int, Session] self._incoming_endpoints = {} # type: Dict[int, Session] @@ -181,7 +170,7 @@ async def _connect(self): raise AMQPConnectionError( ErrorCondition.SocketError, description="Failed to initiate the connection due to exception: " + str(exc), - error=exc + error=exc, ) async def _disconnect(self, *args) -> None: @@ -245,7 +234,7 @@ async def _send_frame(self, channel, frame, timeout=None, **kwargs): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send frame out due to exception: " + str(exc), - error=exc + error=exc, ) else: _LOGGER.warning("Cannot write frame in current state: %r", self.state) @@ -280,7 +269,7 @@ async def _outgoing_empty(self): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send empty frame due to exception: " + str(exc), - error=exc + error=exc, ) async def _outgoing_header(self): @@ -351,8 +340,7 @@ async def _incoming_open(self, channel, frame): _LOGGER.error("OPEN frame received on a channel that is not 0.") await self.close( error=AMQPError( - condition=ErrorCondition.NotAllowed, - description="OPEN frame received on a channel that is not 0." + condition=ErrorCondition.NotAllowed, description="OPEN frame received on a channel that is not 0." ) ) await self._set_state(ConnectionState.END) @@ -360,16 +348,19 @@ async def _incoming_open(self, channel, frame): _LOGGER.error("OPEN frame received in the OPENED state.") await self.close() if frame[4]: - self._remote_idle_timeout = frame[4]/1000 # Convert to seconds + self._remote_idle_timeout = frame[4] / 1000 # Convert to seconds self._remote_idle_timeout_send_frame = self._idle_timeout_empty_frame_send_ratio * self._remote_idle_timeout if frame[2] < 512: - #Close with error - #Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] - #Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] - await self.close(error=AMQPConnectionError( - condition=ErrorCondition.InvalidField, - description="connection_endpoint_frame_received::failed parsing OPEN frame")) + # Close with error + # Codes_S_R_S_CONNECTION_01_143: [If any of the values in the received open frame are invalid then the connection shall be closed.] + # Codes_S_R_S_CONNECTION_01_220: [The error amqp:invalid-field shall be set in the error.condition field of the CLOSE frame.] + await self.close( + error=AMQPConnectionError( + condition=ErrorCondition.InvalidField, + description="connection_endpoint_frame_received::failed parsing OPEN frame", + ) + ) _LOGGER.error("connection_endpoint_frame_received::failed parsing OPEN frame") else: self._remote_max_frame_size = frame[2] @@ -409,7 +400,7 @@ async def _incoming_close(self, channel, frame): ConnectionState.HDR_EXCH, ConnectionState.OPEN_RCVD, ConnectionState.CLOSE_SENT, - ConnectionState.DISCARDING + ConnectionState.DISCARDING, ] if self.state in disconnect_states: await self._disconnect() @@ -427,12 +418,8 @@ async def _incoming_close(self, channel, frame): await self._set_state(ConnectionState.END) if frame[0]: - self._error = AMQPConnectionError( - condition=frame[0][0], - description=frame[0][1], - info=frame[0][2] - ) - _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation + self._error = AMQPConnectionError(condition=frame[0][0], description=frame[0][1], info=frame[0][2]) + _LOGGER.error("Connection error: {}".format(frame[0])) # pylint:disable=logging-format-interpolation async def _incoming_begin(self, channel, frame): # type: (int, Tuple[Any, ...]) -> None @@ -485,7 +472,7 @@ async def _incoming_end(self, channel, frame): _LOGGER.error(f"Invalid channel {channel} ") await self.close(error=end_error) - async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements + async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-many-return-statements # type: (int, Optional[Union[bytes, Tuple[int, Tuple[Any, ...]]]]) -> bool """Process an incoming frame, either directly or by passing to the necessary Session. @@ -534,13 +521,15 @@ async def _process_incoming_frame(self, channel, frame): # pylint:disable=too-ma if performative == 0: await self._incoming_header(channel, fields) return True - if performative == 1: # pylint:disable=no-else-return + if performative == 1: # pylint:disable=no-else-return return False # TODO: incoming EMPTY else: - _LOGGER.error("Unrecognized incoming frame: {}".format(frame)) # pylint:disable=logging-format-interpolation + _LOGGER.error( + "Unrecognized incoming frame: {}".format(frame) + ) # pylint:disable=logging-format-interpolation return True except KeyError: - return True #TODO: channel error + return True # TODO: channel error async def _process_outgoing_frame(self, channel, frame): # type: (int, NamedTuple) -> None @@ -554,14 +543,15 @@ async def _process_outgoing_frame(self, channel, frame): raise ValueError("Connection not open.") now = time.time() if get_local_timeout(now, self._idle_timeout, self._last_frame_received_time) or ( - await self._get_remote_timeout(now)): + await self._get_remote_timeout(now) + ): await self.close( # TODO: check error condition error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, - description="No frame received for the idle timeout." + description="No frame received for the idle timeout.", ), - wait=False + wait=False, ) return await self._send_frame(channel, frame) @@ -634,21 +624,21 @@ async def listen(self, wait=False, batch=1, **kwargs): if self.state not in _CLOSING_STATES: now = time.time() if get_local_timeout(now, self._idle_timeout, self._last_frame_received_time) or ( - await self._get_remote_timeout(now)): + await self._get_remote_timeout(now) + ): # TODO: check error condition await self.close( error=AMQPError( condition=ErrorCondition.ConnectionCloseForced, - description="No frame received for the idle timeout." + description="No frame received for the idle timeout.", ), - wait=False + wait=False, ) return if self.state == ConnectionState.END: # TODO: check error condition self._error = AMQPConnectionError( - condition=ErrorCondition.ConnectionCloseForced, - description="Connection was already closed." + condition=ErrorCondition.ConnectionCloseForced, description="Connection was already closed." ) return for _ in range(batch): @@ -659,7 +649,7 @@ async def listen(self, wait=False, batch=1, **kwargs): self._error = AMQPConnectionError( ErrorCondition.SocketError, description="Can not send frame out due to exception: " + str(exc), - error=exc + error=exc, ) def create_session(self, **kwargs): @@ -684,14 +674,15 @@ def create_session(self, **kwargs): will be logged at the logging.INFO level. Default value is that configured for the connection. """ assigned_channel = self._get_next_outgoing_channel() - kwargs['allow_pipelined_open'] = self._allow_pipelined_open - kwargs['idle_wait_time'] = self._idle_wait_time + kwargs["allow_pipelined_open"] = self._allow_pipelined_open + kwargs["idle_wait_time"] = self._idle_wait_time session = Session( self, assigned_channel, - network_trace=kwargs.pop('network_trace', self._network_trace), + network_trace=kwargs.pop("network_trace", self._network_trace), network_trace_params=dict(self._network_trace_params), - **kwargs) + **kwargs + ) self._outgoing_endpoints[assigned_channel] = session return session @@ -732,9 +723,7 @@ async def close(self, error=None, wait=False): await self._outgoing_close(error=error) if error: self._error = AMQPConnectionError( - condition=error.condition, - description=error.description, - info=error.info + condition=error.condition, description=error.description, info=error.info ) if self.state == ConnectionState.OPEN_PIPE: await self._set_state(ConnectionState.OC_PIPE) @@ -745,7 +734,7 @@ async def close(self, error=None, wait=False): else: await self._set_state(ConnectionState.CLOSE_SENT) await self._wait_for_response(wait, ConnectionState.END) - except Exception as exc: # pylint:disable=broad-except + except Exception as exc: # pylint:disable=broad-except # If error happened during closing, ignore the error and set state to END _LOGGER.info("An error occurred when closing the connection: %r", exc) await self._set_state(ConnectionState.END) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py index b7bfddf7bb8f..9847da7c400a 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py @@ -1,53 +1,29 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- -import threading -import struct from typing import Optional import uuid import logging -import time -from enum import Enum -from io import BytesIO -from urllib.parse import urlparse + import asyncio from ..endpoints import Source, Target -from ..constants import ( - DEFAULT_LINK_CREDIT, - SessionState, - SessionTransferState, - LinkDeliverySettleReason, - LinkState, - Role, - SenderSettleMode, - ReceiverSettleMode -) +from ..constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode from ..performatives import ( AttachFrame, DetachFrame, - TransferFrame, - DispositionFrame, - FlowFrame, ) -from ..error import ( - ErrorCondition, - AMQPLinkError, - AMQPLinkRedirect, - AMQPConnectionError -) +from ..error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError _LOGGER = logging.getLogger(__name__) class Link(object): - """ - - """ + """ """ def __init__(self, session, handle, name, role, **kwargs): self.state = LinkState.DETACHED @@ -55,53 +31,61 @@ def __init__(self, session, handle, name, role, **kwargs): self.handle = handle self.remote_handle = None self.role = role - source_address = kwargs['source_address'] + source_address = kwargs["source_address"] target_address = kwargs["target_address"] - self.source = source_address if isinstance(source_address, Source) else Source( - address=kwargs['source_address'], - durable=kwargs.get('source_durable'), - expiry_policy=kwargs.get('source_expiry_policy'), - timeout=kwargs.get('source_timeout'), - dynamic=kwargs.get('source_dynamic'), - dynamic_node_properties=kwargs.get('source_dynamic_node_properties'), - distribution_mode=kwargs.get('source_distribution_mode'), - filters=kwargs.get('source_filters'), - default_outcome=kwargs.get('source_default_outcome'), - outcomes=kwargs.get('source_outcomes'), - capabilities=kwargs.get('source_capabilities') + self.source = ( + source_address + if isinstance(source_address, Source) + else Source( + address=kwargs["source_address"], + durable=kwargs.get("source_durable"), + expiry_policy=kwargs.get("source_expiry_policy"), + timeout=kwargs.get("source_timeout"), + dynamic=kwargs.get("source_dynamic"), + dynamic_node_properties=kwargs.get("source_dynamic_node_properties"), + distribution_mode=kwargs.get("source_distribution_mode"), + filters=kwargs.get("source_filters"), + default_outcome=kwargs.get("source_default_outcome"), + outcomes=kwargs.get("source_outcomes"), + capabilities=kwargs.get("source_capabilities"), + ) ) - self.target = target_address if isinstance(target_address,Target) else Target( - address=kwargs['target_address'], - durable=kwargs.get('target_durable'), - expiry_policy=kwargs.get('target_expiry_policy'), - timeout=kwargs.get('target_timeout'), - dynamic=kwargs.get('target_dynamic'), - dynamic_node_properties=kwargs.get('target_dynamic_node_properties'), - capabilities=kwargs.get('target_capabilities') + self.target = ( + target_address + if isinstance(target_address, Target) + else Target( + address=kwargs["target_address"], + durable=kwargs.get("target_durable"), + expiry_policy=kwargs.get("target_expiry_policy"), + timeout=kwargs.get("target_timeout"), + dynamic=kwargs.get("target_dynamic"), + dynamic_node_properties=kwargs.get("target_dynamic_node_properties"), + capabilities=kwargs.get("target_capabilities"), + ) ) - self.link_credit = kwargs.pop('link_credit', None) or DEFAULT_LINK_CREDIT + self.link_credit = kwargs.pop("link_credit", None) or DEFAULT_LINK_CREDIT self.current_link_credit = self.link_credit - self.send_settle_mode = kwargs.pop('send_settle_mode', SenderSettleMode.Mixed) - self.rcv_settle_mode = kwargs.pop('rcv_settle_mode', ReceiverSettleMode.First) - self.unsettled = kwargs.pop('unsettled', None) - self.incomplete_unsettled = kwargs.pop('incomplete_unsettled', None) - self.initial_delivery_count = kwargs.pop('initial_delivery_count', 0) + self.send_settle_mode = kwargs.pop("send_settle_mode", SenderSettleMode.Mixed) + self.rcv_settle_mode = kwargs.pop("rcv_settle_mode", ReceiverSettleMode.First) + self.unsettled = kwargs.pop("unsettled", None) + self.incomplete_unsettled = kwargs.pop("incomplete_unsettled", None) + self.initial_delivery_count = kwargs.pop("initial_delivery_count", 0) self.delivery_count = self.initial_delivery_count self.received_delivery_id = None - self.max_message_size = kwargs.pop('max_message_size', None) + self.max_message_size = kwargs.pop("max_message_size", None) self.remote_max_message_size = None - self.available = kwargs.pop('available', None) - self.properties = kwargs.pop('properties', None) + self.available = kwargs.pop("available", None) + self.properties = kwargs.pop("properties", None) self.offered_capabilities = None - self.desired_capabilities = kwargs.pop('desired_capabilities', None) + self.desired_capabilities = kwargs.pop("desired_capabilities", None) - self.network_trace = kwargs['network_trace'] - self.network_trace_params = kwargs['network_trace_params'] - self.network_trace_params['link'] = self.name + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["link"] = self.name self._session = session self._is_closed = False - self._on_link_state_change = kwargs.get('on_link_state_change') - self._on_attach = kwargs.get('on_attach') + self._on_link_state_change = kwargs.get("on_link_state_change") + self._on_attach = kwargs.get("on_attach") self._error = None async def __aenter__(self): @@ -114,7 +98,7 @@ async def __aexit__(self, *args): @classmethod def from_incoming_frame(cls, session, handle, frame): # check link_create_from_endpoint in C lib - raise NotImplementedError('Pending') # TODO: Assuming we establish all links for now... + raise NotImplementedError("Pending") # TODO: Assuming we establish all links for now... def get_state(self): try: @@ -128,10 +112,7 @@ def _check_if_closed(self): try: raise self._error except TypeError: - raise AMQPConnectionError( - condition=ErrorCondition.InternalError, - description="Link already closed." - ) + raise AMQPConnectionError(condition=ErrorCondition.InternalError, description="Link already closed.") async def _set_state(self, new_state): # type: (LinkState) -> None @@ -147,7 +128,7 @@ async def _set_state(self, new_state): pass except Exception as e: # pylint: disable=broad-except _LOGGER.error("Link state change callback failed: '%r'", e, extra=self.network_trace_params) - + async def _on_session_state_change(self): if self._session.state == SessionState.MAPPED: if not self._is_closed and self.state == LinkState.DETACHED: @@ -172,20 +153,20 @@ async def _outgoing_attach(self): max_message_size=self.max_message_size, offered_capabilities=self.offered_capabilities if self.state == LinkState.ATTACH_RCVD else None, desired_capabilities=self.desired_capabilities if self.state == LinkState.DETACHED else None, - properties=self.properties + properties=self.properties, ) if self.network_trace: _LOGGER.info("-> %r", attach_frame, extra=self.network_trace_params) - await self._session._outgoing_attach(attach_frame) + await self._session._outgoing_attach(attach_frame) # pylint: disable=protected-access async def _incoming_attach(self, frame): if self.network_trace: _LOGGER.info("<- %r", AttachFrame(*frame), extra=self.network_trace_params) if self._is_closed: raise ValueError("Invalid link") - elif not frame[5] or not frame[6]: + if not frame[5] or not frame[6]: _LOGGER.info("Cannot get source or target. Detaching link") - await self._set_state(LinkState.DETACHED) + await self._set_state(LinkState.DETACHED) raise ValueError("Invalid link") self.remote_handle = frame[1] # handle self.remote_max_message_size = frame[10] # max_message_size @@ -210,19 +191,19 @@ async def _incoming_attach(self, frame): async def _outgoing_flow(self, **kwargs): flow_frame = { - 'handle': self.handle, - 'delivery_count': self.delivery_count, - 'link_credit': self.current_link_credit, - 'available': kwargs.get('available'), - 'drain': kwargs.get('drain'), - 'echo': kwargs.get('echo'), - 'properties': kwargs.get('properties') + "handle": self.handle, + "delivery_count": self.delivery_count, + "link_credit": self.current_link_credit, + "available": kwargs.get("available"), + "drain": kwargs.get("drain"), + "echo": kwargs.get("echo"), + "properties": kwargs.get("properties"), } - await self._session._outgoing_flow(flow_frame) + await self._session._outgoing_flow(flow_frame) # pylint: disable=protected-access async def _incoming_flow(self, frame): pass - + async def _incoming_disposition(self, frame): pass @@ -230,7 +211,7 @@ async def _outgoing_detach(self, close=False, error=None): detach_frame = DetachFrame(handle=self.handle, closed=close, error=error) if self.network_trace: _LOGGER.info("-> %r", detach_frame, extra=self.network_trace_params) - await self._session._outgoing_detach(detach_frame) + await self._session._outgoing_detach(detach_frame) # pylint: disable=protected-access if close: self._is_closed = True @@ -274,11 +255,6 @@ async def detach(self, close=False, error=None): _LOGGER.info("An error occurred when detaching the link: %r", exc) await self._set_state(LinkState.DETACHED) - async def flow( - self, - *, - link_credit: Optional[int] = None, - **kwargs - ) -> None: + async def flow(self, *, link_credit: Optional[int] = None, **kwargs) -> None: self.current_link_credit = link_credit if link_credit is not None else self.link_credit await self._outgoing_flow(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py index 76b4e01d2c36..d6dfde96e7b1 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_link_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import time import logging @@ -18,20 +18,20 @@ ReceiverSettleMode, ManagementExecuteOperationResult, ManagementOpenResult, - SEND_DISPOSITION_ACCEPT, SEND_DISPOSITION_REJECT, - MessageDeliveryState + MessageDeliveryState, ) -from ..error import ErrorResponse, AMQPException, ErrorCondition -from ..message import Message, Properties, _MessageDelivery +from ..error import AMQPException, ErrorCondition +from ..message import Properties, _MessageDelivery _LOGGER = logging.getLogger(__name__) -class ManagementLink(object): # pylint:disable=too-many-instance-attributes +class ManagementLink(object): # pylint:disable=too-many-instance-attributes """ - # TODO: Fill in docstring + # TODO: Fill in docstring """ + def __init__(self, session, endpoint, **kwargs): self.next_message_id = 0 self.state = ManagementLinkState.IDLE @@ -42,7 +42,7 @@ def __init__(self, session, endpoint, **kwargs): source_address=endpoint, on_link_state_change=self._on_sender_state_change, send_settle_mode=SenderSettleMode.Unsettled, - rcv_settle_mode=ReceiverSettleMode.First + rcv_settle_mode=ReceiverSettleMode.First, ) self._response_link: ReceiverLink = session.create_receiver_link( endpoint, @@ -50,13 +50,13 @@ def __init__(self, session, endpoint, **kwargs): on_link_state_change=self._on_receiver_state_change, on_transfer=self._on_message_received, send_settle_mode=SenderSettleMode.Unsettled, - rcv_settle_mode=ReceiverSettleMode.First + rcv_settle_mode=ReceiverSettleMode.First, ) - self._on_amqp_management_error = kwargs.get('on_amqp_management_error') - self._on_amqp_management_open_complete = kwargs.get('on_amqp_management_open_complete') + self._on_amqp_management_error = kwargs.get("on_amqp_management_error") + self._on_amqp_management_open_complete = kwargs.get("on_amqp_management_open_complete") - self._status_code_field = kwargs.get('status_code_field', b'statusCode') - self._status_description_field = kwargs.get('status_description_field', b'statusDescription') + self._status_code_field = kwargs.get("status_code_field", b"statusCode") + self._status_description_field = kwargs.get("status_description_field", b"statusDescription") self._sender_connected = False self._receiver_connected = False @@ -132,14 +132,13 @@ async def _on_message_received(self, _, message): to_remove_operation = operation break if to_remove_operation: - mgmt_result = ManagementExecuteOperationResult.OK \ - if 200 <= status_code <= 299 else ManagementExecuteOperationResult.FAILED_BAD_STATUS + mgmt_result = ( + ManagementExecuteOperationResult.OK + if 200 <= status_code <= 299 + else ManagementExecuteOperationResult.FAILED_BAD_STATUS + ) await to_remove_operation.on_execute_operation_complete( - mgmt_result, - status_code, - status_description, - message, - response_detail.get(b'error-condition') + mgmt_result, status_code, status_description, message, response_detail.get(b"error-condition") ) self._pending_operations.remove(to_remove_operation) @@ -164,7 +163,7 @@ async def _on_send_complete(self, message_delivery, reason, state): # todo: rea condition=state[SEND_DISPOSITION_REJECT][0][0], # 0 is error condition description=state[SEND_DISPOSITION_REJECT][0][1], # 1 is error description info=state[SEND_DISPOSITION_REJECT][0][2], # 2 is error info - ) + ), ) async def open(self): @@ -174,12 +173,7 @@ async def open(self): await self._response_link.attach() await self._request_link.attach() - async def execute_operation( - self, - message, - on_execute_operation_complete, - **kwargs - ): + async def execute_operation(self, message, on_execute_operation_complete, **kwargs): """Execute a request and wait on a response. :param message: The message to send in the management request. @@ -214,19 +208,11 @@ async def execute_operation( new_properties = Properties(message_id=self.next_message_id) message = message._replace(properties=new_properties) expire_time = (time.time() + timeout) if timeout else None - message_delivery = _MessageDelivery( - message, - MessageDeliveryState.WaitingToBeSent, - expire_time - ) + message_delivery = _MessageDelivery(message, MessageDeliveryState.WaitingToBeSent, expire_time) on_send_complete = partial(self._on_send_complete, message_delivery) - await self._request_link.send_transfer( - message, - on_send_complete=on_send_complete, - timeout=timeout - ) + await self._request_link.send_transfer(message, on_send_complete=on_send_complete, timeout=timeout) self.next_message_id += 1 self._pending_operations.append(PendingManagementOperation(message, on_execute_operation_complete)) @@ -241,7 +227,7 @@ async def close(self): None, None, pending_operation.message, - AMQPException(condition=ErrorCondition.ClientError, description="Management link already closed.") + AMQPException(condition=ErrorCondition.ClientError, description="Management link already closed."), ) self._pending_operations = [] self.state = ManagementLinkState.IDLE diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py index b3fb7a4ac130..d6f225a34427 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_management_operation_async.py @@ -9,10 +9,7 @@ from functools import partial from ._management_link_async import ManagementLink -from ..message import Message from ..error import ( - AMQPException, - AMQPConnectionError, AMQPLinkError, ErrorCondition ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py index bc54577f3215..dea7ac8bda91 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py @@ -1,53 +1,34 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging from typing import Optional, Union from .._decode import decode_payload -from ..endpoints import Target from ._link_async import Link -from ..message import Message, Properties, Header -from ..constants import ( - DEFAULT_LINK_CREDIT, - SessionState, - SessionTransferState, - LinkDeliverySettleReason, - LinkState, - Role -) +from ..constants import LinkState, Role from ..performatives import ( - AttachFrame, - DetachFrame, TransferFrame, DispositionFrame, - FlowFrame, -) -from ..outcomes import ( - Received, - Accepted, - Rejected, - Released, - Modified ) +from ..outcomes import Received, Accepted, Rejected, Released, Modified _LOGGER = logging.getLogger(__name__) class ReceiverLink(Link): - def __init__(self, session, handle, source_address, **kwargs): - name = kwargs.pop('name', None) or str(uuid.uuid4()) + name = kwargs.pop("name", None) or str(uuid.uuid4()) role = Role.Receiver - if 'target_address' not in kwargs: - kwargs['target_address'] = "receiver-link-{}".format(name) + if "target_address" not in kwargs: + kwargs["target_address"] = "receiver-link-{}".format(name) super(ReceiverLink, self).__init__(session, handle, name, role, source_address=source_address, **kwargs) - self._on_transfer = kwargs.pop('on_transfer') + self._on_transfer = kwargs.pop("on_transfer") self._received_payload = bytearray() async def _process_incoming_message(self, frame, message): @@ -89,56 +70,45 @@ async def _incoming_transfer(self, frame): await self._outgoing_disposition(first=frame[1], settled=True, state=delivery_state) async def _wait_for_response(self, wait: Union[bool, float]) -> None: - if wait == True: - await self._session._connection.listen(wait=False) + if wait is True: + await self._session._connection.listen(wait=False) # pylint: disable=protected-access if self.state == LinkState.ERROR: - raise self._error + raise self._error elif wait: - await self._session._connection.listen(wait=wait) + await self._session._connection.listen(wait=wait) # pylint: disable=protected-access if self.state == LinkState.ERROR: - raise self._error + raise self._error async def _outgoing_disposition( - self, - first: int, - last: Optional[int], - settled: Optional[bool], - state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], - batchable: Optional[bool] + self, + first: int, + last: Optional[int], + settled: Optional[bool], + state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], + batchable: Optional[bool], ): disposition_frame = DispositionFrame( - role=self.role, - first=first, - last=last, - settled=settled, - state=state, - batchable=batchable + role=self.role, first=first, last=last, settled=settled, state=state, batchable=batchable ) if self.network_trace: _LOGGER.info("-> %r", DispositionFrame(*disposition_frame), extra=self.network_trace_params) - await self._session._outgoing_disposition(disposition_frame) + await self._session._outgoing_disposition(disposition_frame) # pylint: disable=protected-access async def attach(self): await super().attach() self._received_payload = bytearray() async def send_disposition( - self, - *, - wait: Union[bool, float] = False, - first_delivery_id: int, - last_delivery_id: Optional[int] = None, - settled: Optional[bool] = None, - delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, - batchable: Optional[bool] = None - ): + self, + *, + wait: Union[bool, float] = False, + first_delivery_id: int, + last_delivery_id: Optional[int] = None, + settled: Optional[bool] = None, + delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, + batchable: Optional[bool] = None + ): if self._is_closed: raise ValueError("Link already closed.") - await self._outgoing_disposition( - first_delivery_id, - last_delivery_id, - settled, - delivery_state, - batchable - ) + await self._outgoing_disposition(first_delivery_id, last_delivery_id, settled, delivery_state, batchable) await self._wait_for_response(wait) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py index 6f8c846c8da0..acb1079eae44 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sasl_async.py @@ -4,19 +4,10 @@ # license information. #-------------------------------------------------------------------------- -import struct -from enum import Enum - from ._transport_async import AsyncTransport, WebSocketTransportAsync -from ..types import AMQPTypes, TYPE, VALUE -from ..constants import FIELD, SASLCode, SASL_HEADER_FRAME, TransportType, WEBSOCKET_PORT +from ..constants import SASLCode, SASL_HEADER_FRAME, WEBSOCKET_PORT from .._transport import AMQPS_PORT -from ..performatives import ( - SASLOutcome, - SASLResponse, - SASLChallenge, - SASLInit -) +from ..performatives import SASLInit _SASL_FRAME_TYPE = b'\x01' @@ -96,8 +87,7 @@ async def _negotiate(self): raise NotImplementedError("Unsupported SASL challenge") if fields[0] == SASLCode.Ok: # code return - else: - raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) + raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) class SASLTransport(AsyncTransport, SASLTransportMixinAsync): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py index a72511b639d9..fd75494d78a3 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_sender_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import struct import uuid import logging @@ -11,14 +11,7 @@ from .._encode import encode_payload from ._link_async import Link -from ..constants import ( - SessionTransferState, - LinkDeliverySettleReason, - LinkState, - Role, - SenderSettleMode, - SessionState -) +from ..constants import SessionTransferState, LinkDeliverySettleReason, LinkState, Role, SenderSettleMode, SessionState from ..performatives import ( TransferFrame, ) @@ -28,33 +21,31 @@ class PendingDelivery(object): - def __init__(self, **kwargs): - self.message = kwargs.get('message') + self.message = kwargs.get("message") self.sent = False self.frame = None - self.on_delivery_settled = kwargs.get('on_delivery_settled') + self.on_delivery_settled = kwargs.get("on_delivery_settled") self.start = time.time() self.transfer_state = None - self.timeout = kwargs.get('timeout') - self.settled = kwargs.get('settled', False) + self.timeout = kwargs.get("timeout") + self.settled = kwargs.get("settled", False) async def on_settled(self, reason, state): if self.on_delivery_settled and not self.settled: try: await self.on_delivery_settled(reason, state) - except Exception as e: # pylint:disable=broad-except + except Exception as e: # pylint:disable=broad-except _LOGGER.warning("Message 'on_send_complete' callback failed: %r", e) self.settled = True class SenderLink(Link): - def __init__(self, session, handle, target_address, **kwargs): - name = kwargs.pop('name', None) or str(uuid.uuid4()) + name = kwargs.pop("name", None) or str(uuid.uuid4()) role = Role.Sender - if 'source_address' not in kwargs: - kwargs['source_address'] = "sender-link-{}".format(name) + if "source_address" not in kwargs: + kwargs["source_address"] = "sender-link-{}".format(name) super(SenderLink, self).__init__(session, handle, name, role, target_address=target_address, **kwargs) self._pending_deliveries = [] @@ -91,22 +82,24 @@ async def _outgoing_transfer(self, delivery): encode_payload(output, delivery.message) delivery_count = self.delivery_count + 1 delivery.frame = { - 'handle': self.handle, - 'delivery_tag': struct.pack('>I', abs(delivery_count)), - 'message_format': delivery.message._code, # pylint:disable=protected-access - 'settled': delivery.settled, - 'more': False, - 'rcv_settle_mode': None, - 'state': None, - 'resume': None, - 'aborted': None, - 'batchable': None, - 'payload': output + "handle": self.handle, + "delivery_tag": struct.pack(">I", abs(delivery_count)), + "message_format": delivery.message._code, # pylint:disable=protected-access + "settled": delivery.settled, + "more": False, + "rcv_settle_mode": None, + "state": None, + "resume": None, + "aborted": None, + "batchable": None, + "payload": output, } if self.network_trace: - _LOGGER.info("-> %r", TransferFrame(delivery_id='', **delivery.frame), extra=self.network_trace_params) # pylint:disable=line-to-long + _LOGGER.info( + "-> %r", TransferFrame(delivery_id="", **delivery.frame), extra=self.network_trace_params + ) _LOGGER.info(" %r", delivery.message, extra=self.network_trace_params) - await self._session._outgoing_transfer(delivery) # pylint:disable=protected-access + await self._session._outgoing_transfer(delivery) # pylint:disable=protected-access sent_and_settled = False if delivery.transfer_state == SessionTransferState.OKAY: self.delivery_count = delivery_count @@ -126,7 +119,7 @@ async def _incoming_disposition(self, frame): settled_ids = list(range(frame[1], range_end)) unsettled = [] for delivery in self._pending_deliveries: - if delivery.sent and delivery.frame['delivery_id'] in settled_ids: + if delivery.sent and delivery.frame["delivery_id"] in settled_ids: await delivery.on_settled(LinkDeliverySettleReason.DISPOSITION_RECEIVED, frame[4]) # state continue unsettled.append(delivery) @@ -138,7 +131,7 @@ async def _remove_pending_deliveries(self): futures.append(asyncio.ensure_future(delivery.on_settled(LinkDeliverySettleReason.NOT_DELIVERED, None))) await asyncio.gather(*futures) self._pending_deliveries = [] - + async def _on_session_state_change(self): if self._session.state == SessionState.DISCARDING: await self._remove_pending_deliveries() @@ -166,14 +159,14 @@ async def send_transfer(self, message, *, send_async=False, **kwargs): if self.state != LinkState.ATTACHED: raise AMQPLinkError( # TODO: should we introduce MessageHandler to indicate the handler is in wrong state condition=ErrorCondition.ClientError, # TODO: should this be a ClientError? - description="Link is not attached." + description="Link is not attached.", ) settled = self.send_settle_mode == SenderSettleMode.Settled if self.send_settle_mode == SenderSettleMode.Mixed: - settled = kwargs.pop('settled', True) + settled = kwargs.pop("settled", True) delivery = PendingDelivery( - on_delivery_settled=kwargs.get('on_send_complete'), - timeout=kwargs.get('timeout'), + on_delivery_settled=kwargs.get("on_send_complete"), + timeout=kwargs.get("timeout"), message=message, settled=settled, ) @@ -194,6 +187,7 @@ async def cancel_transfer(self, delivery): if delivery.sent: raise MessageException( ErrorCondition.ClientError, - message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.") + message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.", + ) await delivery.on_settled(LinkDeliverySettleReason.CANCELLED, None) self._pending_deliveries.pop(index) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py index 4fc82112c995..4bb36bc63ae1 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_session_async.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging @@ -22,15 +22,7 @@ from ._sender_async import SenderLink from ._receiver_async import ReceiverLink from ._management_link_async import ManagementLink -from ..performatives import ( - BeginFrame, - EndFrame, - FlowFrame, - AttachFrame, - DetachFrame, - TransferFrame, - DispositionFrame -) +from ..performatives import BeginFrame, EndFrame, FlowFrame, TransferFrame, DispositionFrame from .._encode import encode_frame _LOGGER = logging.getLogger(__name__) @@ -49,27 +41,27 @@ class Session(object): """ def __init__(self, connection, channel, **kwargs): - self.name = kwargs.pop('name', None) or str(uuid.uuid4()) + self.name = kwargs.pop("name", None) or str(uuid.uuid4()) self.state = SessionState.UNMAPPED - self.handle_max = kwargs.get('handle_max', 4294967295) - self.properties = kwargs.pop('properties', None) + self.handle_max = kwargs.get("handle_max", 4294967295) + self.properties = kwargs.pop("properties", None) self.channel = channel self.remote_channel = None - self.next_outgoing_id = kwargs.pop('next_outgoing_id', 0) + self.next_outgoing_id = kwargs.pop("next_outgoing_id", 0) self.next_incoming_id = None - self.incoming_window = kwargs.pop('incoming_window', 1) - self.outgoing_window = kwargs.pop('outgoing_window', 1) + self.incoming_window = kwargs.pop("incoming_window", 1) + self.outgoing_window = kwargs.pop("outgoing_window", 1) self.target_incoming_window = self.incoming_window self.remote_incoming_window = 0 self.remote_outgoing_window = 0 self.offered_capabilities = None - self.desired_capabilities = kwargs.pop('desired_capabilities', None) + self.desired_capabilities = kwargs.pop("desired_capabilities", None) - self.allow_pipelined_open = kwargs.pop('allow_pipelined_open', True) - self.idle_wait_time = kwargs.get('idle_wait_time', 0.1) - self.network_trace = kwargs['network_trace'] - self.network_trace_params = kwargs['network_trace_params'] - self.network_trace_params['session'] = self.name + self.allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) + self.idle_wait_time = kwargs.get("idle_wait_time", 0.1) + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["session"] = self.name self.links = {} self._connection = connection @@ -99,7 +91,7 @@ async def _set_state(self, new_state): _LOGGER.info("Session state changed: %r -> %r", previous_state, new_state, extra=self.network_trace_params) futures = [] for link in self.links.values(): - futures.append(asyncio.ensure_future(link._on_session_state_change())) + futures.append(asyncio.ensure_future(link._on_session_state_change())) # pylint: disable=protected-access await asyncio.gather(*futures) async def _on_connection_state_change(self): @@ -119,7 +111,7 @@ def _get_next_output_handle(self): raise ValueError("Maximum number of handles ({}) has been reached.".format(self.handle_max)) next_handle = next(i for i in range(1, self.handle_max) if i not in self._output_handles) return next_handle - + async def _outgoing_begin(self): begin_frame = BeginFrame( remote_channel=self.remote_channel if self.state == SessionState.BEGIN_RCVD else None, @@ -133,7 +125,7 @@ async def _outgoing_begin(self): ) if self.network_trace: _LOGGER.info("-> %r", begin_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, begin_frame) + await self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access async def _incoming_begin(self, frame): if self.network_trace: @@ -154,7 +146,7 @@ async def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, end_frame) + await self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access async def _incoming_end(self, frame): if self.network_trace: @@ -163,43 +155,45 @@ async def _incoming_end(self, frame): await self._set_state(SessionState.END_RCVD) for _, link in self.links.items(): await link.detach() - # TODO: handling error + # TODO: handling error await self._outgoing_end() await self._set_state(SessionState.UNMAPPED) async def _outgoing_attach(self, frame): - await self._connection._process_outgoing_frame(self.channel, frame) + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access async def _incoming_attach(self, frame): try: - self._input_handles[frame[1]] = self.links[frame[0].decode('utf-8')] # name and handle - await self._input_handles[frame[1]]._incoming_attach(frame) + self._input_handles[frame[1]] = self.links[frame[0].decode("utf-8")] # name and handle + await self._input_handles[frame[1]]._incoming_attach(frame) # pylint: disable=protected-access except KeyError: outgoing_handle = self._get_next_output_handle() # TODO: catch max-handles error if frame[2] == Role.Sender: # role new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) else: new_link = SenderLink.from_incoming_frame(self, outgoing_handle, frame) - await new_link._incoming_attach(frame) + await new_link._incoming_attach(frame) # pylint: disable=protected-access self.links[frame[0]] = new_link self._output_handles[outgoing_handle] = new_link self._input_handles[frame[1]] = new_link except ValueError: # Reject Link await self._input_handles[frame[1]].detach() - + async def _outgoing_flow(self, frame=None): link_flow = frame or {} - link_flow.update({ - 'next_incoming_id': self.next_incoming_id, - 'incoming_window': self.incoming_window, - 'next_outgoing_id': self.next_outgoing_id, - 'outgoing_window': self.outgoing_window - }) + link_flow.update( + { + "next_incoming_id": self.next_incoming_id, + "incoming_window": self.incoming_window, + "next_outgoing_id": self.next_outgoing_id, + "outgoing_window": self.outgoing_window, + } + ) flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - await self._connection._process_outgoing_frame(self.channel, flow_frame) + await self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access async def _incoming_flow(self, frame): if self.network_trace: @@ -209,12 +203,12 @@ async def _incoming_flow(self, frame): self.remote_incoming_window = remote_incoming_id + frame[1] - self.next_outgoing_id # incoming_window self.remote_outgoing_window = frame[3] # outgoing_window if frame[4] is not None: # handle - await self._input_handles[frame[4]]._incoming_flow(frame) + await self._input_handles[frame[4]]._incoming_flow(frame) # pylint: disable=protected-access else: futures = [] for link in self._output_handles.values(): - if self.remote_incoming_window > 0 and not link._is_closed: - futures.append(link._incoming_flow(frame)) + if self.remote_incoming_window > 0 and not link._is_closed: # pylint: disable=protected-access + futures.append(link._incoming_flow(frame)) # pylint: disable=protected-access await asyncio.gather(*futures) async def _outgoing_transfer(self, delivery): @@ -223,58 +217,58 @@ async def _outgoing_transfer(self, delivery): if self.remote_incoming_window <= 0: delivery.transfer_state = SessionTransferState.BUSY else: - payload = delivery.frame['payload'] + payload = delivery.frame["payload"] payload_size = len(payload) - delivery.frame['delivery_id'] = self.next_outgoing_id + delivery.frame["delivery_id"] = self.next_outgoing_id # calculate the transfer frame encoding size excluding the payload - delivery.frame['payload'] = b"" + delivery.frame["payload"] = b"" # TODO: encoding a frame would be expensive, we might want to improve depending on the perf test results encoded_frame = encode_frame(TransferFrame(**delivery.frame))[1] transfer_overhead_size = len(encoded_frame) # available size for payload per frame is calculated as following: # remote max frame size - transfer overhead (calculated) - header (8 bytes) - available_frame_size = self._connection._remote_max_frame_size - transfer_overhead_size - 8 + available_frame_size = self._connection._remote_max_frame_size - transfer_overhead_size - 8 # pylint: disable=protected-access start_idx = 0 remaining_payload_cnt = payload_size # encode n-1 frames if payload_size > available_frame_size while remaining_payload_cnt > available_frame_size: tmp_delivery_frame = { - 'handle': delivery.frame['handle'], - 'delivery_tag': delivery.frame['delivery_tag'], - 'message_format': delivery.frame['message_format'], - 'settled': delivery.frame['settled'], - 'more': True, - 'rcv_settle_mode': delivery.frame['rcv_settle_mode'], - 'state': delivery.frame['state'], - 'resume': delivery.frame['resume'], - 'aborted': delivery.frame['aborted'], - 'batchable': delivery.frame['batchable'], - 'payload': payload[start_idx:start_idx+available_frame_size], - 'delivery_id': self.next_outgoing_id + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": True, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "payload": payload[start_idx : start_idx + available_frame_size], + "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) + await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access start_idx += available_frame_size remaining_payload_cnt -= available_frame_size # encode the last frame tmp_delivery_frame = { - 'handle': delivery.frame['handle'], - 'delivery_tag': delivery.frame['delivery_tag'], - 'message_format': delivery.frame['message_format'], - 'settled': delivery.frame['settled'], - 'more': False, - 'rcv_settle_mode': delivery.frame['rcv_settle_mode'], - 'state': delivery.frame['state'], - 'resume': delivery.frame['resume'], - 'aborted': delivery.frame['aborted'], - 'batchable': delivery.frame['batchable'], - 'payload': payload[start_idx:], - 'delivery_id': self.next_outgoing_id + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": False, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "payload": payload[start_idx:], + "delivery_id": self.next_outgoing_id, } - await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) + await self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -286,31 +280,31 @@ async def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - await self._input_handles[frame[0]]._incoming_transfer(frame) # handle + await self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access except KeyError: - pass #TODO: "unattached handle" + pass # TODO: "unattached handle" if self.incoming_window == 0: self.incoming_window = self.target_incoming_window await self._outgoing_flow() async def _outgoing_disposition(self, frame): - await self._connection._process_outgoing_frame(self.channel, frame) + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access async def _incoming_disposition(self, frame): if self.network_trace: _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) futures = [] for link in self._input_handles.values(): - asyncio.ensure_future(link._incoming_disposition(frame)) + asyncio.ensure_future(link._incoming_disposition(frame)) # pylint: disable=protected-access await asyncio.gather(*futures) async def _outgoing_detach(self, frame): - await self._connection._process_outgoing_frame(self.channel, frame) + await self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access async def _incoming_detach(self, frame): try: link = self._input_handles[frame[0]] # handle - await link._incoming_detach(frame) + await link._incoming_detach(frame) # pylint: disable=protected-access # if link._is_closed: TODO # self.links.pop(link.name, None) # self._input_handles.pop(link.remote_handle, None) @@ -320,7 +314,7 @@ async def _incoming_detach(self, frame): async def _wait_for_response(self, wait, end_state): # type: (Union[bool, float], SessionState) -> None - if wait == True: + if wait is True: await self._connection.listen(wait=False) while self.state != end_state: await asyncio.sleep(self.idle_wait_time) @@ -362,9 +356,10 @@ def create_receiver_link(self, source_address, **kwargs): self, handle=assigned_handle, source_address=source_address, - network_trace=kwargs.pop('network_trace', self.network_trace), + network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs) + **kwargs + ) self.links[link.name] = link self._output_handles[assigned_handle] = link return link @@ -375,16 +370,13 @@ def create_sender_link(self, target_address, **kwargs): self, handle=assigned_handle, target_address=target_address, - network_trace=kwargs.pop('network_trace', self.network_trace), + network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs) + **kwargs + ) self._output_handles[assigned_handle] = link self.links[link.name] = link return link def create_request_response_link_pair(self, endpoint, **kwargs): - return ManagementLink( - self, - endpoint, - network_trace=kwargs.pop('network_trace', self.network_trace), - **kwargs) + return ManagementLink(self, endpoint, network_trace=kwargs.pop("network_trace", self.network_trace), **kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py index 33fe7e812315..cd4e8d49514a 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py @@ -1,4 +1,4 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # This is a fork of the transport.py which was originally written by Barry Pederson and # maintained by the Celery project: https://github.com/celery/py-amqp. # @@ -30,23 +30,20 @@ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF # THE POSSIBILITY OF SUCH DAMAGE. -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- import asyncio import errno -import re import socket import ssl import struct from ssl import SSLError -from contextlib import contextmanager from io import BytesIO import logging -from threading import Lock import certifi -from .._platform import KNOWN_TCP_OPTS, SOL_TCP, pack, unpack +from .._platform import KNOWN_TCP_OPTS, SOL_TCP from .._encode import encode_frame from .._decode import decode_frame, decode_empty_frame from ..constants import TLS_HEADER_FRAME, WEBSOCKET_PORT, AMQP_WS_SUBPROTOCOL @@ -55,13 +52,11 @@ get_errno, to_host_port, DEFAULT_SOCKET_SETTINGS, - IPV6_LITERAL, SIGNED_INT_MAX, _UNAVAIL, set_cloexec, AMQP_PORT, TIMEOUT_INTERVAL, - WebSocketTransport ) @@ -71,20 +66,21 @@ def get_running_loop(): try: import asyncio # pylint: disable=import-error + return asyncio.get_running_loop() except AttributeError: # 3.6 loop = None try: loop = asyncio._get_running_loop() # pylint: disable=protected-access except AttributeError: - _LOGGER.warning('This version of Python is deprecated, please upgrade to >= v3.6') + _LOGGER.warning("This version of Python is deprecated, please upgrade to >= v3.6") if loop is None: - _LOGGER.warning('No running event loop') + _LOGGER.warning("No running event loop") loop = asyncio.get_event_loop() return loop -class AsyncTransportMixin(): +class AsyncTransportMixin: async def receive_frame(self, timeout=None, *args, **kwargs): try: header, channel, payload = await asyncio.wait_for(self.read(**kwargs), timeout=timeout) @@ -104,11 +100,11 @@ async def read(self, verify_frame_type=0, **kwargs): frame_header = memoryview(bytearray(8)) read_frame_buffer.write(await self._read(8, buffer=frame_header, initial=True)) - channel = struct.unpack('>H', frame_header[6:])[0] + channel = struct.unpack(">H", frame_header[6:])[0] size = frame_header[0:4] if size == AMQP_FRAME: # Empty frame or AMQP header negotiation return frame_header, channel, None - size = struct.unpack('>I', size)[0] + size = struct.unpack(">I", size)[0] offset = frame_header[4] frame_type = frame_header[5] @@ -121,7 +117,7 @@ async def read(self, verify_frame_type=0, **kwargs): read_frame_buffer.write(await self._read(size - SIGNED_INT_MAX, buffer=payload[SIGNED_INT_MAX:])) else: read_frame_buffer.write(await self._read(payload_size, buffer=payload)) - except (TimeoutError, socket.timeout, asyncio.IncompleteReadError): + except (TimeoutError, socket.timeout, asyncio.IncompleteReadError): read_frame_buffer.write(self._read_buffer.getvalue()) self._read_buffer = read_frame_buffer self._read_buffer.seek(0) @@ -129,7 +125,7 @@ async def read(self, verify_frame_type=0, **kwargs): except (OSError, IOError, SSLError, socket.error) as exc: # Don't disconnect for ssl read time outs # http://bugs.python.org/issue10272 - if isinstance(exc, SSLError) and 'timed out' in str(exc): + if isinstance(exc, SSLError) and "timed out" in str(exc): raise socket.timeout() if get_errno(exc) not in _UNAVAIL: self.connected = False @@ -142,18 +138,28 @@ async def send_frame(self, channel, frame, **kwargs): if performative is None: data = header else: - encoded_channel = struct.pack('>H', channel) + encoded_channel = struct.pack(">H", channel) data = header + encoded_channel + performative await self.write(data) - #_LOGGER.info("OCH%d -> %r", channel, frame) + # _LOGGER.info("OCH%d -> %r", channel, frame) + class AsyncTransport(AsyncTransportMixin): """Common superclass for TCP and SSL transports.""" - def __init__(self, host, port=AMQP_PORT, connect_timeout=None, - read_timeout=None, write_timeout=None, ssl=False, - socket_settings=None, raise_on_initial_eintr=True, **kwargs): + def __init__( + self, + host, + port=AMQP_PORT, + connect_timeout=None, + read_timeout=None, + write_timeout=None, + ssl=False, + socket_settings=None, + raise_on_initial_eintr=True, + **kwargs + ): self.connected = False self.sock = None self.reader = None @@ -161,7 +167,7 @@ def __init__(self, host, port=AMQP_PORT, connect_timeout=None, self.raise_on_initial_eintr = raise_on_initial_eintr self._read_buffer = BytesIO() self.host, self.port = to_host_port(host, port) - + self.connect_timeout = connect_timeout self.read_timeout = read_timeout self.write_timeout = write_timeout @@ -174,19 +180,23 @@ def _build_ssl_opts(self, sslopts): if sslopts in [True, False, None, {}]: return sslopts try: - if 'context' in sslopts: - return self._build_ssl_context(sslopts, **sslopts.pop('context')) - ssl_version = sslopts.get('ssl_version') + if "context" in sslopts: + return self._build_ssl_context(sslopts, **sslopts.pop("context")) + ssl_version = sslopts.get("ssl_version") if ssl_version is None: ssl_version = ssl.PROTOCOL_TLS # Set SNI headers if supported - server_hostname = sslopts.get('server_hostname') - if (server_hostname is not None) and (hasattr(ssl, 'HAS_SNI') and ssl.HAS_SNI) and (hasattr(ssl, 'SSLContext')): + server_hostname = sslopts.get("server_hostname") + if ( + (server_hostname is not None) + and (hasattr(ssl, "HAS_SNI") and ssl.HAS_SNI) + and (hasattr(ssl, "SSLContext")) + ): context = ssl.SSLContext(ssl_version) - cert_reqs = sslopts.get('cert_reqs', ssl.CERT_REQUIRED) - certfile = sslopts.get('certfile') - keyfile = sslopts.get('keyfile') + cert_reqs = sslopts.get("cert_reqs", ssl.CERT_REQUIRED) + certfile = sslopts.get("certfile") + keyfile = sslopts.get("keyfile") context.verify_mode = cert_reqs if cert_reqs != ssl.CERT_NONE: context.check_hostname = True @@ -195,7 +205,7 @@ def _build_ssl_opts(self, sslopts): return context return True except TypeError: - raise TypeError('SSL configuration must be a dictionary, or the value True.') + raise TypeError("SSL configuration must be a dictionary, or the value True.") def _build_ssl_context(self, sslopts, check_hostname=None, **ctx_options): ctx = ssl.create_default_context(**ctx_options) @@ -211,12 +221,12 @@ async def connect(self): return await self._connect(self.host, self.port, self.connect_timeout) self._init_socket( - self.socket_settings, self.read_timeout, self.write_timeout, + self.socket_settings, + self.read_timeout, + self.write_timeout, ) self.reader, self.writer = await asyncio.open_connection( - sock=self.sock, - ssl=self.sslopts, - server_hostname=self.host if self.sslopts else None + sock=self.sock, ssl=self.sslopts, server_hostname=self.host if self.sslopts else None ) # we've sent the banner; signal connect # EINTR, EAGAIN, EWOULDBLOCK would signal that the banner @@ -245,8 +255,7 @@ async def _connect(self, host, port, timeout): for n, family in enumerate(addr_types): # first, resolve the address for a single address family try: - entries = await self.loop.getaddrinfo( - host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP) + entries = await self.loop.getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM, proto=SOL_TCP) entries_num = len(entries) except socket.gaierror: # we may have depleted all our options @@ -254,10 +263,7 @@ async def _connect(self, host, port, timeout): # if getaddrinfo succeeded before for another address # family, reraise the previous socket.error since it's more # relevant to users - raise (e - if e is not None - else socket.error( - "failed to resolve broker hostname")) + raise e if e is not None else socket.error("failed to resolve broker hostname") continue # pragma: no cover # now that we have address(es) for the hostname, connect to broker @@ -305,7 +311,7 @@ def _get_tcp_socket_defaults(self, sock): tcp_opts = {} for opt in KNOWN_TCP_OPTS: enum = None - if opt == 'TCP_USER_TIMEOUT': + if opt == "TCP_USER_TIMEOUT": try: from socket import TCP_USER_TIMEOUT as enum except ImportError: @@ -318,8 +324,7 @@ def _get_tcp_socket_defaults(self, sock): if opt in DEFAULT_SOCKET_SETTINGS: tcp_opts[enum] = DEFAULT_SOCKET_SETTINGS[opt] elif hasattr(socket, opt): - tcp_opts[enum] = sock.getsockopt( - SOL_TCP, getattr(socket, opt)) + tcp_opts[enum] = sock.getsockopt(SOL_TCP, getattr(socket, opt)) return tcp_opts def _set_socket_options(self, socket_settings): @@ -329,8 +334,7 @@ def _set_socket_options(self, socket_settings): for opt, val in tcp_opts.items(): self.sock.setsockopt(SOL_TCP, opt, val) - async def _read(self, toread, initial=False, buffer=None, - _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)): + async def _read(self, toread, initial=False, buffer=None, _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)): # According to SSL_read(3), it can at most return 16kb of data. # Thus, we use an internal read buffer like TCPTransport._read # to get the exact number of bytes wanted. @@ -342,16 +346,16 @@ async def _read(self, toread, initial=False, buffer=None, try: while toread: try: - view[nbytes:nbytes + toread] = await self.reader.readexactly(toread) + view[nbytes : nbytes + toread] = await self.reader.readexactly(toread) nbytes = toread except asyncio.IncompleteReadError as exc: pbytes = len(exc.partial) - view[nbytes:nbytes + pbytes] = exc.partial + view[nbytes : nbytes + pbytes] = exc.partial nbytes = pbytes except socket.error as exc: # ssl.sock.read may cause a SSLerror without errno # http://bugs.python.org/issue10272 - if isinstance(exc, SSLError) and 'timed out' in str(exc): + if isinstance(exc, SSLError) and "timed out" in str(exc): raise socket.timeout() # ssl.sock.read may cause ENOENT if the # operation couldn't be performed (Issue celery#1414). @@ -361,7 +365,7 @@ async def _read(self, toread, initial=False, buffer=None, continue raise if not nbytes: - raise IOError('Server unexpectedly closed connection') + raise IOError("Server unexpectedly closed connection") length += nbytes toread -= nbytes @@ -412,13 +416,15 @@ async def negotiate(self): await self.write(TLS_HEADER_FRAME) channel, returned_header = await self.receive_frame(verify_frame_type=None) if returned_header[1] == TLS_HEADER_FRAME: - raise ValueError("Mismatching TLS header protocol. Excpected: {}, received: {}".format( - TLS_HEADER_FRAME, returned_header[1])) + raise ValueError( + "Mismatching TLS header protocol. Excpected: {}, received: {}".format( + TLS_HEADER_FRAME, returned_header[1] + ) + ) class WebSocketTransportAsync(AsyncTransportMixin): - def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs - ): + def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs): self._read_buffer = BytesIO() self.loop = get_running_loop() self.socket_lock = asyncio.Lock() @@ -427,19 +433,20 @@ def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, ** self._custom_endpoint = kwargs.get("custom_endpoint") self.host = host self.ws = None - self._http_proxy = kwargs.get('http_proxy', None) + self._http_proxy = kwargs.get("http_proxy", None) async def connect(self): http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None if self._http_proxy: - http_proxy_host = self._http_proxy['proxy_hostname'] - http_proxy_port = self._http_proxy['proxy_port'] - username = self._http_proxy.get('username', None) - password = self._http_proxy.get('password', None) + http_proxy_host = self._http_proxy["proxy_hostname"] + http_proxy_port = self._http_proxy["proxy_port"] + username = self._http_proxy.get("username", None) + password = self._http_proxy.get("password", None) if username or password: http_proxy_auth = (username, password) try: from websocket import create_connection + self.ws = create_connection( url="wss://{}".format(self._custom_endpoint or self.host), subprotocols=[AMQP_WS_SUBPROTOCOL], @@ -448,12 +455,12 @@ async def connect(self): sslopt=self.sslopts, http_proxy_host=http_proxy_host, http_proxy_port=http_proxy_port, - http_proxy_auth=http_proxy_auth + http_proxy_auth=http_proxy_auth, ) except ImportError: raise ValueError("Please install websocket-client library to use websocket transport.") - async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-arguments + async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-arguments """Read exactly n bytes from the peer.""" from websocket import WebSocketTimeoutException @@ -464,20 +471,18 @@ async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argume n -= nbytes try: while n: - data = await self.loop.run_in_executor( - None, self.ws.recv - ) + data = await self.loop.run_in_executor(None, self.ws.recv) if len(data) <= n: - view[length: length + len(data)] = data + view[length : length + len(data)] = data n -= len(data) else: - view[length: length + n] = data[0:n] + view[length : length + n] = data[0:n] self._read_buffer = BytesIO(data[n:]) n = 0 - return view - except WebSocketTimeoutException as wex: + return view + except WebSocketTimeoutException: raise TimeoutError() def close(self): @@ -491,6 +496,4 @@ async def write(self, s): See http://tools.ietf.org/html/rfc5234 http://tools.ietf.org/html/rfc6455#section-5.2 """ - await self.loop.run_in_executor( - None, self.ws.send_binary, s - ) + await self.loop.run_in_executor(None, self.ws.send_binary, s) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py index 6fb937867295..43d7803c87d6 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/authentication.py @@ -5,7 +5,6 @@ #------------------------------------------------------------------------- import time -import urllib from collections import namedtuple from functools import partial @@ -20,12 +19,6 @@ AUTH_TYPE_SASL_PLAIN ) -try: - from urlparse import urlparse - from urllib import quote_plus # type: ignore -except ImportError: - from urllib.parse import urlparse, quote_plus - AccessToken = namedtuple("AccessToken", ["token", "expires_on"]) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py index ced290d75b70..27dfa11b0a9b 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/cbs.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- import logging from datetime import datetime @@ -10,12 +10,7 @@ from .utils import utc_now, utc_from_timestamp from .management_link import ManagementLink from .message import Message, Properties -from .error import ( - AuthenticationException, - ErrorCondition, - TokenAuthFailure, - TokenExpired -) +from .error import AuthenticationException, ErrorCondition, TokenAuthFailure, TokenExpired from .constants import ( CbsState, CbsAuthState, @@ -26,7 +21,7 @@ CBS_OPERATION, ManagementExecuteOperationResult, ManagementOpenResult, - DEFAULT_AUTH_TIMEOUT + DEFAULT_AUTH_TIMEOUT, ) _LOGGER = logging.getLogger(__name__) @@ -42,25 +37,19 @@ def check_expiration_and_refresh_status(expires_on, refresh_window): def check_put_timeout_status(auth_timeout, token_put_time): if auth_timeout > 0: return (int(utc_now().timestamp()) - token_put_time) >= auth_timeout - else: - return False + return False -class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes - def __init__( - self, - session, - auth, - **kwargs - ): +class CBSAuthenticator(object): # pylint:disable=too-many-instance-attributes + def __init__(self, session, auth, **kwargs): self._session = session self._connection = self._session._connection self._mgmt_link = self._session.create_request_response_link_pair( - endpoint='$cbs', + endpoint="$cbs", on_amqp_management_open_complete=self._on_amqp_management_open_complete, on_amqp_management_error=self._on_amqp_management_error, - status_code_field=b'status-code', - status_description_field=b'status-description' + status_code_field=b"status-code", + status_description_field=b"status-description", ) # type: ManagementLink if not auth.get_token or not callable(auth.get_token): @@ -89,15 +78,15 @@ def _put_token(self, token, token_type, audience, expires_on=None): CBS_NAME: audience, CBS_OPERATION: CBS_PUT_TOKEN, CBS_TYPE: token_type, - CBS_EXPIRATION: expires_on - } + CBS_EXPIRATION: expires_on, + }, ) self._mgmt_link.execute_operation( message, self._on_execute_operation_complete, timeout=self._auth_timeout, operation=CBS_PUT_TOKEN, - type=token_type + type=token_type, ) self._mgmt_link.next_message_id += 1 @@ -108,12 +97,15 @@ def _on_amqp_management_open_complete(self, management_open_result): self.state = CbsState.ERROR _LOGGER.info( "Unexpected AMQP management open complete in OPEN, CBS error occurred on connection %r.", - self._connection._container_id # pylint:disable=protected-access + self._connection._container_id, # pylint:disable=protected-access ) elif self.state == CbsState.OPENING: self.state = CbsState.OPEN if management_open_result == ManagementOpenResult.OK else CbsState.CLOSED - _LOGGER.info("CBS for connection %r completed opening with status: %r", - self._connection._container_id, management_open_result) # pylint:disable=protected-access + _LOGGER.info( + "CBS for connection %r completed opening with status: %r", + self._connection._container_id, # pylint: disable=protected-access + management_open_result, + ) # pylint:disable=protected-access def _on_amqp_management_error(self): if self.state == CbsState.CLOSED: @@ -121,22 +113,26 @@ def _on_amqp_management_error(self): elif self.state == CbsState.OPENING: self.state = CbsState.ERROR self._mgmt_link.close() - _LOGGER.info("CBS for connection %r failed to open with status: %r", - self._connection._container_id, ManagementOpenResult.ERROR) # pylint:disable=protected-access + _LOGGER.info( + "CBS for connection %r failed to open with status: %r", + self._connection._container_id, + ManagementOpenResult.ERROR, + ) # pylint:disable=protected-access elif self.state == CbsState.OPEN: self.state = CbsState.ERROR - _LOGGER.info("CBS error occurred on connection %r.", self._connection._container_id) # pylint:disable=protected-access + _LOGGER.info( + "CBS error occurred on connection %r.", self._connection._container_id + ) # pylint:disable=protected-access def _on_execute_operation_complete( - self, + self, execute_operation_result, status_code, status_description, message, error_condition=None + ): # TODO: message and error_condition never used + _LOGGER.info( + "CBS Put token result (%r), status code: %s, status_description: %s.", execute_operation_result, status_code, status_description, - message, - error_condition=None - ): # TODO: message and error_condition never used - _LOGGER.info("CBS Put token result (%r), status code: %s, status_description: %s.", - execute_operation_result, status_code, status_description) + ) self._token_status_code = status_code self._token_status_description = status_description @@ -152,15 +148,17 @@ def _on_execute_operation_complete( def _update_status(self): if self.auth_state == CbsAuthState.OK or self.auth_state == CbsAuthState.REFRESH_REQUIRED: - _LOGGER.debug('update_status In refresh required or OK.') - is_expired, is_refresh_required = check_expiration_and_refresh_status(self._expires_on, self._refresh_window) - _LOGGER.debug('is expired == %r, is refresh required == %r', is_expired, is_refresh_required) + _LOGGER.debug("update_status In refresh required or OK.") + is_expired, is_refresh_required = check_expiration_and_refresh_status( + self._expires_on, self._refresh_window + ) + _LOGGER.debug("is expired == %r, is refresh required == %r", is_expired, is_refresh_required) if is_expired: self.auth_state = CbsAuthState.EXPIRED elif is_refresh_required: self.auth_state = CbsAuthState.REFRESH_REQUIRED elif self.auth_state == CbsAuthState.IN_PROGRESS: - _LOGGER.debug('In update status, in progress. token put time: %r', self._token_put_time) + _LOGGER.debug("In update status, in progress. token put time: %r", self._token_put_time) put_timeout = check_put_timeout_status(self._auth_timeout, self._token_put_time) if put_timeout: self.auth_state = CbsAuthState.TIMEOUT @@ -175,7 +173,7 @@ def _cbs_link_ready(self): # Think how upper layer handle this exception + condition code raise AuthenticationException( condition=ErrorCondition.ClientError, - description="CBS authentication link is in broken status, please recreate the cbs link." + description="CBS authentication link is in broken status, please recreate the cbs link.", ) def open(self): @@ -210,30 +208,27 @@ def handle_token(self): if self.auth_state == CbsAuthState.IDLE: self.update_token() return False - elif self.auth_state == CbsAuthState.IN_PROGRESS: + if self.auth_state == CbsAuthState.IN_PROGRESS: return False - elif self.auth_state == CbsAuthState.OK: + if self.auth_state == CbsAuthState.OK: return True - elif self.auth_state == CbsAuthState.REFRESH_REQUIRED: - _LOGGER.info("Token on connection %r will expire soon - attempting to refresh.", - self._connection._container_id) # pylint:disable=protected-access + if self.auth_state == CbsAuthState.REFRESH_REQUIRED: + _LOGGER.info( + "Token on connection %r will expire soon - attempting to refresh.", self._connection._container_id + ) # pylint:disable=protected-access self.update_token() return False - elif self.auth_state == CbsAuthState.FAILURE: + if self.auth_state == CbsAuthState.FAILURE: raise AuthenticationException( - condition=ErrorCondition.InternalError, - description="Failed to open CBS authentication link." + condition=ErrorCondition.InternalError, description="Failed to open CBS authentication link." ) - elif self.auth_state == CbsAuthState.ERROR: + if self.auth_state == CbsAuthState.ERROR: raise TokenAuthFailure( self._token_status_code, self._token_status_description, - encoding=self._encoding # TODO: drop off all the encodings + encoding=self._encoding, # TODO: drop off all the encodings ) - elif self.auth_state == CbsAuthState.TIMEOUT: + if self.auth_state == CbsAuthState.TIMEOUT: raise TimeoutError("Authentication attempt timed-out.") - elif self.auth_state == CbsAuthState.EXPIRED: - raise TokenExpired( - condition=ErrorCondition.InternalError, - description="CBS Authentication Expired." - ) + if self.auth_state == CbsAuthState.EXPIRED: + raise TokenExpired(condition=ErrorCondition.InternalError, description="CBS Authentication Expired.") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py index db0a6e38d520..e1c088f41c8d 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/client.py @@ -1,18 +1,18 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- # pylint: disable=too-many-lines # TODO: Check types of kwargs (issue exists for this) import logging +import queue import time import uuid -import certifi -import queue from functools import partial from typing import Any, Dict, Optional, Tuple, Union, overload +import certifi from typing_extensions import Literal from ._connection import Connection @@ -59,7 +59,6 @@ from .management_operation import ManagementOperation from .cbs import CBSAuthenticator -from .authentication import _CBSAuth _logger = logging.getLogger(__name__) @@ -187,16 +186,16 @@ def __init__(self, hostname, **kwargs): self._handle_max = kwargs.pop('handle_max', None) # Link settings - self._send_settle_mode = kwargs.pop('send_settle_mode', SenderSettleMode.Unsettled) - self._receive_settle_mode = kwargs.pop('receive_settle_mode', ReceiverSettleMode.Second) - self._desired_capabilities = kwargs.pop('desired_capabilities', None) - self._on_attach = kwargs.pop('on_attach', None) + self._send_settle_mode = kwargs.pop("send_settle_mode", SenderSettleMode.Unsettled) + self._receive_settle_mode = kwargs.pop("receive_settle_mode", ReceiverSettleMode.Second) + self._desired_capabilities = kwargs.pop("desired_capabilities", None) + self._on_attach = kwargs.pop("on_attach", None) # transport - if kwargs.get('transport_type') is TransportType.Amqp and kwargs.get('http_proxy') is not None: + if kwargs.get("transport_type") is TransportType.Amqp and kwargs.get("http_proxy") is not None: raise ValueError("Http proxy settings can't be passed if transport_type is explicitly set to Amqp") - self._transport_type = kwargs.pop('transport_type', TransportType.Amqp) - self._http_proxy = kwargs.pop('http_proxy', None) + self._transport_type = kwargs.pop("transport_type", TransportType.Amqp) + self._http_proxy = kwargs.pop("http_proxy", None) # Custom Endpoint self._custom_endpoint_address = kwargs.get("custom_endpoint_address") @@ -225,7 +224,7 @@ def _client_run(self, **kwargs): self._connection.listen(wait=self._socket_timeout) def _close_link(self, **kwargs): - if self._link and not self._link._is_closed: + if self._link and not self._link._is_closed: # pylint: disable=protected-access self._link.detach(close=True) self._link = None @@ -257,8 +256,8 @@ def _do_retryable_operation(self, operation, *args, **kwargs): finally: end_time = time.time() if absolute_timeout > 0: - absolute_timeout -= (end_time - start_time) - raise retry_settings['history'][-1] + absolute_timeout -= end_time - start_time + raise retry_settings["history"][-1] def open(self, connection=None): """Open the client. The client can create a new Connection @@ -283,7 +282,7 @@ def open(self, connection=None): self._connection = Connection( "amqps://" + self._hostname, sasl_credential=self._auth.sasl, - ssl={'ca_certs':self._connection_verify or certifi.where()}, + ssl={"ca_certs": self._connection_verify or certifi.where()}, container_id=self._name, max_frame_size=self._max_frame_size, channel_max=self._channel_max, @@ -292,20 +291,17 @@ def open(self, connection=None): network_trace=self._network_trace, transport_type=self._transport_type, http_proxy=self._http_proxy, - custom_endpoint_address=self._custom_endpoint_address + custom_endpoint_address=self._custom_endpoint_address, ) self._connection.open() if not self._session: self._session = self._connection.create_session( - incoming_window=self._incoming_window, - outgoing_window=self._outgoing_window + incoming_window=self._incoming_window, outgoing_window=self._outgoing_window ) self._session.begin() if self._auth.auth_type == AUTH_TYPE_CBS: self._cbs_authenticator = CBSAuthenticator( - session=self._session, - auth=self._auth, - auth_timeout=self._auth_timeout + session=self._session, auth=self._auth, auth_timeout=self._auth_timeout ) self._cbs_authenticator.open() self._shutdown = False @@ -401,7 +397,7 @@ def mgmt_request(self, message, **kwargs): operation = kwargs.pop("operation", None) operation_type = kwargs.pop("operation_type", None) node = kwargs.pop("node", "$management") - timeout = kwargs.pop('timeout', 0) + timeout = kwargs.pop("timeout", 0) try: mgmt_link = self._mgmt_links[node] except KeyError: @@ -412,12 +408,9 @@ def mgmt_request(self, message, **kwargs): while not mgmt_link.ready(): self._connection.listen(wait=False) - operation_type = operation_type or b'empty' + operation_type = operation_type or b"empty" status, description, response = mgmt_link.execute( - message, - operation=operation, - operation_type=operation_type, - timeout=timeout + message, operation=operation, operation_type=operation_type, timeout=timeout ) return status, description, response @@ -538,7 +531,8 @@ def _client_ready(self): send_settle_mode=self._send_settle_mode, rcv_settle_mode=self._receive_settle_mode, max_message_size=self._max_message_size, - properties=self._link_properties) + properties=self._link_properties, + ) self._link.attach() return False if self._link.get_state().value != 3: # ATTACHED @@ -566,10 +560,7 @@ def _transfer_message(self, message_delivery, timeout=0): message_delivery.state = MessageDeliveryState.WaitingForSendAck on_send_complete = partial(self._on_send_complete, message_delivery) delivery = self._link.send_transfer( - message_delivery.message, - on_send_complete=on_send_complete, - timeout=timeout, - send_async=True + message_delivery.message, on_send_complete=on_send_complete, timeout=timeout, send_async=True ) return delivery @@ -596,13 +587,10 @@ def _on_send_complete(self, message_delivery, reason, state): message_delivery, condition=error_info[0][0], description=error_info[0][1], - info=error_info[0][2] + info=error_info[0][2], ) except TypeError: - self._process_send_error( - message_delivery, - condition=ErrorCondition.UnknownError - ) + self._process_send_error(message_delivery, condition=ErrorCondition.UnknownError) elif reason == LinkDeliverySettleReason.SETTLED: message_delivery.state = MessageDeliveryState.Ok elif reason == LinkDeliverySettleReason.TIMEOUT: @@ -610,20 +598,13 @@ def _on_send_complete(self, message_delivery, reason, state): message_delivery.error = TimeoutError("Sending message timed out.") else: # NotDelivered and other unknown errors - self._process_send_error( - message_delivery, - condition=ErrorCondition.UnknownError - ) + self._process_send_error(message_delivery, condition=ErrorCondition.UnknownError) def _send_message_impl(self, message, **kwargs): timeout = kwargs.pop("timeout", 0) expire_time = (time.time() + timeout) if timeout else None self.open() - message_delivery = _MessageDelivery( - message, - MessageDeliveryState.WaitingToBeSent, - expire_time - ) + message_delivery = _MessageDelivery(message, MessageDeliveryState.WaitingToBeSent, expire_time) while not self.client_ready(): time.sleep(0.05) @@ -631,9 +612,13 @@ def _send_message_impl(self, message, **kwargs): running = True while running and message_delivery.state not in MESSAGE_DELIVERY_DONE_STATES: running = self.do_work() - if message_delivery.state in (MessageDeliveryState.Error, MessageDeliveryState.Cancelled, MessageDeliveryState.Timeout): + if message_delivery.state in ( + MessageDeliveryState.Error, + MessageDeliveryState.Cancelled, + MessageDeliveryState.Timeout, + ): try: - raise message_delivery.error + raise message_delivery.error # pylint: disable=raising-bad-type except TypeError: # This is a default handler raise MessageException(condition=ErrorCondition.UnknownError, description="Send failed.") @@ -771,7 +756,7 @@ def _client_ready(self): on_transfer=self._message_received, properties=self._link_properties, desired_capabilities=self._desired_capabilities, - on_attach=self._on_attach + on_attach=self._on_attach, ) self._link.attach() return False @@ -879,10 +864,7 @@ def receive_message_batch(self, **kwargs): default is 0. :type timeout: float """ - return self._do_retryable_operation( - self._receive_message_batch_impl, - **kwargs - ) + return self._do_retryable_operation(self._receive_message_batch_impl, **kwargs) @overload def settle_messages( @@ -941,16 +923,16 @@ def settle_messages( ... def settle_messages(self, delivery_id: Union[int, Tuple[int, int]], outcome: str, **kwargs): - batchable = kwargs.pop('batchable', None) - if outcome.lower() == 'accepted': + batchable = kwargs.pop("batchable", None) + if outcome.lower() == "accepted": state = Accepted() - elif outcome.lower() == 'released': + elif outcome.lower() == "released": state = Released() - elif outcome.lower() == 'rejected': + elif outcome.lower() == "rejected": state = Rejected(**kwargs) - elif outcome.lower() == 'modified': + elif outcome.lower() == "modified": state = Modified(**kwargs) - elif outcome.lower() == 'received': + elif outcome.lower() == "received": state = Received(**kwargs) else: raise ValueError("Unrecognized message output: {}".format(outcome)) @@ -965,5 +947,5 @@ def settle_messages(self, delivery_id: Union[int, Tuple[int, int]], outcome: str settled=True, delivery_state=state, batchable=batchable, - wait=True + wait=True, ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py index c68cc05c3d6f..a2d0b4a240e7 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/endpoints.py @@ -143,8 +143,8 @@ class ApacheFilters(object): 'capabilities' ]) Source.__new__.__defaults__ = (None,) * len(Source._fields) -Source._code = 0x00000028 -Source._definition = ( +Source._code = 0x00000028 # pylint: disable=protected-access +Source._definition = ( # pylint: disable=protected-access FIELD("address", AMQPTypes.string, False, None, False), FIELD("durable", AMQPTypes.uint, False, "none", False), FIELD("expiry_policy", AMQPTypes.symbol, False, ExpiryPolicy.SessionEnd, False), @@ -227,9 +227,9 @@ class ApacheFilters(object): 'dynamic_node_properties', 'capabilities' ]) -Target._code = 0x00000029 -Target.__new__.__defaults__ = (None,) * len(Target._fields) -Target._definition = ( +Target._code = 0x00000029 # pylint: disable=protected-access +Target.__new__.__defaults__ = (None,) * len(Target._fields) # pylint: disable=protected-access +Target._definition = ( # pylint: disable=protected-access FIELD("address", AMQPTypes.string, False, None, False), FIELD("durable", AMQPTypes.uint, False, "none", False), FIELD("expiry_policy", AMQPTypes.symbol, False, ExpiryPolicy.SessionEnd, False), diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py index fc2b8cbfe5dc..96c7803fa0bb 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/error.py @@ -183,8 +183,8 @@ def get_backoff_time(self, settings, error): AMQPError = namedtuple('error', ['condition', 'description', 'info']) AMQPError.__new__.__defaults__ = (None,) * len(AMQPError._fields) -AMQPError._code = 0x0000001d -AMQPError._definition = ( +AMQPError._code = 0x0000001d # pylint: disable=protected-access +AMQPError._definition = ( # pylint: disable=protected-access FIELD('condition', AMQPTypes.symbol, True, None, False), FIELD('description', AMQPTypes.string, False, None, False), FIELD('info', FieldDefinition.fields, False, None, False), @@ -254,8 +254,11 @@ class AMQPSessionError(AMQPException): class AMQPLinkError(AMQPException): - """ + """Details of a Link-level error. + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. """ @@ -279,14 +282,20 @@ def __init__(self, condition, description=None, info=None): class AuthenticationException(AMQPException): - """ + """Details of a Authentication error. + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. """ class TokenExpired(AuthenticationException): - """ + """Details of a Token expiration error. + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. """ @@ -308,13 +317,21 @@ def __init__(self, status_code, status_description, **kwargs): class MessageException(AMQPException): - """ + """Details of a Message error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. """ class MessageSendFailed(MessageException): - """ + """Details of a Message send failed error. + + :param bytes condition: The error code. + :keyword str description: A description of the error. + :keyword dict info: A dictionary of additional data associated with the error. """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py index 31e6f395949f..d105486d5f57 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py @@ -1,52 +1,25 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- + -import threading -import struct from typing import Optional import uuid import logging -import time -from enum import Enum -from io import BytesIO -from urllib.parse import urlparse from .endpoints import Source, Target -from .constants import ( - DEFAULT_LINK_CREDIT, - SessionState, - SessionTransferState, - LinkDeliverySettleReason, - LinkState, - Role, - SenderSettleMode, - ReceiverSettleMode -) -from .performatives import ( - AttachFrame, - DetachFrame, - TransferFrame, - DispositionFrame, - FlowFrame, -) +from .constants import DEFAULT_LINK_CREDIT, SessionState, LinkState, Role, SenderSettleMode, ReceiverSettleMode +from .performatives import AttachFrame, DetachFrame -from .error import ( - ErrorCondition, - AMQPLinkError, - AMQPLinkRedirect, - AMQPConnectionError -) +from .error import ErrorCondition, AMQPLinkError, AMQPLinkRedirect, AMQPConnectionError _LOGGER = logging.getLogger(__name__) class Link(object): - """ - - """ + """ """ def __init__(self, session, handle, name, role, **kwargs): self.state = LinkState.DETACHED @@ -54,53 +27,61 @@ def __init__(self, session, handle, name, role, **kwargs): self.handle = handle self.remote_handle = None self.role = role - source_address = kwargs['source_address'] + source_address = kwargs["source_address"] target_address = kwargs["target_address"] - self.source = source_address if isinstance(source_address, Source) else Source( - address=kwargs['source_address'], - durable=kwargs.get('source_durable'), - expiry_policy=kwargs.get('source_expiry_policy'), - timeout=kwargs.get('source_timeout'), - dynamic=kwargs.get('source_dynamic'), - dynamic_node_properties=kwargs.get('source_dynamic_node_properties'), - distribution_mode=kwargs.get('source_distribution_mode'), - filters=kwargs.get('source_filters'), - default_outcome=kwargs.get('source_default_outcome'), - outcomes=kwargs.get('source_outcomes'), - capabilities=kwargs.get('source_capabilities') + self.source = ( + source_address + if isinstance(source_address, Source) + else Source( + address=kwargs["source_address"], + durable=kwargs.get("source_durable"), + expiry_policy=kwargs.get("source_expiry_policy"), + timeout=kwargs.get("source_timeout"), + dynamic=kwargs.get("source_dynamic"), + dynamic_node_properties=kwargs.get("source_dynamic_node_properties"), + distribution_mode=kwargs.get("source_distribution_mode"), + filters=kwargs.get("source_filters"), + default_outcome=kwargs.get("source_default_outcome"), + outcomes=kwargs.get("source_outcomes"), + capabilities=kwargs.get("source_capabilities"), + ) ) - self.target = target_address if isinstance(target_address,Target) else Target( - address=kwargs['target_address'], - durable=kwargs.get('target_durable'), - expiry_policy=kwargs.get('target_expiry_policy'), - timeout=kwargs.get('target_timeout'), - dynamic=kwargs.get('target_dynamic'), - dynamic_node_properties=kwargs.get('target_dynamic_node_properties'), - capabilities=kwargs.get('target_capabilities') + self.target = ( + target_address + if isinstance(target_address, Target) + else Target( + address=kwargs["target_address"], + durable=kwargs.get("target_durable"), + expiry_policy=kwargs.get("target_expiry_policy"), + timeout=kwargs.get("target_timeout"), + dynamic=kwargs.get("target_dynamic"), + dynamic_node_properties=kwargs.get("target_dynamic_node_properties"), + capabilities=kwargs.get("target_capabilities"), + ) ) - self.link_credit = kwargs.pop('link_credit', None) or DEFAULT_LINK_CREDIT + self.link_credit = kwargs.pop("link_credit", None) or DEFAULT_LINK_CREDIT self.current_link_credit = self.link_credit - self.send_settle_mode = kwargs.pop('send_settle_mode', SenderSettleMode.Mixed) - self.rcv_settle_mode = kwargs.pop('rcv_settle_mode', ReceiverSettleMode.First) - self.unsettled = kwargs.pop('unsettled', None) - self.incomplete_unsettled = kwargs.pop('incomplete_unsettled', None) - self.initial_delivery_count = kwargs.pop('initial_delivery_count', 0) + self.send_settle_mode = kwargs.pop("send_settle_mode", SenderSettleMode.Mixed) + self.rcv_settle_mode = kwargs.pop("rcv_settle_mode", ReceiverSettleMode.First) + self.unsettled = kwargs.pop("unsettled", None) + self.incomplete_unsettled = kwargs.pop("incomplete_unsettled", None) + self.initial_delivery_count = kwargs.pop("initial_delivery_count", 0) self.delivery_count = self.initial_delivery_count self.received_delivery_id = None - self.max_message_size = kwargs.pop('max_message_size', None) + self.max_message_size = kwargs.pop("max_message_size", None) self.remote_max_message_size = None - self.available = kwargs.pop('available', None) - self.properties = kwargs.pop('properties', None) + self.available = kwargs.pop("available", None) + self.properties = kwargs.pop("properties", None) self.offered_capabilities = None - self.desired_capabilities = kwargs.pop('desired_capabilities', None) + self.desired_capabilities = kwargs.pop("desired_capabilities", None) - self.network_trace = kwargs['network_trace'] - self.network_trace_params = kwargs['network_trace_params'] - self.network_trace_params['link'] = self.name + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["link"] = self.name self._session = session self._is_closed = False - self._on_link_state_change = kwargs.get('on_link_state_change') - self._on_attach = kwargs.get('on_attach') + self._on_link_state_change = kwargs.get("on_link_state_change") + self._on_attach = kwargs.get("on_attach") self._error = None def __enter__(self): @@ -113,7 +94,7 @@ def __exit__(self, *args): @classmethod def from_incoming_frame(cls, session, handle, frame): # check link_create_from_endpoint in C lib - raise NotImplementedError('Pending') # TODO: Assuming we establish all links for now... + raise NotImplementedError("Pending") # TODO: Assuming we establish all links for now... def get_state(self): try: @@ -127,10 +108,7 @@ def _check_if_closed(self): try: raise self._error except TypeError: - raise AMQPConnectionError( - condition=ErrorCondition.InternalError, - description="Link already closed." - ) + raise AMQPConnectionError(condition=ErrorCondition.InternalError, description="Link already closed.") def _set_state(self, new_state): # type: (LinkState) -> None @@ -146,7 +124,7 @@ def _set_state(self, new_state): pass except Exception as e: # pylint: disable=broad-except _LOGGER.error("Link state change callback failed: '%r'", e, extra=self.network_trace_params) - + def _on_session_state_change(self): if self._session.state == SessionState.MAPPED: if not self._is_closed and self.state == LinkState.DETACHED: @@ -171,20 +149,20 @@ def _outgoing_attach(self): max_message_size=self.max_message_size, offered_capabilities=self.offered_capabilities if self.state == LinkState.ATTACH_RCVD else None, desired_capabilities=self.desired_capabilities if self.state == LinkState.DETACHED else None, - properties=self.properties + properties=self.properties, ) if self.network_trace: _LOGGER.info("-> %r", attach_frame, extra=self.network_trace_params) - self._session._outgoing_attach(attach_frame) + self._session._outgoing_attach(attach_frame) # pylint: disable=protected-access def _incoming_attach(self, frame): if self.network_trace: _LOGGER.info("<- %r", AttachFrame(*frame), extra=self.network_trace_params) if self._is_closed: raise ValueError("Invalid link") - elif not frame[5] or not frame[6]: + if not frame[5] or not frame[6]: _LOGGER.info("Cannot get source or target. Detaching link") - self._set_state(LinkState.DETACHED) + self._set_state(LinkState.DETACHED) raise ValueError("Invalid link") self.remote_handle = frame[1] # handle self.remote_max_message_size = frame[10] # max_message_size @@ -209,19 +187,19 @@ def _incoming_attach(self, frame): def _outgoing_flow(self, **kwargs): flow_frame = { - 'handle': self.handle, - 'delivery_count': self.delivery_count, - 'link_credit': self.current_link_credit, - 'available': kwargs.get('available'), - 'drain': kwargs.get('drain'), - 'echo': kwargs.get('echo'), - 'properties': kwargs.get('properties') + "handle": self.handle, + "delivery_count": self.delivery_count, + "link_credit": self.current_link_credit, + "available": kwargs.get("available"), + "drain": kwargs.get("drain"), + "echo": kwargs.get("echo"), + "properties": kwargs.get("properties"), } - self._session._outgoing_flow(flow_frame) + self._session._outgoing_flow(flow_frame) # pylint: disable=protected-access def _incoming_flow(self, frame): pass - + def _incoming_disposition(self, frame): pass @@ -229,7 +207,7 @@ def _outgoing_detach(self, close=False, error=None): detach_frame = DetachFrame(handle=self.handle, closed=close, error=error) if self.network_trace: _LOGGER.info("-> %r", detach_frame, extra=self.network_trace_params) - self._session._outgoing_detach(detach_frame) + self._session._outgoing_detach(detach_frame) # pylint: disable=protected-access if close: self._is_closed = True @@ -273,11 +251,6 @@ def detach(self, close=False, error=None): _LOGGER.info("An error occurred when detaching the link: %r", exc) self._set_state(LinkState.DETACHED) - def flow( - self, - *, - link_credit: Optional[int] = None, - **kwargs - ) -> None: + def flow(self, *, link_credit: Optional[int] = None, **kwargs) -> None: self.current_link_credit = link_credit if link_credit is not None else self.link_credit self._outgoing_flow(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py index e7e710a28e3c..ac1b7bc08029 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py @@ -18,12 +18,11 @@ ReceiverSettleMode, ManagementExecuteOperationResult, ManagementOpenResult, - SEND_DISPOSITION_ACCEPT, SEND_DISPOSITION_REJECT, MessageDeliveryState ) -from .error import ErrorResponse, AMQPException, ErrorCondition -from .message import Message, Properties, _MessageDelivery +from .error import AMQPException, ErrorCondition +from .message import Properties, _MessageDelivery _LOGGER = logging.getLogger(__name__) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py index 3ccb6544af34..d9e9080ea260 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_operation.py @@ -9,10 +9,7 @@ from functools import partial from .management_link import ManagementLink -from .message import Message from .error import ( - AMQPException, - AMQPConnectionError, AMQPLinkError, ErrorCondition ) @@ -107,7 +104,7 @@ def execute(self, message, operation=None, operation_type=None, timeout=0): if self._mgmt_error: self._responses.pop(operation_id) - raise self._mgmt_error + raise self._mgmt_error # pylint: disable=raising-bad-type response = self._responses.pop(operation_id) return response @@ -118,7 +115,7 @@ def open(self): def ready(self): try: - raise self._mgmt_error + raise self._mgmt_error # pylint: disable=raising-bad-type except TypeError: pass diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py index 072379147a95..2923ddaebc19 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py @@ -1,53 +1,31 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging from typing import Optional, Union from ._decode import decode_payload -from .constants import DEFAULT_LINK_CREDIT, Role -from .endpoints import Target from .link import Link -from .message import Message, Properties, Header -from .constants import ( - DEFAULT_LINK_CREDIT, - SessionState, - SessionTransferState, - LinkDeliverySettleReason, - LinkState -) -from .performatives import ( - AttachFrame, - DetachFrame, - TransferFrame, - DispositionFrame, - FlowFrame, -) -from .outcomes import ( - Received, - Accepted, - Rejected, - Released, - Modified -) +from .constants import LinkState, Role +from .performatives import TransferFrame, DispositionFrame +from .outcomes import Received, Accepted, Rejected, Released, Modified _LOGGER = logging.getLogger(__name__) class ReceiverLink(Link): - def __init__(self, session, handle, source_address, **kwargs): - name = kwargs.pop('name', None) or str(uuid.uuid4()) + name = kwargs.pop("name", None) or str(uuid.uuid4()) role = Role.Receiver - if 'target_address' not in kwargs: - kwargs['target_address'] = "receiver-link-{}".format(name) + if "target_address" not in kwargs: + kwargs["target_address"] = "receiver-link-{}".format(name) super(ReceiverLink, self).__init__(session, handle, name, role, source_address=source_address, **kwargs) - self._on_transfer = kwargs.pop('on_transfer') + self._on_transfer = kwargs.pop("on_transfer") self._received_payload = bytearray() def _process_incoming_message(self, frame, message): @@ -89,56 +67,45 @@ def _incoming_transfer(self, frame): self._outgoing_disposition(first=frame[1], settled=True, state=delivery_state) def _wait_for_response(self, wait: Union[bool, float]) -> None: - if wait == True: - self._session._connection.listen(wait=False) + if wait is True: + self._session._connection.listen(wait=False) # pylint: disable=protected-access if self.state == LinkState.ERROR: - raise self._error + raise self._error elif wait: - self._session._connection.listen(wait=wait) + self._session._connection.listen(wait=wait) # pylint: disable=protected-access if self.state == LinkState.ERROR: - raise self._error + raise self._error def _outgoing_disposition( - self, - first: int, - last: Optional[int], - settled: Optional[bool], - state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], - batchable: Optional[bool] + self, + first: int, + last: Optional[int], + settled: Optional[bool], + state: Optional[Union[Received, Accepted, Rejected, Released, Modified]], + batchable: Optional[bool], ): disposition_frame = DispositionFrame( - role=self.role, - first=first, - last=last, - settled=settled, - state=state, - batchable=batchable + role=self.role, first=first, last=last, settled=settled, state=state, batchable=batchable ) if self.network_trace: _LOGGER.info("-> %r", DispositionFrame(*disposition_frame), extra=self.network_trace_params) - self._session._outgoing_disposition(disposition_frame) + self._session._outgoing_disposition(disposition_frame) # pylint: disable=protected-access def attach(self): super().attach() self._received_payload = bytearray() def send_disposition( - self, - *, - wait: Union[bool, float] = False, - first_delivery_id: int, - last_delivery_id: Optional[int] = None, - settled: Optional[bool] = None, - delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, - batchable: Optional[bool] = None - ): + self, + *, + wait: Union[bool, float] = False, + first_delivery_id: int, + last_delivery_id: Optional[int] = None, + settled: Optional[bool] = None, + delivery_state: Optional[Union[Received, Accepted, Rejected, Released, Modified]] = None, + batchable: Optional[bool] = None + ): if self._is_closed: raise ValueError("Link already closed.") - self._outgoing_disposition( - first_delivery_id, - last_delivery_id, - settled, - delivery_state, - batchable - ) + self._outgoing_disposition(first_delivery_id, last_delivery_id, settled, delivery_state, batchable) self._wait_for_response(wait) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py index d2bbd889e890..b927cf5b627a 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sasl.py @@ -4,18 +4,9 @@ # license information. #-------------------------------------------------------------------------- -import struct -from enum import Enum - from ._transport import SSLTransport, WebSocketTransport, AMQPS_PORT -from .types import AMQPTypes, TYPE, VALUE -from .constants import FIELD, SASLCode, SASL_HEADER_FRAME, TransportType, WEBSOCKET_PORT -from .performatives import ( - SASLOutcome, - SASLResponse, - SASLChallenge, - SASLInit -) +from .constants import SASLCode, SASL_HEADER_FRAME, WEBSOCKET_PORT +from .performatives import SASLInit _SASL_FRAME_TYPE = b'\x01' @@ -92,8 +83,7 @@ def _negotiate(self): raise NotImplementedError("Unsupported SASL challenge") if fields[0] == SASLCode.Ok: # code return - else: - raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) + raise ValueError("SASL negotiation failed.\nOutcome: {}\nDetails: {}".format(*fields)) class SASLTransport(SSLTransport, SASLTransportMixin): diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py index 75a46b3196b9..7ee12e3d7b98 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/sender.py @@ -1,8 +1,8 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import struct import uuid import logging @@ -10,14 +10,7 @@ from ._encode import encode_payload from .link import Link -from .constants import ( - SessionTransferState, - LinkDeliverySettleReason, - LinkState, - Role, - SenderSettleMode, - SessionState -) +from .constants import SessionTransferState, LinkDeliverySettleReason, LinkState, Role, SenderSettleMode, SessionState from .performatives import ( TransferFrame, ) @@ -27,33 +20,31 @@ class PendingDelivery(object): - def __init__(self, **kwargs): - self.message = kwargs.get('message') + self.message = kwargs.get("message") self.sent = False self.frame = None - self.on_delivery_settled = kwargs.get('on_delivery_settled') + self.on_delivery_settled = kwargs.get("on_delivery_settled") self.start = time.time() self.transfer_state = None - self.timeout = kwargs.get('timeout') - self.settled = kwargs.get('settled', False) + self.timeout = kwargs.get("timeout") + self.settled = kwargs.get("settled", False) def on_settled(self, reason, state): if self.on_delivery_settled and not self.settled: try: self.on_delivery_settled(reason, state) - except Exception as e: # pylint:disable=broad-except + except Exception as e: # pylint:disable=broad-except _LOGGER.warning("Message 'on_send_complete' callback failed: %r", e) self.settled = True class SenderLink(Link): - def __init__(self, session, handle, target_address, **kwargs): - name = kwargs.pop('name', None) or str(uuid.uuid4()) + name = kwargs.pop("name", None) or str(uuid.uuid4()) role = Role.Sender - if 'source_address' not in kwargs: - kwargs['source_address'] = "sender-link-{}".format(name) + if "source_address" not in kwargs: + kwargs["source_address"] = "sender-link-{}".format(name) super(SenderLink, self).__init__(session, handle, name, role, target_address=target_address, **kwargs) self._pending_deliveries = [] @@ -90,22 +81,24 @@ def _outgoing_transfer(self, delivery): encode_payload(output, delivery.message) delivery_count = self.delivery_count + 1 delivery.frame = { - 'handle': self.handle, - 'delivery_tag': struct.pack('>I', abs(delivery_count)), - 'message_format': delivery.message._code, # pylint:disable=protected-access - 'settled': delivery.settled, - 'more': False, - 'rcv_settle_mode': None, - 'state': None, - 'resume': None, - 'aborted': None, - 'batchable': None, - 'payload': output + "handle": self.handle, + "delivery_tag": struct.pack(">I", abs(delivery_count)), + "message_format": delivery.message._code, # pylint:disable=protected-access + "settled": delivery.settled, + "more": False, + "rcv_settle_mode": None, + "state": None, + "resume": None, + "aborted": None, + "batchable": None, + "payload": output, } if self.network_trace: - _LOGGER.info("-> %r", TransferFrame(delivery_id='', **delivery.frame), extra=self.network_trace_params) # pylint:disable=line-to-long + _LOGGER.info( + "-> %r", TransferFrame(delivery_id="", **delivery.frame), extra=self.network_trace_params + ) _LOGGER.info(" %r", delivery.message, extra=self.network_trace_params) - self._session._outgoing_transfer(delivery) # pylint:disable=protected-access + self._session._outgoing_transfer(delivery) # pylint:disable=protected-access sent_and_settled = False if delivery.transfer_state == SessionTransferState.OKAY: self.delivery_count = delivery_count @@ -125,7 +118,7 @@ def _incoming_disposition(self, frame): settled_ids = list(range(frame[1], range_end)) unsettled = [] for delivery in self._pending_deliveries: - if delivery.sent and delivery.frame['delivery_id'] in settled_ids: + if delivery.sent and delivery.frame["delivery_id"] in settled_ids: delivery.on_settled(LinkDeliverySettleReason.DISPOSITION_RECEIVED, frame[4]) # state continue unsettled.append(delivery) @@ -135,7 +128,7 @@ def _remove_pending_deliveries(self): for delivery in self._pending_deliveries: delivery.on_settled(LinkDeliverySettleReason.NOT_DELIVERED, None) self._pending_deliveries = [] - + def _on_session_state_change(self): if self._session.state == SessionState.DISCARDING: self._remove_pending_deliveries() @@ -163,14 +156,14 @@ def send_transfer(self, message, *, send_async=False, **kwargs): if self.state != LinkState.ATTACHED: raise AMQPLinkError( # TODO: should we introduce MessageHandler to indicate the handler is in wrong state condition=ErrorCondition.ClientError, # TODO: should this be a ClientError? - description="Link is not attached." + description="Link is not attached.", ) settled = self.send_settle_mode == SenderSettleMode.Settled if self.send_settle_mode == SenderSettleMode.Mixed: - settled = kwargs.pop('settled', True) + settled = kwargs.pop("settled", True) delivery = PendingDelivery( - on_delivery_settled=kwargs.get('on_send_complete'), - timeout=kwargs.get('timeout'), + on_delivery_settled=kwargs.get("on_send_complete"), + timeout=kwargs.get("timeout"), message=message, settled=settled, ) @@ -191,6 +184,7 @@ def cancel_transfer(self, delivery): if delivery.sent: raise MessageException( ErrorCondition.ClientError, - message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.") + message="Transfer cannot be cancelled. Message has already been sent and awaiting disposition.", + ) delivery.on_settled(LinkDeliverySettleReason.CANCELLED, None) self._pending_deliveries.pop(index) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py index 590c45c68d65..7caa2466e8c8 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/session.py @@ -1,12 +1,11 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- +# -------------------------------------------------------------------------- import uuid import logging -from enum import Enum import time from .constants import ( @@ -21,15 +20,7 @@ from .sender import SenderLink from .receiver import ReceiverLink from .management_link import ManagementLink -from .performatives import ( - BeginFrame, - EndFrame, - FlowFrame, - AttachFrame, - DetachFrame, - TransferFrame, - DispositionFrame -) +from .performatives import BeginFrame, EndFrame, FlowFrame, TransferFrame, DispositionFrame from ._encode import encode_frame _LOGGER = logging.getLogger(__name__) @@ -48,27 +39,27 @@ class Session(object): """ def __init__(self, connection, channel, **kwargs): - self.name = kwargs.pop('name', None) or str(uuid.uuid4()) + self.name = kwargs.pop("name", None) or str(uuid.uuid4()) self.state = SessionState.UNMAPPED - self.handle_max = kwargs.get('handle_max', 4294967295) - self.properties = kwargs.pop('properties', None) + self.handle_max = kwargs.get("handle_max", 4294967295) + self.properties = kwargs.pop("properties", None) self.channel = channel self.remote_channel = None - self.next_outgoing_id = kwargs.pop('next_outgoing_id', 0) + self.next_outgoing_id = kwargs.pop("next_outgoing_id", 0) self.next_incoming_id = None - self.incoming_window = kwargs.pop('incoming_window', 1) - self.outgoing_window = kwargs.pop('outgoing_window', 1) + self.incoming_window = kwargs.pop("incoming_window", 1) + self.outgoing_window = kwargs.pop("outgoing_window", 1) self.target_incoming_window = self.incoming_window self.remote_incoming_window = 0 self.remote_outgoing_window = 0 self.offered_capabilities = None - self.desired_capabilities = kwargs.pop('desired_capabilities', None) + self.desired_capabilities = kwargs.pop("desired_capabilities", None) - self.allow_pipelined_open = kwargs.pop('allow_pipelined_open', True) - self.idle_wait_time = kwargs.get('idle_wait_time', 0.1) - self.network_trace = kwargs['network_trace'] - self.network_trace_params = kwargs['network_trace_params'] - self.network_trace_params['session'] = self.name + self.allow_pipelined_open = kwargs.pop("allow_pipelined_open", True) + self.idle_wait_time = kwargs.get("idle_wait_time", 0.1) + self.network_trace = kwargs["network_trace"] + self.network_trace_params = kwargs["network_trace_params"] + self.network_trace_params["session"] = self.name self.links = {} self._connection = connection @@ -97,7 +88,7 @@ def _set_state(self, new_state): self.state = new_state _LOGGER.info("Session state changed: %r -> %r", previous_state, new_state, extra=self.network_trace_params) for link in self.links.values(): - link._on_session_state_change() + link._on_session_state_change() # pylint: disable=protected-access def _on_connection_state_change(self): if self._connection.state in [ConnectionState.CLOSE_RCVD, ConnectionState.END]: @@ -116,7 +107,7 @@ def _get_next_output_handle(self): raise ValueError("Maximum number of handles ({}) has been reached.".format(self.handle_max)) next_handle = next(i for i in range(1, self.handle_max) if i not in self._output_handles) return next_handle - + def _outgoing_begin(self): begin_frame = BeginFrame( remote_channel=self.remote_channel if self.state == SessionState.BEGIN_RCVD else None, @@ -130,7 +121,7 @@ def _outgoing_begin(self): ) if self.network_trace: _LOGGER.info("-> %r", begin_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, begin_frame) + self._connection._process_outgoing_frame(self.channel, begin_frame) # pylint: disable=protected-access def _incoming_begin(self, frame): if self.network_trace: @@ -151,7 +142,7 @@ def _outgoing_end(self, error=None): end_frame = EndFrame(error=error) if self.network_trace: _LOGGER.info("-> %r", end_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, end_frame) + self._connection._process_outgoing_frame(self.channel, end_frame) # pylint: disable=protected-access def _incoming_end(self, frame): if self.network_trace: @@ -165,38 +156,40 @@ def _incoming_end(self, frame): self._set_state(SessionState.UNMAPPED) def _outgoing_attach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) + self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access def _incoming_attach(self, frame): try: - self._input_handles[frame[1]] = self.links[frame[0].decode('utf-8')] # name and handle - self._input_handles[frame[1]]._incoming_attach(frame) + self._input_handles[frame[1]] = self.links[frame[0].decode("utf-8")] # name and handle + self._input_handles[frame[1]]._incoming_attach(frame) # pylint: disable=protected-access except KeyError: outgoing_handle = self._get_next_output_handle() # TODO: catch max-handles error if frame[2] == Role.Sender: # role new_link = ReceiverLink.from_incoming_frame(self, outgoing_handle, frame) else: new_link = SenderLink.from_incoming_frame(self, outgoing_handle, frame) - new_link._incoming_attach(frame) + new_link._incoming_attach(frame) # pylint: disable=protected-access self.links[frame[0]] = new_link self._output_handles[outgoing_handle] = new_link self._input_handles[frame[1]] = new_link except ValueError: # Reject Link self._input_handles[frame[1]].detach() - + def _outgoing_flow(self, frame=None): link_flow = frame or {} - link_flow.update({ - 'next_incoming_id': self.next_incoming_id, - 'incoming_window': self.incoming_window, - 'next_outgoing_id': self.next_outgoing_id, - 'outgoing_window': self.outgoing_window - }) + link_flow.update( + { + "next_incoming_id": self.next_incoming_id, + "incoming_window": self.incoming_window, + "next_outgoing_id": self.next_outgoing_id, + "outgoing_window": self.outgoing_window, + } + ) flow_frame = FlowFrame(**link_flow) if self.network_trace: _LOGGER.info("-> %r", flow_frame, extra=self.network_trace_params) - self._connection._process_outgoing_frame(self.channel, flow_frame) + self._connection._process_outgoing_frame(self.channel, flow_frame) # pylint: disable=protected-access def _incoming_flow(self, frame): if self.network_trace: @@ -206,11 +199,11 @@ def _incoming_flow(self, frame): self.remote_incoming_window = remote_incoming_id + frame[1] - self.next_outgoing_id # incoming_window self.remote_outgoing_window = frame[3] # outgoing_window if frame[4] is not None: # handle - self._input_handles[frame[4]]._incoming_flow(frame) + self._input_handles[frame[4]]._incoming_flow(frame) # pylint: disable=protected-access else: for link in self._output_handles.values(): - if self.remote_incoming_window > 0 and not link._is_closed: - link._incoming_flow(frame) + if self.remote_incoming_window > 0 and not link._is_closed: # pylint: disable=protected-access + link._incoming_flow(frame) # pylint: disable=protected-access def _outgoing_transfer(self, delivery): if self.state != SessionState.MAPPED: @@ -218,58 +211,58 @@ def _outgoing_transfer(self, delivery): if self.remote_incoming_window <= 0: delivery.transfer_state = SessionTransferState.BUSY else: - payload = delivery.frame['payload'] + payload = delivery.frame["payload"] payload_size = len(payload) - delivery.frame['delivery_id'] = self.next_outgoing_id + delivery.frame["delivery_id"] = self.next_outgoing_id # calculate the transfer frame encoding size excluding the payload - delivery.frame['payload'] = b"" + delivery.frame["payload"] = b"" # TODO: encoding a frame would be expensive, we might want to improve depending on the perf test results encoded_frame = encode_frame(TransferFrame(**delivery.frame))[1] transfer_overhead_size = len(encoded_frame) # available size for payload per frame is calculated as following: # remote max frame size - transfer overhead (calculated) - header (8 bytes) - available_frame_size = self._connection._remote_max_frame_size - transfer_overhead_size - 8 + available_frame_size = self._connection._remote_max_frame_size - transfer_overhead_size - 8 # pylint: disable=protected-access start_idx = 0 remaining_payload_cnt = payload_size # encode n-1 frames if payload_size > available_frame_size while remaining_payload_cnt > available_frame_size: tmp_delivery_frame = { - 'handle': delivery.frame['handle'], - 'delivery_tag': delivery.frame['delivery_tag'], - 'message_format': delivery.frame['message_format'], - 'settled': delivery.frame['settled'], - 'more': True, - 'rcv_settle_mode': delivery.frame['rcv_settle_mode'], - 'state': delivery.frame['state'], - 'resume': delivery.frame['resume'], - 'aborted': delivery.frame['aborted'], - 'batchable': delivery.frame['batchable'], - 'payload': payload[start_idx:start_idx+available_frame_size], - 'delivery_id': self.next_outgoing_id + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": True, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "payload": payload[start_idx : start_idx + available_frame_size], + "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) + self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access start_idx += available_frame_size remaining_payload_cnt -= available_frame_size # encode the last frame tmp_delivery_frame = { - 'handle': delivery.frame['handle'], - 'delivery_tag': delivery.frame['delivery_tag'], - 'message_format': delivery.frame['message_format'], - 'settled': delivery.frame['settled'], - 'more': False, - 'rcv_settle_mode': delivery.frame['rcv_settle_mode'], - 'state': delivery.frame['state'], - 'resume': delivery.frame['resume'], - 'aborted': delivery.frame['aborted'], - 'batchable': delivery.frame['batchable'], - 'payload': payload[start_idx:], - 'delivery_id': self.next_outgoing_id + "handle": delivery.frame["handle"], + "delivery_tag": delivery.frame["delivery_tag"], + "message_format": delivery.frame["message_format"], + "settled": delivery.frame["settled"], + "more": False, + "rcv_settle_mode": delivery.frame["rcv_settle_mode"], + "state": delivery.frame["state"], + "resume": delivery.frame["resume"], + "aborted": delivery.frame["aborted"], + "batchable": delivery.frame["batchable"], + "payload": payload[start_idx:], + "delivery_id": self.next_outgoing_id, } - self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) + self._connection._process_outgoing_frame(self.channel, TransferFrame(**tmp_delivery_frame)) # pylint: disable=protected-access self.next_outgoing_id += 1 self.remote_incoming_window -= 1 self.outgoing_window -= 1 @@ -281,29 +274,29 @@ def _incoming_transfer(self, frame): self.remote_outgoing_window -= 1 self.incoming_window -= 1 try: - self._input_handles[frame[0]]._incoming_transfer(frame) # handle + self._input_handles[frame[0]]._incoming_transfer(frame) # pylint: disable=protected-access except KeyError: - pass #TODO: "unattached handle" + pass # TODO: "unattached handle" if self.incoming_window == 0: self.incoming_window = self.target_incoming_window self._outgoing_flow() def _outgoing_disposition(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) + self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access def _incoming_disposition(self, frame): if self.network_trace: _LOGGER.info("<- %r", DispositionFrame(*frame), extra=self.network_trace_params) for link in self._input_handles.values(): - link._incoming_disposition(frame) + link._incoming_disposition(frame) # pylint: disable=protected-access def _outgoing_detach(self, frame): - self._connection._process_outgoing_frame(self.channel, frame) + self._connection._process_outgoing_frame(self.channel, frame) # pylint: disable=protected-access def _incoming_detach(self, frame): try: link = self._input_handles[frame[0]] # handle - link._incoming_detach(frame) + link._incoming_detach(frame) # pylint: disable=protected-access # if link._is_closed: TODO # self.links.pop(link.name, None) # self._input_handles.pop(link.remote_handle, None) @@ -313,7 +306,7 @@ def _incoming_detach(self, frame): def _wait_for_response(self, wait, end_state): # type: (Union[bool, float], SessionState) -> None - if wait == True: + if wait is True: self._connection.listen(wait=False) while self.state != end_state: time.sleep(self.idle_wait_time) @@ -355,9 +348,10 @@ def create_receiver_link(self, source_address, **kwargs): self, handle=assigned_handle, source_address=source_address, - network_trace=kwargs.pop('network_trace', self.network_trace), + network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs) + **kwargs + ) self.links[link.name] = link self._output_handles[assigned_handle] = link return link @@ -368,16 +362,13 @@ def create_sender_link(self, target_address, **kwargs): self, handle=assigned_handle, target_address=target_address, - network_trace=kwargs.pop('network_trace', self.network_trace), + network_trace=kwargs.pop("network_trace", self.network_trace), network_trace_params=dict(self.network_trace_params), - **kwargs) + **kwargs + ) self._output_handles[assigned_handle] = link self.links[link.name] = link return link - + def create_request_response_link_pair(self, endpoint, **kwargs): - return ManagementLink( - self, - endpoint, - network_trace=kwargs.pop('network_trace', self.network_trace), - **kwargs) + return ManagementLink(self, endpoint, network_trace=kwargs.pop("network_trace", self.network_trace), **kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py index 540d4a63d0ea..5baf13992f44 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py @@ -100,15 +100,14 @@ def normalized_data_body(data, **kwargs): encoding = kwargs.get("encoding", "utf-8") if isinstance(data, list): return [encode_str(item, encoding) for item in data] - else: - return [encode_str(data, encoding)] + return [encode_str(data, encoding)] def normalized_sequence_body(sequence): # A helper method to normalize input into AMQP Sequence Body format if isinstance(sequence, list) and all([isinstance(b, list) for b in sequence]): return sequence - elif isinstance(sequence, list): + if isinstance(sequence, list): return [sequence] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/__init__.py new file mode 100644 index 000000000000..34913fb394d7 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/__init__.py @@ -0,0 +1,4 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py new file mode 100644 index 000000000000..d67cceedcd40 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py @@ -0,0 +1,289 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +from typing import Tuple, Union, TYPE_CHECKING +from abc import ABC, abstractmethod + +if TYPE_CHECKING: + from uamqp import types as uamqp_types + +class AmqpTransport(ABC): # pylint: disable=too-many-public-methods + """ + Abstract class that defines a set of common methods needed by producer and consumer. + """ + # define constants + MAX_FRAME_SIZE_BYTES: int + MAX_MESSAGE_LENGTH_BYTES: int + TIMEOUT_FACTOR: int + CONNECTION_CLOSING_STATES: Tuple + + # define symbols + PRODUCT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + VERSION_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + FRAMEWORK_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + PLATFORM_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + USER_AGENT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + + @staticmethod + @abstractmethod + def build_message(**kwargs): + """ + Creates a uamqp.Message or pyamqp.Message with given arguments. + :rtype: uamqp.Message or pyamqp.Message + """ + + @staticmethod + @abstractmethod + def build_batch_message(**kwargs): + """ + Creates a uamqp.BatchMessage or pyamqp.BatchMessage with given arguments. + :rtype: uamqp.BatchMessage or pyamqp.BatchMessage + """ + + @staticmethod + @abstractmethod + def to_outgoing_amqp_message(annotated_message): + """ + Converts an AmqpAnnotatedMessage into an Amqp Message. + :param AmqpAnnotatedMessage annotated_message: AmqpAnnotatedMessage to convert. + :rtype: uamqp.Message or pyamqp.Message + """ + + @staticmethod + @abstractmethod + def get_message_encoded_size(message): + """ + Gets the message encoded size given an underlying Message. + :param uamqp.Message or pyamqp.Message message: Message to get encoded size of. + :rtype: int + """ + + @staticmethod + @abstractmethod + def get_remote_max_message_size(handler): + """ + Returns max peer message size. + :param AMQPClient handler: Client to get remote max message size on link from. + :rtype: int + """ + + @staticmethod + @abstractmethod + def create_retry_policy(config): + """ + Creates the error retry policy. + :param ~azure.eventhub._configuration.Configuration config: Configuration. + """ + + @staticmethod + @abstractmethod + def create_link_properties(link_properties): + """ + Creates and returns the link properties. + :param dict[bytes, int] link_properties: The dict of symbols and corresponding values. + :rtype: dict + """ + + @staticmethod + @abstractmethod + def create_connection(**kwargs): + """ + Creates and returns the uamqp Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + + @staticmethod + @abstractmethod + def close_connection(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + + @staticmethod + @abstractmethod + def get_connection_state(connection): + """ + Gets connection state. + :param connection: uamqp or pyamqp Connection. + """ + + @staticmethod + @abstractmethod + def create_send_client(*, config, **kwargs): + """ + Creates and returns the send client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + + @staticmethod + @abstractmethod + def send_messages(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + + @staticmethod + @abstractmethod + def set_message_partition_key(message, partition_key, **kwargs): + """Set the partition key as an annotation on a uamqp message. + + :param message: The message to update. + :param str partition_key: The partition key value. + :rtype: None + """ + + @staticmethod + @abstractmethod + def add_batch(batch_message, outgoing_event_data, event_data): + """ + Add EventData to the data body of the BatchMessage. + :param batch_message: BatchMessage to add data to. + :param outgoing_event_data: Transformed EventData for sending. + :param event_data: EventData to add to internal batch events. uamqp use only. + :rtype: None + """ + + @staticmethod + @abstractmethod + def create_source(source, offset, selector): + """ + Creates and returns the Source. + + :param str source: Required. + :param int offset: Required. + :param bytes selector: Required. + """ + + @staticmethod + @abstractmethod + def create_receive_client(*, config, **kwargs): + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword Source source: Required. The source. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. Missing in pyamqp. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + @staticmethod + @abstractmethod + def open_receive_client(*, handler, client, auth): + """ + Opens the receive client. + :param ReceiveClient handler: The receive client. + :param ~azure.eventhub.EventHubConsumerClient client: The consumer client. + """ + + @staticmethod + @abstractmethod + def check_link_stolen(consumer, exception): + """ + Checks if link stolen and handles exception. + :param consumer: The EventHubConsumer. + :param exception: Exception to check. + """ + + @staticmethod + @abstractmethod + def create_token_auth(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Whether to update token. If not updating token, + then pass 300 to refresh_window. Only used by uamqp. + """ + + @staticmethod + @abstractmethod + def create_mgmt_client(address, mgmt_auth, config): + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + @staticmethod + @abstractmethod + def get_updated_token(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + + @staticmethod + @abstractmethod + def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQP Client mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + + @staticmethod + @abstractmethod + def get_error(status_code, description): + """ + Gets error corresponding to status code. + :param status_code: Status code. + :param str description: Description of error. + """ + + @staticmethod + @abstractmethod + def check_timeout_exception(base, exception): + """ + Checks if timeout exception. + :param base: ClientBase. + :param exception: Exception to check. + """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py new file mode 100644 index 000000000000..205a3b3419d7 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py @@ -0,0 +1,563 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import logging +import time +from typing import Optional, Union, Any, Tuple + +from .._pyamqp import ( + error as errors, + utils, + SendClient, + constants, + AMQPClient, + ReceiveClient, +) +from .._pyamqp.message import Message, BatchMessage, Header, Properties +from .._pyamqp.authentication import JWTTokenAuth +from .._pyamqp.endpoints import Source, ApacheFilters +from .._pyamqp._connection import Connection, _CLOSING_STATES + +from ._base import AmqpTransport +from .._constants import ( + NO_RETRY_ERRORS, + PROP_PARTITION_KEY, + CUSTOM_CONDITION_BACKOFF, +) + +from ..exceptions import ( + ConnectError, + EventHubError, + AuthenticationError, + ConnectionLostError, + EventDataSendError, +) + +_LOGGER = logging.getLogger(__name__) + + +class PyamqpTransport(AmqpTransport): + """ + Class which defines uamqp-based methods used by the producer and consumer. + """ + + # define constants + MAX_FRAME_SIZE_BYTES = constants.MAX_FRAME_SIZE_BYTES + MAX_MESSAGE_LENGTH_BYTES = constants.MAX_FRAME_SIZE_BYTES # TODO: define actual value in pyamqp + TIMEOUT_FACTOR = 1 + CONNECTION_CLOSING_STATES: Tuple = _CLOSING_STATES + + # define symbols + PRODUCT_SYMBOL = "product" + VERSION_SYMBOL = "version" + FRAMEWORK_SYMBOL = "framework" + PLATFORM_SYMBOL = "platform" + USER_AGENT_SYMBOL = "user-agent" + PROP_PARTITION_KEY_AMQP_SYMBOL = PROP_PARTITION_KEY + + @staticmethod + def build_message(**kwargs): + """ + Creates a pyamqp.Message with given arguments. + :rtype: pyamqp.Message + """ + return Message(**kwargs) + + @staticmethod + def build_batch_message(**kwargs): + """ + Creates a pyamqp.BatchMessage with given arguments. + :rtype: pyamqp.BatchMessage + """ + return BatchMessage(**kwargs) + + @staticmethod + def to_outgoing_amqp_message(annotated_message): + """ + Converts an AmqpAnnotatedMessage into an Amqp Message. + :param AmqpAnnotatedMessage annotated_message: AmqpAnnotatedMessage to convert. + :rtype: pyamqp.Message + """ + message_header = None + if annotated_message.header and any(annotated_message.header.values()): + message_header = Header( + delivery_count=annotated_message.header.delivery_count, + ttl=annotated_message.header.time_to_live, + first_acquirer=annotated_message.header.first_acquirer, + durable=annotated_message.header.durable, + priority=annotated_message.header.priority, + ) + + message_properties = None + if annotated_message.properties and any(annotated_message.properties.values()): + message_properties = Properties( + message_id=annotated_message.properties.message_id, + user_id=annotated_message.properties.user_id, + to=annotated_message.properties.to, + subject=annotated_message.properties.subject, + reply_to=annotated_message.properties.reply_to, + correlation_id=annotated_message.properties.correlation_id, + content_type=annotated_message.properties.content_type, + content_encoding=annotated_message.properties.content_encoding, + creation_time=int(annotated_message.properties.creation_time) + if annotated_message.properties.creation_time + else None, + absolute_expiry_time=int(annotated_message.properties.absolute_expiry_time) + if annotated_message.properties.absolute_expiry_time + else None, + group_id=annotated_message.properties.group_id, + group_sequence=annotated_message.properties.group_sequence, + reply_to_group_id=annotated_message.properties.reply_to_group_id, + ) + + message_dict = { + "header": message_header, + "properties": message_properties, + "application_properties": annotated_message.application_properties, + "message_annotations": annotated_message.annotations, + "delivery_annotations": annotated_message.delivery_annotations, + "data": annotated_message._data_body, # pylint: disable=protected-access + "sequence": annotated_message._sequence_body, # pylint: disable=protected-access + "value": annotated_message._value_body, # pylint: disable=protected-access + "footer": annotated_message.footer, + } + + return Message(**message_dict) + + @staticmethod + def get_batch_message_encoded_size(message): + """ + Gets the batch message encoded size given an underlying Message. + :param pyamqp.BatchMessage message: Message to get encoded size of. + :rtype: int + """ + return utils.get_message_encoded_size(message) + + @staticmethod + def get_message_encoded_size(message): + """ + Gets the message encoded size given an underlying Message. + :param pyamqp.Message: Message to get encoded size of. + :rtype: int + """ + return utils.get_message_encoded_size(message) + + @staticmethod + def get_remote_max_message_size(handler): + """ + Returns max peer message size. + :param AMQPClient handler: Client to get remote max message size on link from. + :rtype: int + """ + return handler._link.remote_max_message_size # pylint: disable=protected-access + + @staticmethod + def create_retry_policy(config): + """ + Creates the error retry policy. + :param ~azure.eventhub._configuration.Configuration config: Configuration. + """ + return errors.RetryPolicy( + retry_total=config.max_retries, # pylint:disable=protected-access + retry_backoff_factor=config.backoff_factor, # pylint:disable=protected-access + retry_backoff_max=config.backoff_max, # pylint:disable=protected-access + retry_mode=config.retry_mode, # pylint:disable=protected-access + no_retry_condition=NO_RETRY_ERRORS, + custom_condition_backoff=CUSTOM_CONDITION_BACKOFF, + ) + + @staticmethod + def create_link_properties(link_properties): + """ + Creates and returns the link properties. + :param dict[bytes, int] link_properties: The dict of symbols and corresponding values. + :rtype: dict + """ + return {symbol: utils.amqp_long_value(value) for (symbol, value) in link_properties.items()} + + @staticmethod + def create_connection(**kwargs): + """ + Creates and returns the uamqp Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") + host = kwargs.pop("host") # pylint:disable=unused-variable + auth = kwargs.pop("auth") # pylint:disable=unused-variable + network_trace = kwargs.pop("debug") + return Connection(endpoint, network_trace=network_trace, **kwargs) + + @staticmethod + def close_connection(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + connection.close() + + @staticmethod + def get_connection_state(connection): + """ + Gets connection state. + :param connection: uamqp or pyamqp Connection. + """ + return connection.state + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the uamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + + target = kwargs.pop("target") + # TODO: extra passed in to pyamqp, but not used. should be used? + msg_timeout = kwargs.pop("msg_timeout") # pylint: disable=unused-variable # TODO: not used by pyamqp? + + return SendClient( + config.hostname, + target, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + **kwargs, + ) + + @staticmethod + def send_messages(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + producer._open() + timeout = timeout_time - time.time() if timeout_time else 0 + producer._handler.send_message(producer._unsent_events[0], timeout=timeout) + # The unsent_events list will always be <= 1. Even for a batch, it gets the underlying singular BatchMessage. + # May want to refactor in the future so that this isn't a list. + producer._unsent_events = None + # TODO: figure out if we want to use below, and see if it affects error story + # try: + # producer._open() + # producer._handler.send_message( + # producer._unsent_events[0], timeout=timeout_time + # ) + # except TimeoutError as exc: + # raise OperationTimeoutError(message=str(exc), details=exc) + # except Exception as exc: + # raise producer._handle_exception(exc) + + @staticmethod + def set_message_partition_key(message, partition_key, **kwargs): + # type: (Message, Optional[Union[bytes, str]], Any) -> Message + """Set the partition key as an annotation on a uamqp message. + :param Message message: The message to update. + :param str partition_key: The partition key value. + :rtype: Message + """ + encoding = kwargs.pop("encoding", "utf-8") + if partition_key: + annotations = message.message_annotations + if annotations is None: + annotations = {} + try: + partition_key = partition_key.decode(encoding) + except AttributeError: + pass + annotations[PROP_PARTITION_KEY] = partition_key # pylint:disable=protected-access + header = Header(durable=True) + return message._replace(message_annotations=annotations, header=header) + return message + + @staticmethod + def add_batch(event_data_batch, outgoing_event_data, event_data): # pylint: disable=unused-argument + """ + Add EventData to the data body of the BatchMessage. + :param event_data_batch: EventDataBatch to add data to. + :param outgoing_event_data: Transformed EventData for sending. + :param event_data: EventData to add to internal batch events. uamqp use only. + :rtype: None + """ + event_data_batch._internal_events.append(event_data) # pylint: disable=protected-access + utils.add_batch(event_data_batch._message, outgoing_event_data._message) # pylint: disable=protected-access + + @staticmethod + def create_source(source, offset, selector): + """ + Creates and returns the Source. + + :param str source: Required. + :param int offset: Required. + :param bytes selector: Required. + """ + source = Source(address=source, filters={}) + if offset is not None: + filter_key = ApacheFilters.selector_filter + source.filters[filter_key] = (filter_key, utils.amqp_string_value(selector)) + return source + + @staticmethod + def create_receive_client(*, config, **kwargs): + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. Missing in pyamqp. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + return ReceiveClient( + config.hostname, + source, + receive_settle_mode=constants.ReceiverSettleMode.First, + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + **kwargs, + ) + + @staticmethod + def open_receive_client(*, handler, client, auth): + """ + Opens the receive client and returns ready status. + :param ReceiveClient handler: The receive client. + :param ~azure.eventhub.EventHubConsumerClient client: The consumer client. + :param auth: Auth. + :rtype: bool + """ + # pylint:disable=protected-access + handler.open(connection=client._conn_manager.get_connection(client._address.hostname, auth)) + + @staticmethod + def check_link_stolen(consumer, exception): + """ + Checks if link stolen and handles exception. + :param consumer: The EventHubConsumer. + :param exception: Exception to check. + """ + + if isinstance(exception, errors.AMQPLinkError) and exception.condition == errors.ErrorCondition.LinkStolen: + raise consumer._handle_exception(exception) # pylint: disable=protected-access + + @staticmethod + def create_token_auth(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Whether to update token. If not updating token, then pass 300 to refresh_window. + """ + # TODO: figure out why we're passing all these args to pyamqp JWTTokenAuth, which aren't being used + update_token = kwargs.pop("update_token") # pylint: disable=unused-variable + if update_token: + # update_token not actually needed by pyamqp + # just using to detect wh + return JWTTokenAuth(auth_uri, auth_uri, get_token) + return JWTTokenAuth( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + ) + # if update_token: + # token_auth.update_token() # TODO: why don't we need to update in pyamqp? + + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): # pylint: disable=unused-argument + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + return AMQPClient( + config.hostname, + auth=mgmt_auth, + network_trace=config.network_tracing, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + ) + + @staticmethod + def get_updated_token(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return mgmt_auth.get_token() + + @staticmethod + def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQPClient mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + return mgmt_client.mgmt_request( + mgmt_msg, operation=operation.decode(), operation_type=operation_type.decode(), **kwargs + ) + + @staticmethod + def get_error(status_code, description): + """ + Gets error and passes in error message, and, if applicable, condition. + :param error: The error to raise. + :param str message: Error message. + :param condition: Optional error condition. Will not be used by uamqp. + """ + if status_code in [401]: + return errors.AuthenticationException( + errors.ErrorCondition.UnauthorizedAccess, + description=f"Management authentication failed. Status code: {status_code}, Description: {description!r}", + ) + if status_code in [404]: + return errors.AMQPConnectionError( + errors.ErrorCondition.NotFound, + description=f"Management connection failed. Status code: {status_code}, Description: {description!r}", + ) + return errors.AMQPConnectionError( + errors.ErrorCondition.UnknownError, + description=f"Management request error. Status code: {status_code}, Description: {description!r}", + ) + + @staticmethod + def check_timeout_exception(base, exception): + """ + Checks if timeout exception. + :param base: ClientBase. + :param exception: Exception to check. + """ + if not base.running and isinstance(exception, TimeoutError): + exception = errors.AuthenticationException( + errors.ErrorCondition.InternalError, description="Authorization timeout." + ) + return exception + + @staticmethod + def _create_eventhub_exception(exception): + if isinstance(exception, errors.AuthenticationException): + error = AuthenticationError(str(exception), exception) + elif isinstance(exception, errors.AMQPLinkError): + error = ConnectError(str(exception), exception) + # TODO: do we need MessageHanlderError in amqp any more + # if connection/session/link error are enough? + # elif isinstance(exception, errors.MessageHandlerError): + # error = ConnectionLostError(str(exception), exception) + elif isinstance(exception, errors.AMQPConnectionError): + error = ConnectError(str(exception), exception) + elif isinstance(exception, TimeoutError): + error = ConnectionLostError(str(exception), exception) + else: + error = EventHubError(str(exception), exception) + return error + + @staticmethod + def _handle_exception(exception, closable): # pylint:disable=too-many-branches, too-many-statements + try: # closable is a producer/consumer object + name = closable._name # pylint: disable=protected-access + except AttributeError: # closable is an client object + name = closable._container_id # pylint: disable=protected-access + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + closable._close_connection() # pylint:disable=protected-access + raise exception + elif isinstance(exception, EventHubError): + closable._close_handler() # pylint:disable=protected-access + raise exception + # TODO: The following errors seem to be useless in EH + # elif isinstance( + # exception, + # ( + # errors.MessageAccepted, + # errors.MessageAlreadySettled, + # errors.MessageModified, + # errors.MessageRejected, + # errors.MessageReleased, + # errors.MessageContentTooLarge, + # ), + # ): + # _LOGGER.info("%r Event data error (%r)", name, exception) + # error = EventDataError(str(exception), exception) + # raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + if isinstance(exception, errors.AuthenticationException): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPLinkError): + if hasattr(closable, "_close_handler"): + closable._close_handler() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPConnectionError): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + # TODO: add MessageHandlerError in amqp? + # elif isinstance(exception, errors.MessageHandlerError): + # if hasattr(closable, "_close_handler"): + # closable._close_handler() # pylint:disable=protected-access + else: # errors.AMQPConnectionError, compat.TimeoutException + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + return PyamqpTransport._create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py new file mode 100644 index 000000000000..478906f27a38 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py @@ -0,0 +1,650 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import time +import logging +from typing import Optional, Union, Any, Tuple + +try: + from uamqp import ( + c_uamqp, + BatchMessage, + constants, + MessageBodyType, + Message, + types, + SendClient, + ReceiveClient, + Source, + utils, + authentication, + AMQPClient, + compat, + errors, + Connection, + ) + from uamqp.message import ( + MessageHeader, + MessageProperties, + ) + uamqp_installed = True +except ImportError: + uamqp_installed = False + +from ._base import AmqpTransport +from ..amqp._constants import AmqpMessageBodyType +from .._constants import ( + NO_RETRY_ERRORS, + PROP_PARTITION_KEY, +) + +from ..exceptions import ( + ConnectError, + OperationTimeoutError, + EventHubError, + AuthenticationError, + ConnectionLostError, + EventDataError, + EventDataSendError, +) + +_LOGGER = logging.getLogger(__name__) + +if uamqp_installed: + def _error_handler(error): + """ + Called internally when an event has failed to send so we + can parse the error to determine whether we should attempt + to retry sending the event again. + Returns the action to take according to error type. + + :param error: The error received in the send attempt. + :type error: Exception + :rtype: ~uamqp.errors.ErrorAction + """ + if error.condition == b"com.microsoft:server-busy": + return errors.ErrorAction(retry=True, backoff=4) + if error.condition == b"com.microsoft:timeout": + return errors.ErrorAction(retry=True, backoff=2) + if error.condition == b"com.microsoft:operation-cancelled": + return errors.ErrorAction(retry=True) + if error.condition == b"com.microsoft:container-close": + return errors.ErrorAction(retry=True, backoff=4) + if error.condition in NO_RETRY_ERRORS: + return errors.ErrorAction(retry=False) + return errors.ErrorAction(retry=True) + + + class UamqpTransport(AmqpTransport): # pylint: disable=too-many-public-methods + """ + Class which defines uamqp-based methods used by the producer and consumer. + """ + # define constants + MAX_FRAME_SIZE_BYTES = constants.MAX_FRAME_SIZE_BYTES + MAX_MESSAGE_LENGTH_BYTES = constants.MAX_MESSAGE_LENGTH_BYTES + TIMEOUT_FACTOR = 1000 + CONNECTION_CLOSING_STATES: Tuple = ( # pylint:disable=protected-access + c_uamqp.ConnectionState.CLOSE_RCVD, # pylint:disable=c-extension-no-member + c_uamqp.ConnectionState.CLOSE_SENT, # pylint:disable=c-extension-no-member + c_uamqp.ConnectionState.DISCARDING, # pylint:disable=c-extension-no-member + c_uamqp.ConnectionState.END, # pylint:disable=c-extension-no-member + ) + + # define symbols + PRODUCT_SYMBOL = types.AMQPSymbol("product") + VERSION_SYMBOL = types.AMQPSymbol("version") + FRAMEWORK_SYMBOL = types.AMQPSymbol("framework") + PLATFORM_SYMBOL = types.AMQPSymbol("platform") + USER_AGENT_SYMBOL = types.AMQPSymbol("user-agent") + PROP_PARTITION_KEY_AMQP_SYMBOL = types.AMQPSymbol(PROP_PARTITION_KEY) + + @staticmethod + def build_message(**kwargs): + """ + Creates a uamqp.Message with given arguments. + :rtype: uamqp.Message + """ + return Message(**kwargs) + + @staticmethod + def build_batch_message(**kwargs): + """ + Creates a uamqp.BatchMessage with given arguments. + :rtype: uamqp.BatchMessage + """ + return BatchMessage(**kwargs) + + @staticmethod + def to_outgoing_amqp_message(annotated_message): + """ + Converts an AmqpAnnotatedMessage into an Amqp Message. + :param AmqpAnnotatedMessage annotated_message: AmqpAnnotatedMessage to convert. + :rtype: uamqp.Message + """ + message_header = None + if annotated_message.header and any(annotated_message.header.values()): + message_header = MessageHeader() + message_header.delivery_count = annotated_message.header.delivery_count + message_header.time_to_live = annotated_message.header.time_to_live + message_header.first_acquirer = annotated_message.header.first_acquirer + message_header.durable = annotated_message.header.durable + message_header.priority = annotated_message.header.priority + + message_properties = None + if annotated_message.properties and any(annotated_message.properties.values()): + message_properties = MessageProperties( + message_id=annotated_message.properties.message_id, + user_id=annotated_message.properties.user_id, + to=annotated_message.properties.to, + subject=annotated_message.properties.subject, + reply_to=annotated_message.properties.reply_to, + correlation_id=annotated_message.properties.correlation_id, + content_type=annotated_message.properties.content_type, + content_encoding=annotated_message.properties.content_encoding, + creation_time=int(annotated_message.properties.creation_time) + if annotated_message.properties.creation_time else None, + absolute_expiry_time=int(annotated_message.properties.absolute_expiry_time) + if annotated_message.properties.absolute_expiry_time else None, + group_id=annotated_message.properties.group_id, + group_sequence=annotated_message.properties.group_sequence, + reply_to_group_id=annotated_message.properties.reply_to_group_id, + encoding=annotated_message._encoding # pylint: disable=protected-access + ) + + # pylint: disable=protected-access + amqp_body_type = annotated_message.body_type + if amqp_body_type == AmqpMessageBodyType.DATA: + amqp_body_type = MessageBodyType.Data + amqp_body = list(annotated_message._data_body) + elif amqp_body_type == AmqpMessageBodyType.SEQUENCE: + amqp_body_type = MessageBodyType.Sequence + amqp_body = list(annotated_message._sequence_body) + else: + amqp_body_type = MessageBodyType.Value + amqp_body = annotated_message._value_body + + return Message( + body=amqp_body, + body_type=amqp_body_type, + header=message_header, + properties=message_properties, + application_properties=annotated_message.application_properties, + annotations=annotated_message.annotations, + delivery_annotations=annotated_message.delivery_annotations, + footer=annotated_message.footer + ) + + @staticmethod + def get_batch_message_encoded_size(message): + """ + Gets the batch message encoded size given an underlying Message. + :param uamqp.BatchMessage message: Message to get encoded size of. + :rtype: int + """ + return message.gather()[0].get_message_encoded_size() + + @staticmethod + def get_message_encoded_size(message): + """ + Gets the message encoded size given an underlying Message. + :param uamqp.Message message: Message to get encoded size of. + :rtype: int + """ + return message.get_message_encoded_size() + + @staticmethod + def get_remote_max_message_size(handler): + """ + Returns max peer message size. + :param AMQPClient handler: Client to get remote max message size on link from. + :rtype: int + """ + return handler.message_handler._link.peer_max_message_size # pylint:disable=protected-access + + @staticmethod + def create_retry_policy(config): + """ + Creates the error retry policy. + :param ~azure.eventhub._configuration.Configuration config: Configuration. + """ + return errors.ErrorPolicy(max_retries=config.max_retries, on_error=_error_handler) + + @staticmethod + def create_link_properties(link_properties): + """ + Creates and returns the link properties. + :param dict[bytes, int] link_properties: The dict of symbols and corresponding values. + :rtype: dict + """ + return {types.AMQPSymbol(symbol): types.AMQPLong(value) for (symbol, value) in link_properties.items()} + + @staticmethod + def create_connection(**kwargs): + """ + Creates and returns the uamqp Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") # pylint:disable=unused-variable + custom_endpoint_address = kwargs.pop("custom_endpoint_address") # pylint:disable=unused-variable + host = kwargs.pop("host") + auth = kwargs.pop("auth") + return Connection( + host, + auth, + **kwargs + ) + + @staticmethod + def close_connection(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + connection.destroy() + + @staticmethod + def get_connection_state(connection): + """ + Gets connection state. + :param connection: uamqp or pyamqp Connection. + """ + return connection._state # pylint:disable=protected-access + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the uamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + target = kwargs.pop("target") + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") + + return SendClient( + target, + debug=network_trace, + error_policy=retry_policy, + **kwargs + ) + + @staticmethod + def _set_msg_timeout(producer, timeout_time, last_exception, logger): + if not timeout_time: + return + remaining_time = timeout_time - time.time() + if remaining_time <= 0.0: + if last_exception: + error = last_exception + else: + error = OperationTimeoutError("Send operation timed out") + logger.info("%r send operation timed out. (%r)", producer._name, error) # pylint: disable=protected-access + raise error + producer._handler._msg_timeout = remaining_time * 1000 # type: ignore # pylint: disable=protected-access + + @staticmethod + def send_messages(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + producer._open() + producer._unsent_events[0].on_send_complete = producer._on_outcome + UamqpTransport._set_msg_timeout(producer, timeout_time, last_exception, logger) + producer._handler.queue_message(*producer._unsent_events) # type: ignore + producer._handler.wait() # type: ignore + producer._unsent_events = producer._handler.pending_messages # type: ignore + if producer._outcome != constants.MessageSendResult.Ok: + if producer._outcome == constants.MessageSendResult.Timeout: + producer._condition = OperationTimeoutError("Send operation timed out") + if producer._condition: + raise producer._condition + + @staticmethod + def set_message_partition_key(message, partition_key, **kwargs): # pylint:disable=unused-argument + # type: (Message, Optional[Union[bytes, str]], Any) -> Message + """Set the partition key as an annotation on a uamqp message. + + :param ~uamqp.Message message: The message to update. + :param str partition_key: The partition key value. + :rtype: Message + """ + if partition_key: + annotations = message.annotations + if annotations is None: + annotations = {} + annotations[ + UamqpTransport.PROP_PARTITION_KEY_AMQP_SYMBOL # TODO: see if setting non-amqp symbol is valid + ] = partition_key + header = MessageHeader() + header.durable = True + message.annotations = annotations + message.header = header + return message + + @staticmethod + def add_batch(batch_message, outgoing_event_data, event_data): + """ + Add EventData to the data body of the BatchMessage. + :param batch_message: BatchMessage to add data to. + :param outgoing_event_data: Transformed EventData for sending. + :param event_data: EventData to add to internal batch events. uamqp use only. + :rtype: None + """ + # pylint: disable=protected-access + batch_message._internal_events.append(event_data) + batch_message._message._body_gen.append( + outgoing_event_data._message + ) + + @staticmethod + def create_source(source, offset, selector): + """ + Creates and returns the Source. + + :param str source: Required. + :param int offset: Required. + :param bytes selector: Required. + """ + source = Source(source) + if offset is not None: + source.set_filter(selector) + return source + + @staticmethod + def create_receive_client(*, config, **kwargs): # pylint: disable=unused-argument + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + symbol_array = kwargs.pop("desired_capabilities") + desired_capabilities = None + if symbol_array: + symbol_array = [types.AMQPSymbol(symbol) for symbol in symbol_array] + desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array)) + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") + link_credit = kwargs.pop("link_credit") + streaming_receive = kwargs.pop("streaming_receive") + message_received_callback = kwargs.pop("message_received_callback") + + client = ReceiveClient( + source, + debug=network_trace, # pylint:disable=protected-access + error_policy=retry_policy, + desired_capabilities=desired_capabilities, + prefetch=link_credit, + receive_settle_mode=constants.ReceiverSettleMode.ReceiveAndDelete, + auto_complete=False, + **kwargs + ) + # pylint:disable=protected-access + client._streaming_receive = streaming_receive + client._message_received_callback = (message_received_callback) + return client + + @staticmethod + def open_receive_client(*, handler, client, auth): + """ + Opens the receive client and returns ready status. + :param ReceiveClient handler: The receive client. + :param ~azure.eventhub.EventHubConsumerClient client: The consumer client. + :param auth: Auth. + :rtype: bool + """ + # pylint:disable=protected-access + handler.open(connection=client._conn_manager.get_connection( + client._address.hostname, auth + )) + + @staticmethod + def check_link_stolen(consumer, exception): + """ + Checks if link stolen and handles exception. + :param consumer: The EventHubConsumer. + :param exception: Exception to check. + """ + if ( + isinstance(exception, errors.LinkDetach) + and exception.condition == constants.ErrorCodes.LinkStolen # pylint: disable=no-member + ): + raise consumer._handle_exception(exception) # pylint: disable=protected-access + + @staticmethod + def create_token_auth(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Required. Whether to update token. If not updating token, + then pass 300 to refresh_window. + """ + update_token = kwargs.pop("update_token") + refresh_window = 300 + if update_token: + refresh_window = 0 + + token_auth = authentication.JWTTokenAuth( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + refresh_window=refresh_window + ) + if update_token: + token_auth.update_token() + return token_auth + + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + mgmt_target = f"amqps://{address.hostname}{address.path}" + return AMQPClient( + mgmt_target, + auth=mgmt_auth, + debug=config.network_tracing + ) + + @staticmethod + def get_updated_token(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return mgmt_auth.token + + @staticmethod + def mgmt_client_request(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQP Client mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + response = mgmt_client.mgmt_request( + mgmt_msg, + operation, + op_type=operation_type, + **kwargs + ) + status_code = response.application_properties[kwargs.get("status_code_field")] + description = response.application_properties.get( + kwargs.get("description_fields") + ) # type: Optional[Union[str, bytes]] + return status_code, description, response + + @staticmethod + def get_error(status_code, description): + """ + Gets error corresponding to status code. + :param status_code: Status code. + :param str description: Description of error. + """ + if status_code in [401]: + return errors.AuthenticationException( + f"Management authentication failed. Status code: {status_code}, Description: {description!r}" + ) + if status_code in [404]: + return ConnectError( + f"Management connection failed. Status code: {status_code}, Description: {description!r}" + ) + return errors.AMQPConnectionError( + f"Management request error. Status code: {status_code}, Description: {description!r}" + ) + + @staticmethod + def check_timeout_exception(base, exception): + """ + Checks if timeout exception. + :param base: ClientBase. + :param exception: Exception to check. + """ + if not base.running and isinstance( + exception, compat.TimeoutException + ): + exception = errors.AuthenticationException( + "Authorization timeout." + ) + return exception + + @staticmethod + def _create_eventhub_exception(exception): + if isinstance(exception, errors.AuthenticationException): + error = AuthenticationError(str(exception), exception) + elif isinstance(exception, errors.VendorLinkDetach): + error = ConnectError(str(exception), exception) + elif isinstance(exception, errors.LinkDetach): + error = ConnectionLostError(str(exception), exception) + elif isinstance(exception, errors.ConnectionClose): + error = ConnectionLostError(str(exception), exception) + elif isinstance(exception, errors.MessageHandlerError): + error = ConnectionLostError(str(exception), exception) + elif isinstance(exception, errors.AMQPConnectionError): + error_type = ( + AuthenticationError + if str(exception).startswith("Unable to open authentication session") + else ConnectError + ) + error = error_type(str(exception), exception) + elif isinstance(exception, compat.TimeoutException): + error = ConnectionLostError(str(exception), exception) + else: + error = EventHubError(str(exception), exception) + return error + + @staticmethod + def _handle_exception( + exception, closable + ): # pylint:disable=too-many-branches, too-many-statements + try: # closable is a producer/consumer object + name = closable._name # pylint: disable=protected-access + except AttributeError: # closable is an client object + name = closable._container_id # pylint: disable=protected-access + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + closable._close_connection() # pylint:disable=protected-access + raise exception + elif isinstance(exception, EventHubError): + closable._close_handler() # pylint:disable=protected-access + raise exception + elif isinstance( + exception, + ( + errors.MessageAccepted, + errors.MessageAlreadySettled, + errors.MessageModified, + errors.MessageRejected, + errors.MessageReleased, + errors.MessageContentTooLarge, + ), + ): + _LOGGER.info("%r Event data error (%r)", name, exception) + error = EventDataError(str(exception), exception) + raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + if isinstance(exception, errors.AuthenticationException): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + elif isinstance(exception, errors.LinkDetach): + if hasattr(closable, "_close_handler"): + closable._close_handler() # pylint:disable=protected-access + elif isinstance(exception, errors.ConnectionClose): + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + elif isinstance(exception, errors.MessageHandlerError): + if hasattr(closable, "_close_handler"): + closable._close_handler() # pylint:disable=protected-access + else: # errors.AMQPConnectionError, compat.TimeoutException + if hasattr(closable, "_close_connection"): + closable._close_connection() # pylint:disable=protected-access + return UamqpTransport._create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py index e9c000e9d5e3..e15f0b884732 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations from contextlib import contextmanager import sys @@ -10,19 +10,28 @@ import datetime import calendar import logging -from typing import TYPE_CHECKING, Type, Optional, Dict, Union, Any, Iterable, Tuple, Mapping +from typing import ( + TYPE_CHECKING, + cast, + Type, + Optional, + Dict, + Union, + Any, + Iterable, + Tuple, + Mapping, + Callable +) import six -from ._pyamqp.message import Header - from azure.core.settings import settings from azure.core.tracing import SpanKind, Link -from .amqp import AmqpAnnotatedMessage +from .amqp import AmqpAnnotatedMessage, AmqpMessageHeader from ._version import VERSION from ._constants import ( - PROP_PARTITION_KEY, MAX_USER_AGENT_LENGTH, USER_AGENT_PREFIX, PROP_LAST_ENQUEUED_SEQUENCE_NUMBER, @@ -32,9 +41,15 @@ PROP_TIMESTAMP, ) + if TYPE_CHECKING: # pylint: disable=ungrouped-imports - from ._pyamqp.message import Message + from ._transport._base import AmqpTransport + try: + from uamqp import types as uamqp_types + except ImportError: + uamqp_types = None + from ._pyamqp import types from azure.core.tracing import AbstractSpan from azure.core.credentials import AzureSasCredential from ._common import EventData @@ -76,66 +91,36 @@ def utc_from_timestamp(timestamp): return datetime.datetime.fromtimestamp(timestamp, tz=TZ_UTC) -def create_properties(user_agent=None): - # type: (Optional[str]) -> Dict[types.AMQPSymbol, str] +def create_properties( + user_agent: Optional[str] = None, *, amqp_transport: AmqpTransport +) -> Union[Dict[uamqp_types.AMQPSymbol, str], Dict[str, str]]: """ Format the properties with which to instantiate the connection. This acts like a user agent over HTTP. :rtype: dict """ - properties = {} - properties["product"] = USER_AGENT_PREFIX - properties["version"] = VERSION - framework = "Python/{}.{}.{}".format( - sys.version_info[0], sys.version_info[1], sys.version_info[2] - ) - properties["framework"] = framework + properties: Dict[Any, str] = {} + properties[amqp_transport.PRODUCT_SYMBOL] = USER_AGENT_PREFIX + properties[amqp_transport.VERSION_SYMBOL] = VERSION + framework = f"Python/{sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}" + properties[amqp_transport.FRAMEWORK_SYMBOL] = framework platform_str = platform.platform() - properties["platform"] = platform_str + properties[amqp_transport.PLATFORM_SYMBOL] = platform_str - final_user_agent = "{}/{} {} ({})".format( - USER_AGENT_PREFIX, VERSION, framework, platform_str - ) + final_user_agent = f"{USER_AGENT_PREFIX}/{VERSION} {framework} ({platform_str})" if user_agent: - final_user_agent = "{} {}".format(user_agent, final_user_agent) + final_user_agent = f"{user_agent} {final_user_agent}" if len(final_user_agent) > MAX_USER_AGENT_LENGTH: raise ValueError( - "The user-agent string cannot be more than {} in length." - "Current user_agent string is: {} with length: {}".format( - MAX_USER_AGENT_LENGTH, final_user_agent, len(final_user_agent) - ) + f"The user-agent string cannot be more than {MAX_USER_AGENT_LENGTH} in length." + f"Current user_agent string is: {final_user_agent} with length: {len(final_user_agent)}" ) - properties["user-agent"] = final_user_agent + properties[amqp_transport.USER_AGENT_SYMBOL] = final_user_agent return properties -def set_message_partition_key(message, partition_key, **kwargs): - # type: (Message, Optional[Union[bytes, str]]) -> Message - """Set the partition key as an annotation on a uamqp message. - - :param ~uamqp.Message message: The message to update. - :param str partition_key: The partition key value. - :rtype: None - """ - encoding = kwargs.pop("encoding", 'utf-8') - if partition_key: - annotations = message.message_annotations - if annotations is None: - annotations = dict() - try: - partition_key = partition_key.decode(encoding) - except AttributeError: - pass - annotations[ - PROP_PARTITION_KEY - ] = partition_key # pylint:disable=protected-access - header = Header(durable=True) - return message._replace(message_annotations=annotations, header=header) - return message - - @contextmanager def send_context_manager(): span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan] @@ -147,6 +132,31 @@ def send_context_manager(): yield None +def set_event_partition_key( + event: Union[AmqpAnnotatedMessage, EventData], + partition_key: Optional[Union[bytes, str]], + amqp_transport: AmqpTransport +) -> None: + if not partition_key: + return + + try: + raw_message = event.raw_amqp_message # type: ignore + except AttributeError: + raw_message = event + + annotations = raw_message.annotations + if annotations is None: + annotations = {} + annotations[ + amqp_transport.PROP_PARTITION_KEY_AMQP_SYMBOL + ] = partition_key # pylint:disable=protected-access + if not raw_message.header: + raw_message.header = AmqpMessageHeader(header=True) + else: + raw_message.header.durable = True + + def trace_message(event, parent_span=None): # type: (EventData, Optional[AbstractSpan]) -> None """Add tracing information to this event. @@ -176,9 +186,7 @@ def trace_message(event, parent_span=None): def get_event_links(events): # pylint:disable=isinstance-second-argument-not-valid-type - trace_events = ( - events if isinstance(events, Iterable) else (events,) - ) + trace_events = events if isinstance(events, Iterable) else (events,) links = [] try: for event in trace_events: # type: ignore @@ -209,15 +217,13 @@ def event_position_selector(value, inclusive=False): value.microsecond / 1000 ) return ( - "amqp.annotation.x-opt-enqueued-time {} '{}'".format( - operator, int(timestamp) - ) + f"amqp.annotation.x-opt-enqueued-time {operator} '{int(timestamp)}'" ).encode("utf-8") elif isinstance(value, six.integer_types): return ( - "amqp.annotation.x-opt-sequence-number {} '{}'".format(operator, value) + f"amqp.annotation.x-opt-sequence-number {operator} '{value}'" ).encode("utf-8") - return ("amqp.annotation.x-opt-offset {} '{}'".format(operator, value)).encode( + return (f"amqp.annotation.x-opt-offset {operator} '{value}'").encode( "utf-8" ) @@ -232,23 +238,23 @@ def get_last_enqueued_event_properties(event_data): if event_data._last_enqueued_event_properties: return event_data._last_enqueued_event_properties - if event_data.message.delivery_annotations: - sequence_number = event_data.message.delivery_annotations.get( + if event_data._message.delivery_annotations: + sequence_number = event_data._message.delivery_annotations.get( PROP_LAST_ENQUEUED_SEQUENCE_NUMBER, None ) - enqueued_time_stamp = event_data.message.delivery_annotations.get( + enqueued_time_stamp = event_data._message.delivery_annotations.get( PROP_LAST_ENQUEUED_TIME_UTC, None ) if enqueued_time_stamp: enqueued_time_stamp = utc_from_timestamp(float(enqueued_time_stamp) / 1000) - retrieval_time_stamp = event_data.message.delivery_annotations.get( + retrieval_time_stamp = event_data._message.delivery_annotations.get( PROP_RUNTIME_INFO_RETRIEVAL_TIME_UTC, None ) if retrieval_time_stamp: retrieval_time_stamp = utc_from_timestamp( float(retrieval_time_stamp) / 1000 ) - offset_bytes = event_data.message.delivery_annotations.get( + offset_bytes = event_data._message.delivery_annotations.get( PROP_LAST_ENQUEUED_OFFSET, None ) offset = offset_bytes.decode("UTF-8") if offset_bytes else None @@ -274,8 +280,8 @@ def parse_sas_credential(credential): return (sas, expiry) -def transform_outbound_single_message(message, message_type): - # type: (Union[AmqpAnnotatedMessage, EventData], Type[EventData]) -> EventData +def transform_outbound_single_message(message, message_type, to_outgoing_amqp_message): + # type: (Union[AmqpAnnotatedMessage, EventData], Type[EventData], Callable) -> EventData """ This method serves multiple goals: 1. update the internal message to reflect any updates to settable properties on EventData @@ -287,14 +293,20 @@ def transform_outbound_single_message(message, message_type): :rtype: EventData """ try: - # EventData # pylint: disable=protected-access - return message._to_outgoing_message() # type: ignore + # If EventData, set EventData._message to uamqp/pyamqp.Message right before sending. + message = cast("EventData", message) + message._message = to_outgoing_amqp_message(message.raw_amqp_message) + return message # type: ignore except AttributeError: - # AmqpAnnotatedMessage # pylint: disable=protected-access + # If AmqpAnnotatedMessage, create EventData object with _from_message. + # event_data._message will be set to outgoing uamqp/pyamqp.Message. + # event_data.raw_amqp_message will be set to AmqpAnnotatedMessage. + message = cast(AmqpAnnotatedMessage, message) + amqp_message = to_outgoing_amqp_message(message) return message_type._from_message( - message=message._to_outgoing_amqp_message(), raw_amqp_message=message # type: ignore + message=amqp_message, raw_amqp_message=message # type: ignore ) @@ -314,14 +326,18 @@ def decode_with_recurse(data, encoding="UTF-8"): return data if isinstance(data, six.binary_type): return data.decode(encoding) - if isinstance(data, Mapping): # pylint:disable=isinstance-second-argument-not-valid-type + if isinstance( + data, Mapping + ): # pylint:disable=isinstance-second-argument-not-valid-type decoded_mapping = {} for k, v in data.items(): decoded_key = decode_with_recurse(k, encoding) decoded_val = decode_with_recurse(v, encoding) decoded_mapping[decoded_key] = decoded_val return decoded_mapping - if isinstance(data, Iterable): # pylint:disable=isinstance-second-argument-not-valid-type + if isinstance( + data, Iterable + ): # pylint:disable=isinstance-second-argument-not-valid-type decoded_list = [] for d in data: decoded_list.append(decode_with_recurse(d, encoding)) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py index b4919bbfcac1..613c2a509829 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/_version.py @@ -3,4 +3,4 @@ # Licensed under the MIT License. # ------------------------------------ -VERSION = "5.8.0b6" +VERSION = "5.8.0b1" diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_async_utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_async_utils.py index 9e604a982b53..a547c83bdba0 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_async_utils.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_async_utils.py @@ -1,15 +1,26 @@ -#------------------------------------------------------------------------- +# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. -#-------------------------------------------------------------------------- - +# -------------------------------------------------------------------------- +import asyncio import sys +from asyncio import Semaphore + def get_dict_with_loop_if_needed(loop): if sys.version_info >= (3, 10): if loop: - raise ValueError("Starting Python 3.10, asyncio no longer supports loop as a parameter.") + raise ValueError( + "Starting Python 3.10, asyncio no longer supports loop as a parameter." + ) elif loop: - return {'loop': loop} + return {"loop": loop} return {} + + +async def semaphore_acquire_with_timeout(semaphore: Semaphore, timeout=None): + try: + return await asyncio.wait_for(semaphore.acquire(), timeout=timeout) + except asyncio.TimeoutError: + return False diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/__init__.py new file mode 100644 index 000000000000..7154a1189cc3 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/__init__.py @@ -0,0 +1,9 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from ._buffered_producer_dispatcher_async import BufferedProducerDispatcher +from ._partition_resolver_async import PartitionResolver +from ._buffered_producer_async import BufferedProducer + +__all__ = ["BufferedProducerDispatcher", "PartitionResolver", "BufferedProducer"] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py new file mode 100644 index 000000000000..d48eee1c358f --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.py @@ -0,0 +1,220 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +import asyncio +import logging +import queue +import time +from asyncio import Lock +from typing import Optional, Callable, Awaitable, TYPE_CHECKING + +from .._producer_async import EventHubProducer +from ..._common import EventDataBatch +from ...exceptions import OperationTimeoutError + +if TYPE_CHECKING: + from .._transport._base_async import AmqpTransportAsync + from ..._producer_client import SendEventTypes + +_LOGGER = logging.getLogger(__name__) + + +class BufferedProducer: + # pylint: disable=too-many-instance-attributes + def __init__( + self, + producer: EventHubProducer, + partition_id: str, + on_success: Callable[["SendEventTypes", Optional[str]], Awaitable[None]], + on_error: Callable[ + ["SendEventTypes", Optional[str], Exception], Awaitable[None] + ], + max_message_size_on_link: int, + *, + amqp_transport: AmqpTransportAsync, + max_buffer_length: int, + max_wait_time: float = 1 + ): + self._buffered_queue: queue.Queue = queue.Queue() + self._max_buffer_len = max_buffer_length + self._cur_buffered_len = 0 + self._producer: EventHubProducer = producer + self._lock = Lock() + self._max_wait_time = max_wait_time + self._on_success = self.failsafe_callback(on_success) + self._on_error = self.failsafe_callback(on_error) + self._last_send_time = None + self._running = False + self._cur_batch: Optional[EventDataBatch] = None + self._max_message_size_on_link = max_message_size_on_link + self._check_max_wait_time_future = None + self.partition_id = partition_id + self._amqp_transport = amqp_transport + + async def start(self): + async with self._lock: + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + self._running = True + if self._max_wait_time: + self._last_send_time = time.time() + self._check_max_wait_time_future = asyncio.ensure_future( + self.check_max_wait_time_worker() + ) + + async def stop(self, flush=True, timeout_time=None, raise_error=False): + self._running = False + if flush: + async with self._lock: + await self._flush(timeout_time=timeout_time, raise_error=raise_error) + else: + if self._cur_buffered_len: + _LOGGER.warning( + "Shutting down Partition %r." + " There are still %r events in the buffer which will be lost", + self.partition_id, + self._cur_buffered_len, + ) + if self._check_max_wait_time_future: + try: + await self._check_max_wait_time_future + except Exception as exc: # pylint: disable=broad-except + _LOGGER.warning( + "Partition %r stopped with error %r", self.partition_id, exc + ) + await self._producer.close() + + async def put_events(self, events, timeout_time=None): + # Put single event or EventDataBatch into the queue. + # This method would raise OperationTimeout if the queue does not have enough space for the input and + # flush cannot finish in timeout. + try: + new_events_len = len(events) + except TypeError: + new_events_len = 1 + if self._max_buffer_len - self._cur_buffered_len < new_events_len: + _LOGGER.info( + "The buffer for partition %r is full. Attempting to flush before adding %r events.", + self.partition_id, + new_events_len, + ) + # flush the buffer + await self.flush(timeout_time=timeout_time) + if timeout_time and time.time() > timeout_time: + raise OperationTimeoutError( + "Failed to enqueue events into buffer due to timeout." + ) + try: + # add single event into current batch + self._cur_batch.add(events) + except AttributeError: # if the input events is a EventDataBatch, put the whole into the buffer + # if there are events in cur_batch, enqueue cur_batch to the buffer + async with self._lock: + if self._cur_batch: + self._buffered_queue.put(self._cur_batch) + self._buffered_queue.put(events) + # create a new batch for incoming events + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + except ValueError: + # add single event exceeds the cur batch size, create new batch + async with self._lock: + self._buffered_queue.put(self._cur_batch) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + self._cur_batch.add(events) + async with self._lock: + self._cur_buffered_len += new_events_len + + def failsafe_callback(self, callback): + async def wrapper_callback(*args, **kwargs): + try: + await callback(*args, **kwargs) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.warning( + "On partition %r, callback %r encountered exception %r", + callback.__name__, + exc, + self.partition_id, + ) + + return wrapper_callback + + async def flush(self, timeout_time=None, raise_error=True): + async with self._lock: + await self._flush(timeout_time, raise_error) + + async def _flush(self, timeout_time=None, raise_error=True): + # pylint: disable=protected-access + # try flushing all the buffered batch within given time + _LOGGER.info("Partition: %r started flushing.", self.partition_id) + if self._cur_batch: # if there is batch, enqueue it to the buffer first + self._buffered_queue.put(self._cur_batch) + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + while self._buffered_queue.qsize() > 0: + remaining_time = timeout_time - time.time() if timeout_time else None + if (remaining_time and remaining_time > 0) or remaining_time is None: + try: + batch = self._buffered_queue.get(block=False) + except queue.Empty: + break + self._buffered_queue.task_done() + try: + _LOGGER.info("Partition %r is sending.", self.partition_id) + await self._producer.send( + batch, + timeout=timeout_time - time.time() if timeout_time else None, + ) + _LOGGER.info( + "Partition %r sending %r events succeeded.", + self.partition_id, + len(batch), + ) + await self._on_success(batch._internal_events, self.partition_id) + except Exception as exc: # pylint: disable=broad-except + _LOGGER.info( + "Partition %r sending %r events failed due to exception: %r", + self.partition_id, + len(batch), + exc, + ) + await self._on_error(batch._internal_events, self.partition_id, exc) + finally: + self._cur_buffered_len -= len(batch) + # If flush could not get the semaphore, we log and raise error if wanted + else: + _LOGGER.info( + "Partition %r fails to flush due to timeout.", self.partition_id + ) + if raise_error: + raise OperationTimeoutError( + "Failed to flush {!r} within {}".format( + self.partition_id, timeout_time + ) + ) + break + # after finishing flushing, reset cur batch and put it into the buffer + self._last_send_time = time.time() + #reset curr_buffered + self._cur_buffered_len = 0 + self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport) + _LOGGER.info("Partition %r finished flushing.", self.partition_id) + + async def check_max_wait_time_worker(self): + while self._running: + if self._cur_buffered_len > 0: + now_time = time.time() + _LOGGER.info( + "Partition %r worker is checking max_wait_time.", self.partition_id + ) + # flush the partition if its beyond the waiting time or the buffer is at max capacity + if (now_time - self._last_send_time > self._max_wait_time) or ( + self._cur_buffered_len >= self._max_buffer_len + ): + # in the worker, not raising error for flush, users can not handle this + async with self._lock: + await self._flush(raise_error=False) + await asyncio.sleep(min(self._max_wait_time, 5)) + + @property + def buffered_event_count(self): + return self._cur_buffered_len diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py new file mode 100644 index 000000000000..64e565944aaf --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.py @@ -0,0 +1,175 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +import asyncio +import logging +from typing import Dict, List, Callable, Optional, Awaitable, TYPE_CHECKING +from asyncio import Lock + +from ._partition_resolver_async import PartitionResolver +from ...aio._producer_async import EventHubProducer +from ._buffered_producer_async import BufferedProducer +from ...exceptions import EventDataSendError, ConnectError, EventHubError + +if TYPE_CHECKING: + from .._transport._base_async import AmqpTransportAsync + from ..._producer_client import SendEventTypes + +_LOGGER = logging.getLogger(__name__) + + +class BufferedProducerDispatcher: + # pylint: disable=too-many-instance-attributes + def __init__( + self, + partitions: List[str], + on_success: Callable[["SendEventTypes", Optional[str]], Awaitable[None]], + on_error: Callable[ + ["SendEventTypes", Optional[str], Exception], Awaitable[None] + ], + create_producer: Callable[..., EventHubProducer], + eventhub_name: str, + max_message_size_on_link: int, + *, + amqp_transport: AmqpTransportAsync, + max_buffer_length: int = 1500, + max_wait_time: float = 1 + ): + self._buffered_producers: Dict[str, BufferedProducer] = {} + self._partition_ids: List[str] = partitions + self._lock = Lock() + self._on_success = on_success + self._on_error = on_error + self._create_producer = create_producer + self._eventhub_name = eventhub_name + self._max_message_size_on_link = max_message_size_on_link + self._partition_resolver = PartitionResolver(self._partition_ids) + self._max_wait_time = max_wait_time + self._max_buffer_length = max_buffer_length + self._amqp_transport = amqp_transport + + async def _get_partition_id(self, partition_id, partition_key): + if partition_id: + if partition_id not in self._partition_ids: + raise ConnectError( + "Invalid partition {} for the event hub {}".format( + partition_id, self._eventhub_name + ) + ) + return partition_id + if isinstance(partition_key, str): + return await self._partition_resolver.get_partition_id_by_partition_key( + partition_key + ) + return await self._partition_resolver.get_next_partition_id() + + async def enqueue_events( + self, events, *, partition_id=None, partition_key=None, timeout_time=None + ): + pid = await self._get_partition_id(partition_id, partition_key) + async with self._lock: + try: + await self._buffered_producers[pid].put_events(events, timeout_time) + except KeyError: + buffered_producer = BufferedProducer( + self._create_producer(partition_id=pid), + pid, + self._on_success, + self._on_error, + self._max_message_size_on_link, + max_wait_time=self._max_wait_time, + max_buffer_length=self._max_buffer_length, + amqp_transport=self._amqp_transport, + ) + await buffered_producer.start() + self._buffered_producers[pid] = buffered_producer + await buffered_producer.put_events(events, timeout_time) + + async def flush(self, timeout_time=None): + # flush all the buffered producer, the method will block until finishes or times out + async with self._lock: + futures = [] + for pid, producer in self._buffered_producers.items(): + # call each producer's flush method + futures.append( + ( + pid, + asyncio.ensure_future( + producer.flush(timeout_time=timeout_time) + ), + ) + ) + + # gather results + exc_results = {} + for pid, future in futures: + try: + await future + except Exception as exc: # pylint: disable=broad-except + exc_results[pid] = exc + + if not exc_results: + _LOGGER.info("Flushing all partitions succeeded") + return + + _LOGGER.warning( + "Flushing all partitions partially failed with result %r.", exc_results + ) + raise EventDataSendError( + message="Flushing all partitions partially failed, failed partitions are {!r}" + " Exception details are {!r}".format(exc_results.keys(), exc_results) + ) + + async def close(self, *, flush=True, timeout_time=None, raise_error=False): + + async with self._lock: + + futures = [] + # stop all buffered producers + for pid, producer in self._buffered_producers.items(): + futures.append( + ( + pid, + asyncio.ensure_future( + producer.stop( + flush=flush, + timeout_time=timeout_time, + raise_error=raise_error, + ) + ), + ) + ) + + exc_results = {} + # gather results + for pid, future in futures: + try: + await future + except Exception as exc: # pylint: disable=broad-except + exc_results[pid] = exc + + if exc_results: + _LOGGER.warning( + "Stopping all partitions failed with result %r.", exc_results + ) + if raise_error: + raise EventHubError( + message="Stopping all partitions partially failed, failed partitions are {!r}" + " Exception details are {!r}".format( + exc_results.keys(), exc_results + ) + ) + + def get_buffered_event_count(self, pid): + try: + return self._buffered_producers[pid].buffered_event_count + except KeyError: + return 0 + + @property + def total_buffered_event_count(self): + return sum( + [self.get_buffered_event_count(pid) for pid in self._buffered_producers] + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_partition_resolver_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_partition_resolver_async.py new file mode 100644 index 000000000000..b741029e9bcb --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_partition_resolver_async.py @@ -0,0 +1,33 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +""" +jenkins-hash lookup3 algorithm implementation +""" + +from asyncio import Lock +from ..._buffered_producer._partition_resolver import ( + generate_hash_code, +) # pylint: disable=protected-access + + +class PartitionResolver: + def __init__(self, partitions): + self._idx = -1 + self._partitions = partitions + self._partitions_cnt = len(self._partitions) + self._lock = Lock() + + async def get_next_partition_id(self): + """ + round-robin partition assignment + """ + async with self._lock: + self._idx += 1 + self._idx %= self._partitions_cnt + return self._partitions[self._idx] + + async def get_partition_id_by_partition_key(self, partition_key): + hash_code = generate_hash_code(partition_key) + return self._partitions[abs(hash_code % self._partitions_cnt)] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py index dc2093c12d38..6f8533a64adb 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py @@ -2,14 +2,13 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -from __future__ import unicode_literals +from __future__ import unicode_literals, annotations import logging import asyncio import time import functools from typing import TYPE_CHECKING, Any, Dict, List, Callable, Optional, Union, cast -import six from azure.core.credentials import ( AccessToken, @@ -23,31 +22,100 @@ _parse_conn_str, _get_backoff_time, ) - -from .._pyamqp.message import Message -from .._pyamqp import constants, error as errors, utils as pyamqp_utils -from .._pyamqp.aio import AMQPClientAsync -from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync from .._utils import utc_from_timestamp, parse_sas_credential -from ..exceptions import ClientClosedError, ConnectError +from ..exceptions import ClientClosedError from .._constants import ( JWT_TOKEN_SCOPE, MGMT_OPERATION, MGMT_PARTITION_OPERATION, MGMT_STATUS_CODE, - MGMT_STATUS_DESC, READ_OPERATION, + MGMT_STATUS_DESC, + READ_OPERATION, ) from ._async_utils import get_dict_with_loop_if_needed -from ._error_async import _handle_exception +from ._connection_manager_async import get_connection_manager +try: + from ._transport._uamqp_transport_async import UamqpTransportAsync +except ImportError: + UamqpTransportAsync = None +from ._transport._pyamqp_transport_async import PyamqpTransportAsync if TYPE_CHECKING: + from .._pyamqp.message import Message + from .._pyamqp.aio import AMQPClientAsync + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + try: + from uamqp import ( + authentication as uamqp_authentication, + Message as uamqp_Message, + AMQPClientAsync as uamqp_AMQPClientAsync, + ) + except ImportError: + uamqp_authentication = None + uamqp_Message = None + uamqp_AMQPClientAsync = None from azure.core.credentials_async import AsyncTokenCredential + CredentialTypes = Union[ + "EventHubSharedKeyCredential", + AsyncTokenCredential, + AzureSasCredential, + AzureNamedKeyCredential, + ] + try: from typing_extensions import Protocol except ImportError: Protocol = object # type: ignore + class AbstractConsumerProducer(Protocol): + @property + def _name(self) -> str: + """Name of the consumer or producer""" + + @_name.setter + def _name(self, value): + pass + + @property + def _client(self) -> ClientBaseAsync: + """The instance of EventHubComsumerClient or EventHubProducerClient""" + + @_client.setter + def _client(self, value): + pass + + @property + def _handler(self) -> Union[uamqp_AMQPClientAsync, AMQPClientAsync]: + """The instance of SendClientAsync or ReceiveClientAsync""" + + @property + def _internal_kwargs(self) -> dict: + """The dict with an event loop that users may pass in to wrap sync calls to async API. + It's furthur passed to uamqp APIs + """ + + @_internal_kwargs.setter + def _internal_kwargs(self, value): + pass + + @property + def running(self): + # type: () -> bool + """Whether the consumer or producer is running""" + + @running.setter + def running(self, value): + pass + + def _create_handler(self, auth: Union[uamqp_authentication.JWTTokenAsync, JWTTokenAuthAsync]) -> None: + pass + + _MIXIN_BASE = AbstractConsumerProducer +else: + _MIXIN_BASE = object + + _LOGGER = logging.getLogger(__name__) @@ -63,7 +131,9 @@ def __init__(self, policy: str, key: str): self.key = key self.token_type = b"servicebus.windows.net:sastoken" - async def get_token(self, *scopes, **kwargs) -> AccessToken: # pylint:disable=unused-argument + async def get_token( + self, *scopes, **kwargs # pylint:disable=unused-argument + ) -> AccessToken: if not scopes: raise ValueError("No token scope provided.") return _generate_sas_token(scopes[0], self.policy, self.key) @@ -101,12 +171,13 @@ class EventhubAzureNamedKeyTokenCredentialAsync(object): :type credential: ~azure.core.credentials.AzureNamedKeyCredential """ - def __init__(self, azure_named_key_credential): - # type: (AzureNamedKeyCredential) -> None + def __init__(self, azure_named_key_credential: AzureNamedKeyCredential) -> None: self._credential = azure_named_key_credential self.token_type = b"servicebus.windows.net:sastoken" - async def get_token(self, *scopes, **kwargs) -> AccessToken: # pylint:disable=unused-argument + async def get_token( + self, *scopes, **kwargs # pylint:disable=unused-argument + ) -> AccessToken: if not scopes: raise ValueError("No token scope provided.") name, key = self._credential.named_key @@ -140,12 +211,14 @@ def __init__( self, fully_qualified_namespace: str, eventhub_name: str, - credential: Union[ - "AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential - ], + credential: "CredentialTypes", **kwargs: Any ) -> None: self._internal_kwargs = get_dict_with_loop_if_needed(kwargs.get("loop", None)) + uamqp_transport = kwargs.get("uamqp_transport", False) + if uamqp_transport and not UamqpTransportAsync: + raise ValueError("To use the uAMQP transport, please install `uamqp>=1.6.0,<2.0.0`.") + self._amqp_transport = UamqpTransportAsync if uamqp_transport else PyamqpTransportAsync if isinstance(credential, AzureSasCredential): self._credential = EventhubAzureSasTokenCredentialAsync(credential) # type: ignore elif isinstance(credential, AzureNamedKeyCredential): @@ -156,6 +229,12 @@ def __init__( fully_qualified_namespace=fully_qualified_namespace, eventhub_name=eventhub_name, credential=self._credential, + amqp_transport=self._amqp_transport, + **kwargs + ) + kwargs["custom_endpoint_address"] = self._config.custom_endpoint_address + self._conn_manager_async = get_connection_manager( + amqp_transport=self._amqp_transport, **kwargs ) @@ -177,7 +256,7 @@ def _from_connection_string(conn_str: str, **kwargs) -> Dict[str, Any]: kwargs["credential"] = EventHubSharedKeyCredential(policy, key) return kwargs - async def _create_auth_async(self) -> JWTTokenAuthAsync: + async def _create_auth_async(self) -> Union[uamqp_authentication.JWTTokenAsync, JWTTokenAuthAsync]: """ Create an ~uamqp.authentication.SASTokenAuthAsync instance to authenticate the session. @@ -189,26 +268,23 @@ async def _create_auth_async(self) -> JWTTokenAuthAsync: except AttributeError: token_type = b"jwt" if token_type == b"servicebus.windows.net:sastoken": - return JWTTokenAuthAsync( - self._auth_uri, + return await self._amqp_transport.create_token_auth_async( self._auth_uri, functools.partial(self._credential.get_token, self._auth_uri), + token_type=token_type, + config=self._config, + update_token=True, ) - return JWTTokenAuthAsync( - self._auth_uri, + return await self._amqp_transport.create_token_auth_async( self._auth_uri, functools.partial(self._credential.get_token, JWT_TOKEN_SCOPE), token_type=token_type, - timeout=self._config.auth_timeout, - http_proxy=self._config.http_proxy, - transport_type=self._config.transport_type, - custom_endpoint_hostname=self._config.custom_endpoint_hostname, - port=self._config.connection_port, - verify=self._config.connection_verify, + config=self._config, + update_token=False, ) async def _close_connection_async(self) -> None: - pass + await self._conn_manager_async.reset_connection_if_broken() async def _backoff_async( self, @@ -241,75 +317,42 @@ async def _backoff_async( ) raise last_exception - async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> Any: + async def _management_request_async(self, mgmt_msg: Union[Message, uamqp_Message], op_type: bytes) -> Any: retried_times = 0 last_exception = None while retried_times <= self._config.max_retries: mgmt_auth = await self._create_auth_async() - hostname = self._address.hostname - custom_endpoint_address = self._config.custom_endpoint_address - if self._config.transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - mgmt_client = AMQPClientAsync( - hostname, - auth=mgmt_auth, - network_trace=self._config.network_tracing, - transport_type=self._config.transport_type, - http_proxy=self._config.http_proxy, - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._config.connection_verify + mgmt_client = self._amqp_transport.create_mgmt_client( + self._address, mgmt_auth=mgmt_auth, config=self._config ) try: - await mgmt_client.open_async() - while not (await mgmt_client.client_ready_async()): + conn = await self._conn_manager_async.get_connection( + host=self._address.hostname, auth=mgmt_auth + ) + await mgmt_client.open_async(connection=conn) + while not await mgmt_client.client_ready_async(): await asyncio.sleep(0.05) - access_token = await mgmt_auth.get_token() - mgmt_msg.application_properties["security_token"] = access_token.token - - if not access_token.token: - _LOGGER.debug("update_token received an empty token") - - status_code, description, response = await mgmt_client.mgmt_request_async( + mgmt_msg.application_properties[ + "security_token" + ] = await self._amqp_transport.get_updated_token_async(mgmt_auth) + status_code, description, response = await self._amqp_transport.mgmt_client_request_async( + mgmt_client, mgmt_msg, - operation=READ_OPERATION.decode(), - operation_type=op_type.decode(), + operation=READ_OPERATION, + operation_type=op_type, status_code_field=MGMT_STATUS_CODE, description_fields=MGMT_STATUS_DESC, ) status_code = int(status_code) - if description and isinstance(description, six.binary_type): + if description and isinstance(description, bytes): description = description.decode("utf-8") if status_code < 400: return response - if status_code in [401]: - raise errors.AuthenticationException( - errors.ErrorCondition.UnauthorizedAccess, - description="Management authentication failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) - if status_code in [404]: - raise errors.AMQPConnectionError( - errors.ErrorCondition.NotFound, - description="Management connection failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) - raise errors.AMQPConnectionError( - errors.ErrorCondition.UnknownError, - description="Management operation failed. Status code: {}, Description: {!r}".format( - status_code, - description - ) - ) + raise self._amqp_transport.get_error(status_code, description) except asyncio.CancelledError: # pylint: disable=try-except-raise raise except Exception as exception: # pylint:disable=broad-except - last_exception = await _handle_exception(exception, self) + last_exception = await self._amqp_transport._handle_exception_async(exception, self) # pylint: disable=protected-access await self._backoff_async( retried_times=retried_times, last_exception=last_exception ) @@ -323,12 +366,14 @@ async def _management_request_async(self, mgmt_msg: Message, op_type: bytes) -> await mgmt_client.close_async() async def _get_eventhub_properties_async(self) -> Dict[str, Any]: - mgmt_msg = Message(application_properties={"name": self.eventhub_name}) + mgmt_msg = self._amqp_transport.build_message( + application_properties={"name": self.eventhub_name} + ) response = await self._management_request_async( mgmt_msg, op_type=MGMT_OPERATION ) output = {} - eh_info = response.value # type: Dict[bytes, Any] + eh_info: Dict[bytes, Any] = response.value if eh_info: output["eventhub_name"] = eh_info[b"name"].decode("utf-8") output["created_at"] = utc_from_timestamp( @@ -345,7 +390,7 @@ async def _get_partition_ids_async(self) -> List[str]: async def _get_partition_properties_async( self, partition_id: str ) -> Dict[str, Any]: - mgmt_msg = Message( + mgmt_msg = self._amqp_transport.build_message( application_properties={ "name": self.eventhub_name, "partition": partition_id, @@ -377,61 +422,7 @@ async def _get_partition_properties_async( return output async def _close_async(self) -> None: - pass - - -if TYPE_CHECKING: - - class AbstractConsumerProducer(Protocol): - @property - def _name(self): - # type: () -> str - """Name of the consumer or producer""" - - @_name.setter - def _name(self, value): - pass - - @property - def _client(self): - # type: () -> ClientBaseAsync - """The instance of EventHubComsumerClient or EventHubProducerClient""" - - @_client.setter - def _client(self, value): - pass - - @property - def _handler(self): - # type: () -> AMQPClientAsync - """The instance of SendClientAsync or ReceiveClientAsync""" - - @property - def _internal_kwargs(self): - # type: () -> dict - """The dict with an event loop that users may pass in to wrap sync calls to async API. - It's furthur passed to uamqp APIs - """ - - @_internal_kwargs.setter - def _internal_kwargs(self, value): - pass - - @property - def running(self): - # type: () -> bool - """Whether the consumer or producer is running""" - - @running.setter - def running(self, value): - pass - - def _create_handler(self, auth: JWTTokenAuthAsync) -> None: - pass - - _MIXIN_BASE = AbstractConsumerProducer -else: - _MIXIN_BASE = object + await self._conn_manager_async.close_connection() class ConsumerProducerMixin(_MIXIN_BASE): @@ -460,12 +451,16 @@ async def _open(self) -> None: await self._handler.close_async() auth = await self._client._create_auth_async() self._create_handler(auth) - await self._handler.open_async() + conn = await self._client._conn_manager_async.get_connection( + host=self._client._address.hostname, auth=auth + ) + await self._handler.open_async(connection=conn) while not await self._handler.client_ready_async(): await asyncio.sleep(0.05, **self._internal_kwargs) + # pylint: disable=protected-access self._max_message_size_on_link = ( - self._handler._link.remote_max_message_size - or constants.MAX_FRAME_SIZE_BYTES + self._client._amqp_transport.get_remote_max_message_size(self._handler) + or self._client._amqp_transport.MAX_MESSAGE_LENGTH_BYTES ) self.running = True @@ -480,12 +475,11 @@ async def _close_connection_async(self) -> None: await self._client._conn_manager_async.reset_connection_if_broken() # pylint:disable=protected-access async def _handle_exception(self, exception: Exception) -> Exception: - if not self.running and isinstance(exception, TimeoutError): - exception = errors.AuthenticationException( - errors.ErrorCondition.InternalError, - description="Authorization timeout." - ) - return await _handle_exception(exception, self) + # pylint: disable=protected-access + exception = self._client._amqp_transport.check_timeout_exception(self, exception) + return await self._client._amqp_transport._handle_exception_async( + exception, self + ) async def _do_retryable_operation( self, diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py new file mode 100644 index 000000000000..e1be06565dca --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py @@ -0,0 +1,125 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from __future__ import annotations +from typing import TYPE_CHECKING, Optional, Union +from asyncio import Lock + +from .._connection_manager import _ConnectionMode +from .._constants import TransportType + +if TYPE_CHECKING: + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + from .._pyamqp.aio._connection_async import Connection as ConnectionAsync + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAuthAsync + from uamqp.async_ops import ConnectionAsync as uamqp_ConnectionAsync + + try: + from typing_extensions import Protocol + except ImportError: + Protocol = object # type: ignore + + class ConnectionManager(Protocol): + async def get_connection( + self, + *, + host: Optional[str] = None, + auth: Optional[Union[uamqp_JWTTokenAuthAsync, JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, + ) -> Union[ConnectionAsync, uamqp_ConnectionAsync]: + pass + + async def close_connection(self) -> None: + pass + + async def reset_connection_if_broken(self) -> None: + pass + + +class _SharedConnectionManager(object): # pylint:disable=too-many-instance-attributes + def __init__(self, **kwargs) -> None: + self._loop = kwargs.get("loop") + self._lock = Lock(loop=self._loop) + self._conn = None + + self._container_id = kwargs.get("container_id") + self._custom_endpoint_address = kwargs.get("custom_endpoint_address") + self._debug = kwargs.get("debug") + self._error_policy = kwargs.get("error_policy") + self._properties = kwargs.get("properties") + self._encoding = kwargs.get("encoding") or "UTF-8" + self._transport_type = kwargs.get("transport_type") or TransportType.Amqp + self._http_proxy = kwargs.get("http_proxy") + self._max_frame_size = kwargs.get("max_frame_size") + self._channel_max = kwargs.get("channel_max") + self._idle_timeout = kwargs.get("idle_timeout") + self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get("remote_idle_timeout_empty_frame_send_ratio") + self._amqp_transport = kwargs.get("amqp_transport") + + async def get_connection( + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuthAsync, uamqp_JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, + ) -> Union[ConnectionAsync, uamqp_ConnectionAsync]: + async with self._lock: + if self._conn is None: + self._conn = self._amqp_transport.create_connection_async( + host=host, + auth=auth, + endpoint=endpoint, + custom_endpoint_address=self._custom_endpoint_address, + container_id=self._container_id, + max_frame_size=self._max_frame_size, + channel_max=self._channel_max, + idle_timeout=self._idle_timeout, + properties=self._properties, + remote_idle_timeout_empty_frame_send_ratio=self._remote_idle_timeout_empty_frame_send_ratio, + error_policy=self._error_policy, + debug=self._debug, + loop=self._loop, + encoding=self._encoding, + ) + return self._conn + + async def close_connection(self) -> None: + async with self._lock: + if self._conn: + await self._amqp_transport.close_connection_async(self._conn) + self._conn = None + + async def reset_connection_if_broken(self) -> None: + async with self._lock: + conn_state = self._amqp_transport.get_connection_state(self._conn) + if self._conn and conn_state in self._amqp_transport.CONNECTION_CLOSING_STATES: + self._conn = None + + +class _SeparateConnectionManager(object): + def __init__(self, **kwargs) -> None: + pass + + async def get_connection( + self, + *, + host: Optional[str] = None, + auth: Optional[Union[JWTTokenAuthAsync, uamqp_JWTTokenAuthAsync]] = None, + endpoint: Optional[str] = None, + ) -> None: + pass # return None + + async def close_connection(self) -> None: + pass + + async def reset_connection_if_broken(self) -> None: + pass + + +def get_connection_manager(**kwargs) -> "ConnectionManager": + connection_mode = kwargs.get("connection_mode", _ConnectionMode.SeparateConnection) # type: ignore + if connection_mode == _ConnectionMode.ShareConnection: + return _SharedConnectionManager(**kwargs) + return _SeparateConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py index 670bfcfabae9..c1e4fec6d275 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py @@ -2,42 +2,43 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -import time -import asyncio +from __future__ import annotations import uuid +import asyncio import logging from collections import deque -from typing import TYPE_CHECKING, Callable, Awaitable, cast, Dict, Optional, Union, List -from urllib.parse import urlparse - -from .._pyamqp import ( - types, - utils as pyamqp_utils, - error, - constants as pyamqp_constants -) -from .._pyamqp.endpoints import Source, ApacheFilters -from .._pyamqp.message import Message -from .._pyamqp.aio import ReceiveClientAsync +from typing import TYPE_CHECKING, Callable, Awaitable, Dict, Optional, Union, List +from functools import partial from ._client_base_async import ConsumerProducerMixin from ._async_utils import get_dict_with_loop_if_needed from .._common import EventData from .._utils import create_properties, event_position_selector -from .._constants import EPOCH_SYMBOL, TIMEOUT_SYMBOL, RECEIVER_RUNTIME_METRIC_SYMBOL, NO_RETRY_ERRORS, \ - CUSTOM_CONDITION_BACKOFF +from .._constants import EPOCH_SYMBOL, TIMEOUT_SYMBOL, RECEIVER_RUNTIME_METRIC_SYMBOL if TYPE_CHECKING: from typing import Deque + + try: + from uamqp import ReceiveClientAsync as uamqp_ReceiveClientAsync, Message as uamqp_Message + from uamqp.types import AMQPType as uamqp_AMQPType + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAsync + except ImportError: + uamqp_Message = None + uamqp_ReceiveClientAsync = None + uamqp_AMQPType = None + uamqp_JWTTokenAsync = None + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync + from .._pyamqp.aio._client_async import ReceiveClientAsync + from .._pyamqp import types + from ._consumer_client_async import EventHubConsumerClient _LOGGER = logging.getLogger(__name__) -class EventHubConsumer( - ConsumerProducerMixin -): # pylint:disable=too-many-instance-attributes +class EventHubConsumer(ConsumerProducerMixin): # pylint:disable=too-many-instance-attributes """ A consumer responsible for reading EventData from a specific Event Hub partition and as a member of a specific consumer group. @@ -78,17 +79,16 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N owner_level = kwargs.get("owner_level", None) keep_alive = kwargs.get("keep_alive", None) auto_reconnect = kwargs.get("auto_reconnect", True) - track_last_enqueued_event_properties = kwargs.get( - "track_last_enqueued_event_properties", False - ) + track_last_enqueued_event_properties = kwargs.get("track_last_enqueued_event_properties", False) idle_timeout = kwargs.get("idle_timeout", None) self.running = False self.closed = False - self._on_event_received = kwargs[ + self._amqp_transport = kwargs.pop("amqp_transport") + self._on_event_received: Callable[[Union[Optional[EventData], List[EventData]]], Awaitable[None]] = kwargs[ "on_event_received" - ] # type: Callable[[Union[Optional[EventData], List[EventData]]], Awaitable[None]] + ] self._internal_kwargs = get_dict_with_loop_if_needed(kwargs.get("loop", None)) self._client = client self._source = source @@ -98,87 +98,56 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs) -> N self._owner_level = owner_level self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect - self._retry_policy = error.RetryPolicy( - retry_total=self._client._config.max_retries, # pylint:disable=protected-access - retry_backoff_factor=self._client._config.backoff_factor, # pylint:disable=protected-access - retry_backoff_max=self._client._config.backoff_max, # pylint:disable=protected-access - retry_mode=self._client._config.retry_mode, # pylint:disable=protected-access - no_retry_condition=NO_RETRY_ERRORS, - custom_condition_backoff=CUSTOM_CONDITION_BACKOFF, - ) + self._retry_policy = self._amqp_transport.create_retry_policy(self._client._config) self._reconnect_backoff = 1 self._timeout = 0 - self._idle_timeout = idle_timeout - self._link_properties = {} - partition = self._source.split("/")[-1] - self._partition = partition - self._name = "EHReceiver-{}-partition{}".format(uuid.uuid4(), partition) + self._idle_timeout = (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) if idle_timeout else None + link_properties: Union[Dict[uamqp_AMQPType, uamqp_AMQPType], Dict[types.AMQPTypes, types.AMQPTypes]] = {} + self._partition = self._source.split("/")[-1] + self._name = f"EHReceiver-{uuid.uuid4()}-partition{self._partition}" if owner_level is not None: - self._link_properties[EPOCH_SYMBOL] = pyamqp_utils.amqp_long_value(int(owner_level)) + link_properties[EPOCH_SYMBOL] = int(owner_level) link_property_timeout_ms = ( - self._client._config.receive_timeout or self._timeout # pylint:disable=protected-access - ) * 1000 - self._link_properties[TIMEOUT_SYMBOL] = pyamqp_utils.amqp_long_value(int(link_property_timeout_ms)) - self._handler = None # type: Optional[ReceiveClientAsync] - self._track_last_enqueued_event_properties = ( - track_last_enqueued_event_properties - ) - self._message_buffer = deque() # type: Deque[Message] - self._last_received_event = None # type: Optional[EventData] + self._client._config.receive_timeout or self._timeout # pylint:disable=protected-access + ) * self._amqp_transport.TIMEOUT_FACTOR + link_properties[TIMEOUT_SYMBOL] = int(link_property_timeout_ms) + self._link_properties = self._amqp_transport.create_link_properties(link_properties) + self._handler: Optional[ReceiveClientAsync] = None + self._track_last_enqueued_event_properties = track_last_enqueued_event_properties + self._message_buffer: Deque[uamqp_Message] = deque() + self._last_received_event: Optional[EventData] = None self._message_buffer_lock = asyncio.Lock() self._last_callback_called_time = None self._callback_task_run = None - def _create_handler(self, auth: "JWTTokenAuthAsync") -> None: - source = Source(self._source, filters={}) - if self._offset is not None: - filter_key = ApacheFilters.selector_filter - source.filters[filter_key] = ( - filter_key, - pyamqp_utils.amqp_string_value( - event_position_selector( - self._offset, - self._offset_inclusive - ) - ) - ) + def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) -> None: + source = self._amqp_transport.create_source( + self._source, self._offset, event_position_selector(self._offset, self._offset_inclusive) + ) desired_capabilities = [RECEIVER_RUNTIME_METRIC_SYMBOL] if self._track_last_enqueued_event_properties else None - custom_endpoint_address = self._client._config.custom_endpoint_address - transport_type = self._client._config.transport_type # pylint:disable=protected-access - hostname = urlparse(source.address).hostname - if transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - self._handler = ReceiveClientAsync( - hostname, - source, + self._handler = self._amqp_transport.create_receive_client( + config=self._client._config, # pylint:disable=protected-access + source=source, auth=auth, - idle_timeout=self._idle_timeout, network_trace=self._client._config.network_tracing, # pylint:disable=protected-access link_credit=self._prefetch, link_properties=self._link_properties, - transport_type=transport_type, - http_proxy=self._client._config.http_proxy, # pylint:disable=protected-access + idle_timeout=self._idle_timeout, retry_policy=self._retry_policy, + keep_alive_interval=self._keep_alive, client_name=self._name, - receive_settle_mode=pyamqp_constants.ReceiverSettleMode.First, - properties=create_properties(self._client._config.user_agent), # pylint:disable=protected-access + properties=create_properties( + self._client._config.user_agent, amqp_transport=self._amqp_transport # pylint:disable=protected-access + ), desired_capabilities=desired_capabilities, streaming_receive=True, - message_received_callback=self._message_received, - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._client._config.connection_verify, + message_received_callback=partial(self._amqp_transport.message_received_async, self), ) async def _open_with_retry(self) -> None: await self._do_retryable_operation(self._open, operation_need_param=False) - async def _message_received(self, message: Message) -> None: - async with self._message_buffer_lock: - self._message_buffer.append(message) - def _next_message_in_buffer(self): # pylint:disable=protected-access message = self._message_buffer.popleft() @@ -186,64 +155,5 @@ def _next_message_in_buffer(self): self._last_received_event = event_data return event_data - async def _callback_task(self, batch, max_batch_size, max_wait_time): - while self._callback_task_run: - async with self._message_buffer_lock: - messages = [ - self._message_buffer.popleft() for _ in range(min(max_batch_size, len(self._message_buffer))) - ] - events = [EventData._from_message(message) for message in messages] - now_time = time.time() - if len(events) > 0: - await self._on_event_received(events if batch else events[0]) - self._last_callback_called_time = now_time - else: - if max_wait_time and (now_time - self._last_callback_called_time) > max_wait_time: - # no events received, and need to callback - await self._on_event_received([] if batch else None) - self._last_callback_called_time = now_time - # backoff a bit to avoid throttling CPU when no events are coming - await asyncio.sleep(0.05) - - async def _receive_task(self): - max_retries = ( - self._client._config.max_retries # pylint:disable=protected-access - ) - retried_times = 0 - while retried_times <= max_retries: - try: - await self._open() - await cast(ReceiveClientAsync, self._handler).do_work_async(batch=self._prefetch) - except asyncio.CancelledError: # pylint: disable=try-except-raise - raise - except Exception as exception: # pylint: disable=broad-except - if ( - isinstance(exception, error.AMQPLinkError) - and exception.condition == error.ErrorCondition.LinkStolen # pylint: disable=no-member - ): - raise await self._handle_exception(exception) - if not self.running: # exit by close - return - if self._last_received_event: - self._offset = self._last_received_event.offset - last_exception = await self._handle_exception(exception) - retried_times += 1 - if retried_times > max_retries: - _LOGGER.info( - "%r operation has exhausted retry. Last exception: %r.", - self._name, - last_exception, - ) - raise last_exception - async def receive(self, batch=False, max_batch_size=300, max_wait_time=None) -> None: - self._callback_task_run = True - self._last_callback_called_time = time.time() - callback_task = asyncio.ensure_future(self._callback_task(batch, max_batch_size, max_wait_time)) - receive_task = asyncio.ensure_future(self._receive_task()) - - try: - await receive_task - finally: - self._callback_task_run = False - await callback_task + await self._amqp_transport.receive_messages_async(self, batch, max_batch_size, max_wait_time) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py index 1671d4c8cdcd..73a91416b1a4 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from __future__ import annotations import asyncio import logging import datetime @@ -21,12 +22,12 @@ from ._eventprocessor.event_processor import EventProcessor from ._consumer_async import EventHubConsumer from ._client_base_async import ClientBaseAsync -from .._constants import ALL_PARTITIONS +from .._constants import ALL_PARTITIONS, TransportType from .._eventprocessor.common import LoadBalancingStrategy if TYPE_CHECKING: - from azure.core.credentials_async import AsyncTokenCredential + from ._client_base_async import CredentialTypes from ._eventprocessor.partition_context import PartitionContext from ._eventprocessor.checkpoint_store import CheckpointStore from .._common import EventData @@ -35,7 +36,9 @@ _LOGGER = logging.getLogger(__name__) -class EventHubConsumerClient(ClientBaseAsync): # pylint: disable=client-accepts-api-version-keyword +class EventHubConsumerClient( + ClientBaseAsync +): # pylint: disable=client-accepts-api-version-keyword """The EventHubConsumerClient class defines a high level interface for receiving events from the Azure Event Hubs service. @@ -127,7 +130,9 @@ class EventHubConsumerClient(ClientBaseAsync): # pylint: disable=client-accepts If port is not specified in the `custom_endpoint_address`, by default port 443 will be used. :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. - Default is None in which case `certifi.where()` will be used. + Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -155,10 +160,17 @@ def __init__( "partition_ownership_expiration_interval", None ) if self._partition_ownership_expiration_interval is None: - self._partition_ownership_expiration_interval = 6 * self._load_balancing_interval - load_balancing_strategy = kwargs.pop("load_balancing_strategy", None) or LoadBalancingStrategy.GREEDY - self._load_balancing_strategy = LoadBalancingStrategy(load_balancing_strategy) if load_balancing_strategy \ + self._partition_ownership_expiration_interval = ( + 6 * self._load_balancing_interval + ) + load_balancing_strategy = ( + kwargs.pop("load_balancing_strategy", None) or LoadBalancingStrategy.GREEDY + ) + self._load_balancing_strategy = ( + LoadBalancingStrategy(load_balancing_strategy) + if load_balancing_strategy else LoadBalancingStrategy.GREEDY + ) self._consumer_group = consumer_group network_tracing = kwargs.pop("logging_enable", False) super(EventHubConsumerClient, self).__init__( @@ -166,7 +178,7 @@ def __init__( eventhub_name=eventhub_name, credential=credential, network_tracing=network_tracing, - **kwargs + **kwargs, ) self._lock = asyncio.Lock(**self._internal_kwargs) self._event_processors = dict() # type: Dict[Tuple[str, str], EventProcessor] @@ -205,7 +217,8 @@ def _create_consumer( prefetch=prefetch, idle_timeout=self._idle_timeout, track_last_enqueued_event_properties=track_last_enqueued_event_properties, - **self._internal_kwargs + amqp_transport=self._amqp_transport, + **self._internal_kwargs, ) return handler @@ -221,6 +234,7 @@ def from_connection_string( auth_timeout: float = 60, user_agent: Optional[str] = None, retry_total: int = 3, + transport_type: TransportType = TransportType.Amqp, checkpoint_store: Optional["CheckpointStore"] = None, load_balancing_interval: float = 10, **kwargs: Any @@ -296,7 +310,6 @@ def from_connection_string( Default is None in which case `certifi.where()` will be used. :rtype: ~azure.eventhub.aio.EventHubConsumerClient - .. admonition:: Example: .. literalinclude:: ../samples/async_samples/sample_code_eventhub_async.py @@ -316,36 +329,37 @@ def from_connection_string( auth_timeout=auth_timeout, user_agent=user_agent, retry_total=retry_total, + transport_type=transport_type, checkpoint_store=checkpoint_store, load_balancing_interval=load_balancing_interval, - **kwargs + **kwargs, ) return cls(**constructor_args) async def _receive( - self, - on_event, - batch=False, - *, - max_batch_size: int = 300, - max_wait_time: Optional[float] = None, - partition_id: Optional[str] = None, - owner_level: Optional[int] = None, - prefetch: int = 300, - track_last_enqueued_event_properties: bool = False, - starting_position: Optional[ - Union[str, int, datetime.datetime, Dict[str, Any]] - ] = None, - starting_position_inclusive: Union[bool, Dict[str, bool]] = False, - on_error: Optional[ - Callable[["PartitionContext", Exception], Awaitable[None]] - ] = None, - on_partition_initialize: Optional[ - Callable[["PartitionContext"], Awaitable[None]] - ] = None, - on_partition_close: Optional[ - Callable[["PartitionContext", "CloseReason"], Awaitable[None]] - ] = None + self, + on_event, + batch=False, + *, + max_batch_size: int = 300, + max_wait_time: Optional[float] = None, + partition_id: Optional[str] = None, + owner_level: Optional[int] = None, + prefetch: int = 300, + track_last_enqueued_event_properties: bool = False, + starting_position: Optional[ + Union[str, int, datetime.datetime, Dict[str, Any]] + ] = None, + starting_position_inclusive: Union[bool, Dict[str, bool]] = False, + on_error: Optional[ + Callable[["PartitionContext", Exception], Awaitable[None]] + ] = None, + on_partition_initialize: Optional[ + Callable[["PartitionContext"], Awaitable[None]] + ] = None, + on_partition_close: Optional[ + Callable[["PartitionContext", "CloseReason"], Awaitable[None]] + ] = None ): async with self._lock: error = None @@ -357,7 +371,7 @@ async def _receive( ) ) elif partition_id is None and any( - x[0] == self._consumer_group for x in self._event_processors + x[0] == self._consumer_group for x in self._event_processors ): error = ( "This consumer client is already receiving events " @@ -389,12 +403,14 @@ async def _receive( load_balancing_interval=self._load_balancing_interval, load_balancing_strategy=self._load_balancing_strategy, partition_ownership_expiration_interval=self._partition_ownership_expiration_interval, - initial_event_position=starting_position if starting_position is not None else "@latest", + initial_event_position=starting_position + if starting_position is not None + else "@latest", initial_event_position_inclusive=starting_position_inclusive or False, owner_level=owner_level, prefetch=prefetch, track_last_enqueued_event_properties=track_last_enqueued_event_properties, - **self._internal_kwargs + **self._internal_kwargs, ) self._event_processors[ (self._consumer_group, partition_id or ALL_PARTITIONS) @@ -413,7 +429,9 @@ async def _receive( async def receive( self, - on_event: Callable[["PartitionContext", Optional["EventData"]], Awaitable[None]], + on_event: Callable[ + ["PartitionContext", Optional["EventData"]], Awaitable[None] + ], *, max_wait_time: Optional[float] = None, partition_id: Optional[str] = None, @@ -517,12 +535,14 @@ async def receive( starting_position_inclusive=starting_position_inclusive, on_error=on_error, on_partition_initialize=on_partition_initialize, - on_partition_close=on_partition_close + on_partition_close=on_partition_close, ) async def receive_batch( self, - on_event_batch: Callable[["PartitionContext", List["EventData"]], Awaitable[None]], + on_event_batch: Callable[ + ["PartitionContext", List["EventData"]], Awaitable[None] + ], *, max_batch_size: int = 300, max_wait_time: Optional[float] = None, @@ -635,7 +655,7 @@ async def receive_batch( starting_position_inclusive=starting_position_inclusive, on_error=on_error, on_partition_initialize=on_partition_initialize, - on_partition_close=on_partition_close + on_partition_close=on_partition_close, ) async def get_eventhub_properties(self) -> Dict[str, Any]: diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_error_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_error_async.py deleted file mode 100644 index ebcca5b12860..000000000000 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_error_async.py +++ /dev/null @@ -1,78 +0,0 @@ -# -------------------------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See License.txt in the project root for license information. -# -------------------------------------------------------------------------------------------- -import asyncio -import logging -from typing import TYPE_CHECKING, Union, cast - -from ..exceptions import ( - _create_eventhub_exception, - EventHubError, - EventDataSendError, - EventDataError, -) - -from .._pyamqp import error as errors - -if TYPE_CHECKING: - from ._client_base_async import ClientBaseAsync, ConsumerProducerMixin - -_LOGGER = logging.getLogger(__name__) - - -async def _handle_exception( # pylint:disable=too-many-branches, too-many-statements - exception: Exception, closable: Union["ClientBaseAsync", "ConsumerProducerMixin"] -) -> Exception: - # pylint: disable=protected-access - if isinstance(exception, asyncio.CancelledError): - raise exception - error = exception - try: - name = cast("ConsumerProducerMixin", closable)._name - except AttributeError: - name = cast("ClientBaseAsync", closable)._container_id - if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise - _LOGGER.info("%r stops due to keyboard interrupt", name) - await cast("ConsumerProducerMixin", closable)._close_connection_async() - raise error - elif isinstance(exception, EventHubError): - await cast("ConsumerProducerMixin", closable)._close_handler_async() - raise error - # TODO: The following errors seem to be useless in EH - # elif isinstance( - # exception, - # ( - # errors.MessageAccepted, - # errors.MessageAlreadySettled, - # errors.MessageModified, - # errors.MessageRejected, - # errors.MessageReleased, - # errors.MessageContentTooLarge, - # ), - # ): - # _LOGGER.info("%r Event data error (%r)", name, exception) - # error = EventDataError(str(exception), exception) - # raise error - elif isinstance(exception, errors.MessageException): - _LOGGER.info("%r Event data send error (%r)", name, exception) - error = EventDataSendError(str(exception), exception) - raise error - else: - try: - if isinstance(exception, errors.AuthenticationException): - await closable._close_connection_async() # pylint:disable=protected-access - elif isinstance(exception, errors.AMQPLinkError): - await closable._close_handler_async() # pylint:disable=protected-access - elif isinstance(exception, errors.AMQPConnectionError): - await closable._close_connection_async() # pylint:disable=protected-access - # TODO: add MessageHandlerError in amqp? - # elif isinstance(exception, errors.MessageHandlerError): - # if hasattr(closable, "_close_handler"): - # closable._close_handler() # pylint:disable=protected-access - else: # errors.AMQPConnectionError, compat.TimeoutException - await closable._close_connection_async() # pylint:disable=protected-access - return _create_eventhub_exception(exception) - except AttributeError: - pass - return _create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/_ownership_manager.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/_ownership_manager.py index ab973d9879e1..01227e2e2524 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/_ownership_manager.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/_ownership_manager.py @@ -52,8 +52,7 @@ def __init__( self.partition_id = partition_id async def claim_ownership(self) -> List[str]: - """Claims ownership for this EventProcessor - """ + """Claims ownership for this EventProcessor""" if not self.cached_parition_ids: await self._retrieve_partition_ids() @@ -100,15 +99,13 @@ async def release_ownership(self, partition_id: str) -> None: await self.checkpoint_store.claim_ownership(partition_ownership) async def _retrieve_partition_ids(self) -> None: - """List all partition ids of the event hub that the EventProcessor is working on. - """ + """List all partition ids of the event hub that the EventProcessor is working on.""" self.cached_parition_ids = await self.eventhub_client.get_partition_ids() def _balance_ownership( # pylint:disable=too-many-locals self, ownership_list: Iterable[Dict[str, Any]], all_partition_ids: List[str] ) -> List[Dict[str, Any]]: - """Balances and claims ownership of partitions for this EventProcessor. - """ + """Balances and claims ownership of partitions for this EventProcessor.""" now = time.time() ownership_dict = { x["partition_id"]: x for x in ownership_list @@ -147,7 +144,7 @@ def _balance_ownership( # pylint:disable=too-many-locals # Py2 math.ceil() returns float, a/b return int if not divisable. # Py3 math.ceil() returns int, a/b return float if not divisable. # Even though this is py3 code, make it the same as the py2-compatible sync code. - max_count_per_owner = int(math.ceil(all_partition_count*1.0 / owners_count)) + max_count_per_owner = int(math.ceil(all_partition_count * 1.0 / owners_count)) # end of calculating expected count per owner to_claim = active_ownership_self @@ -156,9 +153,11 @@ def _balance_ownership( # pylint:disable=too-many-locals if self.load_balancing_strategy is LoadBalancingStrategy.GREEDY: # Greedily claim more partitions if there are claimable partitions to_greedy_claim_ids = random.sample( - claimable_partition_ids, k=min( - max_count_per_owner - len(active_ownership_self), len(claimable_partition_ids) - ) + claimable_partition_ids, + k=min( + max_count_per_owner - len(active_ownership_self), + len(claimable_partition_ids), + ), ) if to_greedy_claim_ids: for pid in to_greedy_claim_ids: @@ -194,13 +193,9 @@ def _balance_ownership( # pylint:disable=too-many-locals active_ownership_count_group_by_owner = Counter( dict((x, len(y)) for x, y in active_ownership_by_owner.items()) ) - most_frequent_owner_id = active_ownership_count_group_by_owner.most_common( - 1 - )[ - 0 - ][ - 0 - ] + most_frequent_owner_id = ( + active_ownership_count_group_by_owner.most_common(1)[0][0] + ) # randomly choose a partition to steal from the most_frequent_owner to_steal_partition = random.choice( active_ownership_by_owner[most_frequent_owner_id] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/checkpoint_store.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/checkpoint_store.py index 0f951813b16f..b56c3578fec2 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/checkpoint_store.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/checkpoint_store.py @@ -15,7 +15,11 @@ class CheckpointStore(ABC): @abstractmethod async def list_ownership( - self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any + self, + fully_qualified_namespace: str, + eventhub_name: str, + consumer_group: str, + **kwargs: Any ) -> Iterable[Dict[str, Any]]: """Retrieves a complete ownership list from the chosen storage service. @@ -87,7 +91,11 @@ async def update_checkpoint( @abstractmethod async def list_checkpoints( - self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any + self, + fully_qualified_namespace: str, + eventhub_name: str, + consumer_group: str, + **kwargs: Any ) -> Iterable[Dict[str, Any]]: """List the updated checkpoints from the chosen storage service. diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py index 0fdb7f3368ae..69e8567b7a86 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/event_processor.py @@ -52,7 +52,10 @@ def __init__( self, eventhub_client: "EventHubConsumerClient", consumer_group: str, - event_handler: Callable[[PartitionContext, Union[Optional[EventData], List[EventData]]], Awaitable[None]], + event_handler: Callable[ + [PartitionContext, Union[Optional[EventData], List[EventData]]], + Awaitable[None], + ], *, batch: Optional[bool] = False, max_batch_size: Optional[int] = 300, @@ -96,10 +99,14 @@ def __init__( self._initial_event_position = initial_event_position self._initial_event_position_inclusive = initial_event_position_inclusive self._load_balancing_interval = load_balancing_interval - self._ownership_timeout = partition_ownership_expiration_interval \ - if partition_ownership_expiration_interval is not None \ + self._ownership_timeout = ( + partition_ownership_expiration_interval + if partition_ownership_expiration_interval is not None else self._load_balancing_interval * 6 - self._load_balancing_strategy = load_balancing_strategy or LoadBalancingStrategy.GREEDY + ) + self._load_balancing_strategy = ( + load_balancing_strategy or LoadBalancingStrategy.GREEDY + ) self._tasks = {} # type: Dict[str, asyncio.Task] self._partition_contexts = {} # type: Dict[str, PartitionContext] self._owner_level = owner_level @@ -133,7 +140,7 @@ async def _cancel_tasks_for_partitions( _LOGGER.debug( "EventProcessor %r tries to cancel partitions %r", self._id, - to_cancel_partitions + to_cancel_partitions, ) for partition_id in to_cancel_partitions: task = self._tasks.get(partition_id) @@ -142,9 +149,11 @@ async def _cancel_tasks_for_partitions( _LOGGER.info( "EventProcessor %r has cancelled partition %r", self._id, - partition_id + partition_id, ) - if partition_id not in self._consumers: # task is cancelled before the consumer is created + if ( + partition_id not in self._consumers + ): # task is cancelled before the consumer is created del self._tasks[partition_id] def _create_tasks_for_claimed_ownership( @@ -155,7 +164,7 @@ def _create_tasks_for_claimed_ownership( _LOGGER.debug( "EventProcessor %r tries to claim partition %r", self._id, - claimed_partitions + claimed_partitions, ) for partition_id in claimed_partitions: if partition_id not in self._tasks or self._tasks[partition_id].done(): @@ -167,7 +176,7 @@ def _create_tasks_for_claimed_ownership( _LOGGER.info( "EventProcessor %r has claimed partition %r", self._id, - partition_id + partition_id, ) async def _process_error( @@ -215,7 +224,9 @@ async def _close_partition( await self._process_error(partition_context, err) async def _on_event_received( - self, partition_context: PartitionContext, event: Union[Optional[EventData], List[EventData]] + self, + partition_context: PartitionContext, + event: Union[Optional[EventData], List[EventData]], ) -> None: if event: try: @@ -253,7 +264,9 @@ async def _receive( ) = self.get_init_event_position(partition_id, checkpoint) if partition_id in self._partition_contexts: partition_context = self._partition_contexts[partition_id] - partition_context._last_received_event = None # pylint:disable=protected-access + partition_context._last_received_event = ( # pylint:disable=protected-access + None + ) else: partition_context = PartitionContext( self._namespace, @@ -379,7 +392,7 @@ async def start(self) -> None: self._eventhub_name, self._consumer_group, err, - load_balancing_interval + load_balancing_interval, ) await self._process_error(None, err) # type: ignore diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/in_memory_checkpoint_store.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/in_memory_checkpoint_store.py index 22ef721c0ee0..f38af0274eb2 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/in_memory_checkpoint_store.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/in_memory_checkpoint_store.py @@ -3,7 +3,9 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # ----------------------------------------------------------------------------------- from typing import Dict, Any, Iterable, Optional, Union -from azure.eventhub._eventprocessor.in_memory_checkpoint_store import InMemoryCheckpointStore as CheckPointStoreImpl +from azure.eventhub._eventprocessor.in_memory_checkpoint_store import ( + InMemoryCheckpointStore as CheckPointStoreImpl, +) from .checkpoint_store import CheckpointStore @@ -12,9 +14,15 @@ def __init__(self): self._checkpoint_store_impl = CheckPointStoreImpl() async def list_ownership( - self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any + self, + fully_qualified_namespace: str, + eventhub_name: str, + consumer_group: str, + **kwargs: Any ) -> Iterable[Dict[str, Any]]: - return self._checkpoint_store_impl.list_ownership(fully_qualified_namespace, eventhub_name, consumer_group) + return self._checkpoint_store_impl.list_ownership( + fully_qualified_namespace, eventhub_name, consumer_group + ) async def claim_ownership( self, ownership_list: Iterable[Dict[str, Any]], **kwargs: Any @@ -27,6 +35,12 @@ async def update_checkpoint( self._checkpoint_store_impl.update_checkpoint(checkpoint) async def list_checkpoints( - self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any + self, + fully_qualified_namespace: str, + eventhub_name: str, + consumer_group: str, + **kwargs: Any ) -> Iterable[Dict[str, Any]]: - return self._checkpoint_store_impl.list_checkpoints(fully_qualified_namespace, eventhub_name, consumer_group) + return self._checkpoint_store_impl.list_checkpoints( + fully_qualified_namespace, eventhub_name, consumer_group + ) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/partition_context.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/partition_context.py index 3f45981d6252..3c49318c63e3 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/partition_context.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/partition_context.py @@ -56,7 +56,9 @@ def last_enqueued_event_properties(self) -> Optional[Dict[str, Any]]: return get_last_enqueued_event_properties(self._last_received_event) return None - async def update_checkpoint(self, event: Optional["EventData"] = None, **kwargs: Any) -> None: + async def update_checkpoint( + self, event: Optional["EventData"] = None, **kwargs: Any + ) -> None: """Updates the receive checkpoint to the given events offset. :param ~azure.eventhub.EventData event: The EventData instance which contains the offset and diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/utils.py index 140e1da3e677..d341442be704 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/utils.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_eventprocessor/utils.py @@ -11,7 +11,7 @@ def get_running_loop() -> asyncio.AbstractEventLoop: return asyncio.get_running_loop() except AttributeError: # 3.5 / 3.6 loop = ( - asyncio._get_running_loop() # pylint: disable=protected-access, no-member + asyncio._get_running_loop() # pylint: disable=protected-access, no-member ) if loop is None: raise RuntimeError("No running event loop") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py index 819eef1b072b..a5bb2ee2c68f 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_async.py @@ -2,43 +2,45 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- +from __future__ import annotations import uuid import asyncio import logging from typing import Iterable, Union, Optional, Any, AnyStr, List, TYPE_CHECKING -import time from azure.core.tracing import AbstractSpan from .._common import EventData, EventDataBatch -from ..exceptions import OperationTimeoutError from .._producer import _set_partition_key, _set_trace_message from .._utils import ( create_properties, - set_message_partition_key, trace_message, send_context_manager, transform_outbound_single_message, ) -from .._constants import TIMEOUT_SYMBOL, NO_RETRY_ERRORS, CUSTOM_CONDITION_BACKOFF +from .._constants import TIMEOUT_SYMBOL +from ..amqp import AmqpAnnotatedMessage from ._client_base_async import ConsumerProducerMixin from ._async_utils import get_dict_with_loop_if_needed -from .._pyamqp import ( - error, - utils as pyamqp_utils, -) -from .._pyamqp.aio import SendClientAsync if TYPE_CHECKING: - from uamqp.authentication import JWTTokenAsync # pylint: disable=ungrouped-imports + try: + from uamqp import constants, SendClientAsync as uamqp_SendClientAsync + from uamqp.constants import MessageSendResult as uamqp_MessageSendResult + from uamqp.authentication import JWTTokenAsync as uamqp_JWTTokenAsync + except ImportError: + uamqp_MessageSendResult = None + uamqp_SendClientAsync = None + uamqp_JWTTokenAsync = None + + from .._pyamqp.aio._client_async import SendClientAsync + from .._pyamqp.aio._authentication_async import JWTTokenAuthAsync from ._producer_client_async import EventHubProducerClient _LOGGER = logging.getLogger(__name__) -class EventHubProducer( - ConsumerProducerMixin -): # pylint: disable=too-many-instance-attributes +class EventHubProducer(ConsumerProducerMixin): # pylint: disable=too-many-instance-attributes """A producer responsible for transmitting batches of EventData to a specific Event Hub. Depending on the options specified at creation, the producer may @@ -61,8 +63,9 @@ class EventHubProducer( Default value is `True`. """ - def __init__(self, client: "EventHubProducerClient", target: str, **kwargs) -> None: + def __init__(self, client: EventHubProducerClient, target: str, **kwargs) -> None: super().__init__() + self._amqp_transport = kwargs.pop("amqp_transport") partition = kwargs.get("partition", None) send_timeout = kwargs.get("send_timeout", 60) keep_alive = kwargs.get("keep_alive", None) @@ -80,12 +83,9 @@ def __init__(self, client: "EventHubProducerClient", target: str, **kwargs) -> N self._keep_alive = keep_alive self._auto_reconnect = auto_reconnect self._timeout = send_timeout - self._idle_timeout = idle_timeout - self._retry_policy = error.RetryPolicy( - retry_total=self._client._config.max_retries, # pylint: disable=protected-access - no_retry_condition=NO_RETRY_ERRORS, - custom_condition_backoff=CUSTOM_CONDITION_BACKOFF - ) + self._idle_timeout = (idle_timeout * self._amqp_transport.TIMEOUT_FACTOR) if idle_timeout else None + + self._retry_policy = self._amqp_transport.create_retry_policy(config=self._client._config) self._reconnect_backoff = 1 self._name = "EHProducer-{}".format(uuid.uuid4()) self._unsent_events = [] # type: List[Any] @@ -93,97 +93,108 @@ def __init__(self, client: "EventHubProducerClient", target: str, **kwargs) -> N if partition: self._target += "/Partitions/" + partition self._name += "-partition{}".format(partition) - self._handler = None # type: Optional[SendClientAsync] - self._condition = None # type: Optional[Exception] + self._handler: Optional[Union[uamqp_SendClientAsync, SendClientAsync]] = None + self._outcome: Optional[uamqp_MessageSendResult] = None + self._condition: Optional[Exception] = None self._lock = asyncio.Lock(**self._internal_kwargs) - self._link_properties = {TIMEOUT_SYMBOL: pyamqp_utils.amqp_long_value(int(self._timeout * 1000))} - - def _create_handler(self, auth: "JWTTokenAsync") -> None: - transport_type = self._client._config.transport_type # pylint:disable=protected-access - custom_endpoint_address = self._client._config.custom_endpoint_address # pylint: disable=protected-access - hostname = self._client._address.hostname # pylint: disable=protected-access - if transport_type.name == 'AmqpOverWebsocket': - hostname += '/$servicebus/websocket/' - if custom_endpoint_address: - custom_endpoint_address += '/$servicebus/websocket/' - self._handler = SendClientAsync( - hostname, - self._target, + self._link_properties = self._amqp_transport.create_link_properties( + {TIMEOUT_SYMBOL: int(self._timeout * self._amqp_transport.TIMEOUT_FACTOR)} + ) + + def _create_handler(self, auth: Union[uamqp_JWTTokenAsync, JWTTokenAuthAsync]) -> None: + self._handler = self._amqp_transport.create_send_client( + config=self._client._config, # pylint:disable=protected-access + target=self._target, auth=auth, + network_trace=self._client._config.network_tracing, # pylint:disable=protected-access idle_timeout=self._idle_timeout, - network_trace=self._client._config.network_tracing, # pylint: disable=protected-access retry_policy=self._retry_policy, keep_alive_interval=self._keep_alive, - transport_type=transport_type, - http_proxy=self._client._config.http_proxy, # pylint:disable=protected-access client_name=self._name, link_properties=self._link_properties, - properties=create_properties(self._client._config.user_agent), # pylint: disable=protected-access - custom_endpoint_address=custom_endpoint_address, - connection_verify=self._client._config.connection_verify, - **self._internal_kwargs + properties=create_properties( + self._client._config.user_agent, # pylint: disable=protected-access + amqp_transport=self._amqp_transport, + ), + msg_timeout=self._timeout * self._amqp_transport.TIMEOUT_FACTOR, ) async def _open_with_retry(self) -> Any: - return await self._do_retryable_operation( - self._open, operation_need_param=False - ) + return await self._do_retryable_operation(self._open, operation_need_param=False) async def _send_event_data( self, timeout_time: Optional[float] = None, last_exception: Optional[Exception] = None, ) -> None: - # TODO: Correct uAMQP type hints if self._unsent_events: - await self._open() - timeout = timeout_time - time.time() if timeout_time else 0 - await self._handler.send_message_async(self._unsent_events[0], timeout=timeout) - self._unsent_events = None + await self._amqp_transport.send_messages_async(self, timeout_time, last_exception, _LOGGER) - async def _send_event_data_with_retry( - self, timeout: Optional[float] = None - ) -> None: + async def _send_event_data_with_retry(self, timeout: Optional[float] = None) -> None: await self._do_retryable_operation(self._send_event_data, timeout=timeout) + def _on_outcome(self, outcome: uamqp_MessageSendResult, condition: Optional[Exception]) -> None: + """ + ONLY USED FOR uamqp_transport=True. Called when the outcome is received for a delivery. + + :param outcome: The outcome of the message delivery - success or failure. + :type outcome: ~uamqp.constants.MessageSendResult + :param condition: Detail information of the outcome. + + """ + self._outcome = outcome + self._condition = condition + def _wrap_eventdata( self, - event_data: Union[EventData, EventDataBatch, Iterable[EventData]], + event_data: Union[EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData]], span: Optional[AbstractSpan], partition_key: Optional[AnyStr], ) -> Union[EventData, EventDataBatch]: - if isinstance(event_data, EventData): - outgoing_event_data = transform_outbound_single_message(event_data, EventData) + if isinstance(event_data, (EventData, AmqpAnnotatedMessage)): + outgoing_event_data = transform_outbound_single_message( + event_data, EventData, self._amqp_transport.to_outgoing_amqp_message + ) if partition_key: - set_message_partition_key(outgoing_event_data.message, partition_key) + self._amqp_transport.set_message_partition_key( + outgoing_event_data._message, partition_key # pylint: disable=protected-access + ) wrapper_event_data = outgoing_event_data trace_message(wrapper_event_data, span) else: - if isinstance( - event_data, EventDataBatch - ): # The partition_key in the param will be omitted. - if ( - partition_key and partition_key != event_data._partition_key # pylint: disable=protected-access - ): - raise ValueError( - "The partition_key does not match the one of the EventDataBatch" + if isinstance(event_data, EventDataBatch): # The partition_key in the param will be omitted. + if not event_data: + return event_data + # If AmqpTransports are not the same, create batch with correct BatchMessage. + if self._amqp_transport.TIMEOUT_FACTOR != event_data._amqp_transport.TIMEOUT_FACTOR: # pylint: disable=protected-access + # pylint: disable=protected-access + event_data = EventDataBatch._from_batch( + event_data._internal_events, + amqp_transport=self._amqp_transport, + partition_key=event_data._partition_key, + partition_id=event_data._partition_id, + max_size_in_bytes=event_data.max_size_in_bytes, ) - for event in event_data.message.data: # pylint: disable=protected-access + if partition_key and partition_key != event_data._partition_key: # pylint: disable=protected-access + raise ValueError("The partition_key does not match the one of the EventDataBatch") + for event in event_data._message.data: # pylint: disable=protected-access trace_message(event, span) wrapper_event_data = event_data # type:ignore else: if partition_key: - event_data = _set_partition_key(event_data, partition_key) + event_data = _set_partition_key(event_data, partition_key, self._amqp_transport) event_data = _set_trace_message(event_data, span) - wrapper_event_data = EventDataBatch._from_batch(event_data, partition_key) # type: ignore # pylint: disable=protected-access + wrapper_event_data = EventDataBatch._from_batch( # type: ignore # pylint: disable=protected-access + event_data, self._amqp_transport, partition_key + ) return wrapper_event_data async def send( self, - event_data: Union[EventData, EventDataBatch, Iterable[EventData]], + event_data: Union[EventData, AmqpAnnotatedMessage, EventDataBatch, Iterable[EventData]], *, partition_key: Optional[AnyStr] = None, - timeout: Optional[float] = None + timeout: Optional[float] = None, ) -> None: """ Sends an event data and blocks until acknowledgement is @@ -213,20 +224,16 @@ async def send( with send_context_manager() as child: self._check_closed() wrapper_event_data = self._wrap_eventdata(event_data, child, partition_key) - self._unsent_events = [wrapper_event_data.message] + + if not wrapper_event_data: + return + + self._unsent_events = [wrapper_event_data._message] # pylint: disable=protected-access if child: - self._client._add_span_request_attributes( # pylint: disable=protected-access - child - ) + self._client._add_span_request_attributes(child) # pylint: disable=protected-access - try: - await self._open() - await self._handler.send_message_async(wrapper_event_data.message, timeout=timeout) - except TimeoutError as exception: - raise OperationTimeoutError(message=str(exception), details=exception) - except Exception as exception: # pylint:disable=broad-except - raise (await self._handle_exception(exception)) + await self._send_event_data_with_retry(timeout=timeout) # pylint:disable=unexpected-keyword-arg async def close(self) -> None: """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py index 9868d114b79f..098833c43f7e 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py @@ -4,27 +4,32 @@ # -------------------------------------------------------------------------------------------- import asyncio import logging +import time -from typing import Any, Union, TYPE_CHECKING, List, Optional, Dict, cast - -from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential +from typing import Any, Union, List, Optional, Dict, Callable, cast +from typing_extensions import TYPE_CHECKING, Literal, Awaitable, overload from ..exceptions import ConnectError, EventHubError from ..amqp import AmqpAnnotatedMessage from ._client_base_async import ClientBaseAsync from ._producer_async import EventHubProducer -from .._constants import ALL_PARTITIONS, MAX_MESSAGE_LENGTH_BYTES +from ._buffered_producer import BufferedProducerDispatcher +from .._utils import set_event_partition_key +from .._constants import ALL_PARTITIONS, TransportType from .._common import EventDataBatch, EventData if TYPE_CHECKING: - from azure.core.credentials_async import AsyncTokenCredential + from ._client_base_async import CredentialTypes SendEventTypes = List[Union[EventData, AmqpAnnotatedMessage]] _LOGGER = logging.getLogger(__name__) -class EventHubProducerClient(ClientBaseAsync): +class EventHubProducerClient( + ClientBaseAsync +): # pylint: disable=client-accepts-api-version-keyword + # pylint: disable=too-many-instance-attributes """ The EventHubProducerClient class defines a high level interface for sending events to the Azure Event Hubs service. @@ -38,6 +43,36 @@ class EventHubProducerClient(ClientBaseAsync): generated by the azure-identity library and objects that implement the `get_token(self, *scopes)` method. :type credential: ~azure.core.credentials_async.AsyncTokenCredential or ~azure.core.credentials.AzureSasCredential or ~azure.core.credentials.AzureNamedKeyCredential + :keyword bool buffered_mode: If True, the producer client will collect events in a buffer, efficiently batch, + then publish. Default is False. + :keyword on_success: The callback to be called once a batch has been successfully published. + The callback takes two parameters: + - `events`: The list of events that have been successfully published + - `partition_id`: The partition id that the events in the list have been published to. + The callback function should be defined like: `on_success(events, partition_id)`. + Required when `buffered_mode` is True while optional if `buffered_mode` is False. + :paramtype on_success: Optional[Callable[[SendEventTypes, Optional[str]], Awaitable[None]]] + :keyword on_error: The callback to be called once a batch has failed to be published. + Required when in `buffered_mode` is True while optional if `buffered_mode` is False. + The callback function should be defined like: `on_error(events, partition_id, error)`, where: + - `events`: The list of events that failed to be published, + - `partition_id`: The partition id that the events in the list have been tried to be published to and + - `error`: The exception related to the sending failure. + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + :paramtype on_error: Optional[Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]]] + :keyword int max_buffer_length: Buffered mode only. + The total number of events per partition that can be buffered before a flush will be triggered. + The default value is 1500 in buffered mode. + :keyword Optional[float] max_wait_time: Buffered mode only. + The amount of time to wait for a batch to be built with events in the buffer before publishing. + The default value is 1 in buffered mode. :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`. :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service. The default value is 60 seconds. If set to 0, no timeout will be enforced from the client. @@ -62,7 +97,7 @@ class EventHubProducerClient(ClientBaseAsync): If the port 5671 is unavailable/blocked in the network environment, `TransportType.AmqpOverWebsocket` could be used instead which uses port 443 for communication. :paramtype transport_type: ~azure.eventhub.TransportType - :keyword Dict http_proxy: HTTP proxy settings. This must be a dictionary with the following + :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'`. :keyword str custom_endpoint_address: The custom endpoint address to use for establishing a connection to @@ -73,6 +108,8 @@ class EventHubProducerClient(ClientBaseAsync): :keyword str connection_verify: Path to the custom CA_BUNDLE file of the SSL certificate which is used to authenticate the identity of the connection endpoint. Default is None in which case `certifi.where()` will be used. + :keyword bool uamqp_transport: Whether to use the `uamqp` library as the underlying transport. The default value is + False and the Pure Python AMQP library will be used as the underlying transport. .. admonition:: Example: @@ -84,14 +121,50 @@ class EventHubProducerClient(ClientBaseAsync): :caption: Create a new instance of the EventHubProducerClient. """ + @overload + def __init__( + self, + fully_qualified_namespace: str, + eventhub_name: str, + credential: "CredentialTypes", + *, + buffered_mode: Literal[False] = False, + **kwargs: Any + ) -> None: + ... + + @overload def __init__( self, fully_qualified_namespace: str, eventhub_name: str, - credential: Union[ - "AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential - ], - **kwargs + credential: "CredentialTypes", + *, + buffered_mode: Literal[True], + on_error: Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]], + on_success: Callable[[SendEventTypes, Optional[str]], Awaitable[None]], + max_buffer_length: int = 1500, + max_wait_time: float = 1, + **kwargs: Any + ) -> None: + ... + + def __init__( + self, + fully_qualified_namespace: str, + eventhub_name: str, + credential: "CredentialTypes", + *, + buffered_mode: bool = False, + on_error: Optional[ + Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]] + ] = None, + on_success: Optional[ + Callable[[SendEventTypes, Optional[str]], Awaitable[None]] + ] = None, + max_buffer_length: Optional[int] = None, + max_wait_time: Optional[float] = None, + **kwargs: Any ) -> None: super(EventHubProducerClient, self).__init__( fully_qualified_namespace=fully_qualified_namespace, @@ -108,6 +181,35 @@ def __init__( ) # sync the creation of self._producers self._max_message_size_on_link = 0 self._partition_ids = None # Optional[List[str]] + self._buffered_mode = buffered_mode + self._on_success = on_success + self._on_error = on_error + self._buffered_producer_dispatcher = None + self._max_buffer_length = max_buffer_length + self._max_wait_time = max_wait_time + if self._buffered_mode: + setattr(self, "send_batch", self._buffered_send_batch) + setattr(self, "send_event", self._buffered_send_event) + if not self._on_error: + raise TypeError( + "EventHubProducerClient in buffered mode missing 1 required keyword argument: 'on_error'" + ) + if not self._on_success: + raise TypeError( + "EventHubProducerClient in buffered mode missing 1 required keyword argument: 'on_success'" + ) + if self._max_wait_time is None: + self._max_wait_time = 1 + if self._max_wait_time <= 0: + raise ValueError( + "'max_wait_time' must be a float greater than 0 in buffered mode" + ) + if self._max_buffer_length is None: + self._max_buffer_length = 1500 + if self._max_buffer_length <= 0: + raise ValueError( + "'max_buffer_length' must be an integer greater than 0 in buffered mode" + ) async def __aenter__(self): return self @@ -115,6 +217,77 @@ async def __aenter__(self): async def __aexit__(self, *args): await self.close() + async def _buffered_send(self, events, **kwargs): + try: + await self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) + except AttributeError: + await self._get_partitions() + await self._get_max_message_size() + self._buffered_producer_dispatcher = BufferedProducerDispatcher( + self._partition_ids, + self._on_success, + self._on_error, + self._create_producer, + self.eventhub_name, + self._max_message_size_on_link, + max_wait_time=self._max_wait_time, + max_buffer_length=self._max_buffer_length, + amqp_transport=self._amqp_transport + ) + await self._buffered_producer_dispatcher.enqueue_events(events, **kwargs) + + async def _batch_preparer(self, event_data_batch, **kwargs): + partition_id = kwargs.pop("partition_id", None) + partition_key = kwargs.pop("partition_key", None) + + if isinstance(event_data_batch, EventDataBatch): + if partition_id or partition_key: + raise TypeError( + "partition_id and partition_key should be None when sending an EventDataBatch " + "because type EventDataBatch itself may have partition_id or partition_key" + ) + to_send_batch = event_data_batch + else: + to_send_batch = await self.create_batch( + partition_id=partition_id, partition_key=partition_key + ) + to_send_batch._load_events( # pylint:disable=protected-access + event_data_batch + ) + + return ( + to_send_batch, + to_send_batch._partition_id, # pylint:disable=protected-access + partition_key, + ) + + async def _buffered_send_batch(self, event_data_batch, **kwargs): + batch, pid, pkey = await self._batch_preparer(event_data_batch, **kwargs) + + if len(batch) == 0: + return + + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + await self._buffered_send( + event_data_batch, + partition_id=pid, + partition_key=pkey, + timeout_time=timeout_time, + ) + + async def _buffered_send_event(self, event, **kwargs): + partition_key = kwargs.get("partition_key") + set_event_partition_key(event, partition_key, self._amqp_transport) + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + await self._buffered_send( + event, + partition_id=kwargs.get("partition_id"), + partition_key=partition_key, + timeout_time=timeout_time, + ) + async def _get_partitions(self) -> None: if not self._partition_ids: self._partition_ids = await self.get_partition_ids() # type: ignore @@ -129,10 +302,12 @@ async def _get_max_message_size(self) -> None: EventHubProducer, self._producers[ALL_PARTITIONS] )._open_with_retry() self._max_message_size_on_link = ( - cast( # type: ignore + self._amqp_transport.get_remote_max_message_size( + cast( # type: ignore EventHubProducer, self._producers[ALL_PARTITIONS] - )._handler._link.remote_max_message_size - or MAX_MESSAGE_LENGTH_BYTES + )._handler + ) + or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES ) async def _start_producer( @@ -178,28 +353,97 @@ def _create_producer( partition=partition_id, send_timeout=send_timeout, idle_timeout=self._idle_timeout, + amqp_transport = self._amqp_transport, **self._internal_kwargs ) return handler @classmethod + @overload + def from_connection_string( + cls, + conn_str: str, + *, + eventhub_name: Optional[str] = None, + buffered_mode: Literal[False] = False, + **kwargs: Any + ) -> "EventHubProducerClient": + ... + + @classmethod + @overload def from_connection_string( cls, conn_str: str, *, eventhub_name: Optional[str] = None, + buffered_mode: Literal[True], + on_error: Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]], + on_success: Callable[[SendEventTypes, Optional[str]], Awaitable[None]], + max_buffer_length: int = 1500, + max_wait_time: float = 1, + **kwargs: Any + ) -> "EventHubProducerClient": + ... + + @classmethod + def from_connection_string( + cls, + conn_str: str, + *, + eventhub_name: Optional[str] = None, + buffered_mode: bool = False, + on_error: Optional[ + Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]] + ] = None, + on_success: Optional[ + Callable[[SendEventTypes, Optional[str]], Awaitable[None]] + ] = None, + max_buffer_length: Optional[int] = None, + max_wait_time: Optional[float] = None, logging_enable: bool = False, + http_proxy: Optional[Dict[str, Union[str, int]]] = None, auth_timeout: float = 60, user_agent: Optional[str] = None, retry_total: int = 3, + transport_type: TransportType = TransportType.Amqp, **kwargs: Any ) -> "EventHubProducerClient": """Create an EventHubProducerClient from a connection string. :param str conn_str: The connection string of an Event Hub. :keyword str eventhub_name: The path of the specific Event Hub to connect the client to. + :keyword bool buffered_mode: If True, the producer client will collect events in a buffer, efficiently batch, + then publish. Default is False. + :keyword on_success: The callback to be called once a batch has been successfully published. + The callback takes two parameters: + - `events`: The list of events that have been successfully published + - `partition_id`: The partition id that the events in the list have been published to. + The callback function should be defined like: `on_success(events, partition_id)`. + It is required when `buffered_mode` is True while optional if `buffered_mode` is False. + :paramtype on_success: Optional[Callable[[SendEventTypes, Optional[str]], Awaitable[None]]] + :keyword on_error: The callback to be called once a batch has failed to be published. + The callback function should be defined like: `on_error(events, partition_id, error)`, where: + - `events`: The list of events that failed to be published, + - `partition_id`: The partition id that the events in the list have been tried to be published to and + - `error`: The exception related to the sending failure. + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + :paramtype on_error: Optional[Callable[[SendEventTypes, Optional[str], Exception], Awaitable[None]]] + :keyword int max_buffer_length: Buffered mode only. + The total number of events per partition that can be buffered before a flush will be triggered. + The default value is 1500 in buffered mode. + :keyword Optional[float] max_wait_time: Buffered mode only. + The amount of time to wait for a batch to be built with events in the buffer before publishing. + The default value is 1 in buffered mode. :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`. - :keyword Dict http_proxy: HTTP proxy settings. This must be a dictionary with the following + :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'`. :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service. @@ -247,34 +491,127 @@ def from_connection_string( constructor_args = cls._from_connection_string( conn_str, eventhub_name=eventhub_name, + buffered_mode=buffered_mode, + on_success=on_success, + on_error=on_error, + max_buffer_length=max_buffer_length, + max_wait_time=max_wait_time, logging_enable=logging_enable, + http_proxy=http_proxy, auth_timeout=auth_timeout, user_agent=user_agent, retry_total=retry_total, + transport_type=transport_type, **kwargs ) return cls(**constructor_args) - async def send_batch( - self, - event_data_batch: Union[EventDataBatch, SendEventTypes], - *, - timeout: Optional[Union[int, float]] = None, - **kwargs + async def send_event( + self, event_data: Union[EventData, AmqpAnnotatedMessage], **kwargs: Any ) -> None: - """Sends event data and blocks until acknowledgement is received or operation times out. + """ + Sends an event data. + By default, the method will block until acknowledgement is received or operation times out. + If the `EventHubProducerClient` is configured to run in buffered mode, the method will enqueue the event + into local buffer and return. The producer will do automatic batching and sending in the background. + + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + + :param event_data: The `EventData` object to be sent. + :type event_data: Union[~azure.eventhub.EventData, ~azure.eventhub.amqp.AmqpAnnotatedMessage] + :keyword float timeout: The maximum wait time to send the event data in non-buffered mode or the + maximum wait time to enqueue the event data into the buffer in buffered mode. + In non-buffered mode, the default wait time specified when the producer + was created will be used. In buffered mode, the default wait time is None. + :keyword str partition_id: The specific partition ID to send to. Default is None, in which case the service + will assign to all partitions using round-robin. + A `TypeError` will be raised if partition_id is specified and event_data_batch is an `EventDataBatch` because + `EventDataBatch` itself has partition_id. + :keyword str partition_key: With the given partition_key, event data will be sent to + a particular partition of the Event Hub decided by the service. + A `TypeError` will be raised if partition_key is specified and event_data_batch is an `EventDataBatch` because + `EventDataBatch` itself has partition_key. + If both partition_id and partition_key are provided, the partition_id will take precedence. + **WARNING: Setting partition_key of non-string value on the events to be sent is discouraged + as the partition_key will be ignored by the Event Hub service and events will be assigned + to all partitions using round-robin. Furthermore, there are SDKs for consuming events which expect + partition_key to only be string type, they might fail to parse the non-string value.** + :rtype: None + :raises: :class:`AuthenticationError` + :class:`ConnectError` + :class:`ConnectionLostError` + :class:`EventDataError` + :class:`EventDataSendError` + :class:`EventHubError` + :raises OperationTimeoutError: If the value specified by the timeout parameter elapses before the event can be + sent in non-buffered mode or the events can not be enqueued into the buffered in buffered mode. + """ + input_pid = kwargs.get("partition_id") + pid = input_pid or ALL_PARTITIONS + partition_key = kwargs.get("partition_key") + timeout = kwargs.get("timeout") + try: + try: + await cast(EventHubProducer, self._producers[pid]).send( + event_data, partition_key=partition_key, timeout=timeout + ) + except (KeyError, AttributeError, EventHubError): + await self._start_producer(pid, timeout) + await cast(EventHubProducer, self._producers[pid]).send( + event_data, partition_key=partition_key, timeout=timeout + ) + if self._on_success: + await self._on_success([event_data], input_pid) + except Exception as exc: # pylint: disable=broad-except + if self._on_error: + await self._on_error([event_data], input_pid, exc) + else: + raise - If you're sending a finite list of `EventData` or `AmqpAnnotatedMessage` and you know it's within the event hub - frame size limit, you can send them with a `send_batch` call. Otherwise, use :meth:`create_batch` + async def send_batch( + self, event_data_batch: Union[EventDataBatch, SendEventTypes], **kwargs: Any + ) -> None: + # pylint: disable=protected-access + """ + Sends a batch of event data. + By default, the method will block until acknowledgement is received or operation times out. + If the `EventHubProducerClient` is configured to run in buffered mode, the method will enqueue the events + into local buffer and return. The producer will do automatic sending in the background. + + If `buffered_mode` is False, `on_error` callback is optional and errors will be handled as follows: + - If an `on_error` callback is passed during the producer client instantiation, + then error information will be passed to the `on_error` callback, which will then be called. + - If an `on_error` callback is not passed in during client instantiation, + then the error will be raised by default. + + If `buffered_mode` is True, `on_error` callback is required and errors will be handled as follows: + - If events fail to enqueue within the given timeout, then an error will be directly raised. + - If events fail to send after enqueuing successfully, the `on_error` callback will be called. + + In buffered mode, sending a batch will remain intact and sent as a single unit. + The batch will not be rearranged. This may result in inefficiency of sending events. + + If you're sending a finite list of `EventData` or `AmqpAnnotatedMessage` and you know it's within the + event hub frame size limit, you can send them with a `send_batch` call. Otherwise, use :meth:`create_batch` to create `EventDataBatch` and add either `EventData` or `AmqpAnnotatedMessage` into the batch one by one until the size limit, and then call this method to send out the batch. - :param event_data_batch: The `EventDataBatch` object to be sent or a list of `EventData` to be sent - in a batch. All `EventData` in the list or `EventDataBatch` will land on the same partition. + :param event_data_batch: The `EventDataBatch` object to be sent or a list of `EventData` to be sent in a batch. + All `EventData` or `AmqpAnnotatedMessage` in the list or `EventDataBatch` will land on the same partition. :type event_data_batch: Union[~azure.eventhub.EventDataBatch, List[Union[~azure.eventhub.EventData, ~azure.eventhub.amqp.AmqpAnnotatedMessage]] - :keyword float timeout: The maximum wait time to send the event data. - If not specified, the default wait time specified when the producer was created will be used. + :keyword float timeout: The maximum wait time to send the event data in non-buffered mode or the + maximum wait time to enqueue the event data into the buffer in buffered mode. + In non-buffered mode, the default wait time specified when the producer + was created will be used. In buffered mode, the default wait time is None. :keyword str partition_id: The specific partition ID to send to. Default is None, in which case the service will assign to all partitions using round-robin. A `TypeError` will be raised if partition_id is specified and event_data_batch is an `EventDataBatch` because @@ -297,6 +634,8 @@ async def send_batch( :class:`EventHubError` :class:`ValueError` :class:`TypeError` + :raises OperationTimeoutError: If the value specified by the timeout parameter elapses before the event can be + sent in non-buffered mode or the events can be enqueued into the buffered in buffered mode. .. admonition:: Example: @@ -308,40 +647,33 @@ async def send_batch( :caption: Asynchronously sends event data """ - partition_id = kwargs.get("partition_id") - partition_key = kwargs.get("partition_key") - - if isinstance(event_data_batch, EventDataBatch): - if partition_id or partition_key: - raise TypeError( - "partition_id and partition_key should be None when sending an EventDataBatch " - "because type EventDataBatch itself may have partition_id or partition_key" - ) - to_send_batch = event_data_batch - else: - to_send_batch = await self.create_batch( - partition_id=partition_id, partition_key=partition_key - ) - to_send_batch._load_events( # pylint:disable=protected-access - event_data_batch - ) + batch, pid, pkey = await self._batch_preparer(event_data_batch, **kwargs) - if len(to_send_batch) == 0: + if len(batch) == 0: return - partition_id = ( - to_send_batch._partition_id # pylint:disable=protected-access - or ALL_PARTITIONS - ) + partition_id = pid or ALL_PARTITIONS + timeout = kwargs.pop("timeout", None) + try: - await cast(EventHubProducer, self._producers[partition_id]).send( - to_send_batch, timeout=timeout - ) - except (KeyError, AttributeError, EventHubError): - await self._start_producer(partition_id, timeout) - await cast(EventHubProducer, self._producers[partition_id]).send( - to_send_batch, timeout=timeout - ) + try: + await cast(EventHubProducer, self._producers[partition_id]).send( + batch, partition_key=pkey, timeout=timeout + ) + if self._on_success: + await self._on_success(batch._internal_events, pid) + except (KeyError, AttributeError, EventHubError): + await self._start_producer(partition_id, timeout) + await cast(EventHubProducer, self._producers[partition_id]).send( + batch, partition_key=pkey, timeout=timeout + ) + if self._on_success: + await self._on_success(batch._internal_events, pid) + except Exception as exc: # pylint: disable=broad-except + if self._on_error: + await self._on_error(batch._internal_events, pid, exc) + else: + raise async def create_batch( self, @@ -391,6 +723,7 @@ async def create_batch( max_size_in_bytes=(max_size_in_bytes or self._max_message_size_on_link), partition_id=partition_id, partition_key=partition_key, + amqp_transport=self._amqp_transport ) return event_data_batch @@ -441,10 +774,34 @@ async def get_partition_properties(self, partition_id: str) -> Dict[str, Any]: EventHubProducerClient, self )._get_partition_properties_async(partition_id) - async def close(self) -> None: + async def flush(self, **kwargs: Any) -> None: + """ + Buffered mode only. + Flush events in the buffer to be sent immediately if the client is working in buffered mode. + + :keyword Optional[float] timeout: Timeout to flush the buffered events, default is None which means no timeout. + :rtype: None + :raises EventDataSendError: If the producer fails to flush the buffer within the given timeout + in buffered mode. + """ + async with self._lock: + if self._buffered_mode and self._buffered_producer_dispatcher: + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + await self._buffered_producer_dispatcher.flush( + timeout_time=timeout_time + ) + + async def close(self, *, flush: bool = True, **kwargs: Any) -> None: """Close the Producer client underlying AMQP connection and links. + :keyword bool flush: Buffered mode only. If set to True, events in the buffer will be sent + immediately. Default is True. + :keyword Optional[float] timeout: Buffered mode only. Timeout to close the producer. + Default is None which means no timeout. :rtype: None + :raises EventHubError: If an error occurred when flushing the buffer if `flush` is set to True or closing the + underlying AMQP connections in buffered mode. .. admonition:: Example: @@ -457,9 +814,60 @@ async def close(self) -> None: """ async with self._lock: + if self._buffered_mode and self._buffered_producer_dispatcher: + timeout = kwargs.get("timeout") + timeout_time = time.time() + timeout if timeout else None + await self._buffered_producer_dispatcher.close( + flush=flush, timeout_time=timeout_time, raise_error=True + ) + self._buffered_producer_dispatcher = None + for pid in self._producers: if self._producers[pid] is not None: await self._producers[pid].close() # type: ignore self._producers[pid] = None await super(EventHubProducerClient, self)._close_async() + + def get_buffered_event_count(self, partition_id: str) -> Optional[int]: + """ + The number of events that are buffered and waiting to be published for a given partition. + Returns None in non-buffered mode. **NOTE: The event buffer is processed in a background + coroutine, therefore the number of events in the buffer reported by this API should be + considered only an approximation and is only recommend for use in debugging. For a + partition ID that has no events buffered, 0 will be returned regardless of whether that + partition ID actually exists within the Event Hub.** + + :param str partition_id: The target partition ID. + :rtype: int or None + """ + if not self._buffered_mode: + return None + + try: + return cast( + BufferedProducerDispatcher, self._buffered_producer_dispatcher + ).get_buffered_event_count(partition_id) + except AttributeError: + return 0 + + @property + def total_buffered_event_count(self) -> Optional[int]: + """ + The total number of events that are currently buffered and waiting to be published, + across all partitions. Returns None in non-buffered mode. **NOTE: The event buffer + is processed in a background coroutine, therefore the number of events in the buffer + reported by this API should be considered only an approximation and is only recommend + for use in debugging.** + + :rtype: int or None + """ + if not self._buffered_mode: + return None + + try: + return cast( + BufferedProducerDispatcher, self._buffered_producer_dispatcher + ).total_buffered_event_count + except AttributeError: + return 0 diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/__init__.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/__init__.py new file mode 100644 index 000000000000..34913fb394d7 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/__init__.py @@ -0,0 +1,4 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py new file mode 100644 index 000000000000..81c654c88019 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py @@ -0,0 +1,272 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +from __future__ import annotations +from typing import Tuple, Union, TYPE_CHECKING +from abc import ABC, abstractmethod + +if TYPE_CHECKING: + from uamqp import types as uamqp_types + +class AmqpTransportAsync(ABC): # pylint: disable=too-many-public-methods + """ + Abstract class that defines a set of common methods needed by producer and consumer. + """ + # define constants + MAX_FRAME_SIZE_BYTES: int + MAX_MESSAGE_LENGTH_BYTES: int + TIMEOUT_FACTOR: int + CONNECTION_CLOSING_STATES: Tuple + + # define symbols + PRODUCT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + VERSION_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + FRAMEWORK_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + PLATFORM_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + USER_AGENT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes] + + + @staticmethod + @abstractmethod + def build_message(**kwargs): + """ + Creates a uamqp.Message or pyamqp.Message with given arguments. + :rtype: uamqp.Message or pyamqp.Message + """ + + @staticmethod + @abstractmethod + def build_batch_message(**kwargs): + """ + Creates a uamqp.BatchMessage or pyamqp.BatchMessage with given arguments. + :rtype: uamqp.BatchMessage or pyamqp.BatchMessage + """ + + @staticmethod + @abstractmethod + def to_outgoing_amqp_message(annotated_message): + """ + Converts an AmqpAnnotatedMessage into an Amqp Message. + :param AmqpAnnotatedMessage annotated_message: AmqpAnnotatedMessage to convert. + :rtype: uamqp.Message or pyamqp.Message + """ + + @staticmethod + @abstractmethod + def get_batch_message_encoded_size(message): + """ + Gets the batch message encoded size given an underlying Message. + :param uamqp.BatchMessage message: Message to get encoded size of. + :rtype: int + """ + + @staticmethod + @abstractmethod + def get_remote_max_message_size(handler): + """ + Returns max peer message size. + :param AMQPClient handler: Client to get remote max message size on link from. + :rtype: int + """ + + @staticmethod + @abstractmethod + def create_retry_policy(config): + """ + Creates the error retry policy. + :param ~azure.eventhub._configuration.Configuration config: Configuration. + """ + + @staticmethod + @abstractmethod + def create_link_properties(link_properties): + """ + Creates and returns the link properties. + :param dict[bytes, int] link_properties: The dict of symbols and corresponding values. + :rtype: dict + """ + + @staticmethod + @abstractmethod + async def create_connection_async(**kwargs): + """ + Creates and returns the uamqp async Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + + @staticmethod + @abstractmethod + async def close_connection_async(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + + @staticmethod + @abstractmethod + def get_connection_state(connection): + """ + Gets connection state. + :param connection: uamqp or pyamqp Connection. + """ + + @staticmethod + @abstractmethod + def create_send_client(*, config, **kwargs): + """ + Creates and returns the send client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + + @staticmethod + @abstractmethod + async def send_messages_async(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + + @staticmethod + @abstractmethod + def set_message_partition_key(message, partition_key, **kwargs): + """Set the partition key as an annotation on a uamqp message. + + :param message: The message to update. + :param str partition_key: The partition key value. + :rtype: None + """ + + @staticmethod + @abstractmethod + def create_source(source, offset, selector): + """ + Creates and returns the Source. + + :param str source: Required. + :param int offset: Required. + :param bytes selector: Required. + """ + + @staticmethod + @abstractmethod + def create_receive_client(*, config, **kwargs): + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword Source source: Required. The source. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. Missing in pyamqp. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + @staticmethod + @abstractmethod + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): + """ + Receives messages, creates events, and returns them by calling the on received callback. + :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. + :param bool batch: If receive batch or single event. + :param int max_batch_size: Max batch size. + :param int or None max_wait_time: Max wait time. + """ + + @staticmethod + @abstractmethod + async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Whether to update token. If not updating token, + then pass 300 to refresh_window. Only used by uamqp. + """ + + @staticmethod + @abstractmethod + def create_mgmt_client(address, mgmt_auth, config): + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + @staticmethod + @abstractmethod + async def get_updated_token_async(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + + @staticmethod + @abstractmethod + async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQP Client mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + + @staticmethod + @abstractmethod + def get_error(status_code, description): + """ + Gets error corresponding to status code. + :param status_code: Status code. + :param str description: Description of error. + """ + + @staticmethod + @abstractmethod + def check_timeout_exception(base, exception): + """ + Checks if timeout exception. + :param base: ClientBase. + :param exception: Exception to check. + """ diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py new file mode 100644 index 000000000000..adca8c81b736 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py @@ -0,0 +1,359 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from __future__ import annotations +import asyncio +import time +import logging +from typing import Union, cast, TYPE_CHECKING, List + +from ..._pyamqp import constants, error as errors +from ..._pyamqp.aio import AMQPClientAsync, SendClientAsync, ReceiveClientAsync +from ..._pyamqp.aio._authentication_async import JWTTokenAuthAsync +from ..._pyamqp.aio._connection_async import Connection as ConnectionAsync + +from ._base_async import AmqpTransportAsync +from ..._transport._pyamqp_transport import PyamqpTransport +from ...exceptions import ( + EventHubError, + EventDataSendError, +) +from ..._common import EventData + +if TYPE_CHECKING: + from .._client_base_async import ClientBaseAsync, ConsumerProducerMixin + from ..._pyamqp.message import Message + +_LOGGER = logging.getLogger(__name__) + + +class PyamqpTransportAsync(PyamqpTransport, AmqpTransportAsync): + """ + Class which defines pyamqp-based methods used by the producer and consumer. + """ + + @staticmethod + async def create_connection_async(**kwargs): + """ + Creates and returns the pyamqp Connection object. + :keyword str host: The hostname, used by pyamqp. + :keyword JWTTokenAuthAsync auth: The auth, used by pyamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") + host = kwargs.pop("host") # pylint:disable=unused-variable + auth = kwargs.pop("auth") # pylint:disable=unused-variable + network_trace = kwargs.pop("debug") + return ConnectionAsync(endpoint, network_trace=network_trace, **kwargs) + + @staticmethod + async def close_connection(connection): + """ + Closes existing connection. + :param connection: pyamqp Connection. + """ + await connection.close() + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the pyamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + target = kwargs.pop("target") + # TODO: extra passed in to pyamqp, but not used. should be used? + msg_timeout = kwargs.pop("msg_timeout") # pylint: disable=unused-variable # TODO: not used by pyamqp? + + return SendClientAsync( + config.hostname, + target, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + **kwargs, + ) + + @staticmethod + async def send_messages_async(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by pyamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + await producer._open() + timeout = timeout_time - time.time() if timeout_time else 0 + await producer._handler.send_message_async(producer._unsent_events[0], timeout=timeout) + producer._unsent_events = None + + @staticmethod + def create_receive_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + return ReceiveClientAsync( + config.hostname, + source, + receive_settle_mode=constants.ReceiverSettleMode.First, # TODO: make more descriptive in pyamqp? + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + **kwargs, + ) + + @staticmethod + async def _callback_task(consumer, batch, max_batch_size, max_wait_time): + while consumer._callback_task_run: # pylint: disable=protected-access + async with consumer._message_buffer_lock: # pylint: disable=protected-access + messages = [ + consumer._message_buffer.popleft() # pylint: disable=protected-access + for _ in range(min(max_batch_size, len(consumer._message_buffer))) # pylint: disable=protected-access + ] + events = [EventData._from_message(message) for message in messages] # pylint: disable=protected-access + now_time = time.time() + if len(events) > 0: + await consumer._on_event_received(events if batch else events[0]) # pylint: disable=protected-access + consumer._last_callback_called_time = now_time # pylint: disable=protected-access + else: + if max_wait_time and (now_time - consumer._last_callback_called_time) > max_wait_time: # pylint: disable=protected-access + # no events received, and need to callback + await consumer._on_event_received([] if batch else None) # pylint: disable=protected-access + consumer._last_callback_called_time = now_time # pylint: disable=protected-access + # backoff a bit to avoid throttling CPU when no events are coming + await asyncio.sleep(0.05) + + @staticmethod + async def _receive_task(consumer): + max_retries = consumer._client._config.max_retries # pylint:disable=protected-access + retried_times = 0 + while retried_times <= max_retries: + try: + await consumer._open() # pylint: disable=protected-access + await cast(ReceiveClientAsync, consumer._handler).do_work_async(batch=consumer._prefetch) # pylint: disable=protected-access + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as exception: # pylint: disable=broad-except + if ( + isinstance(exception, errors.AMQPLinkError) + and exception.condition == errors.ErrorCondition.LinkStolen # pylint: disable=no-member + ): + raise await consumer._handle_exception(exception) # pylint: disable=protected-access + if not consumer.running: # exit by close + return + if consumer._last_received_event: # pylint: disable=protected-access + consumer._offset = consumer._last_received_event.offset # pylint: disable=protected-access + last_exception = await consumer._handle_exception(exception) # pylint: disable=protected-access + retried_times += 1 + if retried_times > max_retries: + _LOGGER.info( + "%r operation has exhausted retry. Last exception: %r.", + consumer._name, # pylint: disable=protected-access + last_exception, + ) + raise last_exception + + @staticmethod + async def message_received_async(consumer, message: Message) -> None: + async with consumer._message_buffer_lock: # pylint: disable=protected-access + consumer._message_buffer.append(message) # pylint: disable=protected-access + + @staticmethod + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): + """ + Receives messages, creates events, and returns them by calling the on received callback. + :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. + :param bool batch: If receive batch or single event. + :param int max_batch_size: Max batch size. + :param int or None max_wait_time: Max wait time. + """ + # pylint:disable=protected-access + consumer._callback_task_run = True + consumer._last_callback_called_time = time.time() + callback_task = asyncio.ensure_future( + PyamqpTransportAsync._callback_task(consumer, batch, max_batch_size, max_wait_time) + ) + receive_task = asyncio.ensure_future(PyamqpTransportAsync._receive_task(consumer)) + + try: + await receive_task + finally: + consumer._callback_task_run = False + await callback_task + + @staticmethod + async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Required. Whether to update token. If not updating token, + then pass 300 to refresh_window. + """ + # TODO: figure out why we're passing all these args to pyamqp JWTTokenAuth, which aren't being used + update_token = kwargs.pop("update_token") # pylint: disable=unused-variable + if update_token: + # update_token not actually needed by pyamqp + # just using to detect wh + return JWTTokenAuthAsync(auth_uri, auth_uri, get_token) + return JWTTokenAuthAsync( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + ) + # if update_token: + # token_auth.update_token() # TODO: why don't we need to update in pyamqp? + + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): # pylint: disable=unused-argument + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + return AMQPClientAsync( + config.hostname, + auth=mgmt_auth, + network_trace=config.network_tracing, + transport_type=config.transport_type, + http_proxy=config.http_proxy, + custom_endpoint_address=config.custom_endpoint_address, + connection_verify=config.connection_verify, + ) + + @staticmethod + async def get_updated_token_async(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return await mgmt_auth.get_token() + + @staticmethod + async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQPClientAsync mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + return await mgmt_client.mgmt_request_async( + mgmt_msg, operation=operation.decode(), operation_type=operation_type.decode(), **kwargs + ) + + @staticmethod + async def _handle_exception_async( # pylint:disable=too-many-branches, too-many-statements + exception: Exception, closable: Union["ClientBaseAsync", "ConsumerProducerMixin"] + ) -> Exception: + # pylint: disable=protected-access + if isinstance(exception, asyncio.CancelledError): + raise exception + error = exception + try: + name = cast("ConsumerProducerMixin", closable)._name + except AttributeError: + name = cast("ClientBaseAsync", closable)._container_id + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + await cast("ConsumerProducerMixin", closable)._close_connection_async() + raise error + elif isinstance(exception, EventHubError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + raise error + # TODO: The following errors seem to be useless in EH + # elif isinstance( + # exception, + # ( + # errors.MessageAccepted, + # errors.MessageAlreadySettled, + # errors.MessageModified, + # errors.MessageRejected, + # errors.MessageReleased, + # errors.MessageContentTooLarge, + # ), + # ): + # _LOGGER.info("%r Event data error (%r)", name, exception) + # error = EventDataError(str(exception), exception) + # raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + try: + if isinstance(exception, errors.AuthenticationException): + await closable._close_connection_async() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPLinkError): + await closable._close_handler_async() # pylint:disable=protected-access + elif isinstance(exception, errors.AMQPConnectionError): + await closable._close_connection_async() # pylint:disable=protected-access + # TODO: add MessageHandlerError in amqp? + # elif isinstance(exception, errors.MessageHandlerError): + # if hasattr(closable, "_close_handler"): + # closable._close_handler() # pylint:disable=protected-access + else: # errors.AMQPConnectionError, compat.TimeoutException + await closable._close_connection_async() # pylint:disable=protected-access + return PyamqpTransportAsync._create_eventhub_exception(exception) + except AttributeError: + pass + return PyamqpTransportAsync._create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py new file mode 100644 index 000000000000..29b388d122d6 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py @@ -0,0 +1,390 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from __future__ import annotations +import asyncio +import time +import logging +from typing import Union, cast, TYPE_CHECKING, List, Optional + +try: + from uamqp import ( + constants, + types, + SendClientAsync, + ReceiveClientAsync, + utils, + authentication, + AMQPClientAsync, + errors, + ) + from uamqp.async_ops import ConnectionAsync + from ..._transport._uamqp_transport import UamqpTransport + uamqp_installed = True +except ImportError: + uamqp_installed = False + +from ._base_async import AmqpTransportAsync +from ...exceptions import ( + OperationTimeoutError, + EventHubError, + EventDataError, + EventDataSendError, +) + +if TYPE_CHECKING: + from .._client_base_async import ClientBaseAsync, ConsumerProducerMixin + from ..._common import EventData + try: + from uamqp import Message + except ImportError: + Message = None + +_LOGGER = logging.getLogger(__name__) + +if uamqp_installed: + class UamqpTransportAsync(UamqpTransport, AmqpTransportAsync): + """ + Class which defines uamqp-based methods used by the producer and consumer. + """ + + @staticmethod + async def create_connection_async(**kwargs): + """ + Creates and returns the uamqp async Connection object. + :keyword str host: The hostname, used by uamqp. + :keyword JWTTokenAuth auth: The auth, used by uamqp. + :keyword str endpoint: The endpoint, used by pyamqp. + :keyword str container_id: Required. + :keyword int max_frame_size: Required. + :keyword int channel_max: Required. + :keyword int idle_timeout: Required. + :keyword Dict properties: Required. + :keyword int remote_idle_timeout_empty_frame_send_ratio: Required. + :keyword error_policy: Required. + :keyword bool debug: Required. + :keyword str encoding: Required. + """ + endpoint = kwargs.pop("endpoint") # pylint:disable=unused-variable + host = kwargs.pop("host") + auth = kwargs.pop("auth") + return ConnectionAsync( + host, + auth, + **kwargs + ) + + @staticmethod + async def close_connection_async(connection): + """ + Closes existing connection. + :param connection: uamqp or pyamqp Connection. + """ + await connection.destroy_async() + + @staticmethod + def create_send_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the uamqp SendClient. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str target: Required. The target. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword keep_alive_interval: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + """ + target = kwargs.pop("target") + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") + + return SendClientAsync( + target, + debug=network_trace, # pylint:disable=protected-access + error_policy=retry_policy, + **kwargs + ) + + @staticmethod + async def send_messages_async(producer, timeout_time, last_exception, logger): + """ + Handles sending of event data messages. + :param ~azure.eventhub._producer.EventHubProducer producer: The producer with handler to send messages. + :param int timeout_time: Timeout time. + :param last_exception: Exception to raise if message timed out. Only used by uamqp transport. + :param logger: Logger. + """ + # pylint: disable=protected-access + await producer._open() + producer._unsent_events[0].on_send_complete = producer._on_outcome + UamqpTransportAsync._set_msg_timeout(producer, timeout_time, last_exception, logger) + producer._handler.queue_message(*producer._unsent_events) # type: ignore + await producer._handler.wait_async() # type: ignore + producer._unsent_events = producer._handler.pending_messages # type: ignore + if producer._outcome != constants.MessageSendResult.Ok: + if producer._outcome == constants.MessageSendResult.Timeout: + producer._condition = OperationTimeoutError("Send operation timed out") + if producer._condition: + raise producer._condition + + @staticmethod + def create_receive_client(*, config, **kwargs): # pylint:disable=unused-argument + """ + Creates and returns the receive client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + + :keyword str source: Required. The source. + :keyword str offset: Required. + :keyword str offset_inclusive: Required. + :keyword JWTTokenAuth auth: Required. + :keyword int idle_timeout: Required. + :keyword network_trace: Required. + :keyword retry_policy: Required. + :keyword str client_name: Required. + :keyword dict link_properties: Required. + :keyword properties: Required. + :keyword link_credit: Required. The prefetch. + :keyword keep_alive_interval: Required. + :keyword desired_capabilities: Required. + :keyword streaming_receive: Required. + :keyword message_received_callback: Required. + :keyword timeout: Required. + """ + + source = kwargs.pop("source") + symbol_array = kwargs.pop("desired_capabilities") + desired_capabilities = None + if symbol_array: + symbol_array = [types.AMQPSymbol(symbol) for symbol in symbol_array] + desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array)) + retry_policy = kwargs.pop("retry_policy") + network_trace = kwargs.pop("network_trace") + link_credit = kwargs.pop("link_credit") + streaming_receive = kwargs.pop("streaming_receive") + message_received_callback = kwargs.pop("message_received_callback") + + client = ReceiveClientAsync( + source, + debug=network_trace, # pylint:disable=protected-access + error_policy=retry_policy, + desired_capabilities=desired_capabilities, + prefetch=link_credit, + receive_settle_mode=constants.ReceiverSettleMode.ReceiveAndDelete, + auto_complete=False, + **kwargs + ) + # pylint:disable=protected-access + client._streaming_receive = streaming_receive + client._message_received_callback = (message_received_callback) + return client + + @staticmethod + def message_received_async(consumer, message: Message) -> None: + consumer._message_buffer.append(message) # pylint: disable=protected-access + + @staticmethod + async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time): + """ + Receives messages, creates events, and returns them by calling the on received callback. + :param ~azure.eventhub.aio.EventHubConsumer consumer: The EventHubConsumer. + :param bool batch: If receive batch or single event. + :param int max_batch_size: Max batch size. + :param int or None max_wait_time: Max wait time. + """ + # pylint:disable=protected-access + max_retries = ( + consumer._client._config.max_retries # pylint:disable=protected-access + ) + has_not_fetched_once = True # ensure one trip when max_wait_time is very small + deadline = time.time() + (max_wait_time or 0) # max_wait_time can be None + while len(consumer._message_buffer) < max_batch_size and ( + time.time() < deadline or has_not_fetched_once + ): + retried_times = 0 + has_not_fetched_once = False + while retried_times <= max_retries: + try: + await consumer._open() + await cast( + ReceiveClientAsync, consumer._handler + ).do_work_async() # uamqp sleeps 0.05 if none received + break + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as exception: # pylint: disable=broad-except + if ( + isinstance(exception, errors.LinkDetach) + and exception.condition == constants.ErrorCodes.LinkStolen # pylint: disable=no-member + ): + raise await consumer._handle_exception(exception) + if not consumer.running: # exit by close + return + if consumer._last_received_event: + consumer._offset = consumer._last_received_event.offset + last_exception = await consumer._handle_exception(exception) + retried_times += 1 + if retried_times > max_retries: + _LOGGER.info( + "%r operation has exhausted retry. Last exception: %r.", + consumer._name, + last_exception, + ) + raise last_exception + + if consumer._message_buffer: + while consumer._message_buffer: + if batch: + events_for_callback: List[EventData] = [] + for _ in range(min(max_batch_size, len(consumer._message_buffer))): + events_for_callback.append(consumer._next_message_in_buffer()) + await consumer._on_event_received(events_for_callback) + else: + await consumer._on_event_received(consumer._next_message_in_buffer()) + elif max_wait_time: + if batch: + await consumer._on_event_received([]) + else: + await consumer._on_event_received(None) + + @staticmethod + async def create_token_auth_async(auth_uri, get_token, token_type, config, **kwargs): + """ + Creates the JWTTokenAuth. + :param str auth_uri: The auth uri to pass to JWTTokenAuth. + :param get_token: The callback function used for getting and refreshing + tokens. It should return a valid jwt token each time it is called. + :param bytes token_type: Token type. + :param ~azure.eventhub._configuration.Configuration config: EH config. + + :keyword bool update_token: Required. Whether to update token. If not updating token, + then pass 300 to refresh_window. + """ + update_token = kwargs.pop("update_token") + refresh_window = 300 + if update_token: + refresh_window = 0 + + token_auth = authentication.JWTTokenAsync( + auth_uri, + auth_uri, + get_token, + token_type=token_type, + timeout=config.auth_timeout, + http_proxy=config.http_proxy, + transport_type=config.transport_type, + custom_endpoint_hostname=config.custom_endpoint_hostname, + port=config.connection_port, + verify=config.connection_verify, + refresh_window=refresh_window + ) + if update_token: + await token_auth.update_token() + return token_auth + + @staticmethod + def create_mgmt_client(address, mgmt_auth, config): + """ + Creates and returns the mgmt AMQP client. + :param _Address address: Required. The Address. + :param JWTTokenAuth mgmt_auth: Auth for client. + :param ~azure.eventhub._configuration.Configuration config: The configuration. + """ + + mgmt_target = f"amqps://{address.hostname}{address.path}" + return AMQPClientAsync( + mgmt_target, + auth=mgmt_auth, + debug=config.network_tracing + ) + + @staticmethod + async def get_updated_token_async(mgmt_auth): + """ + Return updated auth token. + :param mgmt_auth: Auth. + """ + return mgmt_auth.token + + @staticmethod + async def mgmt_client_request_async(mgmt_client, mgmt_msg, **kwargs): + """ + Send mgmt request. + :param AMQP Client mgmt_client: Client to send request with. + :param str mgmt_msg: Message. + :keyword bytes operation: Operation. + :keyword operation_type: Op type. + :keyword status_code_field: mgmt status code. + :keyword description_fields: mgmt status desc. + """ + operation_type = kwargs.pop("operation_type") + operation = kwargs.pop("operation") + response = await mgmt_client.mgmt_request_async( + mgmt_msg, + operation, + op_type=operation_type, + **kwargs + ) + status_code = response.application_properties[kwargs.get("status_code_field")] + description = response.application_properties.get( + kwargs.get("description_fields") + ) # type: Optional[Union[str, bytes]] + return status_code, description, response + + @staticmethod + async def _handle_exception_async( # pylint:disable=too-many-branches, too-many-statements + exception: Exception, closable: Union["ClientBaseAsync", "ConsumerProducerMixin"] + ) -> Exception: + # pylint: disable=protected-access + if isinstance(exception, asyncio.CancelledError): + raise exception + error = exception + try: + name = cast("ConsumerProducerMixin", closable)._name + except AttributeError: + name = cast("ClientBaseAsync", closable)._container_id + if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise + _LOGGER.info("%r stops due to keyboard interrupt", name) + await cast("ConsumerProducerMixin", closable)._close_connection_async() + raise error + elif isinstance(exception, EventHubError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + raise error + elif isinstance( + exception, + ( + errors.MessageAccepted, + errors.MessageAlreadySettled, + errors.MessageModified, + errors.MessageRejected, + errors.MessageReleased, + errors.MessageContentTooLarge, + ), + ): + _LOGGER.info("%r Event data error (%r)", name, exception) + error = EventDataError(str(exception), exception) + raise error + elif isinstance(exception, errors.MessageException): + _LOGGER.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + else: + try: + if isinstance(exception, errors.AuthenticationException): + await closable._close_connection_async() + elif isinstance(exception, errors.LinkDetach): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + elif isinstance(exception, errors.ConnectionClose): + await closable._close_connection_async() + elif isinstance(exception, errors.MessageHandlerError): + await cast("ConsumerProducerMixin", closable)._close_handler_async() + else: # errors.AMQPConnectionError, compat.TimeoutException, and any other errors + await closable._close_connection_async() + except AttributeError: + pass + return UamqpTransport._create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py index 21b56b8ba91c..d6cc5937a11d 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_message.py @@ -4,75 +4,12 @@ # license information. # ------------------------------------------------------------------------- -from typing import Optional, Any, cast, Mapping, Dict +from __future__ import annotations +from typing import Optional, Any, cast, Mapping, Dict, Union, List +from ._amqp_utils import normalized_data_body, normalized_sequence_body from ._constants import AmqpMessageBodyType -from .._pyamqp.message import Message, Header, Properties -from .._pyamqp import utils as pyamqp_utils - - -class DictMixin(object): - def __setitem__(self, key, item): - # type: (Any, Any) -> None - self.__dict__[key] = item - - def __getitem__(self, key): - # type: (Any) -> Any - return self.__dict__[key] - - def __repr__(self): - # type: () -> str - return str(self) - - def __len__(self): - # type: () -> int - return len(self.keys()) - - def __delitem__(self, key): - # type: (Any) -> None - self.__dict__[key] = None - - def __eq__(self, other): - # type: (Any) -> bool - """Compare objects by comparing all attributes.""" - if isinstance(other, self.__class__): - return self.__dict__ == other.__dict__ - return False - - def __ne__(self, other): - # type: (Any) -> bool - """Compare objects by comparing all attributes.""" - return not self.__eq__(other) - - def __str__(self): - # type: () -> str - return str({k: v for k, v in self.__dict__.items() if not k.startswith("_")}) - - def has_key(self, k): - # type: (Any) -> bool - return k in self.__dict__ - - def update(self, *args, **kwargs): - # type: (Any, Any) -> None - return self.__dict__.update(*args, **kwargs) - - def keys(self): - # type: () -> list - return [k for k in self.__dict__ if not k.startswith("_")] - - def values(self): - # type: () -> list - return [v for k, v in self.__dict__.items() if not k.startswith("_")] - - def items(self): - # type: () -> list - return [(k, v) for k, v in self.__dict__.items() if not k.startswith("_")] - - def get(self, key, default=None): - # type: (Any, Optional[Any]) -> Any - if key in self.__dict__: - return self.__dict__[key] - return default +from .._mixin import DictMixin class AmqpAnnotatedMessage(object): @@ -108,79 +45,65 @@ class AmqpAnnotatedMessage(object): def __init__(self, **kwargs): # type: (Any) -> None - self._message = kwargs.pop("message", None) self._encoding = kwargs.pop("encoding", "UTF-8") + self._data_body: Optional[Union[str, bytes, List[Union[str, bytes]]]] = None + self._sequence_body: Optional[List[Any]] = None + self._value_body: Any = None # internal usage only for Event Hub received message - if self._message: - self._from_amqp_message(self._message) + message = kwargs.pop("message", None) + if message: + self._from_amqp_message(message) return # manually constructed AMQPAnnotatedMessage - input_count_validation = len([key for key in ("data_body", "sequence_body", "value_body") if key in kwargs]) + input_count_validation = len( + [ + key + for key in ("data_body", "sequence_body", "value_body") + if key in kwargs + ] + ) if input_count_validation != 1: raise ValueError( "There should be one and only one of either data_body, sequence_body " "or value_body being set as the body of the AmqpAnnotatedMessage." ) - self._body = None - self._body_type = None + self._body_type: AmqpMessageBodyType = None # type: ignore if "data_body" in kwargs: - self._body = pyamqp_utils.normalized_data_body(kwargs.get("data_body")) - self._message = Message(data=self._body) + self._data_body = normalized_data_body(kwargs.get("data_body")) self._body_type = AmqpMessageBodyType.DATA elif "sequence_body" in kwargs: - self._body = pyamqp_utils.normalized_sequence_body(kwargs.get("sequence_body")) + self._sequence_body = normalized_sequence_body(kwargs.get("sequence_body")) self._body_type = AmqpMessageBodyType.SEQUENCE - self._message = Message(sequence=self._body) elif "value_body" in kwargs: - self._body = kwargs.get("value_body") + self._value_body = kwargs.get("value_body") self._body_type = AmqpMessageBodyType.VALUE - self._message = Message(value=self._body) - #self._message = uamqp.message.Message(body=self._body, body_type=self._body_type) header_dict = cast(Mapping, kwargs.get("header")) self._header = AmqpMessageHeader(**header_dict) if "header" in kwargs else None self._footer = kwargs.get("footer") properties_dict = cast(Mapping, kwargs.get("properties")) - self._properties = AmqpMessageProperties(**properties_dict) if "properties" in kwargs else None + self._properties = ( + AmqpMessageProperties(**properties_dict) if "properties" in kwargs else None + ) self._application_properties = kwargs.get("application_properties") self._annotations = kwargs.get("annotations") self._delivery_annotations = kwargs.get("delivery_annotations") - def __str__(self): + def __str__(self) -> str: if self._body_type == AmqpMessageBodyType.DATA: - output_str = "" - for data_section in self.body: - try: - output_str += data_section.decode(self._encoding) - except AttributeError: - output_str += str(data_section) - return output_str - elif self._body_type == AmqpMessageBodyType.SEQUENCE: - output_str = "" - for sequence_section in self.body: - for d in sequence_section: - try: - output_str += d.decode(self._encoding) - except AttributeError: - output_str += str(d) - return output_str - else: - if not self.body: - return "" - try: - return self.body.decode(self._encoding) - except AttributeError: - return str(self.body) - - def __repr__(self): - # type: () -> str + return "".join(d.decode(self._encoding) for d in self._data_body) # type: ignore + if self._body_type == AmqpMessageBodyType.SEQUENCE: + return str(self._sequence_body) + if self._body_type == AmqpMessageBodyType.VALUE: + return str(self._value_body) + return "" + + def __repr__(self) -> str: # pylint: disable=bare-except - message_repr = "body={}".format( - str(self) - ) + message_repr = "body={}".format(str(self)) message_repr += ", body_type={}".format(self.body_type) try: message_repr += ", header={}".format(self.header) @@ -195,11 +118,15 @@ def __repr__(self): except: message_repr += ", properties=" try: - message_repr += ", application_properties={}".format(self.application_properties) + message_repr += ", application_properties={}".format( + self.application_properties + ) except: message_repr += ", application_properties=" try: - message_repr += ", delivery_annotations={}".format(self.delivery_annotations) + message_repr += ", delivery_annotations={}".format( + self.delivery_annotations + ) except: message_repr += ", delivery_annotations=" try: @@ -209,8 +136,6 @@ def __repr__(self): return "AmqpAnnotatedMessage({})".format(message_repr)[:1024] def _from_amqp_message(self, message): - # populate the properties from an uamqp message - # TODO: message.properties should not be a list self._properties = AmqpMessageProperties( message_id=message.properties.message_id, user_id=message.properties.user_id, @@ -228,7 +153,7 @@ def _from_amqp_message(self, message): ) if message.properties else None self._header = AmqpMessageHeader( delivery_count=message.header.delivery_count, - time_to_live=message.header.time_to_live, + time_to_live=message.header.ttl, first_acquirer=message.header.first_acquirer, durable=message.header.durable, priority=message.header.priority @@ -237,82 +162,41 @@ def _from_amqp_message(self, message): self._annotations = message.message_annotations if message.message_annotations else {} self._delivery_annotations = message.delivery_annotations if message.delivery_annotations else {} self._application_properties = message.application_properties if message.application_properties else {} - - def _to_outgoing_amqp_message(self): - message_header = None - if self.header and any(self.header.values()): - message_header = Header( - delivery_count=self.header.delivery_count, - ttl=self.header.time_to_live, - first_acquirer=self.header.first_acquirer, - durable=self.header.durable, - priority=self.header.priority - ) - - message_properties = None - if self.properties and any(self.properties.values()): - message_properties = Properties( - message_id=self.properties.message_id, - user_id=self.properties.user_id, - to=self.properties.to, - subject=self.properties.subject, - reply_to=self.properties.reply_to, - correlation_id=self.properties.correlation_id, - content_type=self.properties.content_type, - content_encoding=self.properties.content_encoding, - creation_time=int(self.properties.creation_time) if self.properties.creation_time else None, - absolute_expiry_time=int(self.properties.absolute_expiry_time) - if self.properties.absolute_expiry_time else None, - group_id=self.properties.group_id, - group_sequence=self.properties.group_sequence, - reply_to_group_id=self.properties.reply_to_group_id - ) - - dict = { - "header": message_header, - "properties": message_properties, - "application_properties": self.application_properties, - "message_annotations": self.annotations, - "delivery_annotations": self.delivery_annotations, - "footer": self.footer - } - - if self.body_type == AmqpMessageBodyType.DATA: - dict["data"] = self._body - elif self.body_type == AmqpMessageBodyType.SEQUENCE: - dict["sequence"] = self._body + if message.data: + self._data_body = cast(List, list(message.data)) + self._body_type = AmqpMessageBodyType.DATA + elif message.sequence: + self._sequence_body = cast(List, list(message.sequence)) + self._body_type = AmqpMessageBodyType.SEQUENCE else: - dict["value"] = self._body - - return Message(**dict) + self._value_body = message.value + self._body_type = AmqpMessageBodyType.VALUE @property - def body(self): - # type: () -> Any + def body(self) -> Any: """The body of the Message. The format may vary depending on the body type: For ~azure.eventhub.AmqpMessageBodyType.DATA, the body could be bytes or Iterable[bytes] For ~azure.eventhub.AmqpMessageBodyType.SEQUENCE, the body could be List or Iterable[List] For ~azure.eventhub.AmqpMessageBodyType.VALUE, the body could be any type. :rtype: Any """ - return self._message.data or self._message.sequence or self._message.value + if self._body_type == AmqpMessageBodyType.DATA: # pylint:disable=no-else-return + return (i for i in cast(List, self._data_body)) # type: ignore + elif self._body_type == AmqpMessageBodyType.SEQUENCE: + return (i for i in cast(List, self._sequence_body)) + elif self._body_type == AmqpMessageBodyType.VALUE: + return self._value_body + return None @property - def body_type(self): - # type: () -> AmqpMessageBodyType + def body_type(self) -> AmqpMessageBodyType: """The body type of the underlying AMQP message. rtype: ~azure.eventhub.amqp.AmqpMessageBodyType """ - if self._message.data: - return AmqpMessageBodyType.DATA - elif self._message.sequence: - return AmqpMessageBodyType.SEQUENCE - else: - return AmqpMessageBodyType.VALUE + return self._body_type @property - def properties(self): - # type: () -> Optional[AmqpMessageProperties] + def properties(self) -> Optional[AmqpMessageProperties]: """ Properties to add to the message. :rtype: Optional[~azure.eventhub.amqp.AmqpMessageProperties] @@ -320,13 +204,11 @@ def properties(self): return self._properties @properties.setter - def properties(self, value): - # type: (AmqpMessageProperties) -> None + def properties(self, value: AmqpMessageProperties) -> None: self._properties = value @property - def application_properties(self): - # type: () -> Optional[Dict] + def application_properties(self) -> Optional[Dict[Union[str, bytes], Any]]: """ Service specific application properties. @@ -335,13 +217,11 @@ def application_properties(self): return self._application_properties @application_properties.setter - def application_properties(self, value): - # type: (Dict) -> None + def application_properties(self, value: Optional[Dict[Union[str, bytes], Any]]) -> None: self._application_properties = value @property - def annotations(self): - # type: () -> Optional[Dict] + def annotations(self) -> Optional[Dict[Union[str, bytes], Any]]: """ Service specific message annotations. @@ -350,13 +230,11 @@ def annotations(self): return self._annotations @annotations.setter - def annotations(self, value): - # type: (Dict) -> None + def annotations(self, value: Optional[Dict[Union[str, bytes], Any]]) -> None: self._annotations = value @property - def delivery_annotations(self): - # type: () -> Optional[Dict] + def delivery_annotations(self) -> Optional[Dict[Union[str, bytes], Any]]: """ Delivery-specific non-standard properties at the head of the message. Delivery annotations convey information from the sending peer to the receiving peer. @@ -366,13 +244,11 @@ def delivery_annotations(self): return self._delivery_annotations @delivery_annotations.setter - def delivery_annotations(self, value): - # type: (Dict) -> None + def delivery_annotations(self, value: Optional[Dict[Union[str, bytes], Any]]) -> None: self._delivery_annotations = value @property - def header(self): - # type: () -> Optional[AmqpMessageHeader] + def header(self) -> Optional[AmqpMessageHeader]: """ The message header. :rtype: Optional[~azure.eventhub.amqp.AmqpMessageHeader] @@ -380,13 +256,11 @@ def header(self): return self._header @header.setter - def header(self, value): - # type: (AmqpMessageHeader) -> None + def header(self, value: AmqpMessageHeader) -> None: self._header = value @property - def footer(self): - # type: () -> Optional[Dict] + def footer(self) -> Optional[Dict[Any, Any]]: """ The message footer. @@ -395,10 +269,8 @@ def footer(self): return self._footer @footer.setter - def footer(self, value): - # type: (Dict) -> None + def footer(self, value: Optional[Dict[Any, Any]]) -> None: self._footer = value - # self._message.footer = value class AmqpMessageHeader(DictMixin): @@ -464,6 +336,7 @@ class AmqpMessageHeader(DictMixin): priority messages. Messages with higher priorities MAY be delivered before those with lower priorities. :vartype priority: Optional[int] """ + def __init__(self, **kwargs): self.delivery_count = kwargs.get("delivery_count") self.time_to_live = kwargs.get("time_to_live") @@ -546,6 +419,7 @@ class AmqpMessageProperties(DictMixin): to this message to a specific group. :vartype reply_to_group_id: Optional[bytes] """ + def __init__(self, **kwargs): self.message_id = kwargs.get("message_id") self.user_id = kwargs.get("user_id") diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_utils.py b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_utils.py new file mode 100644 index 000000000000..c620c149ea5e --- /dev/null +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/amqp/_amqp_utils.py @@ -0,0 +1,25 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + +def encode_str(data, encoding='utf-8'): + try: + return data.encode(encoding) + except AttributeError: + return data + +def normalized_data_body(data, **kwargs): + # A helper method to normalize input into AMQP Data Body format + encoding = kwargs.get("encoding", "utf-8") + if isinstance(data, list): + return [encode_str(item, encoding) for item in data] + return [encode_str(data, encoding)] + +def normalized_sequence_body(sequence): + # A helper method to normalize input into AMQP Sequence Body format + if isinstance(sequence, list) and all([isinstance(b, list) for b in sequence]): + return sequence + if isinstance(sequence, list): + return [sequence] diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/exceptions.py b/sdk/eventhub/azure-eventhub/azure/eventhub/exceptions.py index d6b0258a04d0..f686251e6e95 100644 --- a/sdk/eventhub/azure-eventhub/azure/eventhub/exceptions.py +++ b/sdk/eventhub/azure-eventhub/azure/eventhub/exceptions.py @@ -2,15 +2,8 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -import logging import six -from ._constants import NO_RETRY_ERRORS -from ._pyamqp import error as errors - -_LOGGER = logging.getLogger(__name__) - - class EventHubError(Exception): """Represents an error occurred in the client. @@ -100,74 +93,3 @@ class OperationTimeoutError(EventHubError): class OwnershipLostError(Exception): """Raised when `update_checkpoint` detects the ownership to a partition has been lost.""" - - -def _create_eventhub_exception(exception): - if isinstance(exception, errors.AuthenticationException): - error = AuthenticationError(str(exception), exception) - elif isinstance(exception, errors.AMQPLinkError): - error = ConnectError(str(exception), exception) - # TODO: do we need MessageHanlderError in amqp any more - # if connection/session/link error are enough? - # elif isinstance(exception, errors.MessageHandlerError): - # error = ConnectionLostError(str(exception), exception) - elif isinstance(exception, errors.AMQPConnectionError): - error = ConnectError(str(exception), exception) - elif isinstance(exception, TimeoutError): - error = ConnectionLostError(str(exception), exception) - else: - error = EventHubError(str(exception), exception) - return error - - -def _handle_exception( - exception, closable -): # pylint:disable=too-many-branches, too-many-statements - try: # closable is a producer/consumer object - name = closable._name # pylint: disable=protected-access - except AttributeError: # closable is an client object - name = closable._container_id # pylint: disable=protected-access - if isinstance(exception, KeyboardInterrupt): # pylint:disable=no-else-raise - _LOGGER.info("%r stops due to keyboard interrupt", name) - closable._close_connection() # pylint:disable=protected-access - raise exception - elif isinstance(exception, EventHubError): - closable._close_handler() # pylint:disable=protected-access - raise exception - # TODO: The following errors seem to be useless in EH - # elif isinstance( - # exception, - # ( - # errors.MessageAccepted, - # errors.MessageAlreadySettled, - # errors.MessageModified, - # errors.MessageRejected, - # errors.MessageReleased, - # errors.MessageContentTooLarge, - # ), - # ): - # _LOGGER.info("%r Event data error (%r)", name, exception) - # error = EventDataError(str(exception), exception) - # raise error - elif isinstance(exception, errors.MessageException): - _LOGGER.info("%r Event data send error (%r)", name, exception) - error = EventDataSendError(str(exception), exception) - raise error - else: - if isinstance(exception, errors.AuthenticationException): - if hasattr(closable, "_close_connection"): - closable._close_connection() # pylint:disable=protected-access - elif isinstance(exception, errors.AMQPLinkError): - if hasattr(closable, "_close_handler"): - closable._close_handler() # pylint:disable=protected-access - elif isinstance(exception, errors.AMQPConnectionError): - if hasattr(closable, "_close_connection"): - closable._close_connection() # pylint:disable=protected-access - # TODO: add MessageHandlerError in amqp? - # elif isinstance(exception, errors.MessageHandlerError): - # if hasattr(closable, "_close_handler"): - # closable._close_handler() # pylint:disable=protected-access - else: # errors.AMQPConnectionError, compat.TimeoutException - if hasattr(closable, "_close_connection"): - closable._close_connection() # pylint:disable=protected-access - return _create_eventhub_exception(exception) diff --git a/sdk/eventhub/azure-eventhub/conftest.py b/sdk/eventhub/azure-eventhub/conftest.py index 836f79850f3d..f2b5ae7c0443 100644 --- a/sdk/eventhub/azure-eventhub/conftest.py +++ b/sdk/eventhub/azure-eventhub/conftest.py @@ -18,6 +18,11 @@ from azure.eventhub import EventHubProducerClient from azure.eventhub._pyamqp import ReceiveClient from azure.eventhub._pyamqp.authentication import SASTokenAuth +try: + import uamqp + uamqp_transport_params = [True, False] +except ImportError: + uamqp_transport_params = [False] from devtools_testutils import get_region_override @@ -42,6 +47,9 @@ def sleep(request): sleep = request.config.getoption("--sleep") return sleep.lower() in ('true', 'yes', '1', 'y') +@pytest.fixture(scope="session", params=uamqp_transport_params) +def uamqp_transport(request): + return request.param def get_logger(filename, level=logging.INFO): azure_logger = logging.getLogger("azure.eventhub") @@ -68,6 +76,13 @@ def get_logger(filename, level=logging.INFO): log = get_logger(None, logging.DEBUG) +@pytest.fixture(scope="session") +def timeout_factor(uamqp_transport): + if uamqp_transport: + return 1000 + else: + return 1 + @pytest.fixture(scope="session") def resource_group(): try: @@ -188,22 +203,26 @@ def invalid_policy(live_eventhub): @pytest.fixture() -def connstr_receivers(live_eventhub): +def connstr_receivers(live_eventhub, uamqp_transport): connection_str = live_eventhub["connection_str"] partitions = [str(i) for i in range(PARTITION_COUNT)] receivers = [] for p in partitions: uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = SASTokenAuth( - uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] - ) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], p) - receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, debug=False, timeout=0, link_credit=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + receiver = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) receiver.open() receivers.append(receiver) yield connection_str, receivers @@ -212,9 +231,9 @@ def connstr_receivers(live_eventhub): @pytest.fixture() -def connstr_senders(live_eventhub): +def connstr_senders(live_eventhub, uamqp_transport): connection_str = live_eventhub["connection_str"] - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) partitions = client.get_partition_ids() senders = [] diff --git a/sdk/eventhub/azure-eventhub/dev_requirements.txt b/sdk/eventhub/azure-eventhub/dev_requirements.txt index 9c91833e14d8..269b8da868a9 100644 --- a/sdk/eventhub/azure-eventhub/dev_requirements.txt +++ b/sdk/eventhub/azure-eventhub/dev_requirements.txt @@ -6,4 +6,4 @@ azure-mgmt-resource==20.0.0 aiohttp>=3.0 websocket-client -e ../../../tools/azure-devtools --e ../../servicebus/azure-servicebus \ No newline at end of file +uamqp>=1.6.0,<2.0.0 diff --git a/sdk/eventhub/azure-eventhub/samples/README.md b/sdk/eventhub/azure-eventhub/samples/README.md index 437f14852f8d..bb2adc1044a1 100644 --- a/sdk/eventhub/azure-eventhub/samples/README.md +++ b/sdk/eventhub/azure-eventhub/samples/README.md @@ -78,8 +78,13 @@ The [sync version](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/e - Send AMQPAnnotatedMessage of different body types. - Receive messages and parse the body according to the body type. +- [send_buffered_mode.py](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/eventhub/azure-eventhub/samples/sync_samples/send_buffered_mode.py) ([async_version](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/eventhub/azure-eventhub/samples/async_samples/send_buffered_mode_async.py)) - Examples to send events in buffered mode: + - Send single events, which will be automatically batched. + - Send a batch of events by enqueuing an EventDataBatch object to the buffer. + - Send events in buffer immediately by calling `flush`. + ## Prerequisites -- Python 3.6 or later. +- Python 3.7 or later. - **Microsoft Azure Subscription:** To use Azure services, including Azure Event Hubs, you'll need a subscription. If you do not have an existing Azure account, you may sign up for a free trial or use your MSDN subscriber benefits when you [create an account](https://account.windowsazure.com/Home/Index). diff --git a/sdk/eventhub/azure-eventhub/samples/async_samples/send_buffered_mode_async.py b/sdk/eventhub/azure-eventhub/samples/async_samples/send_buffered_mode_async.py new file mode 100644 index 000000000000..33a5f4f4ad97 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/samples/async_samples/send_buffered_mode_async.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python + +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +""" +Examples to show sending events in buffered mode to an Event Hub asynchronously. +""" + +import time +import asyncio +import os + +from azure.eventhub.aio import EventHubProducerClient +from azure.eventhub import EventData + +CONNECTION_STR = os.environ['EVENT_HUB_CONN_STR'] +EVENTHUB_NAME = os.environ['EVENT_HUB_NAME'] + + +async def on_success(events, pid): + # sending succeeded + print(events, pid) + + +async def on_error(events, pid, error): + # sending failed + print(events, pid, error) + + +async def run(): + + producer = EventHubProducerClient.from_connection_string( + conn_str=CONNECTION_STR, + eventhub_name=EVENTHUB_NAME, + buffered_mode=True, + on_success=on_success, + on_error=on_error + ) + + # exiting the context manager will automatically call flush + async with producer: + # single events will be batched automatically + for i in range(10): + # the method returning indicates the event has been enqueued to the buffer + await producer.send_event(EventData('Single data {}'.format(i))) + + batch = await producer.create_batch() + for i in range(10): + batch.add(EventData('Single data in batch {}'.format(i))) + # alternatively, you can enqueue an EventDataBatch object to the buffer + await producer.send_batch(batch) + + # calling flush sends out the events in the buffer immediately + await producer.flush() + +start_time = time.time() +asyncio.run(run()) +print("Send messages in {} seconds.".format(time.time() - start_time)) diff --git a/sdk/eventhub/azure-eventhub/samples/sync_samples/send_buffered_mode.py b/sdk/eventhub/azure-eventhub/samples/sync_samples/send_buffered_mode.py new file mode 100644 index 000000000000..2939e3445b76 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/samples/sync_samples/send_buffered_mode.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python + +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +""" +Examples to show sending events in buffered mode to an Event Hub. +""" + +import time +import os +from azure.eventhub import EventHubProducerClient, EventData + +CONNECTION_STR = os.environ['EVENT_HUB_CONN_STR'] +EVENTHUB_NAME = os.environ['EVENT_HUB_NAME'] + + +def on_success(events, pid): + # sending succeeded + print(events, pid) + + +def on_error(events, pid, error): + # sending failed + print(events, pid, error) + + +producer = EventHubProducerClient.from_connection_string( + conn_str=CONNECTION_STR, + eventhub_name=EVENTHUB_NAME, + buffered_mode=True, + on_success=on_success, + on_error=on_error +) + +start_time = time.time() + +# exiting the context manager will automatically call flush +with producer: + # single events will be batched automatically + for i in range(10): + # the method returning indicates the event has been enqueued to the buffer + producer.send_event(EventData('Single data {}'.format(i))) + + batch = producer.create_batch() + for i in range(10): + batch.add(EventData('Single data in batch {}'.format(i))) + # alternatively, you can enqueue an EventDataBatch object to the buffer + producer.send_batch(batch) + + # calling flush sends out the events in the buffer immediately + producer.flush() + +print("Send messages in {} seconds.".format(time.time() - start_time)) diff --git a/sdk/eventhub/azure-eventhub/setup.py b/sdk/eventhub/azure-eventhub/setup.py index 8730981bc8ea..592f926a2288 100644 --- a/sdk/eventhub/azure-eventhub/setup.py +++ b/sdk/eventhub/azure-eventhub/setup.py @@ -58,14 +58,13 @@ "Development Status :: 4 - Beta", 'Programming Language :: Python', 'Programming Language :: Python :: 3 :: Only', - 'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', 'License :: OSI Approved :: MIT License', ], - python_requires=">=3.6", + python_requires=">=3.7", zip_safe=False, packages=find_packages(exclude=exclude_packages), install_requires=[ diff --git a/sdk/eventhub/azure-eventhub/stress/Dockerfile b/sdk/eventhub/azure-eventhub/stress/Dockerfile index 05c55e17da28..033b01c3fa0c 100644 --- a/sdk/eventhub/azure-eventhub/stress/Dockerfile +++ b/sdk/eventhub/azure-eventhub/stress/Dockerfile @@ -1,6 +1,7 @@ -# syntax=docker/dockerfile:1 - -FROM python:3.8-slim-buster +# internal users should provide MCR registry to build via 'docker build . --build-arg REGISTRY="mcr.microsoft.com/mirror/docker/library/"' +# public OSS users should simply leave this argument blank or ignore its presence entirely +ARG REGISTRY="mcr.microsoft.com/mirror/docker/library/" +FROM ${REGISTRY}python:3.8-slim-buster WORKDIR /app diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py index 273ece110966..cbd2b5131426 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_auth_async.py @@ -16,17 +16,21 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_secret_credential_async(live_eventhub): +async def test_client_secret_credential_async(live_eventhub, uamqp_transport): credential = EnvironmentCredential() producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], credential=credential, - user_agent='customized information') + user_agent='customized information', + uamqp_transport=uamqp_transport + ) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + user_agent='customized information', + uamqp_transport=uamqp_transport + ) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -49,11 +53,11 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_sas_credential_async(live_eventhub): +async def test_client_sas_credential_async(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], - eventhub_name=live_eventhub['event_hub']) + eventhub_name=live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -66,7 +70,8 @@ async def test_client_sas_credential_async(live_eventhub): token = (await credential.get_token(auth_uri)).token producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=EventHubSASTokenCredential(token, time.time() + 3000)) + credential=EventHubSASTokenCredential(token, time.time() + 3000), + uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -76,7 +81,7 @@ async def test_client_sas_credential_async(live_eventhub): # Finally let's do it with SAS token + conn str token_conn_str = "Endpoint=sb://{}/;SharedAccessSignature={};".format(hostname, token.decode()) conn_str_producer_client = EventHubProducerClient.from_connection_string(token_conn_str, - eventhub_name=live_eventhub['event_hub']) + eventhub_name=live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with conn_str_producer_client: batch = await conn_str_producer_client.create_batch(partition_id='0') @@ -86,10 +91,10 @@ async def test_client_sas_credential_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_azure_sas_credential_async(live_eventhub): +async def test_client_azure_sas_credential_async(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] - producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub']) + producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub'], uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -101,7 +106,7 @@ async def test_client_azure_sas_credential_async(live_eventhub): token = (await credential.get_token(auth_uri)).token.decode() producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=AzureSasCredential(token)) + credential=AzureSasCredential(token), uamqp_transport=uamqp_transport) async with producer_client: batch = await producer_client.create_batch(partition_id='0') @@ -111,14 +116,14 @@ async def test_client_azure_sas_credential_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_client_azure_named_key_credential_async(live_eventhub): +async def test_client_azure_named_key_credential_async(live_eventhub, uamqp_transport): credential = AzureNamedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + user_agent='customized information', uamqp_transport=uamqp_transport) assert (await consumer_client.get_eventhub_properties()) is not None diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py new file mode 100644 index 000000000000..b3b155f6780c --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_buffered_producer_async.py @@ -0,0 +1,550 @@ +#!/usr/bin/env python + +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- +import asyncio +from collections import defaultdict +from uuid import uuid4 + +import pytest + +from azure.eventhub import EventData +from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient +from azure.eventhub.aio._buffered_producer import PartitionResolver +from azure.eventhub.amqp import ( + AmqpAnnotatedMessage, +) +from azure.eventhub.exceptions import EventDataSendError, OperationTimeoutError, EventHubError + + +async def random_pkey_generation(partitions): + pr = PartitionResolver(partitions) + total = len(partitions) + dic = {} + + while total: + key = str(uuid4()) + pid = await pr.get_partition_id_by_partition_key(key) + if pid in dic: + continue + else: + dic[pid] = key + total -= 1 + + return dic + + +@pytest.mark.liveTest() +@pytest.mark.asyncio +async def test_producer_client_constructor(connection_str, uamqp_transport): + async def on_success(events, pid): + pass + + async def on_error(events, error, pid): + pass + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, uamqp_transport=uamqp_transport) + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_success=on_success, uamqp_transport=uamqp_transport) + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_error=on_error, uamqp_transport=uamqp_transport) + with pytest.raises(ValueError): + EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + max_wait_time=0, + uamqp_transport=uamqp_transport + ) + with pytest.raises(ValueError): + EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + max_buffer_length=0, + uamqp_transport=uamqp_transport + ) + + +@pytest.mark.liveTest +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flush_after_sending, close_after_sending", + [ + (False, False), + (True, False), + (False, True) + ] +) +async def test_basic_send_single_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + asyncio.sleep(5) + sent_events = defaultdict(list) + + async def on_success(events, pid): + if len(events) > 1: + on_success.batching = True + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + on_success.batching = False # ensure batching happened + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + partitions = await producer.get_partition_ids() + partitions_cnt = len(partitions) + # perform single sending round-robin + total_single_event_cnt = 100 + eventdata_set, amqpannoated_set = set(), set() + for i in range(total_single_event_cnt // 2): + event = EventData("test:{}".format(i)) + event.properties = {"event_idx": i} + await producer.send_event(event) + eventdata_set.add(i) + for i in range(total_single_event_cnt // 2, total_single_event_cnt): + event = AmqpAnnotatedMessage(data_body="test:{}".format(i)) + event.application_properties = {"event_idx": i} + amqpannoated_set.add(i) + await producer.send_event(event) + + for pid in partitions: + assert producer.get_buffered_event_count(pid) > 0 + assert producer.total_buffered_event_count > 0 + + if not flush_after_sending and not close_after_sending: + # ensure it's buffered sending + for pid in partitions: + assert len(sent_events[pid]) < total_single_event_cnt // partitions_cnt + assert sum([len(sent_events[pid]) for pid in partitions]) < total_single_event_cnt + else: + if flush_after_sending: + await producer.flush() + if close_after_sending: + await producer.close() + # ensure all events are sent after calling flush + assert sum([len(sent_events[pid]) for pid in partitions]) == total_single_event_cnt + + # give some time for producer to complete sending and consumer to complete receiving + await asyncio.sleep(10) + assert len(sent_events) == len(received_events) == partitions_cnt + + for pid in partitions: + assert producer.get_buffered_event_count(pid) == 0 + assert producer.total_buffered_event_count == 0 + assert not on_error.err + + # ensure all events are received in the correct partition + for pid in partitions: + assert len(sent_events[pid]) >= total_single_event_cnt // partitions_cnt + assert len(sent_events[pid]) == len(received_events[pid]) + for i in range(len(sent_events[pid])): + event = sent_events[pid][i] + try: # amqp annotated message + event_idx = event.application_properties["event_idx"] + amqpannoated_set.remove(event_idx) + except AttributeError: # event data + event_idx = event.properties["event_idx"] + eventdata_set.remove(event_idx) + assert received_events[pid][i].properties[b"event_idx"] == event_idx + assert partitions[event_idx % partitions_cnt] == pid + + assert on_success.batching + assert not eventdata_set + assert not amqpannoated_set + + await consumer.close() + await receive_thread + + +@pytest.mark.liveTest +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flush_after_sending, close_after_sending", + [ + (True, False), + (False, True), + (False, False) + ] +) +async def test_basic_send_batch_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + asyncio.sleep(5) + sent_events = defaultdict(list) + + async def on_success(events, pid): + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + partitions = await producer.get_partition_ids() + partitions_cnt = len(partitions) + # perform batch sending round-robin + total_events_cnt = 100 + batch_cnt = partitions_cnt * 2 - 1 + each_partition_cnt = total_events_cnt // batch_cnt + remain_events = total_events_cnt % batch_cnt + batches = [] + event_idx = 0 + eventdata_set, amqpannoated_set = set(), set() + for i in range(batch_cnt): + batch = await producer.create_batch() + for j in range(each_partition_cnt // 2): + event = EventData("test{}:{}".format(i, event_idx)) + event.properties = {'batch_idx': i, 'event_idx': event_idx} + batch.add(event) + eventdata_set.add(event_idx) + event_idx += 1 + for j in range(each_partition_cnt // 2, each_partition_cnt): + event = AmqpAnnotatedMessage(data_body="test{}:{}".format(i, event_idx)) + event.application_properties = {'batch_idx': i, 'event_idx': event_idx} + batch.add(event) + amqpannoated_set.add(event_idx) + event_idx += 1 + batches.append(batch) + + # put remain_events in the last batch + last_batch = await producer.create_batch() + for i in range(remain_events): + event = EventData("test:{}:{}".format(len(batches), event_idx)) + event.properties = {'batch_idx': len(batches), 'event_idx': event_idx} + last_batch.add(event) + eventdata_set.add(event_idx) + event_idx += 1 + batches.append(last_batch) + + for batch in batches: + await producer.send_batch(batch) + + if not flush_after_sending and not close_after_sending: + # ensure it's buffered sending + for pid in partitions: + assert len(sent_events[pid]) < each_partition_cnt + assert sum([len(sent_events[pid]) for pid in partitions]) < total_events_cnt + # give some time for producer to complete sending and consumer to complete receiving + else: + if flush_after_sending: + await producer.flush() + if close_after_sending: + await producer.close() + # ensure all events are sent + assert sum([len(sent_events[pid]) for pid in partitions]) == total_events_cnt + + await asyncio.sleep(10) + assert len(sent_events) == len(received_events) == partitions_cnt + + # ensure all events are received in the correct partition + for pid in partitions: + assert len(sent_events[pid]) > 0 + assert len(sent_events[pid]) == len(received_events[pid]) + for i in range(len(sent_events[pid])): + event = sent_events[pid][i] + try: # amqp annotated message + event_idx = event.application_properties["event_idx"] + amqpannoated_set.remove(event_idx) + except AttributeError: # event data + event_idx = event.properties["event_idx"] + eventdata_set.remove(event_idx) + assert received_events[pid][i].properties[b"event_idx"] == event_idx + + assert not amqpannoated_set + assert not eventdata_set + assert not on_error.err + + await consumer.close() + await receive_thread + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_send_with_hybrid_partition_assignment(connection_str, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + sent_events = defaultdict(list) + + async def on_success(events, pid): + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + partitions = await producer.get_partition_ids() + partitions_cnt = len(partitions) + pid_to_pkey = await random_pkey_generation(partitions) + expected_event_idx_to_partition = {} + event_idx = 0 + # 1. send by partition_key, each partition 2 events, two single + one batch containing two + for pid in partitions: + pkey = pid_to_pkey[pid] + await producer.send_event(EventData('{}'.format(event_idx)), partition_key=pkey) + batch = await producer.create_batch(partition_key=pkey) + batch.add(EventData('{}'.format(event_idx + 1))) + await producer.send_batch(batch) + for i in range(2): + expected_event_idx_to_partition[event_idx + i] = pid + event_idx += 2 + + # 2. send by partition_id, each partition 2 events, two single + one batch containing two + for pid in partitions: + await producer.send_event(EventData('{}'.format(event_idx)), partition_id=pid) + batch = await producer.create_batch(partition_id=pid) + batch.add(EventData('{}'.format(event_idx + 1))) + await producer.send_batch(batch) + for i in range(2): + expected_event_idx_to_partition[event_idx + i] = pid + event_idx += 2 + + # 3. send without partition, each partition 2 events, two single + one batch containing two + for _ in partitions: + await producer.send_event(EventData('{}'.format(event_idx))) + batch = await producer.create_batch() + batch.add(EventData('{}'.format(event_idx + 1))) + await producer.send_batch(batch) + event_idx += 2 + + await producer.flush() + assert len(sent_events) == partitions_cnt + + await asyncio.sleep(10) + + visited = set() + for pid in partitions: + assert len(sent_events[pid]) == 2 * 3 + + for sent_event in sent_events[pid]: + if int(sent_event.body_as_str()) in expected_event_idx_to_partition: + assert expected_event_idx_to_partition[int(sent_event.body_as_str())] == pid + + for recv_event in received_events[pid]: + if int(sent_event.body_as_str()) in expected_event_idx_to_partition: + assert expected_event_idx_to_partition[int(sent_event.body_as_str())] == pid + + assert recv_event.body_as_str() not in visited + visited.add(recv_event.body_as_str()) + + assert len(visited) == 2 * 3 * len(partitions) + + assert not on_error.err + await consumer.close() + await receive_thread + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_send_with_timing_configuration(connection_str, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + sent_events = defaultdict(list) + + async def on_success(events, pid): + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + + # test max_wait_time + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + max_wait_time=10, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + partitions = await producer.get_partition_ids() + await producer.send_event(EventData('data')) + await asyncio.sleep(5) + assert not sent_events + await asyncio.sleep(20) + assert sum([len(sent_events[pid]) for pid in partitions]) == 1 + + assert not on_error.err + + # test max_buffer_length per partition + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + max_wait_time=1000, + max_buffer_length=10, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + sent_events.clear() + received_events.clear() + async with producer: + partitions = await producer.get_partition_ids() + for i in range(7): + await producer.send_event(EventData('data'), partition_id="0") + assert not sent_events + batch = await producer.create_batch(partition_id="0") + for i in range(9): + batch.add(EventData('9')) + await producer.send_batch(batch) # will flush 7 events and put the batch in buffer + assert sum([len(sent_events[pid]) for pid in partitions]) == 7 + for i in range(5): + await producer.send_event(EventData('data'), partition_id="0") # will flush batch (9 events) + 1 event, leaving 4 in buffer + assert sum([len(sent_events[pid]) for pid in partitions]) == 17 + await producer.flush() + assert sum([len(sent_events[pid]) for pid in partitions]) == 21 + + await asyncio.sleep(5) + assert sum([len(received_events[pid]) for pid in partitions]) == 21 + assert not on_error.err + await consumer.close() + await receive_thread + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_long_sleep(connection_str, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + sent_events = defaultdict(list) + + async def on_success(events, pid): + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + async with producer: + await producer.send_event(EventData("test"), partition_id="0") + await asyncio.sleep(220) + await producer.send_event(EventData("test"), partition_id="0") + await asyncio.sleep(5) + + assert not on_error.err + assert len(sent_events["0"]) == 2 + assert len(received_events["0"]) == 2 + + await consumer.close() + await receive_thread + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_long_wait_small_buffer(connection_str, uamqp_transport): + received_events = defaultdict(list) + + async def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + + receive_thread = asyncio.ensure_future(consumer.receive(on_event=on_event)) + + sent_events = defaultdict(list) + + async def on_success(events, pid): + sent_events[pid].extend(events) + + async def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + auth_timeout=3, + retry_total=3, + retry_mode='fixed', + retry_backoff_factor=0.01, + max_wait_time=10, + max_buffer_length=100, + uamqp_transport=uamqp_transport + ) + + async with producer: + for i in range(100): + await producer.send_event(EventData("test")) + + await asyncio.sleep(60) + + assert not on_error.err + assert sum([len(sent_events[key]) for key in sent_events]) == 100 + assert sum([len(received_events[key]) for key in received_events]) == 100 + + await consumer.close() + await receive_thread diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py index 0ff6ae711cbf..9831cb9db439 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py @@ -8,11 +8,11 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_no_partition_async(connstr_senders): +async def test_receive_no_partition_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event(partition_context, event): on_event.received += 1 @@ -49,10 +49,10 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_partition_async(connstr_senders): +async def test_receive_partition_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event(partition_context, event): assert partition_context.partition_id == "0" @@ -72,13 +72,13 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_load_balancing_async(connstr_senders): +async def test_receive_load_balancing_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders cs = InMemoryCheckpointStore() client1 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport) client2 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport) async def on_event(partition_context, event): pass @@ -98,13 +98,13 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_no_max_wait_time_async(connstr_senders): +async def test_receive_batch_no_max_wait_time_async(connstr_senders, uamqp_transport): '''Test whether callback is called when max_wait_time is None and max_batch_size has reached ''' connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) @@ -144,10 +144,10 @@ async def on_event_batch(partition_context, event_batch): ]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_empty_with_max_wait_time_async(connection_str, max_wait_time, sleep_time, expected_result): +async def test_receive_batch_empty_with_max_wait_time_async(connection_str, max_wait_time, sleep_time, expected_result, uamqp_transport): '''Test whether event handler is called when max_wait_time > 0 and no event is received ''' - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.event_batch = event_batch @@ -163,13 +163,13 @@ async def on_event_batch(partition_context, event_batch): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_batch_early_callback_async(connstr_senders): +async def test_receive_batch_early_callback_async(connstr_senders, uamqp_transport): ''' Test whether the callback is called once max_batch_size reaches and before max_wait_time reaches. ''' connection_str, senders = connstr_senders for _ in range(10): senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py index 0efb89d6f3eb..43a651072c13 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_negative_async.py @@ -23,24 +23,44 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_hostname_async(invalid_hostname, connstr_receivers): +async def test_send_with_invalid_hostname_async(invalid_hostname, connstr_receivers, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") _, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() batch.add(EventData("test data")) await client.send_batch(batch) + # test setting callback + async def on_error(events, pid, err): + assert len(events) == 1 + assert not pid + on_error.err = err + + on_error.err = None + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) + async with client: + batch = EventDataBatch() + batch.add(EventData("test data")) + await client.send_batch(batch) + assert isinstance(on_error.err, ConnectError) + + on_error.err = None + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) + async with client: + await client.send_event(EventData("test data")) + assert isinstance(on_error.err, ConnectError) + @pytest.mark.parametrize("invalid_place", ["hostname", "key_name", "access_key", "event_hub", "partition"]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_with_invalid_param_async(live_eventhub, invalid_place): +async def test_receive_with_invalid_param_async(live_eventhub, invalid_place, uamqp_transport): eventhub_config = live_eventhub.copy() if invalid_place != "partition": eventhub_config[invalid_place] = "invalid " + invalid_place @@ -50,7 +70,7 @@ async def test_receive_with_invalid_param_async(live_eventhub, invalid_place): eventhub_config['access_key'], eventhub_config['event_hub']) - client = EventHubConsumerClient.from_connection_string(conn_str, consumer_group='$default', retry_total=0) + client = EventHubConsumerClient.from_connection_string(conn_str, consumer_group='$default', retry_total=0, uamqp_transport=uamqp_transport) async def on_event(partition_context, event): pass @@ -69,8 +89,8 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_key_async(invalid_key): - client = EventHubProducerClient.from_connection_string(invalid_key) +async def test_send_with_invalid_key_async(invalid_key, uamqp_transport): + client = EventHubProducerClient.from_connection_string(invalid_key, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -80,8 +100,8 @@ async def test_send_with_invalid_key_async(invalid_key): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_invalid_policy_async(invalid_policy): - client = EventHubProducerClient.from_connection_string(invalid_policy) +async def test_send_with_invalid_policy_async(invalid_policy, uamqp_transport): + client = EventHubProducerClient.from_connection_string(invalid_policy, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -91,8 +111,8 @@ async def test_send_with_invalid_policy_async(invalid_policy): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_non_existing_entity_sender_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str, eventhub_name="nemo") +async def test_non_existing_entity_sender_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, eventhub_name="nemo", uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -102,10 +122,10 @@ async def test_non_existing_entity_sender_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_to_invalid_partitions_async(connection_str): +async def test_send_to_invalid_partitions_async(connection_str, uamqp_transport): partitions = ["XYZ", "-1", "1000", "-"] for p in partitions: - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ConnectError): batch = await client.create_batch(partition_id=p) @@ -117,10 +137,10 @@ async def test_send_to_invalid_partitions_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_too_large_message_async(connection_str): +async def test_send_too_large_message_async(connection_str, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - open issue regarding message size") - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: data = EventData(b"A" * 1100000) with pytest.raises(ValueError): @@ -132,8 +152,8 @@ async def test_send_too_large_message_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_null_body_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_send_null_body_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ValueError): data = EventData(None) @@ -146,11 +166,11 @@ async def test_send_null_body_async(connection_str): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_create_batch_with_invalid_hostname_async(invalid_hostname): +async def test_create_batch_with_invalid_hostname_async(invalid_hostname, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ConnectError): await client.create_batch(max_size_in_bytes=300) @@ -158,21 +178,23 @@ async def test_create_batch_with_invalid_hostname_async(invalid_hostname): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_create_batch_with_too_large_size_async(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_create_batch_with_too_large_size_async(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: with pytest.raises(ValueError): await client.create_batch(max_size_in_bytes=5 * 1024 * 1024) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_invalid_proxy_server(connection_str): +async def test_invalid_proxy_server(connection_str, uamqp_transport): + if uamqp_transport: + pytest.skip() HTTP_PROXY = { 'proxy_hostname': 'fakeproxy', # proxy hostname. 'proxy_port': 3128, # proxy port. } - client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY) + client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY, uamqp_transport=uamqp_transport) async with client: with pytest.raises(EventHubError): batch = await client.create_batch() diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py index fe53764dd8dc..b957d0d197d4 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_properties_async.py @@ -12,9 +12,10 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties(live_eventhub): +async def test_get_properties(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: properties = await client.get_eventhub_properties() @@ -22,16 +23,18 @@ async def test_get_properties(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties_with_auth_error_async(live_eventhub): +async def test_get_properties_with_auth_error_async(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf=") + EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf="), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(AuthenticationError) as e: await client.get_eventhub_properties() client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential("invalid", live_eventhub['access_key']) + EventHubSharedKeyCredential("invalid", live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(AuthenticationError) as e: @@ -39,16 +42,18 @@ async def test_get_properties_with_auth_error_async(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_properties_with_connect_error(live_eventhub): +async def test_get_properties_with_connect_error(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], "invalid", '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(ConnectError) as e: await client.get_eventhub_properties() client = EventHubConsumerClient("invalid.servicebus.windows.net", live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: with pytest.raises(EventHubError) as e: # This can be either ConnectError or ConnectionLostError @@ -56,9 +61,10 @@ async def test_get_properties_with_connect_error(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_partition_ids(live_eventhub): +async def test_get_partition_ids(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: partition_ids = await client.get_partition_ids() @@ -67,9 +73,10 @@ async def test_get_partition_ids(live_eventhub): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_get_partition_properties(live_eventhub): +async def test_get_partition_properties(live_eventhub, uamqp_transport): client = EventHubProducerClient(live_eventhub['hostname'], live_eventhub['event_hub'], - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) async with client: properties = await client.get_partition_properties('0') diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py index 6fb153052023..18279b81cc2d 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py @@ -15,7 +15,7 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_end_of_stream_async(connstr_senders): +async def test_receive_end_of_stream_async(connstr_senders, uamqp_transport): async def on_event(partition_context, event): if partition_context.partition_id == "0": assert event.body_as_str() == "Receiving only a single event" @@ -31,7 +31,7 @@ async def on_event(partition_context, event): on_event.called = False connection_str, senders = connstr_senders # test async producer client - producer_client = EventHubProducerClient.from_connection_string(connection_str) + producer_client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) partitions = await producer_client.get_partition_ids() senders = [] for p in partitions: @@ -61,7 +61,7 @@ async def on_event(partition_context, event): ("enqueued_time", False, "Exclusive")]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_with_event_position_async(connstr_senders, position, inclusive, expected_result): +async def test_receive_with_event_position_async(connstr_senders, position, inclusive, expected_result, uamqp_transport): async def on_event(partition_context, event): assert partition_context.last_enqueued_event_properties.get('sequence_number') == event.sequence_number assert partition_context.last_enqueued_event_properties.get('offset') == event.offset @@ -79,7 +79,7 @@ async def on_event(partition_context, event): on_event.event_position = None connection_str, senders = connstr_senders senders[0].send(EventData(b"Inclusive")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client: task = asyncio.ensure_future(client.receive(on_event, starting_position="-1", @@ -89,7 +89,7 @@ async def on_event(partition_context, event): assert on_event.event_position is not None await task senders[0].send(EventData(expected_result)) - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client2: task = asyncio.ensure_future( client2.receive(on_event, @@ -102,7 +102,7 @@ async def on_event(partition_context, event): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_owner_level_async(connstr_senders): +async def test_receive_owner_level_async(connstr_senders, uamqp_transport): app_prop = {"raw_prop": "raw_value"} async def on_event(partition_context, event): @@ -112,8 +112,8 @@ async def on_error(partition_context, error): on_error.error = None connection_str, senders = connstr_senders - client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) + client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) async with client1, client2: task1 = asyncio.ensure_future(client1.receive(on_event, partition_id="0", starting_position="-1", @@ -136,7 +136,7 @@ async def on_error(partition_context, error): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_over_websocket_async(connstr_senders): +async def test_receive_over_websocket_async(connstr_senders, uamqp_transport): app_prop = {"raw_prop": "raw_value"} content_type = "text/plain" message_id_base = "mess_id_sample_" @@ -149,7 +149,8 @@ async def on_event(partition_context, event): on_event.app_prop = None connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', - transport_type=TransportType.AmqpOverWebsocket) + transport_type=TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) event_list = [] for i in range(5): diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py index 2a1fc55b2749..8e9fc947d3c1 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_reconnect_async.py @@ -12,7 +12,15 @@ from azure.eventhub._pyamqp.aio._authentication_async import SASTokenAuthAsync from azure.eventhub._pyamqp.aio import ReceiveClientAsync from azure.eventhub._pyamqp import error, constants - +from azure.eventhub._utils import transform_outbound_single_message +try: + import uamqp + from uamqp import compat + from azure.eventhub._transport._uamqp_transport import UamqpTransport +except ImportError: + UamqpTransport = None + +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport from azure.eventhub import EventData from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient, EventHubSharedKeyCredential from azure.eventhub.exceptions import OperationTimeoutError @@ -20,10 +28,11 @@ @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_long_interval_async(live_eventhub, sleep): +async def test_send_with_long_interval_async(live_eventhub, sleep, uamqp_transport, timeout_factor): test_partition = "0" sender = EventHubProducerClient(live_eventhub['hostname'], live_eventhub['event_hub'], - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport) async with sender: batch = await sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) @@ -32,31 +41,38 @@ async def test_send_with_long_interval_async(live_eventhub, sleep): if sleep: await asyncio.sleep(250) # EH server side idle timeout is 240 second else: - await sender._producers[test_partition]._handler._connection.close() + if uamqp_transport: + await sender._producers[test_partition]._handler._connection._conn.destroy() + else: + await sender._producers[test_partition]._handler._connection.close() batch = await sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) await sender.send_batch(batch) received = [] uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = SASTokenAuthAsync( - uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] - ) - source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], test_partition) - receiver = ReceiveClientAsync(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAsync.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + receiver = uamqp.async_ops.client_async.ReceiveClientAsync(source, auth=sas_auth, debug=False, timeout=5000, prefetch=500) + else: + sas_auth = SASTokenAuthAsync( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClientAsync(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) try: await receiver.open_async() # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5))]) - received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5))]) + received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5 * timeout_factor))]) + received.extend([EventData._from_message(x) for x in (await receiver.receive_message_batch_async(max_batch_size=1, timeout=5 * timeout_factor))]) finally: await receiver.close_async() assert len(received) == 2 @@ -65,51 +81,71 @@ async def test_send_with_long_interval_async(live_eventhub, sleep): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_connection_idle_timeout_and_reconnect_async(connstr_receivers): +async def test_send_connection_idle_timeout_and_reconnect_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - # no retry, should just raise error - client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10, retry_total=0) + if uamqp_transport: + amqp_transport = UamqpTransport + retry_total = 3 + timeout_exc = compat.TimeoutException + else: + amqp_transport = PyamqpTransport + retry_total = 0 + timeout_exc = TimeoutError + + client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport) async with client: ed = EventData('data') sender = client._create_producer(partition_id='0') async with sender: await sender._open_with_retry() await asyncio.sleep(11) - sender._unsent_events = [ed.message] - with pytest.raises(error.AMQPConnectionError): - await sender._send_event_data() + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + if uamqp_transport: + sender._unsent_events[0].on_send_complete = sender._on_outcome + with pytest.raises((uamqp.errors.ConnectionClose, + uamqp.errors.MessageHandlerError, OperationTimeoutError)): + await sender._send_event_data() + else: + with pytest.raises(error.AMQPConnectionError): + await sender._send_event_data() + if uamqp_transport: + await sender._send_event_data_with_retry() # with retry, should work - client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10) - async with client: - ed = EventData('data') - sender = client._create_producer(partition_id='0') - async with sender: - await sender._open_with_retry() - await asyncio.sleep(11) - sender._unsent_events = [ed.message] - await sender._send_event_data() + if not uamqp_transport: + client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10) + async with client: + ed = EventData('data') + sender = client._create_producer(partition_id='0') + async with sender: + await sender._open_with_retry() + await asyncio.sleep(11) + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + await sender._send_event_data() retry = 0 while retry < 3: try: - messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10) + messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10 * timeout_factor) if messages: received_ed1 = EventData._from_message(messages[0]) assert received_ed1.body_as_str() == 'data' break - except TimeoutError: + except timeout_exc: retry += 1 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_receive_connection_idle_timeout_and_reconnect_async(connstr_senders): +async def test_receive_connection_idle_timeout_and_reconnect_async(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string( conn_str=connection_str, consumer_group='$default', - idle_timeout=10 + idle_timeout=10, + uamqp_transport=uamqp_transport ) async def on_event_received(event): @@ -118,19 +154,26 @@ async def on_event_received(event): async with client: consumer = client._create_consumer("$default", "0", "-1", on_event_received) async with consumer: - await consumer._open() - await asyncio.sleep(11) + if uamqp_transport: + await consumer._open_with_retry() + else: + await consumer._open() + await asyncio.sleep(11) ed = EventData("Event") senders[0].send(ed) - with pytest.raises(error.AMQPConnectionError): + if uamqp_transport: await consumer._handler.do_work_async() - assert consumer._handler._connection.state == constants.ConnectionState.END - - try: - await asyncio.wait_for(consumer.receive(), timeout=10) - except asyncio.TimeoutError: - pass + assert consumer._handler._connection._state == uamqp.c_uamqp.ConnectionState.DISCARDING + await consumer.receive(batch=False, max_batch_size=1, max_wait_time=10) + else: + with pytest.raises(error.AMQPConnectionError): + await consumer._handler.do_work_async() + assert consumer._handler._connection.state == constants.ConnectionState.END + try: + await asyncio.wait_for(consumer.receive(), timeout=10) + except asyncio.TimeoutError: + pass assert on_event_received.event.body_as_str() == "Event" diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py index 6f8b8f57e6de..8fedec30abad 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py @@ -10,22 +10,35 @@ import pytest import time import json +import uamqp from azure.eventhub import EventData, TransportType, EventDataBatch from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient -from azure.eventhub.exceptions import EventDataSendError +from azure.eventhub.exceptions import EventDataSendError, OperationTimeoutError from azure.eventhub.amqp import ( AmqpMessageHeader, AmqpMessageBodyType, AmqpAnnotatedMessage, AmqpMessageProperties, ) +try: + import uamqp + from uamqp.constants import TransportType as uamqp_TransportType, MessageState + from uamqp.message import MessageProperties +except ImportError: + uamqp_TransportType = TransportType + MessageProperties = None +from azure.eventhub._pyamqp.message import Properties +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp.client import ReceiveClient +from azure.eventhub._pyamqp.error import AMQPConnectionError + @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_amqp_annotated_message(connstr_receivers): +async def test_send_amqp_annotated_message(connstr_receivers, uamqp_transport): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: sequence_body = [b'message', 123.456, True] footer = {'footer_key': 'footer_value'} @@ -77,6 +90,10 @@ async def test_send_amqp_annotated_message(connstr_receivers): batch.add(event_data) await client.send_batch(batch) await client.send_batch([data_message, value_message, sequence_message, event_data]) + await client.send_event(data_message) + await client.send_event(value_message) + await client.send_event(sequence_message) + await client.send_event(event_data) received_count = {} received_count["data_msg"] = 0 @@ -121,7 +138,7 @@ async def on_event(partition_context, event): on_event.received = [] client = EventHubConsumerClient.from_connection_string(connection_str, - consumer_group='$default') + consumer_group='$default', uamqp_transport=uamqp_transport) async with client: task = asyncio.ensure_future(client.receive(on_event, starting_position="-1")) await asyncio.sleep(15) @@ -130,23 +147,23 @@ async def on_event(partition_context, event): await task - assert len(on_event.received) == 8 - assert received_count["data_msg"] == 2 - assert received_count["seq_msg"] == 2 - assert received_count["value_msg"] == 2 - assert received_count["normal_msg"] == 2 + assert len(on_event.received) == 12 + assert received_count["data_msg"] == 3 + assert received_count["seq_msg"] == 3 + assert received_count["value_msg"] == 3 + assert received_count["normal_msg"] == 3 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_partition_key_async(connstr_receivers): +async def test_send_with_partition_key_async(connstr_receivers, live_eventhub, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: data_val = 0 for partition in [b"a", b"b", b"c", b"d", b"e", b"f"]: partition_key = b"test_partition_" + partition - for i in range(50): + for i in range(10): batch = await client.create_batch(partition_key=partition_key) batch.add(EventData(str(data_val))) data_val += 1 @@ -154,151 +171,222 @@ async def test_send_with_partition_key_async(connstr_receivers): await client.send_batch(await client.create_batch()) + for partition in [b"a", b"b", b"c", b"d", b"e", b"f"]: + partition_key = b"test_partition_" + partition + for i in range(10): + event_data = EventData(str(data_val)) + event_data.properties = {'is_single': True} + data_val += 1 + await client.send_event(event_data, partition_key=partition_key) + + batch_cnt = 0 + single_cnt = 0 found_partition_keys = {} + reconnect_receivers = [] for index, partition in enumerate(receivers): - received = partition.receive_message_batch(timeout=5) - for message in received: + retry_total = 0 + while retry_total < 3: + timeout = (5 * retry_total) * timeout_factor try: - event_data = EventData._from_message(message) - existing = found_partition_keys[event_data.partition_key] - assert existing == index - except KeyError: - found_partition_keys[event_data.partition_key] = index + received = partition.receive_message_batch(timeout=timeout) + for message in received: + try: + event_data = EventData._from_message(message) + if event_data.properties and event_data.properties[b'is_single']: + single_cnt += 1 + else: + batch_cnt += 1 + existing = found_partition_keys[event_data.partition_key] + assert existing == index + except KeyError: + found_partition_keys[event_data.partition_key] = index + if received: + break + retry_total += 1 + except AMQPConnectionError: + for r in reconnect_receivers: + r.close() + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + index) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) + partition.open() + reconnect_receivers.append(partition) + retry_total += 1 + if retry_total == 3: + raise OperationTimeoutError(f"Exhausted retries for receiving from {live_eventhub['hostname']}.") + for r in reconnect_receivers: + r.close() + + assert single_cnt == 60 + assert batch_cnt == 60 + assert len(found_partition_keys) == 6 @pytest.mark.parametrize("payload", [b"", b"A single event"]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_and_receive_small_body_async(connstr_receivers, payload): +async def test_send_and_receive_small_body_async(connstr_receivers, payload, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch() batch.add(EventData(payload)) await client.send_batch(batch) + await client.send_event(EventData(payload)) received = [] for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=5)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=5 * timeout_factor)]) - assert len(received) == 1 + assert len(received) == 2 assert list(received[0].body)[0] == payload + assert list(received[1].body)[0] == payload @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_partition_async(connstr_receivers): +async def test_send_partition_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch() batch.add(EventData(b"Data")) await client.send_batch(batch) + await client.send_event(EventData(b"Data")) + async with client: batch = await client.create_batch(partition_id="1") batch.add(EventData(b"Data")) await client.send_batch(batch) + await client.send_event(EventData(b"Data"), partition_id="1") - partition_0 = receivers[0].receive_message_batch(timeout=10) - partition_1 = receivers[1].receive_message_batch(timeout=10) - assert len(partition_0) + len(partition_1) == 2 + partition_0 = receivers[0].receive_message_batch(timeout=10 * timeout_factor) + partition_1 = receivers[1].receive_message_batch(timeout=10 * timeout_factor) + assert len(partition_1) >= 2 + assert len(partition_0) + len(partition_1) == 4 async with client: batch = await client.create_batch() batch.add(EventData(b"Data")) await client.send_batch(batch) - + await client.send_event(EventData(b"Data")) async with client: - batch = await client.create_batch(partition_id="1") + batch = await client.create_batch(partition_id="0") + batch.add(EventData(b"Data")) + await client.send_batch(batch) + await client.send_event(EventData(b"Data"), partition_id="0") + async with client: + batch = EventDataBatch(partition_id="0") batch.add(EventData(b"Data")) await client.send_batch(batch) time.sleep(5) - partition_0 = receivers[0].receive_message_batch(timeout=10) - partition_1 = receivers[1].receive_message_batch(timeout=10) - assert len(partition_0) + len(partition_1) == 2 + partition_0 = receivers[0].receive_message_batch(timeout=10 * timeout_factor) + partition_1 = receivers[1].receive_message_batch(timeout=10 * timeout_factor) + assert len(partition_0) >= 3 + assert len(partition_0) + len(partition_1) == 5 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_non_ascii_async(connstr_receivers): +async def test_send_non_ascii_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch(partition_id="0") batch.add(EventData(u"é,è,à,ù,â,ê,î,ô,û")) batch.add(EventData(json.dumps({"foo": u"漢字"}))) await client.send_batch(batch) + await client.send_event(EventData(u"é,è,à,ù,â,ê,î,ô,û"), partition_id="0") + await client.send_event(EventData(json.dumps({"foo": u"漢字"})), partition_id="0") await asyncio.sleep(1) # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] + \ - [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] + \ + [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] - assert len(partition_0) == 2 + assert len(partition_0) == 4 assert partition_0[0].body_as_str() == u"é,è,à,ù,â,ê,î,ô,û" assert partition_0[1].body_as_json() == {"foo": u"漢字"} + assert partition_0[2].body_as_str() == u"é,è,à,ù,â,ê,î,ô,û" + assert partition_0[3].body_as_json() == {"foo": u"漢字"} @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_multiple_partition_with_app_prop_async(connstr_receivers): +async def test_send_multiple_partition_with_app_prop_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: ed0 = EventData(b"Message 0") ed0.properties = app_prop batch = await client.create_batch(partition_id="0") batch.add(ed0) await client.send_batch(batch) + await client.send_event(ed0, partition_id="0") ed1 = EventData(b"Message 1") ed1.properties = app_prop batch = await client.create_batch(partition_id="1") batch.add(ed1) await client.send_batch(batch) - - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] - assert len(partition_0) == 1 + await client.send_event(ed1, partition_id="1") + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5 * timeout_factor)] + assert len(partition_0) == 2 assert partition_0[0].properties[b"raw_prop"] == b"raw_value" - partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=50)] - assert len(partition_1) == 1 + assert partition_0[1].properties[b"raw_prop"] == b"raw_value" + partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=5 * timeout_factor)] + assert len(partition_0) == 2 assert partition_1[0].properties[b"raw_prop"] == b"raw_value" + assert partition_0[1].properties[b"raw_prop"] == b"raw_value" @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_over_websocket_async(connstr_receivers): - pytest.skip("websocket unsupported") +async def test_send_over_websocket_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers client = EventHubProducerClient.from_connection_string(connection_str, - transport_type=TransportType.AmqpOverWebsocket) + transport_type=uamqp_TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) async with client: batch = await client.create_batch(partition_id="0") batch.add(EventData("Event Data")) await client.send_batch(batch) + await client.send_event(EventData("Event Data"), partition_id="0") time.sleep(1) received = [] - received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=10)) - assert len(received) == 1 + received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=10 * timeout_factor)) + assert len(received) == 2 @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_with_create_event_batch_async(connstr_receivers): +async def test_send_with_create_event_batch_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} client = EventHubProducerClient.from_connection_string(connection_str, - transport_type=TransportType.AmqpOverWebsocket) + transport_type=TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) async with client: event_data_batch = await client.create_batch(max_size_in_bytes=100000) while True: @@ -311,23 +399,22 @@ async def test_send_with_create_event_batch_async(connstr_receivers): await client.send_batch(event_data_batch) received = [] for r in receivers: - received.extend(r.receive_message_batch(timeout=10)) + received.extend(r.receive_message_batch(timeout=10 * timeout_factor)) assert len(received) >= 1 assert EventData._from_message(received[0]).properties[b"raw_prop"] == b"raw_value" - @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_async(connstr_receivers): +async def test_send_list_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) payload = "A1" async with client: await client.send_batch([EventData(payload)]) received = [] for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10 * timeout_factor)]) assert len(received) == 1 assert received[0].body_as_str() == payload @@ -335,13 +422,13 @@ async def test_send_list_async(connstr_receivers): @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_partition_async(connstr_receivers): +async def test_send_list_partition_async(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) payload = "A1" async with client: await client.send_batch([EventData(payload)], partition_id="0") - message = receivers[0].receive_message_batch(timeout=10)[0] + message = receivers[0].receive_message_batch(timeout=10 * timeout_factor)[0] received = EventData._from_message(message) assert received.body_as_str() == payload @@ -352,8 +439,8 @@ async def test_send_list_partition_async(connstr_receivers): ]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_list_wrong_data_async(connection_str, to_send, exception_type): - client = EventHubProducerClient.from_connection_string(connection_str) +async def test_send_list_wrong_data_async(connection_str, to_send, exception_type, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) async with client: with pytest.raises(exception_type): await client.send_batch(to_send) @@ -362,10 +449,59 @@ async def test_send_list_wrong_data_async(connection_str, to_send, exception_typ @pytest.mark.parametrize("partition_id, partition_key", [("0", None), (None, "pk")]) @pytest.mark.liveTest @pytest.mark.asyncio -async def test_send_batch_pid_pk_async(invalid_hostname, partition_id, partition_key): +async def test_send_batch_pid_pk_async(invalid_hostname, partition_id, partition_key, uamqp_transport): # Use invalid_hostname because this is not a live test. - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) batch = EventDataBatch(partition_id=partition_id, partition_key=partition_key) async with client: with pytest.raises(TypeError): await client.send_batch(batch, partition_id=partition_id, partition_key=partition_key) + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_send_with_callback_async(connstr_receivers, uamqp_transport): + + async def on_error(events, pid, err): + on_error.err = err + + async def on_success(events, pid): + sent_events.append((events, pid)) + + sent_events = [] + on_error.err = None + connection_str, receivers = connstr_receivers + client = EventHubProducerClient.from_connection_string(connection_str, on_success=on_success, on_error=on_error, uamqp_transport=uamqp_transport) + + async with client: + batch = await client.create_batch() + batch.add(EventData(b"Data")) + batch.add(EventData(b"Data")) + await client.send_batch(batch) + assert len(sent_events[-1][0]) == 2 + assert not sent_events[-1][1] + await client.send_event(EventData(b"Data")) + assert len(sent_events[-1][0]) == 1 + assert not sent_events[-1][1] + + batch = await client.create_batch(partition_key='key') + batch.add(EventData(b"Data")) + batch.add(EventData(b"Data")) + await client.send_batch(batch) + assert len(sent_events[-1][0]) == 2 + assert not sent_events[-1][1] + await client.send_event(EventData(b"Data"), partition_key='key') + assert len(sent_events[-1][0]) == 1 + assert not sent_events[-1][1] + + batch = await client.create_batch(partition_id="0") + batch.add(EventData(b"Data")) + await client.send_batch(batch) + batch.add(EventData(b"Data")) + assert len(sent_events[-1][0]) == 2 + assert sent_events[-1][1] == "0" + await client.send_event(EventData(b"Data"), partition_id="0") + assert len(sent_events[-1][0]) == 1 + assert sent_events[-1][1] == "0" + + assert not on_error.err diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py index c00ea84067ea..f6dbfa3771c1 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_auth.py @@ -14,17 +14,20 @@ @pytest.mark.liveTest -def test_client_secret_credential(live_eventhub): +def test_client_secret_credential(live_eventhub, uamqp_transport): credential = EnvironmentCredential() producer_client = EventHubProducerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], credential=credential, - user_agent='customized information') + user_agent='customized information', + uamqp_transport=uamqp_transport) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + user_agent='customized information', + uamqp_transport=uamqp_transport + ) with producer_client: batch = producer_client.create_batch(partition_id='0') batch.add(EventData(body='A single message')) @@ -51,10 +54,12 @@ def on_event(partition_context, event): @pytest.mark.liveTest -def test_client_sas_credential(live_eventhub): +def test_client_sas_credential(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] - producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub']) + producer_client = EventHubProducerClient.from_connection_string( + live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub'], uamqp_transport=uamqp_transport + ) with producer_client: batch = producer_client.create_batch(partition_id='0') @@ -67,7 +72,8 @@ def test_client_sas_credential(live_eventhub): token = credential.get_token(auth_uri).token producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=EventHubSASTokenCredential(token, time.time() + 3000)) + credential=EventHubSASTokenCredential(token, time.time() + 3000), + uamqp_transport=uamqp_transport) with producer_client: batch = producer_client.create_batch(partition_id='0') @@ -77,7 +83,8 @@ def test_client_sas_credential(live_eventhub): # Finally let's do it with SAS token + conn str token_conn_str = "Endpoint=sb://{}/;SharedAccessSignature={};".format(hostname, token.decode()) conn_str_producer_client = EventHubProducerClient.from_connection_string(token_conn_str, - eventhub_name=live_eventhub['event_hub']) + eventhub_name=live_eventhub['event_hub'], + uamqp_transport=uamqp_transport) with conn_str_producer_client: batch = conn_str_producer_client.create_batch(partition_id='0') @@ -86,10 +93,12 @@ def test_client_sas_credential(live_eventhub): @pytest.mark.liveTest -def test_client_azure_sas_credential(live_eventhub): +def test_client_azure_sas_credential(live_eventhub, uamqp_transport): # This should "just work" to validate known-good. hostname = live_eventhub['hostname'] - producer_client = EventHubProducerClient.from_connection_string(live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub']) + producer_client = EventHubProducerClient.from_connection_string( + live_eventhub['connection_str'], eventhub_name = live_eventhub['event_hub'], uamqp_transport=uamqp_transport + ) with producer_client: batch = producer_client.create_batch(partition_id='0') @@ -102,7 +111,8 @@ def test_client_azure_sas_credential(live_eventhub): token = credential.get_token(auth_uri).token.decode() producer_client = EventHubProducerClient(fully_qualified_namespace=hostname, eventhub_name=live_eventhub['event_hub'], - credential=AzureSasCredential(token)) + credential=AzureSasCredential(token), + uamqp_transport=uamqp_transport) with producer_client: batch = producer_client.create_batch(partition_id='0') @@ -111,13 +121,14 @@ def test_client_azure_sas_credential(live_eventhub): @pytest.mark.liveTest -def test_client_azure_named_key_credential(live_eventhub): +def test_client_azure_named_key_credential(live_eventhub, uamqp_transport): credential = AzureNamedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) consumer_client = EventHubConsumerClient(fully_qualified_namespace=live_eventhub['hostname'], eventhub_name=live_eventhub['event_hub'], consumer_group='$default', credential=credential, - user_agent='customized information') + user_agent='customized information', + uamqp_transport=uamqp_transport) assert consumer_client.get_eventhub_properties() is not None diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py new file mode 100644 index 000000000000..dede600e7405 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_buffered_producer.py @@ -0,0 +1,560 @@ +#!/usr/bin/env python + +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import time +from collections import defaultdict +from threading import Thread +from uuid import uuid4 +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from azure.eventhub import EventData +from azure.eventhub import EventHubProducerClient, EventHubConsumerClient +from azure.eventhub._buffered_producer import PartitionResolver +from azure.eventhub.amqp import ( + AmqpAnnotatedMessage, +) +from azure.eventhub.exceptions import EventDataSendError, OperationTimeoutError, EventHubError + + +def random_pkey_generation(partitions): + pr = PartitionResolver(partitions) + total = len(partitions) + dic = {} + + while total: + key = str(uuid4()) + pid = pr.get_partition_id_by_partition_key(key) + if pid in dic: + continue + else: + dic[pid] = key + total -= 1 + + return dic + + +@pytest.mark.liveTest() +def test_producer_client_constructor(connection_str, uamqp_transport): + def on_success(events, pid): + pass + + def on_error(events, error, pid): + pass + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, uamqp_transport=uamqp_transport) + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_success=on_success, uamqp_transport=uamqp_transport) + with pytest.raises(TypeError): + EventHubProducerClient.from_connection_string(connection_str, buffered_mode=True, on_error=on_error, uamqp_transport=uamqp_transport) + with pytest.raises(ValueError): + EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + max_wait_time=0, + uamqp_transport=uamqp_transport + ) + with pytest.raises(ValueError): + EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + max_buffer_length=0, + uamqp_transport=uamqp_transport + ) + + +@pytest.mark.liveTest +@pytest.mark.parametrize( + "flush_after_sending, close_after_sending", + [ + (False, False), + (True, False), + (False, True) + ] +) +@pytest.mark.liveTest +def test_basic_send_single_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + + time.sleep(10) + sent_events = defaultdict(list) + + def on_success(events, pid): + if len(events) > 1: + on_success.batching = True + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + on_success.batching = False # ensure batching happened + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + partitions = producer.get_partition_ids() + partitions_cnt = len(partitions) + # perform single sending round-robin + total_single_event_cnt = 100 + eventdata_set, amqpannoated_set = set(), set() + for i in range(total_single_event_cnt // 2): + event = EventData("test:{}".format(i)) + event.properties = {"event_idx": i} + producer.send_event(event) + eventdata_set.add(i) + for i in range(total_single_event_cnt // 2, total_single_event_cnt): + event = AmqpAnnotatedMessage(data_body="test:{}".format(i)) + event.application_properties = {"event_idx": i} + amqpannoated_set.add(i) + producer.send_event(event) + + for pid in partitions: + assert producer.get_buffered_event_count(pid) > 0 + assert producer.total_buffered_event_count > 0 + + if not flush_after_sending and not close_after_sending: + # ensure it's buffered sending + for pid in partitions: + assert len(sent_events[pid]) < total_single_event_cnt // partitions_cnt + assert sum([len(sent_events[pid]) for pid in partitions]) < total_single_event_cnt + else: + if flush_after_sending: + producer.flush() + if close_after_sending: + producer.close() + # ensure all events are sent after calling flush + assert sum([len(sent_events[pid]) for pid in partitions]) == total_single_event_cnt + + # give some time for producer to complete sending and consumer to complete receiving + time.sleep(10) + assert len(sent_events) == len(received_events) == partitions_cnt + + for pid in partitions: + assert producer.get_buffered_event_count(pid) == 0 + assert producer.total_buffered_event_count == 0 + assert not on_error.err + + # ensure all events are received in the correct partition + for pid in partitions: + assert len(sent_events[pid]) >= total_single_event_cnt // partitions_cnt + assert len(sent_events[pid]) == len(received_events[pid]) + for i in range(len(sent_events[pid])): + event = sent_events[pid][i] + try: # amqp annotated message + event_idx = event.application_properties["event_idx"] + amqpannoated_set.remove(event_idx) + except AttributeError: # event data + event_idx = event.properties["event_idx"] + eventdata_set.remove(event_idx) + assert received_events[pid][i].properties[b"event_idx"] == event_idx + assert partitions[event_idx % partitions_cnt] == pid + + assert on_success.batching + assert not eventdata_set + assert not amqpannoated_set + + consumer.close() + receive_thread.join() + + +@pytest.mark.liveTest +@pytest.mark.parametrize( + "flush_after_sending, close_after_sending", + [ + (False, False), + (True, False), + (False, True) + ] +) +def test_basic_send_batch_events_round_robin(connection_str, flush_after_sending, close_after_sending, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + + time.sleep(10) + sent_events = defaultdict(list) + + def on_success(events, pid): + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + partitions = producer.get_partition_ids() + partitions_cnt = len(partitions) + # perform batch sending round-robin + total_events_cnt = 100 + batch_cnt = partitions_cnt * 2 - 1 + each_partition_cnt = total_events_cnt // batch_cnt + remain_events = total_events_cnt % batch_cnt + batches = [] + event_idx = 0 + eventdata_set, amqpannoated_set = set(), set() + for i in range(batch_cnt): + batch = producer.create_batch() + for j in range(each_partition_cnt // 2): + event = EventData("test{}:{}".format(i, event_idx)) + event.properties = {'batch_idx': i, 'event_idx': event_idx} + batch.add(event) + eventdata_set.add(event_idx) + event_idx += 1 + for j in range(each_partition_cnt // 2, each_partition_cnt): + event = AmqpAnnotatedMessage(data_body="test{}:{}".format(i, event_idx)) + event.application_properties = {'batch_idx': i, 'event_idx': event_idx} + batch.add(event) + amqpannoated_set.add(event_idx) + event_idx += 1 + batches.append(batch) + + # put remain_events in the last batch + last_batch = producer.create_batch() + for i in range(remain_events): + event = EventData("test:{}:{}".format(len(batches), event_idx)) + event.properties = {'batch_idx': len(batches), 'event_idx': event_idx} + last_batch.add(event) + eventdata_set.add(event_idx) + event_idx += 1 + batches.append(last_batch) + + for batch in batches: + producer.send_batch(batch) + + if not flush_after_sending and not close_after_sending: + # ensure it's buffered sending + for pid in partitions: + assert len(sent_events[pid]) < each_partition_cnt + assert sum([len(sent_events[pid]) for pid in partitions]) < total_events_cnt + # give some time for producer to complete sending and consumer to complete receiving + else: + if flush_after_sending: + producer.flush() + if close_after_sending: + producer.close() + # ensure all events are sent + assert sum([len(sent_events[pid]) for pid in partitions]) == total_events_cnt + + time.sleep(20) + assert len(sent_events) == len(received_events) == partitions_cnt + + # ensure all events are received in the correct partition + for pid in partitions: + assert len(sent_events[pid]) > 0 + assert len(sent_events[pid]) == len(received_events[pid]) + for i in range(len(sent_events[pid])): + event = sent_events[pid][i] + try: # amqp annotated message + event_idx = event.application_properties["event_idx"] + amqpannoated_set.remove(event_idx) + except AttributeError: # event data + event_idx = event.properties["event_idx"] + eventdata_set.remove(event_idx) + assert received_events[pid][i].properties[b"event_idx"] == event_idx + + assert not amqpannoated_set + assert not eventdata_set + assert not on_error.err + + consumer.close() + receive_thread.join() + + +@pytest.mark.liveTest +def test_send_with_hybrid_partition_assignment(connection_str, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + + time.sleep(5) + sent_events = defaultdict(list) + + def on_success(events, pid): + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + partitions = producer.get_partition_ids() + partitions_cnt = len(partitions) + pid_to_pkey = random_pkey_generation(partitions) + expected_event_idx_to_partition = {} + event_idx = 0 + # 1. send by partition_key, each partition 2 events, two single + one batch containing two + for pid in partitions: + pkey = pid_to_pkey[pid] + producer.send_event(EventData('{}'.format(event_idx)), partition_key=pkey) + batch = producer.create_batch(partition_key=pkey) + batch.add(EventData('{}'.format(event_idx + 1))) + producer.send_batch(batch) + for i in range(2): + expected_event_idx_to_partition[event_idx + i] = pid + event_idx += 2 + + # 2. send by partition_id, each partition 2 events, two single + one batch containing two + for pid in partitions: + producer.send_event(EventData('{}'.format(event_idx)), partition_id=pid) + batch = producer.create_batch(partition_id=pid) + batch.add(EventData('{}'.format(event_idx + 1))) + producer.send_batch(batch) + for i in range(2): + expected_event_idx_to_partition[event_idx + i] = pid + event_idx += 2 + + # 3. send without partition, each partition 2 events, two single + one batch containing two + for _ in partitions: + producer.send_event(EventData('{}'.format(event_idx))) + batch = producer.create_batch() + batch.add(EventData('{}'.format(event_idx + 1))) + producer.send_batch(batch) + event_idx += 2 + + producer.flush() + assert len(sent_events) == partitions_cnt + + time.sleep(10) + + visited = set() + for pid in partitions: + assert len(sent_events[pid]) == 2 * 3 + + for sent_event in sent_events[pid]: + if int(sent_event.body_as_str()) in expected_event_idx_to_partition: + assert expected_event_idx_to_partition[int(sent_event.body_as_str())] == pid + + for recv_event in received_events[pid]: + if int(sent_event.body_as_str()) in expected_event_idx_to_partition: + assert expected_event_idx_to_partition[int(sent_event.body_as_str())] == pid + + assert recv_event.body_as_str() not in visited + visited.add(recv_event.body_as_str()) + + assert len(visited) == 2 * 3 * len(partitions) + + assert not on_error.err + consumer.close() + receive_thread.join() + + +def test_send_with_timing_configuration(connection_str, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + + time.sleep(5) + sent_events = defaultdict(list) + + def on_success(events, pid): + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None + + # test max_wait_time + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + max_wait_time=10, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + partitions = producer.get_partition_ids() + producer.send_event(EventData('data')) + time.sleep(5) + assert not sent_events + time.sleep(10) + assert sum([len(sent_events[pid]) for pid in partitions]) == 1 + + assert not on_error.err + + # test max_buffer_length per partition + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + max_wait_time=1000, + max_buffer_length=10, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + sent_events.clear() + received_events.clear() + with producer: + partitions = producer.get_partition_ids() + for i in range(7): + producer.send_event(EventData('data'), partition_id="0") + assert not sent_events + batch = producer.create_batch(partition_id="0") + for i in range(9): + batch.add(EventData('9')) + producer.send_batch(batch) # will flush 7 events and put the batch in buffer + assert sum([len(sent_events[pid]) for pid in partitions]) == 7 + for i in range(5): + producer.send_event(EventData('data'), partition_id="0") # will flush batch (9 events) + 1 event, leaving 4 in buffer + assert sum([len(sent_events[pid]) for pid in partitions]) == 17 + producer.flush() + assert sum([len(sent_events[pid]) for pid in partitions]) == 21 + + time.sleep(5) + assert sum([len(received_events[pid]) for pid in partitions]) == 21 + assert not on_error.err + consumer.close() + receive_thread.join() + + +@pytest.mark.liveTest +def test_long_sleep(connection_str, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + + time.sleep(5) + sent_events = defaultdict(list) + + def on_success(events, pid): + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + uamqp_transport=uamqp_transport + ) + + with producer: + producer.send_event(EventData("test"), partition_id="0") + time.sleep(220) + producer.send_event(EventData("test"), partition_id="0") + time.sleep(5) + + assert not on_error.err + assert len(sent_events["0"]) == 2 + assert len(received_events["0"]) == 2 + + consumer.close() + receive_thread.join() + +@pytest.mark.liveTest +def test_long_wait_small_buffer(connection_str, uamqp_transport): + received_events = defaultdict(list) + + def on_event(partition_context, event): + received_events[partition_context.partition_id].append(event) + + consumer = EventHubConsumerClient.from_connection_string(connection_str, consumer_group="$default", uamqp_transport=uamqp_transport) + receive_thread = Thread(target=consumer.receive, args=(on_event,)) + receive_thread.daemon = True + receive_thread.start() + time.sleep(10) + + sent_events = defaultdict(list) + + def on_success(events, pid): + sent_events[pid].extend(events) + + def on_error(events, pid, err): + on_error.err = err + + on_error.err = None # ensure no error + producer = EventHubProducerClient.from_connection_string( + connection_str, + buffered_mode=True, + on_success=on_success, + on_error=on_error, + auth_timeout=3, + retry_total=3, + retry_mode='fixed', + retry_backoff_factor=0.01, + max_wait_time=10, + max_buffer_length=100, + uamqp_transport=uamqp_transport + ) + + with producer: + for i in range(100): + producer.send_event(EventData("test")) + + time.sleep(60) + + assert not on_error.err + assert sum([len(sent_events[key]) for key in sent_events]) == 100 + assert sum([len(received_events[key]) for key in received_events]) == 100 + + consumer.close() + receive_thread.join() diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_consumer_client.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_consumer_client.py index 9e09cd156ef8..8b7420075ef4 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_consumer_client.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_consumer_client.py @@ -9,11 +9,16 @@ @pytest.mark.liveTest -def test_receive_no_partition(connstr_senders): +def test_receive_no_partition(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', receive_timeout=1) + client = EventHubConsumerClient.from_connection_string( + connection_str, + consumer_group='$default', + receive_timeout=1, + uamqp_transport=uamqp_transport + ) def on_event(partition_context, event): on_event.received += 1 @@ -46,10 +51,12 @@ def on_event(partition_context, event): @pytest.mark.liveTest -def test_receive_partition(connstr_senders): +def test_receive_partition(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) def on_event(partition_context, event): on_event.received += 1 @@ -74,16 +81,18 @@ def on_event(partition_context, event): @pytest.mark.liveTest -def test_receive_load_balancing(connstr_senders): +def test_receive_load_balancing(connstr_senders, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - test code using multiple threads. Sometimes OSX aborts python process") connection_str, senders = connstr_senders cs = InMemoryCheckpointStore() client1 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport + ) client2 = EventHubConsumerClient.from_connection_string( - connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1) + connection_str, consumer_group='$default', checkpoint_store=cs, load_balancing_interval=1, uamqp_transport=uamqp_transport + ) def on_event(partition_context, event): pass @@ -105,13 +114,15 @@ def on_event(partition_context, event): assert len(client2._event_processors[("$default", ALL_PARTITIONS)]._consumers) == 1 -def test_receive_batch_no_max_wait_time(connstr_senders): +def test_receive_batch_no_max_wait_time(connstr_senders, uamqp_transport): '''Test whether callback is called when max_wait_time is None and max_batch_size has reached ''' connection_str, senders = connstr_senders senders[0].send(EventData("Test EventData")) senders[1].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) @@ -146,14 +157,14 @@ def on_event_batch(partition_context, event_batch): worker.join() + @pytest.mark.parametrize("max_wait_time, sleep_time, expected_result", [(3, 10, []), - (3, 2, None), - ]) -def test_receive_batch_empty_with_max_wait_time(connection_str, max_wait_time, sleep_time, expected_result): + (3, 2, None)]) +def test_receive_batch_empty_with_max_wait_time(uamqp_transport, connection_str, max_wait_time, sleep_time, expected_result): '''Test whether event handler is called when max_wait_time > 0 and no event is received ''' - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) def on_event_batch(partition_context, event_batch): on_event_batch.event_batch = event_batch @@ -168,13 +179,15 @@ def on_event_batch(partition_context, event_batch): worker.join() -def test_receive_batch_early_callback(connstr_senders): +def test_receive_batch_early_callback(connstr_senders, uamqp_transport): ''' Test whether the callback is called once max_batch_size reaches and before max_wait_time reaches. ''' connection_str, senders = connstr_senders for _ in range(10): senders[0].send(EventData("Test EventData")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) def on_event_batch(partition_context, event_batch): on_event_batch.received += len(event_batch) diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py index 35722dfbf635..cd9d60ae7372 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_negative.py @@ -23,24 +23,46 @@ @pytest.mark.liveTest -def test_send_batch_with_invalid_hostname(invalid_hostname): +def test_send_batch_with_invalid_hostname(invalid_hostname, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) with client: with pytest.raises(ConnectError): batch = EventDataBatch() batch.add(EventData("test data")) client.send_batch(batch) + # test setting callback + def on_error(events, pid, err): + assert len(events) == 1 + assert not pid + on_error.err = err + + on_error.err = None + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) + with client: + batch = EventDataBatch() + batch.add(EventData("test data")) + client.send_batch(batch) + assert isinstance(on_error.err, ConnectError) + + on_error.err = None + client = EventHubProducerClient.from_connection_string(invalid_hostname, on_error=on_error, uamqp_transport=uamqp_transport) + with client: + client.send_event(EventData("test data")) + assert isinstance(on_error.err, ConnectError) + @pytest.mark.liveTest -def test_receive_with_invalid_hostname_sync(invalid_hostname): +def test_receive_with_invalid_hostname_sync(invalid_hostname, uamqp_transport): def on_event(partition_context, event): pass - client = EventHubConsumerClient.from_connection_string(invalid_hostname, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + invalid_hostname, consumer_group='$default', uamqp_transport=uamqp_transport + ) with client: thread = threading.Thread(target=client.receive, args=(on_event, )) @@ -51,8 +73,8 @@ def on_event(partition_context, event): @pytest.mark.liveTest -def test_send_batch_with_invalid_key(invalid_key): - client = EventHubProducerClient.from_connection_string(invalid_key) +def test_send_batch_with_invalid_key(invalid_key, uamqp_transport): + client = EventHubProducerClient.from_connection_string(invalid_key, uamqp_transport=uamqp_transport) try: with pytest.raises(ConnectError): batch = EventDataBatch() @@ -63,10 +85,10 @@ def test_send_batch_with_invalid_key(invalid_key): @pytest.mark.liveTest -def test_send_batch_to_invalid_partitions(connection_str): +def test_send_batch_to_invalid_partitions(connection_str, uamqp_transport): partitions = ["XYZ", "-1", "1000", "-"] for p in partitions: - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ConnectError): batch = client.create_batch(partition_id=p) @@ -77,10 +99,10 @@ def test_send_batch_to_invalid_partitions(connection_str): @pytest.mark.liveTest -def test_send_batch_too_large_message(connection_str): +def test_send_batch_too_large_message(connection_str, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - open issue regarding message size") - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: data = EventData(b"A" * 1100000) batch = client.create_batch() @@ -91,8 +113,8 @@ def test_send_batch_too_large_message(connection_str): @pytest.mark.liveTest -def test_send_batch_null_body(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +def test_send_batch_null_body(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) try: with pytest.raises(ValueError): data = EventData(None) @@ -104,32 +126,33 @@ def test_send_batch_null_body(connection_str): @pytest.mark.liveTest -def test_create_batch_with_invalid_hostname_sync(invalid_hostname): +def test_create_batch_with_invalid_hostname_sync(invalid_hostname, uamqp_transport): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - it keeps reporting 'Unable to set external certificates' " "and blocking other tests") - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) with client: with pytest.raises(ConnectError): client.create_batch(max_size_in_bytes=300) @pytest.mark.liveTest -def test_create_batch_with_too_large_size_sync(connection_str): - client = EventHubProducerClient.from_connection_string(connection_str) +def test_create_batch_with_too_large_size_sync(connection_str, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: with pytest.raises(ValueError): client.create_batch(max_size_in_bytes=5 * 1024 * 1024) @pytest.mark.liveTest -def test_invalid_proxy_server(connection_str): +def test_invalid_proxy_server(connection_str, uamqp_transport): + if uamqp_transport: + pytest.skip() HTTP_PROXY = { 'proxy_hostname': 'fakeproxy', # proxy hostname. 'proxy_port': 3128, # proxy port. } - client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY) + client = EventHubProducerClient.from_connection_string(connection_str, http_proxy=HTTP_PROXY, uamqp_transport=uamqp_transport) with client: with pytest.raises(EventHubError): batch = client.create_batch() - diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py index eb197eec44b0..6a4cb8b6eccf 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_properties.py @@ -12,57 +12,67 @@ @pytest.mark.liveTest -def test_get_properties(live_eventhub): +def test_get_properties(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport + ) with client: properties = client.get_eventhub_properties() assert properties['eventhub_name'] == live_eventhub['event_hub'] and properties['partition_ids'] == ['0', '1'] @pytest.mark.liveTest -def test_get_properties_with_auth_error_sync(live_eventhub): +def test_get_properties_with_auth_error_sync(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf=")) + EventHubSharedKeyCredential(live_eventhub['key_name'], "AaBbCcDdEeFf="), + uamqp_transport=uamqp_transport + ) with client: with pytest.raises(AuthenticationError) as e: client.get_eventhub_properties() client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential("invalid", live_eventhub['access_key']) + EventHubSharedKeyCredential("invalid", live_eventhub['access_key']), uamqp_transport=uamqp_transport ) with client: with pytest.raises(AuthenticationError) as e: client.get_eventhub_properties() @pytest.mark.liveTest -def test_get_properties_with_connect_error(live_eventhub): +def test_get_properties_with_connect_error(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], "invalid", '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) with client: with pytest.raises(ConnectError) as e: client.get_eventhub_properties() client = EventHubConsumerClient("invalid.servicebus.windows.net", live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport ) with client: with pytest.raises(EventHubError) as e: # This can be either ConnectError or ConnectionLostError client.get_eventhub_properties() @pytest.mark.liveTest -def test_get_partition_ids(live_eventhub): +def test_get_partition_ids(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport + ) with client: partition_ids = client.get_partition_ids() assert partition_ids == ['0', '1'] @pytest.mark.liveTest -def test_get_partition_properties(live_eventhub): +def test_get_partition_properties(live_eventhub, uamqp_transport): client = EventHubConsumerClient(live_eventhub['hostname'], live_eventhub['event_hub'], '$default', - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key']), + uamqp_transport=uamqp_transport + ) with client: properties = client.get_partition_properties('0') assert properties['eventhub_name'] == live_eventhub['event_hub'] \ diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py index 21d6e249581e..7597531904af 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py @@ -4,18 +4,16 @@ # license information. #-------------------------------------------------------------------------- -import os import threading import pytest import time -import datetime from azure.eventhub import EventData, TransportType, EventHubConsumerClient from azure.eventhub.exceptions import EventHubError @pytest.mark.liveTest -def test_receive_end_of_stream(connstr_senders): +def test_receive_end_of_stream(connstr_senders, uamqp_transport): def on_event(partition_context, event): if partition_context.partition_id == "0": assert event.body_as_str() == "Receiving only a single event" @@ -29,7 +27,9 @@ def on_event(partition_context, event): assert ", partition_key: 0" in event_str on_event.called = False connection_str, senders = connstr_senders - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) with client: thread = threading.Thread(target=client.receive, args=(on_event,), kwargs={"partition_id": "0", "starting_position": "@latest"}) @@ -50,7 +50,7 @@ def on_event(partition_context, event): ("sequence", True, "Inclusive"), ("enqueued_time", False, "Exclusive")]) @pytest.mark.liveTest -def test_receive_with_event_position_sync(connstr_senders, position, inclusive, expected_result): +def test_receive_with_event_position_sync(uamqp_transport, connstr_senders, position, inclusive, expected_result): def on_event(partition_context, event): assert partition_context.last_enqueued_event_properties.get('sequence_number') == event.sequence_number assert partition_context.last_enqueued_event_properties.get('offset') == event.offset @@ -69,7 +69,9 @@ def on_event(partition_context, event): connection_str, senders = connstr_senders senders[0].send(EventData(b"Inclusive")) senders[1].send(EventData(b"Inclusive")) - client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) with client: thread = threading.Thread(target=client.receive, args=(on_event,), kwargs={"starting_position": "-1", @@ -82,7 +84,9 @@ def on_event(partition_context, event): thread.join() senders[0].send(EventData(expected_result)) senders[1].send(EventData(expected_result)) - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client2 = EventHubConsumerClient.from_connection_string( + connection_str, consumer_group='$default', uamqp_transport=uamqp_transport + ) with client2: thread = threading.Thread(target=client2.receive, args=(on_event,), kwargs={"starting_position": on_event.event_position, @@ -90,14 +94,14 @@ def on_event(partition_context, event): "track_last_enqueued_event_properties": True}) thread.daemon = True thread.start() - time.sleep(10) + time.sleep(15) assert on_event.event.body_as_str() == expected_result thread.join() @pytest.mark.liveTest -def test_receive_owner_level(connstr_senders): +def test_receive_owner_level(connstr_senders, uamqp_transport): def on_event(partition_context, event): pass def on_error(partition_context, error): @@ -105,8 +109,8 @@ def on_error(partition_context, error): on_error.error = None connection_str, senders = connstr_senders - client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') - client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default') + client1 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) + client2 = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', uamqp_transport=uamqp_transport) with client1, client2: thread1 = threading.Thread(target=client1.receive, args=(on_event,), kwargs={"partition_id": "0", "starting_position": "-1", @@ -129,7 +133,7 @@ def on_error(partition_context, error): @pytest.mark.liveTest -def test_receive_over_websocket_sync(connstr_senders): +def test_receive_over_websocket_sync(connstr_senders, uamqp_transport): app_prop = {"raw_prop": "raw_value"} content_type = "text/plain" message_id_base = "mess_id_sample_" @@ -143,7 +147,8 @@ def on_event(partition_context, event): connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', - transport_type=TransportType.AmqpOverWebsocket) + transport_type=TransportType.AmqpOverWebsocket, + uamqp_transport=uamqp_transport) event_list = [] for i in range(5): diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py index 003247f7f235..2488abad148e 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_reconnect.py @@ -7,25 +7,32 @@ import time import pytest -from azure.eventhub._pyamqp.authentication import SASTokenAuth -from azure.eventhub._pyamqp.client import ReceiveClient -from azure.eventhub._pyamqp import error, constants - - from azure.eventhub import ( EventData, EventHubSharedKeyCredential, EventHubProducerClient, - EventHubConsumerClient + EventHubConsumerClient, ) from azure.eventhub.exceptions import OperationTimeoutError +from azure.eventhub._utils import transform_outbound_single_message +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp import ReceiveClient, error, constants +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport +try: + import uamqp + from uamqp import compat + from azure.eventhub._transport._uamqp_transport import UamqpTransport +except ImportError: + UamqpTransport = None @pytest.mark.liveTest -def test_send_with_long_interval_sync(live_eventhub, sleep): +def test_send_with_long_interval_sync(live_eventhub, sleep, uamqp_transport, timeout_factor): test_partition = "0" sender = EventHubProducerClient(live_eventhub['hostname'], live_eventhub['event_hub'], - EventHubSharedKeyCredential(live_eventhub['key_name'], live_eventhub['access_key'])) + EventHubSharedKeyCredential(live_eventhub['key_name'], + live_eventhub['access_key']), uamqp_transport=uamqp_transport + ) with sender: batch = sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) @@ -33,7 +40,10 @@ def test_send_with_long_interval_sync(live_eventhub, sleep): if sleep: time.sleep(250) else: - sender._producers[test_partition]._handler._connection.close() + if uamqp_transport: + sender._producers[test_partition]._handler._connection._conn.destroy() + else: + sender._producers[test_partition]._handler._connection.close() batch = sender.create_batch(partition_id=test_partition) batch.add(EventData(b"A single event")) sender.send_batch(batch) @@ -41,22 +51,28 @@ def test_send_with_long_interval_sync(live_eventhub, sleep): received = [] uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) - sas_auth = SASTokenAuth( - uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] - ) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub['hostname'], live_eventhub['event_hub'], live_eventhub['consumer_group'], test_partition) - receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + receiver = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=5000, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key'] + ) + receiver = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, debug=False, link_credit=500) + try: receiver.open() # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5)]) - received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5)]) + received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5 * timeout_factor)]) + received.extend([EventData._from_message(x) for x in receiver.receive_message_batch(max_batch_size=1, timeout=5 * timeout_factor)]) finally: receiver.close() assert len(received) == 2 @@ -64,50 +80,72 @@ def test_send_with_long_interval_sync(live_eventhub, sleep): @pytest.mark.liveTest -def test_send_connection_idle_timeout_and_reconnect_sync(connstr_receivers): +def test_send_connection_idle_timeout_and_reconnect_sync(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - # no retry, should just raise error - client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10, retry_total=0) - with client: - ed = EventData('data') - sender = client._create_producer(partition_id='0') - with sender: - sender._open_with_retry() - time.sleep(11) - sender._unsent_events = [ed.message] - with pytest.raises(error.AMQPConnectionError): - sender._send_event_data() - - # with retry, should work - client = EventHubProducerClient.from_connection_string(conn_str=connection_str, idle_timeout=10) + if uamqp_transport: + amqp_transport = UamqpTransport + retry_total = 3 + timeout_exc = compat.TimeoutException + else: + amqp_transport = PyamqpTransport + retry_total = 0 + timeout_exc = TimeoutError + client = EventHubProducerClient.from_connection_string( + conn_str=connection_str, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport + ) with client: ed = EventData('data') sender = client._create_producer(partition_id='0') with sender: sender._open_with_retry() time.sleep(11) - sender._unsent_events = [ed.message] - sender._send_event_data() + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + if uamqp_transport: + sender._unsent_events[0].on_send_complete = sender._on_outcome + with pytest.raises((uamqp.errors.ConnectionClose, + uamqp.errors.MessageHandlerError, OperationTimeoutError)): + sender._send_event_data() + else: + with pytest.raises(error.AMQPConnectionError): + sender._send_event_data() + if uamqp_transport: + sender._send_event_data_with_retry() + + if not uamqp_transport: + client = EventHubProducerClient.from_connection_string( + conn_str=connection_str, idle_timeout=10, uamqp_transport=uamqp_transport + ) + with client: + ed = EventData('data') + sender = client._create_producer(partition_id='0') + with sender: + sender._open_with_retry() + time.sleep(11) + ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message) + sender._unsent_events = [ed._message] + sender._send_event_data() retry = 0 while retry < 3: try: - messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10) + messages = receivers[0].receive_message_batch(max_batch_size=10, timeout=10 * timeout_factor) if messages: received_ed1 = EventData._from_message(messages[0]) assert received_ed1.body_as_str() == 'data' break - except TimeoutError: + except timeout_exc: retry += 1 @pytest.mark.liveTest -def test_receive_connection_idle_timeout_and_reconnect_sync(connstr_senders): +def test_receive_connection_idle_timeout_and_reconnect_sync(connstr_senders, uamqp_transport): connection_str, senders = connstr_senders client = EventHubConsumerClient.from_connection_string( conn_str=connection_str, consumer_group='$default', - idle_timeout=10 + idle_timeout=10, + uamqp_transport=uamqp_transport ) def on_event_received(event): @@ -121,9 +159,13 @@ def on_event_received(event): ed = EventData("Event") senders[0].send(ed) - with pytest.raises(error.AMQPConnectionError): + if uamqp_transport: consumer._handler.do_work() - assert consumer._handler._connection.state == constants.ConnectionState.END + assert consumer._handler._connection._state == uamqp.c_uamqp.ConnectionState.DISCARDING + else: + with pytest.raises(error.AMQPConnectionError): + consumer._handler.do_work() + assert consumer._handler._connection.state == constants.ConnectionState.END duration = 10 now_time = time.time() diff --git a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py index f04cb602e90b..9bbc6ae60697 100644 --- a/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py +++ b/sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py @@ -13,23 +13,35 @@ from azure.eventhub import EventData, TransportType, EventDataBatch from azure.eventhub import EventHubProducerClient, EventHubConsumerClient -from azure.eventhub.exceptions import EventDataSendError +from azure.eventhub.exceptions import EventDataSendError, OperationTimeoutError from azure.eventhub.amqp import ( AmqpMessageHeader, AmqpMessageBodyType, AmqpAnnotatedMessage, AmqpMessageProperties, ) +try: + import uamqp + from uamqp.constants import TransportType as uamqp_TransportType, MessageState + from uamqp.message import MessageProperties +except ImportError: + uamqp_TransportType = TransportType + MessageProperties = None +from azure.eventhub._pyamqp.message import Properties +from azure.eventhub._pyamqp.authentication import SASTokenAuth +from azure.eventhub._pyamqp import ReceiveClient +from azure.eventhub._pyamqp.error import AMQPConnectionError + @pytest.mark.liveTest -def test_send_with_partition_key(connstr_receivers): +def test_send_with_partition_key(connstr_receivers, live_eventhub, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: data_val = 0 for partition in [b"a", b"b", b"c", b"d", b"e", b"f"]: partition_key = b"test_partition_" + partition - for i in range(50): + for i in range(10): batch = client.create_batch(partition_key=partition_key) batch.add(EventData(str(data_val))) data_val += 1 @@ -37,42 +49,96 @@ def test_send_with_partition_key(connstr_receivers): client.send_batch(client.create_batch()) + for partition in [b"a", b"b", b"c", b"d", b"e", b"f"]: + partition_key = b"test_partition_" + partition + for i in range(10): + event_data = EventData(str(data_val)) + event_data.properties = {'is_single': True} + data_val += 1 + client.send_event(event_data, partition_key=partition_key) + + batch_cnt = 0 + single_cnt = 0 found_partition_keys = {} + reconnect_receivers = [] for index, partition in enumerate(receivers): - received = partition.receive_message_batch(timeout=5) - for message in received: + retry_total = 0 + while retry_total < 3: + timeout = (5 + retry_total) * timeout_factor try: - event_data = EventData._from_message(message) - existing = found_partition_keys[event_data.partition_key] - assert existing == index - except KeyError: - found_partition_keys[event_data.partition_key] = index + received = partition.receive_message_batch(timeout=timeout) + for message in received: + try: + event_data = EventData._from_message(message) + if event_data.properties and event_data.properties[b'is_single']: + single_cnt += 1 + else: + batch_cnt += 1 + existing = found_partition_keys[event_data.partition_key] + assert existing == index + except KeyError: + found_partition_keys[event_data.partition_key] = index + if received: + break + retry_total += 1 + except AMQPConnectionError: + for r in reconnect_receivers: + r.close() + uri = "sb://{}/{}".format(live_eventhub['hostname'], live_eventhub['event_hub']) + source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( + live_eventhub['hostname'], + live_eventhub['event_hub'], + live_eventhub['consumer_group'], + index) + if uamqp_transport: + sas_auth = uamqp.authentication.SASTokenAuth.from_shared_access_key( + uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = uamqp.ReceiveClient(source, auth=sas_auth, debug=False, timeout=0, prefetch=500) + else: + sas_auth = SASTokenAuth( + uri, uri, live_eventhub['key_name'], live_eventhub['access_key']) + partition = ReceiveClient(live_eventhub['hostname'], source, auth=sas_auth, network_trace=False, timeout=0, link_credit=500) + partition.open() + reconnect_receivers.append(partition) + retry_total += 1 + if retry_total == 3: + raise OperationTimeoutError(f"Exhausted retries for receiving from {live_eventhub['hostname']}.") + + for r in reconnect_receivers: + r.close() + + assert single_cnt == 60 + assert batch_cnt == 60 + assert len(found_partition_keys) == 6 @pytest.mark.liveTest -def test_send_and_receive_large_body_size(connstr_receivers): +def test_send_and_receive_large_body_size(connstr_receivers, uamqp_transport, timeout_factor): if sys.platform.startswith('darwin'): pytest.skip("Skipping on OSX - open issue regarding message size") connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: payload = 250 * 1024 batch = client.create_batch() batch.add(EventData("A" * payload)) client.send_batch(batch) + client.send_event(EventData("A" * payload)) received = [] + timeout = 10 * timeout_factor for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) - assert len(received) == 1 + assert len(received) == 2 assert len(list(received[0].body)[0]) == payload + assert len(list(received[1].body)[0]) == payload @pytest.mark.liveTest -def test_send_amqp_annotated_message(connstr_receivers): +def test_send_amqp_annotated_message(connstr_receivers, uamqp_transport): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: sequence_body = [b'message', 123.456, True] footer = {'footer_key': 'footer_value'} @@ -108,7 +174,7 @@ def test_send_amqp_annotated_message(connstr_receivers): ) body_ed = """{"json_key": "json_val"}""" - prop_ed = {"raw_prop": "raw_value"} + prop_ed = {b"raw_prop": b"raw_value"} cont_type_ed = "text/plain" corr_id_ed = "corr_id" mess_id_ed = "mess_id" @@ -116,6 +182,7 @@ def test_send_amqp_annotated_message(connstr_receivers): event_data.content_type = cont_type_ed event_data.correlation_id = corr_id_ed event_data.message_id = mess_id_ed + event_data.properties = prop_ed batch = client.create_batch() batch.add(data_message) @@ -124,6 +191,10 @@ def test_send_amqp_annotated_message(connstr_receivers): batch.add(event_data) client.send_batch(batch) client.send_batch([data_message, value_message, sequence_message, event_data]) + client.send_event(data_message) + client.send_event(value_message) + client.send_event(sequence_message) + client.send_event(event_data) received_count = {} received_count["data_msg"] = 0 @@ -146,6 +217,7 @@ def check_values(event): assert event.correlation_id == corr_id_ed assert event.message_id == mess_id_ed assert event.content_type == cont_type_ed + assert event.properties == prop_ed assert event.body_type == AmqpMessageBodyType.DATA received_count["normal_msg"] += 1 elif raw_amqp_message.body_type == AmqpMessageBodyType.SEQUENCE: @@ -168,7 +240,8 @@ def on_event(partition_context, event): on_event.received = [] client = EventHubConsumerClient.from_connection_string(connection_str, - consumer_group='$default') + consumer_group='$default', + uamqp_transport=uamqp_transport) with client: thread = threading.Thread(target=client.receive, args=(on_event,), kwargs={"starting_position": "-1"}) @@ -177,138 +250,169 @@ def on_event(partition_context, event): for event in on_event.received: check_values(event) - assert len(on_event.received) == 8 - assert received_count["data_msg"] == 2 - assert received_count["seq_msg"] == 2 - assert received_count["value_msg"] == 2 - assert received_count["normal_msg"] == 2 + assert len(on_event.received) == 12 + assert received_count["data_msg"] == 3 + assert received_count["seq_msg"] == 3 + assert received_count["value_msg"] == 3 + assert received_count["normal_msg"] == 3 @pytest.mark.parametrize("payload", [b"", b"A single event"]) @pytest.mark.liveTest -def test_send_and_receive_small_body(connstr_receivers, payload): +def test_send_and_receive_small_body(connstr_receivers, payload, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: batch = client.create_batch() batch.add(EventData(payload)) client.send_batch(batch) + client.send_event(EventData(payload)) received = [] + timeout = 5 * timeout_factor for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=5)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) - assert len(received) == 1 + assert len(received) == 2 assert list(received[0].body)[0] == payload + assert list(received[1].body)[0] == payload @pytest.mark.liveTest -def test_send_partition(connstr_receivers): +def test_send_partition(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + timeout = 10 * timeout_factor + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: batch = client.create_batch() batch.add(EventData(b"Data")) client.send_batch(batch) + client.send_event(EventData(b"Data")) with client: batch = client.create_batch(partition_id="1") batch.add(EventData(b"Data")) client.send_batch(batch) + client.send_event(EventData(b"Data"), partition_id="1") - partition_0 = receivers[0].receive_message_batch(timeout=10) - partition_1 = receivers[1].receive_message_batch(timeout=10) - assert len(partition_0) + len(partition_1) == 2 + partition_0 = receivers[0].receive_message_batch(timeout=timeout) + partition_1 = receivers[1].receive_message_batch(timeout=timeout) + assert len(partition_1) >= 2 + assert len(partition_0) + len(partition_1) == 4 with client: batch = client.create_batch() batch.add(EventData(b"Data")) client.send_batch(batch) + client.send_event(EventData(b"Data")) with client: - batch = client.create_batch(partition_id="1") + batch = client.create_batch(partition_id="0") + batch.add(EventData(b"Data")) + client.send_batch(batch) + client.send_event(EventData(b"Data"), partition_id="0") + + with client: + batch = EventDataBatch(partition_id="0") batch.add(EventData(b"Data")) client.send_batch(batch) time.sleep(5) - partition_0 = receivers[0].receive_message_batch(timeout=10) - partition_1 = receivers[1].receive_message_batch(timeout=10) - assert len(partition_0) + len(partition_1) == 2 + partition_0 = receivers[0].receive_message_batch(timeout=timeout) + partition_1 = receivers[1].receive_message_batch(timeout=timeout) + assert len(partition_0) >= 3 + assert len(partition_0) + len(partition_1) == 5 @pytest.mark.liveTest -def test_send_non_ascii(connstr_receivers): +def test_send_non_ascii(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + timeout = 5 * timeout_factor + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: batch = client.create_batch(partition_id="0") batch.add(EventData(u"é,è,à,ù,â,ê,î,ô,û")) batch.add(EventData(json.dumps({"foo": u"漢字"}))) client.send_batch(batch) + client.send_event(EventData(u"é,è,à,ù,â,ê,î,ô,û"), partition_id="0") + client.send_event(EventData(json.dumps({"foo": u"漢字"})), partition_id="0") time.sleep(1) # receive_message_batch() returns immediately once it receives any messages before the max_batch_size # and timeout reach. Could be 1, 2, or any number between 1 and max_batch_size. # So call it twice to ensure the two events are received. - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] + \ - [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] - assert len(partition_0) == 2 + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=timeout)] + \ + [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=timeout)] + assert len(partition_0) == 4 assert partition_0[0].body_as_str() == u"é,è,à,ù,â,ê,î,ô,û" assert partition_0[1].body_as_json() == {"foo": u"漢字"} + assert partition_0[2].body_as_str() == u"é,è,à,ù,â,ê,î,ô,û" + assert partition_0[3].body_as_json() == {"foo": u"漢字"} @pytest.mark.liveTest -def test_send_multiple_partitions_with_app_prop(connstr_receivers): +def test_send_multiple_partitions_with_app_prop(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers + timeout = 5 * timeout_factor app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} - client = EventHubProducerClient.from_connection_string(connection_str) + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: ed0 = EventData(b"Message 0") ed0.properties = app_prop batch = client.create_batch(partition_id="0") batch.add(ed0) client.send_batch(batch) + client.send_event(ed0, partition_id="0") ed1 = EventData(b"Message 1") ed1.properties = app_prop batch = client.create_batch(partition_id="1") batch.add(ed1) client.send_batch(batch) + client.send_event(ed1, partition_id="1") - partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=5)] - assert len(partition_0) == 1 + partition_0 = [EventData._from_message(x) for x in receivers[0].receive_message_batch(timeout=timeout)] + assert len(partition_0) == 2 assert partition_0[0].properties[b"raw_prop"] == b"raw_value" - partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=5)] - assert len(partition_1) == 1 + assert partition_0[1].properties[b"raw_prop"] == b"raw_value" + partition_1 = [EventData._from_message(x) for x in receivers[1].receive_message_batch(timeout=timeout)] + assert len(partition_1) == 2 assert partition_1[0].properties[b"raw_prop"] == b"raw_value" + assert partition_1[1].properties[b"raw_prop"] == b"raw_value" @pytest.mark.liveTest -def test_send_over_websocket_sync(connstr_receivers): - pytest.skip("websocket not supported") +def test_send_over_websocket_sync(connstr_receivers, uamqp_transport, timeout_factor): + timeout = 10 * timeout_factor connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str, transport_type=TransportType.AmqpOverWebsocket) + client = EventHubProducerClient.from_connection_string( + connection_str, transport_type=uamqp_TransportType.AmqpOverWebsocket, uamqp_transport=uamqp_transport + ) with client: batch = client.create_batch(partition_id="0") batch.add(EventData("Event Data")) client.send_batch(batch) + client.send_event(EventData("Event Data"), partition_id="0") time.sleep(1) received = [] - received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=10)) - assert len(received) == 1 + received.extend(receivers[0].receive_message_batch(max_batch_size=5, timeout=timeout)) + assert len(received) == 2 @pytest.mark.liveTest -def test_send_with_create_event_batch_with_app_prop_sync(connstr_receivers): +def test_send_with_create_event_batch_with_app_prop_sync(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers + timeout = 5 * timeout_factor app_prop_key = "raw_prop" app_prop_value = "raw_value" app_prop = {app_prop_key: app_prop_value} - client = EventHubProducerClient.from_connection_string(connection_str, transport_type=TransportType.AmqpOverWebsocket) + client = EventHubProducerClient.from_connection_string( + connection_str, transport_type=TransportType.AmqpOverWebsocket, uamqp_transport=uamqp_transport + ) with client: event_data_batch = client.create_batch(max_size_in_bytes=100000) while True: @@ -321,55 +425,104 @@ def test_send_with_create_event_batch_with_app_prop_sync(connstr_receivers): client.send_batch(event_data_batch) received = [] for r in receivers: - received.extend(r.receive_message_batch(timeout=5)) + received.extend(r.receive_message_batch(timeout=timeout)) assert len(received) >= 1 assert EventData._from_message(received[0]).properties[b"raw_prop"] == b"raw_value" @pytest.mark.liveTest -def test_send_list(connstr_receivers): +def test_send_list(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + timeout = 10 * timeout_factor + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) payload = "A1" with client: client.send_batch([EventData(payload)]) received = [] for r in receivers: - received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=10)]) + received.extend([EventData._from_message(x) for x in r.receive_message_batch(timeout=timeout)]) assert len(received) == 1 assert received[0].body_as_str() == payload @pytest.mark.liveTest -def test_send_list_partition(connstr_receivers): +def test_send_list_partition(connstr_receivers, uamqp_transport, timeout_factor): connection_str, receivers = connstr_receivers - client = EventHubProducerClient.from_connection_string(connection_str) + timeout = 10 * timeout_factor + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) payload = "A1" with client: client.send_batch([EventData(payload)], partition_id="0") - message = receivers[0].receive_message_batch(timeout=10)[0] + message = receivers[0].receive_message_batch(timeout=timeout)[0] received = EventData._from_message(message) assert received.body_as_str() == payload @pytest.mark.parametrize("to_send, exception_type", [([EventData("A"*1024)]*1100, ValueError), - ("any str", AttributeError) - ]) + ("any str", AttributeError)]) @pytest.mark.liveTest -def test_send_list_wrong_data(connection_str, to_send, exception_type): - client = EventHubProducerClient.from_connection_string(connection_str) +def test_send_list_wrong_data(connection_str, to_send, exception_type, uamqp_transport): + client = EventHubProducerClient.from_connection_string(connection_str, uamqp_transport=uamqp_transport) with client: with pytest.raises(exception_type): client.send_batch(to_send) @pytest.mark.parametrize("partition_id, partition_key", [("0", None), (None, "pk")]) -def test_send_batch_pid_pk(invalid_hostname, partition_id, partition_key): +def test_send_batch_pid_pk(invalid_hostname, partition_id, partition_key, uamqp_transport): # Use invalid_hostname because this is not a live test. - client = EventHubProducerClient.from_connection_string(invalid_hostname) + client = EventHubProducerClient.from_connection_string(invalid_hostname, uamqp_transport=uamqp_transport) batch = EventDataBatch(partition_id=partition_id, partition_key=partition_key) with client: with pytest.raises(TypeError): client.send_batch(batch, partition_id=partition_id, partition_key=partition_key) + + +@pytest.mark.liveTest +def test_send_with_callback(connstr_receivers, uamqp_transport): + + def on_error(events, pid, err): + on_error.err = err + + def on_success(events, pid): + sent_events.append((events, pid)) + + sent_events = [] + on_error.err = None + connection_str, receivers = connstr_receivers + client = EventHubProducerClient.from_connection_string(connection_str, on_success=on_success, on_error=on_error, uamqp_transport=uamqp_transport) + + with client: + batch = client.create_batch() + batch.add(EventData(b"Data")) + batch.add(EventData(b"Data")) + client.send_batch(batch) + assert len(sent_events[-1][0]) == 2 + assert not sent_events[-1][1] + client.send_event(EventData(b"Data")) + assert len(sent_events[-1][0]) == 1 + assert not sent_events[-1][1] + + batch = client.create_batch(partition_key='key') + batch.add(EventData(b"Data")) + batch.add(EventData(b"Data")) + client.send_batch(batch) + assert len(sent_events[-1][0]) == 2 + assert not sent_events[-1][1] + client.send_event(EventData(b"Data"), partition_key='key') + assert len(sent_events[-1][0]) == 1 + assert not sent_events[-1][1] + + batch = client.create_batch(partition_id="0") + batch.add(EventData(b"Data")) + client.send_batch(batch) + batch.add(EventData(b"Data")) + assert len(sent_events[-1][0]) == 2 + assert sent_events[-1][1] == "0" + client.send_event(EventData(b"Data"), partition_id="0") + assert len(sent_events[-1][0]) == 1 + assert sent_events[-1][1] == "0" + + assert not on_error.err diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_partition_resolver_async.py b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_partition_resolver_async.py new file mode 100644 index 000000000000..195ee9e01315 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/unittest/asynctests/test_partition_resolver_async.py @@ -0,0 +1,44 @@ +# coding=utf-8 +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +import asyncio +from collections import defaultdict +import pytest +from azure.eventhub.aio._buffered_producer._partition_resolver_async import PartitionResolver + + +class TestPartitionResolver: + + @pytest.mark.asyncio + @pytest.mark.parametrize("partition_cnt", [1, 2, 16, 32, 256]) + async def test_basic_round_robin(self, partition_cnt): + partitions = [str(i) for i in range(partition_cnt)] + pr = PartitionResolver(partitions) + for i in range(2*partition_cnt): + expected = str(i % partition_cnt) + real = await pr.get_next_partition_id() + assert expected == real + + @pytest.mark.asyncio + @pytest.mark.parametrize("partition_cnt", [1, 2, 16, 32, 256]) + async def test_concurrent_round_robin_fairly(self, partition_cnt): + partitions = [str(i) for i in range(partition_cnt)] + pr = PartitionResolver(partitions) + dic = defaultdict(int) + lock = asyncio.Lock() + + async def gen_pid(): + pid = await pr.get_next_partition_id() + async with lock: + dic[pid] += 1 + + futures = [asyncio.ensure_future(gen_pid()) for _ in range(5*partition_cnt)] + + for future in futures: + await future + + assert len(dic) == partition_cnt + for i in range(partition_cnt): + assert dic[str(i)] == 5 diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py b/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py index 612f2e8b9605..8963d850de8c 100644 --- a/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py +++ b/sdk/eventhub/azure-eventhub/tests/unittest/test_event_data.py @@ -1,9 +1,24 @@ +# -- coding: utf-8 -- +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +#-------------------------------------------------------------------------- + import platform import pytest from packaging import version -from azure.eventhub.amqp import AmqpAnnotatedMessage +try: + import uamqp + from azure.eventhub._transport._uamqp_transport import UamqpTransport +except ImportError: + UamqpTransport = None + pass +from azure.eventhub._transport._pyamqp_transport import PyamqpTransport +from azure.eventhub._pyamqp.message import Message, Properties, Header +from azure.eventhub.amqp import AmqpAnnotatedMessage, AmqpMessageHeader, AmqpMessageProperties + from azure.eventhub import _common -from azure.eventhub._pyamqp.message import Message, Properties pytestmark = pytest.mark.skipif(platform.python_implementation() == "PyPy", reason="This is ignored for PyPy") @@ -55,24 +70,42 @@ def test_app_properties(): assert event_data.properties["a"] == "b" -def test_sys_properties(): - properties = Properties( - message_id="message_id", - user_id="user_id", - to="to", - subject="subject", - reply_to="reply_to", - correlation_id="correlation_id", - content_type="content_type", - content_encoding="content_encoding", - absolute_expiry_time=1, - creation_time=1, - group_id="group_id", - group_sequence=1, - reply_to_group_id="reply_to_group_id" - ) - message_annotations = {_common.PROP_OFFSET: "@latest"} - message = Message(properties=properties, message_annotations=message_annotations) +def test_sys_properties(uamqp_transport): + if uamqp_transport: + properties = uamqp.message.MessageProperties() + properties.message_id = "message_id" + properties.user_id = "user_id" + properties.to = "to" + properties.subject = "subject" + properties.reply_to = "reply_to" + properties.correlation_id = "correlation_id" + properties.content_type = "content_type" + properties.content_encoding = "content_encoding" + properties.absolute_expiry_time = 1 + properties.creation_time = 1 + properties.group_id = "group_id" + properties.group_sequence = 1 + properties.reply_to_group_id = "reply_to_group_id" + message = uamqp.message.Message(properties=properties) + message.annotations = {_common.PROP_OFFSET: "@latest"} + else: + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(properties=properties, message_annotations=message_annotations) ed = EventData._from_message(message) # type: EventData assert ed.system_properties[_common.PROP_OFFSET] == "@latest" @@ -91,22 +124,35 @@ def test_sys_properties(): assert ed.system_properties[_common.PROP_REPLY_TO_GROUP_ID] == properties.reply_to_group_id -def test_event_data_batch(): - batch = EventDataBatch(max_size_in_bytes=110, partition_key="par") +def test_event_data_batch(uamqp_transport): + if uamqp_transport: + if version.parse(uamqp.__version__) >= version.parse("1.2.8"): + expected_result = 97 + else: + expected_result = 93 + amqp_transport=UamqpTransport + else: + expected_result = 99 + amqp_transport=PyamqpTransport + + batch = EventDataBatch(max_size_in_bytes=110, partition_key="par", amqp_transport=amqp_transport) batch.add(EventData("A")) assert str(batch) == "EventDataBatch(max_size_in_bytes=110, partition_id=None, partition_key='par', event_count=1)" assert repr(batch) == "EventDataBatch(max_size_in_bytes=110, partition_id=None, partition_key='par', event_count=1)" - # TODO: uamqp uses 93 bytes for encode, while python amqp uses 99 bytes - # we should understand why extra bytes are needed to encode the content and how it could be improved - assert batch.size_in_bytes == 99 and len(batch) == 1 + assert batch.size_in_bytes == expected_result and len(batch) == 1 with pytest.raises(ValueError): batch.add(EventData("A")) -def test_event_data_from_message(): - message = Message(data=b'A') +def test_event_data_from_message(uamqp_transport): + if uamqp_transport: + amqp_transport = UamqpTransport + else: + amqp_transport = PyamqpTransport + annotated_message = AmqpAnnotatedMessage(data_body=b'A') + message = amqp_transport.to_outgoing_amqp_message(annotated_message) event = EventData._from_message(message) assert event.content_type is None assert event.correlation_id is None @@ -118,7 +164,7 @@ def test_event_data_from_message(): assert event.content_type == 'content_type' assert event.correlation_id == 'correlation_id' assert event.message_id == 'message_id' - assert event.body == b'A' + assert list(event.body) == [b'A'] def test_amqp_message_str_repr(): @@ -126,3 +172,163 @@ def test_amqp_message_str_repr(): message = AmqpAnnotatedMessage(data_body=data_body) assert str(message) == 'A' assert 'AmqpAnnotatedMessage(body=A, body_type=data' in repr(message) + + +def test_amqp_message_from_message(uamqp_transport): + if uamqp_transport: + header = uamqp.message.MessageHeader() + header.delivery_count = 1 + header.time_to_live = 10000 + header.first_acquirer = True + header.durable = True + header.priority = 1 + properties = uamqp.message.MessageProperties() + properties.message_id = "message_id" + properties.user_id = "user_id" + properties.to = "to" + properties.subject = "subject" + properties.reply_to = "reply_to" + properties.correlation_id = "correlation_id" + properties.content_type = "content_type" + properties.content_encoding = "content_encoding" + properties.absolute_expiry_time = 1 + properties.creation_time = 1 + properties.group_id = "group_id" + properties.group_sequence = 1 + properties.reply_to_group_id = "reply_to_group_id" + message = uamqp.message.Message(header=header, properties=properties) + message.annotations = {_common.PROP_OFFSET: "@latest"} + else: + header = Header( + delivery_count=1, + ttl=10000, + first_acquirer=True, + durable=True, + priority=1 + ) + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(properties=properties, header=header, message_annotations=message_annotations) + + amqp_message = AmqpAnnotatedMessage(message=message) + assert amqp_message.properties.message_id == message.properties.message_id + assert amqp_message.properties.user_id == message.properties.user_id + assert amqp_message.properties.to == message.properties.to + assert amqp_message.properties.subject == message.properties.subject + assert amqp_message.properties.reply_to == message.properties.reply_to + assert amqp_message.properties.correlation_id == message.properties.correlation_id + assert amqp_message.properties.content_type == message.properties.content_type + assert amqp_message.properties.absolute_expiry_time == message.properties.absolute_expiry_time + assert amqp_message.properties.creation_time == message.properties.creation_time + assert amqp_message.properties.group_id == message.properties.group_id + assert amqp_message.properties.group_sequence == message.properties.group_sequence + assert amqp_message.properties.reply_to_group_id == message.properties.reply_to_group_id + assert amqp_message.header.time_to_live == message.header.ttl + assert amqp_message.header.delivery_count == message.header.delivery_count + assert amqp_message.header.first_acquirer == message.header.first_acquirer + assert amqp_message.header.durable == message.header.durable + assert amqp_message.header.priority == message.header.priority + assert amqp_message.annotations == message.message_annotations + +def test_legacy_message(uamqp_transport): + if uamqp_transport: + header = uamqp.message.MessageHeader() + header.delivery_count = 1 + header.time_to_live = 10000 + header.first_acquirer = True + header.durable = True + header.priority = 1 + properties = uamqp.message.MessageProperties() + properties.message_id = "message_id" + properties.user_id = "user_id" + properties.to = "to" + properties.subject = "subject" + properties.reply_to = "reply_to" + properties.correlation_id = "correlation_id" + properties.content_type = "content_type" + properties.content_encoding = "content_encoding" + properties.absolute_expiry_time = 1 + properties.creation_time = 1 + properties.group_id = "group_id" + properties.group_sequence = 1 + properties.reply_to_group_id = "reply_to_group_id" + message = uamqp.message.Message(body=b'abc', header=header, properties=properties) + message.annotations = {_common.PROP_OFFSET: "@latest"} + amqp_transport = UamqpTransport + else: + header = Header( + delivery_count=1, + ttl=10000, + first_acquirer=True, + durable=True, + priority=1 + ) + properties = Properties( + message_id="message_id", + user_id="user_id", + to="to", + subject="subject", + reply_to="reply_to", + correlation_id="correlation_id", + content_type="content_type", + content_encoding="content_encoding", + absolute_expiry_time=1, + creation_time=1, + group_id="group_id", + group_sequence=1, + reply_to_group_id="reply_to_group_id" + ) + message_annotations = {_common.PROP_OFFSET: "@latest"} + message = Message(data=b'abc', properties=properties, header=header, message_annotations=message_annotations) + amqp_transport = PyamqpTransport + event_data = EventData._from_message(message=message) + assert event_data.message.properties.user_id == b'user_id' + assert event_data.message.properties.message_id == b'message_id' + assert event_data.message.properties.to == b'to' + assert event_data.message.properties.subject == b'subject' + assert event_data.message.properties.reply_to == b"reply_to" + assert event_data.message.properties.correlation_id == b"correlation_id" + assert event_data.message.properties.content_type == b"content_type" + assert event_data.message.properties.content_encoding == b"content_encoding" + assert event_data.message.properties.absolute_expiry_time == 1 + assert event_data.message.properties.creation_time == 1 + assert event_data.message.properties.group_id == b"group_id" + assert event_data.message.properties.group_sequence == 1 + assert event_data.message.properties.reply_to_group_id == b"reply_to_group_id" + assert event_data.message.state.value == 2 + assert event_data.message.delivery_annotations == {} + assert event_data.message.delivery_no is None + assert event_data.message.delivery_tag is None + assert event_data.message.on_send_complete is None + assert event_data.message.footer == {} + assert event_data.message.retries == 0 + assert event_data.message.idle_time == 0 + + event_data_batch = EventDataBatch(partition_key=b'par', partition_id='1', amqp_transport=amqp_transport) + event_data_batch.add(event_data) + assert event_data_batch.message.max_message_length == 1024 * 1024 + assert event_data_batch.message.size_offset == 0 + assert event_data_batch.message.batch_format == 0x80013700 + assert len(event_data_batch.message.annotations) == 1 + assert event_data_batch.message.application_properties is None + assert event_data_batch.message.header.delivery_count == 0 + assert event_data_batch.message.header.time_to_live is None + assert event_data_batch.message.header.first_acquirer is None + assert event_data_batch.message.header.durable is True + assert event_data_batch.message.header.priority is None + assert event_data_batch.message.on_send_complete is None + assert event_data_batch.message.properties is None diff --git a/sdk/eventhub/azure-eventhub/tests/unittest/test_partition_resolver.py b/sdk/eventhub/azure-eventhub/tests/unittest/test_partition_resolver.py new file mode 100644 index 000000000000..1ff803127823 --- /dev/null +++ b/sdk/eventhub/azure-eventhub/tests/unittest/test_partition_resolver.py @@ -0,0 +1,61 @@ +# coding=utf-8 +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +import pytest +from threading import Lock +from azure.eventhub._buffered_producer._partition_resolver import generate_hash_code, PartitionResolver + + +class TestPartitionResolver: + def test_partition_key(self): + input = { + "7": -15263, + "7149583486996073602": 12977, + "FWfAT": -22341, + "sOdeEAsyQoEuEFPGerWO": -6503, + "FAyAIctPeCgmiwLKbJcyswoHglHVjQdvtBowLACDNORsYvOcLddNJYDmhAVkbyLOrHTKLneMNcbgWVlasVywOByANjs": 5226, + "1XYM6!(7(lF5wq4k4m*e$Nc!1ezLJv*1YK1Y-C^*&B$O)lq^iUkG(TNzXG;Zi#z2Og*Qq0#^*k):vXh$3,C7We7%W0meJ;b3,rQCg^J;^twXgs5E$$hWKxqp": 23950, + "E(x;RRIaQcJs*P;D&jTPau-4K04oqr:lF6Z):ERpo&;9040qyV@G1_c9mgOs-8_8/10Fwa-7b7-yP!T-!IH&968)FWuI;(^g$2fN;)HJ^^yTn:": -29304, + "!c*_!I@1^c": 15372, + "p4*!jioeO/z-!-;w:dh": -3104, + "$0cb": 26269, + "-4189260826195535198": 453 + } + + for k, v in input.items(): + assert generate_hash_code(k) == v + + @pytest.mark.parametrize("partition_cnt", [1, 2, 16, 32, 256]) + def test_basic_round_robin(self, partition_cnt): + partitions = [str(i) for i in range(partition_cnt)] + pr = PartitionResolver(partitions) + for i in range(2*partition_cnt): + expected = str(i % partition_cnt) + real = pr.get_next_partition_id() + assert expected == real + + @pytest.mark.parametrize("partition_cnt", [1, 2, 16, 32, 256]) + def test_concurrent_round_robin_fairly(self, partition_cnt): + partitions = [str(i) for i in range(partition_cnt)] + pr = PartitionResolver(partitions) + exc = ThreadPoolExecutor() + + dic = defaultdict(int) + lock = Lock() + + def gen_pid(): + pid = pr.get_next_partition_id() + with lock: + dic[pid] += 1 + + for i in range(5*partition_cnt): + exc.submit(gen_pid) + + exc.shutdown() + assert len(dic) == partition_cnt + for i in range(partition_cnt): + assert dic[str(i)] == 5 diff --git a/sdk/eventhub/tests.yml b/sdk/eventhub/tests.yml index 7dc641b87cef..25d75cbf03a7 100644 --- a/sdk/eventhub/tests.yml +++ b/sdk/eventhub/tests.yml @@ -4,6 +4,7 @@ stages: - template: ../../eng/pipelines/templates/stages/archetype-sdk-tests.yml parameters: ServiceDirectory: eventhub + TestTimeoutInMinutes: 240 BuildTargetingString: azure-eventhub MatrixReplace: - TestSamples=.*/true diff --git a/shared_requirements.txt b/shared_requirements.txt index 4ce74d1923f7..c08aaac23e5b 100644 --- a/shared_requirements.txt +++ b/shared_requirements.txt @@ -176,7 +176,6 @@ opentelemetry-sdk<2.0.0,>=1.5.0,!=1.10a0 #override azure-eventhub-checkpointstoreblob-aio azure-core<2.0.0,>=1.20.1 #override azure-eventhub-checkpointstoreblob-aio aiohttp<4.0,>=3.0 #override azure-eventhub-checkpointstoretable azure-core<2.0.0,>=1.14.0 -#override azure-eventhub uamqp>=1.6.0,<2.0.0 #override azure-appconfiguration msrest>=0.6.10 #override azure-servicebus uamqp>=1.5.1,<2.0.0 #override azure-servicebus msrest>=0.6.17,<2.0.0