add switch to async

sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

@@ -194,7 +194,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)
 
 
@@ -291,7 +290,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         self._uamqp_transport = kwargs.pop("uamqp_transport", True)
-        self._amqp_transport = UamqpTransport if self._uamqp_transport else None
+        self._amqp_transport = kwargs.pop("amqp_transport", UamqpTransport)
 
         self.eventhub_name = eventhub_name
         if not eventhub_name:
@@ -302,7 +301,7 @@ def __init__(
         if isinstance(credential, AzureSasCredential):
             self._credential = EventhubAzureSasTokenCredential(credential)
         elif isinstance(credential, AzureNamedKeyCredential):
-            self._credential = EventhubAzureNamedKeyTokenCredential(credential)
+            self._credential = EventhubAzureNamedKeyTokenCredential(credential)  # type: ignore
         else:
             self._credential = credential  # type: ignore
         self._keep_alive = kwargs.get("keep_alive", 30)

sdk/eventhub/azure-eventhub/azure/eventhub/_common.py

@@ -511,8 +511,9 @@ def __init__(
         **kwargs,
     ) -> None:
         # TODO: this changes API, check with Anna if valid -
-        # If possible, move out message creation to right before sending.
+        # Need move out message creation to right before sending.
         # Might take more time to loop through events and add them all to batch in `send` than in `add` here
+        # Default async vs sync might cause issues.
         self._amqp_transport = kwargs.pop("amqp_transport", UamqpTransport)
 
 

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py

@@ -22,7 +22,7 @@
 
 if TYPE_CHECKING:
     from typing import Deque
-    from uamqp import ReceiveClient as uamqp_ReceiveClient, Message as uamqp_Message
+    from uamqp import ReceiveClient as uamqp_ReceiveClient, Message as uamqp_Message, types as uamqp_types
     from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth
 
     from ._consumer_client import EventHubConsumerClient
@@ -97,7 +97,7 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs: Any)
         self._auto_reconnect = auto_reconnect
         self._retry_policy = self._amqp_transport.create_retry_policy(self._client._config)
         self._reconnect_backoff = 1
-        link_properties: Dict[bytes, int] = {}
+        link_properties: Dict[uamqp_types.AMQPTypes, uamqp_types.AMQPType] = {}
         self._error = None
         self._timeout = 0
         self._idle_timeout = (idle_timeout * self._amqp_transport.IDLE_TIMEOUT_FACTOR) if idle_timeout else None

sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py

@@ -19,8 +19,6 @@
 )
 from typing_extensions import Literal
 
-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

sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py

@@ -43,7 +43,6 @@
 
 # TODO: remove after fixing up async
 from uamqp import types
-from uamqp.message import MessageHeader
 PROP_PARTITION_KEY_AMQP_SYMBOL = types.AMQPSymbol(PROP_PARTITION_KEY)
 
 
@@ -132,7 +131,7 @@ def send_context_manager():
     else:
         yield None
 
-# TODO: delete after async unit tests have been refactored
+
 def set_event_partition_key(event, partition_key):
     # type: (Union[AmqpAnnotatedMessage, EventData], Optional[Union[bytes, str]]) -> None
     if not partition_key:
@@ -155,27 +154,6 @@ def set_event_partition_key(event, partition_key):
         raw_message.header.durable = True
 
 
-def set_message_partition_key(message, partition_key):
-    # type: (Message, Optional[Union[bytes, str]]) -> None
-    """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
-    """
-    if partition_key:
-        annotations = message.annotations
-        if annotations is None:
-            annotations = dict()
-        annotations[
-            PROP_PARTITION_KEY_AMQP_SYMBOL
-        ] = partition_key  # pylint:disable=protected-access
-        header = MessageHeader()
-        header.durable = True
-        message.annotations = annotations
-        message.header = header
-
-
 @contextmanager
 def send_context_manager():
     span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_async.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 asyncio
 import logging
 import queue
@@ -14,6 +15,7 @@
 from ...exceptions import OperationTimeoutError
 
 if TYPE_CHECKING:
+    from .._transport._base_async import AmqpTransportAsync
     from ..._producer_client import SendEventTypes
 
 _LOGGER = logging.getLogger(__name__)
@@ -32,7 +34,8 @@ def __init__(
         max_message_size_on_link: int,
         *,
         max_wait_time: float = 1,
-        max_buffer_length: int
+        max_buffer_length: int,
+        amqp_transport: AmqpTransportAsync
     ):
         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
 
     async def start(self):
         async 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()
@@ -113,11 +117,11 @@ async 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
 
@@ -145,7 +149,7 @@ async def _flush(self, timeout_time=None, raise_error=True):
         _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)
+            self._cur_batch = EventDataBatch(self._max_message_size_on_link, amqp_transport=self._amqp_transport)
         while self._cur_buffered_len:
             remaining_time = timeout_time - time.time() if timeout_time else None
             if (remaining_time and remaining_time > 0) or remaining_time is None:
@@ -187,7 +191,7 @@ async 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)
 
     async def check_max_wait_time_worker(self):

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_buffered_producer/_buffered_producer_dispatcher_async.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 asyncio
 import logging
 from typing import Dict, List, Callable, Optional, Awaitable, TYPE_CHECKING
@@ -13,6 +14,7 @@
 from ...exceptions import EventDataSendError, ConnectError, EventHubError
 
 if TYPE_CHECKING:
+    from .._transport._base_async import AmqpTransportAsync
     from ..._producer_client import SendEventTypes
 
 _LOGGER = logging.getLogger(__name__)
@@ -33,6 +35,7 @@ def __init__(
         *,
         max_buffer_length: int = 1500,
         max_wait_time: float = 1,
+        amqp_transport: AmqpTransportAsync,
     ):
         self._buffered_producers: Dict[str, BufferedProducer] = {}
         self._partition_ids: List[str] = partitions
@@ -45,6 +48,7 @@ def __init__(
         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:
@@ -77,6 +81,7 @@ async def enqueue_events(
                     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