[EventHubs] uamqp switch support

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Release History
 
+## 5.10.1 (Unreleased)
+
+### Bugs Fixed
+
+### Other Changes
+
+- Internal refactoring to support upcoming Pure Python AMQP-based release.
+
 ## 5.10.0 (2022-06-08)
 
 ### Features Added

sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py

@@ -2,12 +2,12 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
-from uamqp import constants
 from ._common import EventData, EventDataBatch
 from ._version import VERSION
 
 __version__ = VERSION
 
+from ._constants import TransportType
 from ._producer_client import EventHubProducerClient
 from ._consumer_client import EventHubConsumerClient
 from ._client_base import EventHubSharedKeyCredential
@@ -19,8 +19,6 @@
     EventHubConnectionStringProperties,
 )
 
-TransportType = constants.TransportType
-
 __all__ = [
     "EventData",
     "EventDataBatch",

sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer.py

@@ -2,6 +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 annotations
 import time
 import queue
 import logging
@@ -14,6 +15,7 @@
 from ..exceptions import OperationTimeoutError
 
 if TYPE_CHECKING:
+    from .._transport._base import AmqpTransport
     from .._producer_client import SendEventTypes
 
 _LOGGER = logging.getLogger(__name__)
@@ -31,7 +33,8 @@ def __init__(
         executor: ThreadPoolExecutor,
         *,
         max_wait_time: float = 1,
-        max_buffer_length: int
+        max_buffer_length: int,
+        amqp_transport: AmqpTransport
     ):
         self._buffered_queue: queue.Queue = queue.Queue()
         self._max_buffer_len = max_buffer_length
@@ -47,11 +50,12 @@ def __init__(
         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._amqp_transport = amqp_transport
         self.partition_id = partition_id
 
     def start(self):
         with self._lock:
-            self._cur_batch = EventDataBatch(self._max_message_size_on_link)
+            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()
@@ -111,11 +115,11 @@ def put_events(self, events, timeout_time=None):
                 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)
+            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
             self._buffered_queue.put(self._cur_batch)
-            self._cur_batch = EventDataBatch(self._max_message_size_on_link)
+            self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport)
             self._cur_batch.add(events)
         self._cur_buffered_len += new_events_len
 
@@ -182,7 +186,7 @@ def flush(self, timeout_time=None, raise_error=True):
                     break
             # after finishing flushing, reset cur batch and put it into the buffer
             self._last_send_time = time.time()
-            self._cur_batch = EventDataBatch(self._max_message_size_on_link)
+            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):

sdk/eventhub/azure-eventhub/azure/eventhub/_buffered_producer/_buffered_producer_dispatcher.py

@@ -2,6 +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 annotations
 import logging
 from threading import Lock
 from concurrent.futures import ThreadPoolExecutor
@@ -14,6 +15,7 @@
 
 if TYPE_CHECKING:
     from .._producer_client import SendEventTypes
+    from .._transport._base import AmqpTransport
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -31,7 +33,8 @@ def __init__(
         *,
         max_buffer_length: int = 1500,
         max_wait_time: float = 1,
-        executor: Optional[Union[ThreadPoolExecutor, int]] = None
+        executor: Optional[Union[ThreadPoolExecutor, int]] = None,
+        amqp_transport: AmqpTransport
     ):
         self._buffered_producers: Dict[str, BufferedProducer] = {}
         self._partition_ids: List[str] = partitions
@@ -45,6 +48,7 @@ def __init__(
         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()
@@ -86,6 +90,7 @@ def enqueue_events(
                     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