[EventHubs] add amqp switch support

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 }}
 

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,13 +23,66 @@
 
 - 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
 
 - 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`.

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
 

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",

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",
+]

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