Merge branch 'eventhubs_preview2' into eventhubs_eph

sdk/eventhub/azure-eventhubs/azure/eventhub/_connection_manager.py

@@ -74,4 +74,4 @@ def reset_connection_if_broken(self):
 
 
 def get_connection_manager(**kwargs):
-    return _SharedConnectionManager(**kwargs)
+    return _SeparateConnectionManager(**kwargs)

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py

@@ -57,6 +57,13 @@ def __init__(  # pylint: disable=super-init-not-called
         :type owner_level: int
         :param loop: An event loop.
         """
+        event_position = kwargs.get("event_position", None)
+        prefetch = kwargs.get("prefetch", 300)
+        owner_level = kwargs.get("owner_level", None)
+        keep_alive = kwargs.get("keep_alive", None)
+        auto_reconnect = kwargs.get("auto_reconnect", True)
+        loop = kwargs.get("loop", None)
+
         super(EventHubConsumer, self).__init__()
         self.loop = loop or asyncio.get_event_loop()
         self.running = False
@@ -76,7 +83,9 @@ def __init__(  # pylint: disable=super-init-not-called
         partition = self.source.split('/')[-1]
         self.name = "EHReceiver-{}-partition{}".format(uuid.uuid4(), partition)
         if owner_level:
-            self.properties = {types.AMQPSymbol(self._epoch): types.AMQPLong(int(owner_level))}
+            self._link_properties[types.AMQPSymbol(self._epoch)] = types.AMQPLong(int(owner_level))
+        link_property_timeout_ms = (self.client.config.receive_timeout or self.timeout) * 1000
+        self._link_properties[types.AMQPSymbol(self._timeout)] = types.AMQPLong(int(link_property_timeout_ms))
         self._handler = None
 
     def __aiter__(self):
@@ -110,7 +119,7 @@ def _create_handler(self):
             auth=self.client.get_auth(**alt_creds),
             debug=self.client.config.network_tracing,
             prefetch=self.prefetch,
-            link_properties=self.properties,
+            link_properties=self._link_properties,
             timeout=self.timeout,
             error_policy=self.retry_policy,
             keep_alive_interval=self.keep_alive,
@@ -177,6 +186,9 @@ async def receive(self, **kwargs):
                 :caption: Receives events asynchronously
 
         """
+        max_batch_size = kwargs.get("max_batch_size", None)
+        timeout = kwargs.get("timeout", None)
+
         self._check_closed()
         max_batch_size = min(self.client.config.max_batch_size, self.prefetch) if max_batch_size is None else max_batch_size
         timeout = self.client.config.receive_timeout if timeout is None else timeout

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py

@@ -8,7 +8,7 @@
 from typing import Iterable, Union
 import time
 
-from uamqp import constants, errors
+from uamqp import types, constants, errors
 from uamqp import SendClientAsync
 
 from azure.eventhub.common import EventData, _BatchSendEventData
@@ -54,6 +54,12 @@ def __init__(  # pylint: disable=super-init-not-called
         :type auto_reconnect: bool
         :param loop: An event loop. If not specified the default event loop will be used.
         """
+        partition = kwargs.get("partition", None)
+        send_timeout = kwargs.get("send_timeout", 60)
+        keep_alive = kwargs.get("keep_alive", None)
+        auto_reconnect = kwargs.get("auto_reconnect", True)
+        loop = kwargs.get("loop", None)
+
         super(EventHubProducer, self).__init__()
         self.loop = loop or asyncio.get_event_loop()
         self.running = False
@@ -75,6 +81,7 @@ def __init__(  # pylint: disable=super-init-not-called
         self._handler = None
         self._outcome = None
         self._condition = None
+        self._link_properties = {types.AMQPSymbol(self._timeout): types.AMQPLong(int(self.timeout * 1000))}
 
     def _create_handler(self):
         self._handler = SendClientAsync(
@@ -85,6 +92,7 @@ def _create_handler(self):
             error_policy=self.retry_policy,
             keep_alive_interval=self.keep_alive,
             client_name=self.name,
+            link_properties=self._link_properties,
             properties=self.client._create_properties(
                 self.client.config.user_agent),  # pylint: disable=protected-access
             loop=self.loop)
@@ -122,11 +130,14 @@ async def _send_event_data(self, timeout=None):
                             error = OperationTimeoutError("send operation timed out")
                         log.info("%r send operation timed out. (%r)", self.name, error)
                         raise error
+                    self._handler._msg_timeout = remaining_time  # pylint: disable=protected-access
                     self._handler.queue_message(*self.unsent_events)
                     await self._handler.wait_async()
                     self.unsent_events = self._handler.pending_messages
-                if self._outcome != constants.MessageSendResult.Ok:
-                    _error(self._outcome, self._condition)
+                    if self._outcome != constants.MessageSendResult.Ok:
+                        if self._outcome == constants.MessageSendResult.Timeout:
+                            self._condition = OperationTimeoutError("send operation timed out")
+                        _error(self._outcome, self._condition)
                 return
             except Exception as exception:
                 last_exception = await self._handle_exception(exception, retry_count, max_retries, timeout_time)
@@ -144,7 +155,7 @@ def _on_outcome(self, outcome, condition):
         self._outcome = outcome
         self._condition = condition
 
-    async def send(self, event_data, partition_key=None, timeout=None):
+    async def send(self, event_data, **kwargs):
         # type:(Union[EventData, Iterable[EventData]], Union[str, bytes]) -> None
         """
         Sends an event data and blocks until acknowledgement is
@@ -209,4 +220,5 @@ async def close(self, **kwargs):
                 :caption: Close down the handler.
 
         """
+        exception = kwargs.get("exception", None)
         await super(EventHubProducer, self).close(exception)

sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py

@@ -53,6 +53,12 @@ def __init__(self, client, source, **kwargs):
          consumer if owner_level is set.
         :type owner_level: int
         """
+        event_position = kwargs.get("event_position", None)
+        prefetch = kwargs.get("prefetch", 300)
+        owner_level = kwargs.get("owner_level", None)
+        keep_alive = kwargs.get("keep_alive", None)
+        auto_reconnect = kwargs.get("auto_reconnect", True)
+
         super(EventHubConsumer, self).__init__()
         self.running = False
         self.client = client
@@ -71,7 +77,9 @@ def __init__(self, client, source, **kwargs):
         partition = self.source.split('/')[-1]
         self.name = "EHConsumer-{}-partition{}".format(uuid.uuid4(), partition)
         if owner_level:
-            self.properties = {types.AMQPSymbol(self._epoch): types.AMQPLong(int(owner_level))}
+            self._link_properties[types.AMQPSymbol(self._epoch)] = types.AMQPLong(int(owner_level))
+        link_property_timeout_ms = (self.client.config.receive_timeout or self.timeout) * 1000
+        self._link_properties[types.AMQPSymbol(self._timeout)] = types.AMQPLong(int(link_property_timeout_ms))
         self._handler = None
 
     def __iter__(self):
@@ -105,7 +113,7 @@ def _create_handler(self):
             auth=self.client.get_auth(**alt_creds),
             debug=self.client.config.network_tracing,
             prefetch=self.prefetch,
-            link_properties=self.properties,
+            link_properties=self._link_properties,
             timeout=self.timeout,
             error_policy=self.retry_policy,
             keep_alive_interval=self.keep_alive,
@@ -170,7 +178,8 @@ def receive(self, **kwargs):
                 :caption: Receive events from the EventHub.
 
         """
-        self._check_closed()
+        max_batch_size = kwargs.get("max_batch_size", None)
+        timeout = kwargs.get("timeout", None)
 
         self._check_closed()
         max_batch_size = min(self.client.config.max_batch_size, self.prefetch) if max_batch_size is None else max_batch_size

sdk/eventhub/azure-eventhubs/azure/eventhub/error.py

@@ -19,6 +19,7 @@
 
 log = logging.getLogger(__name__)
 
+
 def _error_handler(error):
     """
     Called internally when an event has failed to send so we

sdk/eventhub/azure-eventhubs/azure/eventhub/producer.py

@@ -64,6 +64,11 @@ def __init__(self, client, target, **kwargs):
          Default value is `True`.
         :type auto_reconnect: bool
         """
+        partition = kwargs.get("partition", None)
+        send_timeout = kwargs.get("send_timeout", 60)
+        keep_alive = kwargs.get("keep_alive", None)
+        auto_reconnect = kwargs.get("auto_reconnect", True)
+
         super(EventHubProducer, self).__init__()
         self.running = False
         self.client = client
@@ -84,6 +89,7 @@ def __init__(self, client, target, **kwargs):
         self._handler = None
         self._outcome = None
         self._condition = None
+        self._link_properties = {types.AMQPSymbol(self._timeout): types.AMQPLong(int(self.timeout * 1000))}
 
     def _create_handler(self):
         self._handler = SendClient(
@@ -131,10 +137,13 @@ def _send_event_data(self, timeout=None):
                             error = OperationTimeoutError("send operation timed out")
                         log.info("%r send operation timed out. (%r)", self.name, error)
                         raise error
+                    self._handler._msg_timeout = remaining_time  # pylint: disable=protected-access
                     self._handler.queue_message(*self.unsent_events)
                     self._handler.wait()
                     self.unsent_events = self._handler.pending_messages
                     if self._outcome != constants.MessageSendResult.Ok:
+                        if self._outcome == constants.MessageSendResult.Timeout:
+                            self._condition = OperationTimeoutError("send operation timed out")
                         _error(self._outcome, self._condition)
                 return
             except Exception as exception:
@@ -153,7 +162,7 @@ def _on_outcome(self, outcome, condition):
         self._outcome = outcome
         self._condition = condition
 
-    def send(self, event_data, partition_key=None, timeout=None):
+    def send(self, event_data, **kwargs):
         # type:(Union[EventData, Iterable[EventData]], Union[str, bytes], float) -> None
         """
         Sends an event data and blocks until acknowledgement is
@@ -219,4 +228,5 @@ def close(self, **kwargs):
                 :caption: Close down the handler.
 
         """
+        exception = kwargs.get("exception", None)
         super(EventHubProducer, self).close(exception)