diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/_connection_manager.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/_connection_manager.py new file mode 100644 index 000000000000..166703d698e7 --- /dev/null +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/_connection_manager.py @@ -0,0 +1,77 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from threading import RLock +from uamqp import Connection, TransportType, c_uamqp + + +class _SharedConnectionManager(object): + def __init__(self, **kwargs): + self._lock = RLock() + self._conn = None # type: Connection + + self._container_id = kwargs.get("container_id") + 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") + + def get_connection(self, host, auth): + # type: (...) -> Connection + with self._lock: + if self._conn is None: + self._conn = Connection( + host, + auth, + 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, + encoding=self._encoding) + return self._conn + + def close_connection(self): + with self._lock: + if self._conn: + self._conn.destroy() + self._conn = None + + def reset_connection_if_broken(self): + with self._lock: + if self._conn and self._conn._state in ( + c_uamqp.ConnectionState.CLOSE_RCVD, + c_uamqp.ConnectionState.CLOSE_SENT, + c_uamqp.ConnectionState.DISCARDING, + c_uamqp.ConnectionState.END, + ): + self._conn = None + + +class _SeparateConnectionManager(object): + def __init__(self, **kwargs): + pass + + def get_connection(self, host, auth): + return None + + def close_connection(self): + pass + + def reset_connection_if_broken(self): + pass + + +def get_connection_manager(**kwargs): + return _SeparateConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py new file mode 100644 index 000000000000..9ac6fb468945 --- /dev/null +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py @@ -0,0 +1,109 @@ +# -------------------------------------------------------------------------------------------- +# 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 + +import logging +import time + +from uamqp import errors +from azure.eventhub.error import EventHubError, _handle_exception + +log = logging.getLogger(__name__) + + +class ConsumerProducerMixin(object): + def __init__(self): + self.client = None + self._handler = None + self.name = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close(exc_val) + + def _check_closed(self): + if self.error: + raise EventHubError("{} has been closed. Please create a new consumer to receive event data.".format(self.name)) + + def _create_handler(self): + pass + + def _redirect(self, redirect): + self.redirected = redirect + self.running = False + self._close_connection() + + def _open(self, timeout_time=None): + """ + Open the EventHubConsumer using the supplied connection. + If the handler has previously been redirected, the redirect + context will be used to create a new handler before opening it. + + """ + # pylint: disable=protected-access + if not self.running: + if self.redirected: + alt_creds = { + "username": self.client._auth_config.get("iot_username"), + "password": self.client._auth_config.get("iot_password")} + else: + alt_creds = {} + self._create_handler() + self._handler.open(connection=self.client._conn_manager.get_connection( + self.client.address.hostname, + self.client.get_auth(**alt_creds) + )) + while not self._handler.client_ready(): + if timeout_time and time.time() >= timeout_time: + return + time.sleep(0.05) + self.running = True + + def _close_handler(self): + self._handler.close() # close the link (sharing connection) or connection (not sharing) + self.running = False + + def _close_connection(self): + self._close_handler() + self.client._conn_manager.reset_connection_if_broken() + + def _handle_exception(self, exception, retry_count, max_retries, timeout_time): + _handle_exception(exception, retry_count, max_retries, self, timeout_time) + + def close(self, exception=None): + # type:(Exception) -> None + """ + Close down the handler. If the handler has already closed, + this will be a no op. An optional exception can be passed in to + indicate that the handler was shutdown due to error. + + :param exception: An optional exception if the handler is closing + due to an error. + :type exception: Exception + + Example: + .. literalinclude:: ../examples/test_examples_eventhub.py + :start-after: [START eventhub_client_receiver_close] + :end-before: [END eventhub_client_receiver_close] + :language: python + :dedent: 4 + :caption: Close down the handler. + + """ + self.running = False + if self.error: + return + if isinstance(exception, errors.LinkRedirect): + self.redirected = exception + elif isinstance(exception, EventHubError): + self.error = exception + elif exception: + self.error = EventHubError(str(exception)) + else: + self.error = EventHubError("{} handler is closed.".format(self.name)) + if self._handler: + self._handler.close() # this will close link if sharing connection. Otherwise close connection diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_connection_manager_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_connection_manager_async.py new file mode 100644 index 000000000000..3178e1fb72a7 --- /dev/null +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_connection_manager_async.py @@ -0,0 +1,78 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from asyncio import Lock +from uamqp import TransportType, c_uamqp +from uamqp.async_ops import ConnectionAsync + + +class _SharedConnectionManager(object): + def __init__(self, **kwargs): + self._lock = Lock() + self._conn = None + + self._container_id = kwargs.get("container_id") + 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") + + async def get_connection(self, host, auth): + # type: (...) -> ConnectionAsync + async with self._lock: + if self._conn is None: + self._conn = ConnectionAsync( + host, + auth, + 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, + encoding=self._encoding) + return self._conn + + async def close_connection(self): + async with self._lock: + if self._conn: + await self._conn.destroy_async() + self._conn = None + + async def reset_connection_if_broken(self): + async with self._lock: + if self._conn and self._conn._state in ( + c_uamqp.ConnectionState.CLOSE_RCVD, + c_uamqp.ConnectionState.CLOSE_SENT, + c_uamqp.ConnectionState.DISCARDING, + c_uamqp.ConnectionState.END, + ): + self._conn = None + + +class _SeparateConnectionManager(object): + def __init__(self, **kwargs): + pass + + async def get_connection(self, host, auth): + pass # return None + + async def close_connection(self): + pass + + def reset_connection_if_broken(self): + pass + + +def get_connection_manager(**kwargs): + return _SharedConnectionManager(**kwargs) diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py new file mode 100644 index 000000000000..5a0f0d9eaa4d --- /dev/null +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py @@ -0,0 +1,112 @@ +# -------------------------------------------------------------------------------------------- +# 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 +import time + +from uamqp import errors +from azure.eventhub.error import EventHubError, ConnectError +from ..aio.error_async import _handle_exception + +log = logging.getLogger(__name__) + + +class ConsumerProducerMixin(object): + + def __init__(self): + self.client = None + self._handler = None + self.name = None + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close(exc_val) + + def _check_closed(self): + if self.error: + raise EventHubError("{} has been closed. Please create a new consumer to receive event data.".format(self.name)) + + def _create_handler(self): + pass + + async def _redirect(self, redirect): + self.redirected = redirect + self.running = False + await self._close_connection() + + async def _open(self, timeout_time=None): + """ + Open the EventHubConsumer using the supplied connection. + If the handler has previously been redirected, the redirect + context will be used to create a new handler before opening it. + + """ + # pylint: disable=protected-access + if not self.running: + if self.redirected: + alt_creds = { + "username": self.client._auth_config.get("iot_username"), + "password": self.client._auth_config.get("iot_password")} + else: + alt_creds = {} + self._create_handler() + await self._handler.open_async(connection=await self.client._conn_manager.get_connection( + self.client.address.hostname, + self.client.get_auth(**alt_creds) + )) + while not await self._handler.client_ready_async(): + if timeout_time and time.time() >= timeout_time: + return + await asyncio.sleep(0.05) + self.running = True + + async def _close_handler(self): + await self._handler.close_async() # close the link (sharing connection) or connection (not sharing) + self.running = False + + async def _close_connection(self): + await self._close_handler() + await self.client._conn_manager.reset_connection_if_broken() + + async def _handle_exception(self, exception, retry_count, max_retries, timeout_time): + await _handle_exception(exception, retry_count, max_retries, self, timeout_time) + + async def close(self, exception=None): + # type: (Exception) -> None + """ + Close down the handler. If the handler has already closed, + this will be a no op. An optional exception can be passed in to + indicate that the handler was shutdown due to error. + + :param exception: An optional exception if the handler is closing + due to an error. + :type exception: Exception + + Example: + .. literalinclude:: ../examples/async_examples/test_examples_eventhub_async.py + :start-after: [START eventhub_client_async_receiver_close] + :end-before: [END eventhub_client_async_receiver_close] + :language: python + :dedent: 4 + :caption: Close down the handler. + + """ + self.running = False + if self.error: + return + if isinstance(exception, errors.LinkRedirect): + self.redirected = exception + elif isinstance(exception, EventHubError): + self.error = exception + elif isinstance(exception, (errors.LinkDetach, errors.ConnectionClose)): + self.error = ConnectError(str(exception), exception) + elif exception: + self.error = EventHubError(str(exception)) + else: + self.error = EventHubError("This receive handler is now closed.") + if self._handler: + await self._handler.close_async() \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py index f552cb0a167b..17479119ccad 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py @@ -22,6 +22,8 @@ from .producer_async import EventHubProducer from .consumer_async import EventHubConsumer +from ._connection_manager_async import get_connection_manager +from .error_async import _handle_exception log = logging.getLogger(__name__) @@ -42,6 +44,16 @@ class EventHubClient(EventHubClientAbstract): """ + def __init__(self, host, event_hub_path, credential, **kwargs): + super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs) + self._conn_manager = get_connection_manager(**kwargs) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + def _create_auth(self, username=None, password=None): """ Create an ~uamqp.authentication.cbs_auth_async.SASTokenAuthAsync instance to authenticate @@ -85,17 +97,21 @@ def _create_auth(self, username=None, password=None): get_jwt_token, http_proxy=http_proxy, transport_type=transport_type) + async def _handle_exception(self, exception, retry_count, max_retries): + await _handle_exception(exception, retry_count, max_retries, self) + + async def _close_connection(self): + await self._conn_manager.reset_connection_if_broken() + async def _management_request(self, mgmt_msg, op_type): - alt_creds = { - "username": self._auth_config.get("iot_username"), - "password": self._auth_config.get("iot_password")} - connect_count = 0 + max_retries = self.config.max_retries + retry_count = 0 while True: - connect_count += 1 - mgmt_auth = self._create_auth(**alt_creds) + mgmt_auth = self._create_auth() mgmt_client = AMQPClientAsync(self.mgmt_target, auth=mgmt_auth, debug=self.config.network_tracing) try: - await mgmt_client.open_async() + conn = await self._conn_manager.get_connection(self.host, mgmt_auth) + await mgmt_client.open_async(connection=conn) response = await mgmt_client.mgmt_request_async( mgmt_msg, constants.READ_OPERATION, @@ -103,15 +119,9 @@ async def _management_request(self, mgmt_msg, op_type): status_code_field=b'status-code', description_fields=b'status-description') return response - except (errors.AMQPConnectionError, errors.TokenAuthFailure, compat.TimeoutException) as failure: - if connect_count >= self.config.max_retries: - err = ConnectError( - "Can not connect to EventHubs or get management info from the service. " - "Please make sure the connection string or token is correct and retry. " - "Besides, this method doesn't work if you use an IoT connection string.", - failure - ) - raise err + except Exception as exception: + await self._handle_exception(exception, retry_count, max_retries) + retry_count += 1 finally: await mgmt_client.close_async() @@ -184,8 +194,7 @@ async def get_partition_properties(self, partition): return output def create_consumer( - self, consumer_group, partition_id, event_position, owner_level=None, - operation=None, prefetch=None, loop=None): + self, consumer_group, partition_id, event_position, **kwargs): # type: (str, str, EventPosition, int, str, int, asyncio.AbstractEventLoop) -> EventHubConsumer """ Create an async consumer to the client for a particular consumer group and partition. @@ -217,8 +226,12 @@ def create_consumer( :caption: Add an async consumer to the client for a particular consumer group and partition. """ - prefetch = self.config.prefetch if prefetch is None else prefetch + owner_level = kwargs.get("owner_level", None) + operation = kwargs.get("operation", None) + prefetch = kwargs.get("prefetch", None) + loop = kwargs.get("loop", None) + prefetch = prefetch or self.config.prefetch path = self.address.path + operation if operation else self.address.path source_url = "amqps://{}{}/ConsumerGroups/{}/Partitions/{}".format( self.address.hostname, path, consumer_group, partition_id) @@ -228,7 +241,7 @@ def create_consumer( return handler def create_producer( - self, partition_id=None, operation=None, send_timeout=None, loop=None): + self, **kwargs): # type: (str, str, float, asyncio.AbstractEventLoop) -> EventHubProducer """ Create an async producer to send EventData object to an EventHub. @@ -255,6 +268,11 @@ def create_producer( :caption: Add an async producer to the client to send EventData. """ + partition_id = kwargs.get("partition_id", None) + operation = kwargs.get("operation", None) + send_timeout = kwargs.get("send_timeout", None) + loop = kwargs.get("loop", None) + target = "amqps://{}{}".format(self.address.hostname, self.address.path) if operation: target = target + operation @@ -263,3 +281,6 @@ def create_producer( handler = EventHubProducer( self, target, partition=partition_id, send_timeout=send_timeout, loop=loop) return handler + + async def close(self): + await self._conn_manager.close_connection() \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py index 6cf020176d96..d4d4143810af 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py @@ -6,36 +6,39 @@ import uuid import logging from typing import List +import time from uamqp import errors, types, compat from uamqp import ReceiveClientAsync, Source from azure.eventhub import EventData, EventPosition from azure.eventhub.error import EventHubError, AuthenticationError, ConnectError, ConnectionLostError, _error_handler +from ..aio.error_async import _handle_exception +from ._consumer_producer_mixin_async import ConsumerProducerMixin log = logging.getLogger(__name__) -class EventHubConsumer(object): +class EventHubConsumer(ConsumerProducerMixin): """ A consumer responsible for reading EventData from a specific Event Hub - partition and as a member of a specific consumer group. + partition and as a member of a specific consumer group. A consumer may be exclusive, which asserts ownership over the partition for the consumer - group to ensure that only one consumer from that group is reading the from the partition. - These exclusive consumers are sometimes referred to as "Epoch Consumers." + group to ensure that only one consumer from that group is reading the from the partition. + These exclusive consumers are sometimes referred to as "Epoch Consumers." A consumer may also be non-exclusive, allowing multiple consumers from the same consumer - group to be actively reading events from the partition. These non-exclusive consumers are - sometimes referred to as "Non-Epoch Consumers." + group to be actively reading events from the partition. These non-exclusive consumers are + sometimes referred to as "Non-Epoch Consumers." """ timeout = 0 _epoch = b'com.microsoft:epoch' + _timeout = b'com.microsoft:timeout' def __init__( # pylint: disable=super-init-not-called - self, client, source, event_position=None, prefetch=300, owner_level=None, - keep_alive=None, auto_reconnect=True, loop=None): + self, client, source, **kwargs): """ Instantiate an async consumer. EventHubConsumer should be instantiated by calling the `create_consumer` method in EventHubClient. @@ -54,6 +57,14 @@ 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 self.client = client @@ -68,107 +79,61 @@ def __init__( # pylint: disable=super-init-not-called self.reconnect_backoff = 1 self.redirected = None self.error = None - self.properties = None + self._link_properties = {} partition = self.source.split('/')[-1] self.name = "EHReceiver-{}-partition{}".format(uuid.uuid4(), partition) - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) # pylint: disable=protected-access if owner_level: - self.properties = {types.AMQPSymbol(self._epoch): types.AMQPLong(int(owner_level))} - self._handler = ReceiveClientAsync( - source, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), # pylint: disable=protected-access - loop=self.loop) - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close(exc_val) + 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): return self async def __anext__(self): - await self._open() max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 while True: - connecting_count += 1 try: + await self._open() if not self.messages_iter: self.messages_iter = self._handler.receive_messages_iter_async() message = await self.messages_iter.__anext__() event_data = EventData(message=message) self.offset = EventPosition(event_data.offset, inclusive=False) return event_data - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubConsumer disconnected due to token error. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - await self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry and self.auto_reconnect: - log.info("EventHubConsumer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer connection lost. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer timed out receiving event data. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer timed out. Shutting down.") - await self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except StopAsyncIteration: - raise - except Exception as e: - log.error("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Receive failed: {}".format(e), e) - await self.close(exception=error) - raise error + except Exception as exception: + await self._handle_exception(exception, retry_count, max_retries) + retry_count += 1 + + def _create_handler(self): + alt_creds = { + "username": self.client._auth_config.get("iot_username"), + "password": self.client._auth_config.get("iot_password")} + source = Source(self.source) + if self.offset is not None: + source.set_filter(self.offset._selector()) + self._handler = ReceiveClientAsync( + source, + auth=self.client.get_auth(**alt_creds), + debug=self.client.config.network_tracing, + prefetch=self.prefetch, + link_properties=self._link_properties, + timeout=self.timeout, + error_policy=self.retry_policy, + keep_alive_interval=self.keep_alive, + client_name=self.name, + properties=self.client._create_properties( + self.client.config.user_agent), # pylint: disable=protected-access + loop=self.loop) + self.messages_iter = None - def _check_closed(self): - if self.error: - raise EventHubError("This consumer has been closed. Please create a new consumer to receive event data.", - self.error) + async def _redirect(self, redirect): + self.messages_iter = None + await super(EventHubConsumer, self)._redirect(redirect) - async def _open(self): + async def _open(self, timeout_time=None): """ Open the EventHubConsumer using the supplied connection. If the handler has previously been redirected, the redirect @@ -176,121 +141,10 @@ async def _open(self): """ # pylint: disable=protected-access - self._check_closed() - if self.redirected: + if not self.running and self.redirected: + self.client._process_redirect_uri(self.redirected) self.source = self.redirected.address - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) # pylint: disable=protected-access - alt_creds = { - "username": self.client._auth_config.get("iot_username"), - "password":self.client._auth_config.get("iot_password")} - self._handler = ReceiveClientAsync( - source, - auth=self.client.get_auth(**alt_creds), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), # pylint: disable=protected-access - loop=self.loop) - if not self.running: - await self._connect() - self.running = True - - async def _connect(self): - connected = await self._build_connection() - if not connected: - await asyncio.sleep(self.reconnect_backoff) - while not await self._build_connection(is_reconnect=True): - await asyncio.sleep(self.reconnect_backoff) - - async def _build_connection(self, is_reconnect=False): # pylint: disable=too-many-statements - # pylint: disable=protected-access - if is_reconnect: - alt_creds = { - "username": self.client._auth_config.get("iot_username"), - "password":self.client._auth_config.get("iot_password")} - await self._handler.close_async() - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) # pylint: disable=protected-access - self._handler = ReceiveClientAsync( - source, - auth=self.client.get_auth(**alt_creds), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), # pylint: disable=protected-access - loop=self.loop) - self.messages_iter = None - try: - await self._handler.open_async() - while not await self._handler.client_ready_async(): - await asyncio.sleep(0.05) - return True - except errors.AuthenticationException as shutdown: - if is_reconnect: - log.info("EventHubConsumer couldn't authenticate. Shutting down. (%r)", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubConsumer couldn't authenticate. Attempting reconnect.") - return False - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubConsumer detached. Attempting reconnect.") - return False - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if is_reconnect: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubConsumer detached. Attempting reconnect.") - return False - except errors.AMQPConnectionError as shutdown: - if is_reconnect: - log.info("EventHubConsumer connection error (%r). Shutting down.", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubConsumer couldn't authenticate. Attempting reconnect.") - return False - except compat.TimeoutException as shutdown: - if is_reconnect: - log.info("EventHubConsumer authentication timed out. Shutting down.") - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubConsumer authentication timed out. Attempting reconnect.") - return False - except Exception as e: - log.error("Unexpected error occurred when building connection (%r). Shutting down.", e) - error = EventHubError("Unexpected error occurred when building connection", e) - await self.close(exception=error) - raise error - - async def _reconnect(self): - """If the EventHubConsumer was disconnected from the service with - a retryable error - attempt to reconnect.""" - return await self._build_connection(is_reconnect=True) + await super(EventHubConsumer, self)._open(timeout_time) @property def queue_size(self): @@ -305,7 +159,7 @@ def queue_size(self): return self._handler._received_messages.qsize() return 0 - async def receive(self, max_batch_size=None, timeout=None): + async def receive(self, **kwargs): # type: (int, float) -> List[EventData] """ Receive events asynchronously from the EventHub. @@ -332,78 +186,47 @@ async def receive(self, max_batch_size=None, timeout=None): :caption: Receives events asynchronously """ - self._check_closed() - await self._open() + 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 + if not timeout: + timeout = 100_000 # timeout None or 0 mean no timeout. 100000 seconds is equivalent to no timeout data_batch = [] + start_time = time.time() + timeout_time = start_time + timeout max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 + last_exception = None while True: - connecting_count += 1 try: - timeout_ms = 1000 * timeout if timeout else 0 + await self._open(timeout_time) + remaining_time = timeout_time - time.time() + if remaining_time <= 0.0: + if last_exception: + log.info("%r receive operation timed out. (%r)", self.name, last_exception) + raise last_exception + return data_batch + + remaining_time_ms = 1000 * remaining_time message_batch = await self._handler.receive_message_batch_async( max_batch_size=max_batch_size, - timeout=timeout_ms) + timeout=remaining_time_ms) for message in message_batch: event_data = EventData(message=message) self.offset = EventPosition(event_data.offset) data_batch.append(event_data) return data_batch - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubConsumer disconnected due to token error. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - await self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry and self.auto_reconnect: - log.info("EventHubConsumer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer connection lost. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer timed out receiving event data. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubConsumer timed out. Shutting down.") - await self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except Exception as e: - log.info("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Receive failed: {}".format(e), e) - await self.close(exception=error) - raise error + except EventHubError: + raise + except Exception as exception: + last_exception = await self._handle_exception(exception, retry_count, max_retries, timeout_time) + retry_count += 1 - async def close(self, exception=None): + async def close(self, **kwargs): # type: (Exception) -> None """ Close down the handler. If the handler has already closed, @@ -423,6 +246,7 @@ async def close(self, exception=None): :caption: Close down the handler. """ + exception = kwargs.get("exception", None) self.running = False if self.error: return @@ -436,4 +260,4 @@ async def close(self, exception=None): self.error = EventHubError(str(exception)) else: self.error = EventHubError("This receive handler is now closed.") - await self._handler.close_async() \ No newline at end of file + await self._handler.close_async() diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/error_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/error_async.py new file mode 100644 index 000000000000..957a3005662e --- /dev/null +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/error_async.py @@ -0,0 +1,108 @@ +import asyncio +import time +import logging + +from uamqp import errors, compat +from ..error import EventHubError, EventDataSendError, \ + EventDataError, ConnectError, ConnectionLostError, AuthenticationError + + +log = logging.getLogger(__name__) + + +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 + + +async def _handle_exception(exception, retry_count, max_retries, closable, timeout_time=None): + try: + name = closable.name + except AttributeError: + name = closable.container_id + if isinstance(exception, KeyboardInterrupt): + log.info("%r stops due to keyboard interrupt", name) + closable.close() + raise + elif isinstance(exception, EventHubError): + closable.close() + raise + elif isinstance(exception, ( + errors.MessageAccepted, + errors.MessageAlreadySettled, + errors.MessageModified, + errors.MessageRejected, + errors.MessageReleased, + errors.MessageContentTooLarge) + ): + log.info("%r Event data error (%r)", name, exception) + error = EventDataError(str(exception), exception) + raise error + elif isinstance(exception, errors.MessageException): + log.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + elif retry_count >= max_retries: + error = _create_eventhub_exception(exception) + log.info("%r has exhausted retry. Exception still occurs (%r)", name, exception) + raise error + else: + if isinstance(exception, errors.AuthenticationException): + if hasattr(closable, "_close_connection"): + await closable._close_connection() + elif isinstance(exception, errors.LinkRedirect): + log.info("%r link redirect received. Redirecting...", name) + redirect = exception + if hasattr(closable, "_redirect"): + await closable._redirect(redirect) + elif isinstance(exception, errors.LinkDetach): + if hasattr(closable, "_close_handler"): + await closable._close_handler() + elif isinstance(exception, errors.ConnectionClose): + if hasattr(closable, "_close_connection"): + await closable._close_connection() + elif isinstance(exception, errors.MessageHandlerError): + if hasattr(closable, "_close_handler"): + await closable._close_handler() + elif isinstance(exception, errors.AMQPConnectionError): + if hasattr(closable, "_close_connection"): + await closable._close_connection() + elif isinstance(exception, compat.TimeoutException): + pass # Timeout doesn't need to recreate link or connection to retry + else: + if hasattr(closable, "_close_connection"): + await closable._close_connection() + # start processing retry delay + try: + backoff_factor = closable.client.config.backoff_factor + backoff_max = closable.client.config.backoff_max + except AttributeError: + backoff_factor = closable.config.backoff_factor + backoff_max = closable.config.backoff_max + backoff = backoff_factor * 2 ** retry_count + if backoff <= backoff_max and (timeout_time is None or time.time() + backoff <= timeout_time): + await asyncio.sleep(backoff) + log.info("%r has an exception (%r). Retrying...", format(name), exception) + return _create_eventhub_exception(exception) + else: + error = _create_eventhub_exception(exception) + log.info("%r operation has timed out. Last exception before timeout is (%r)", name, error) + raise error + # end of processing retry delay \ No newline at end of file diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py index aef8dc50ff02..9612b4156327 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py @@ -5,30 +5,33 @@ import uuid import asyncio import logging -from typing import Iterator, Generator, List, Union +from typing import Iterable, Union +import time -from uamqp import constants, errors, compat +from uamqp import types, constants, errors from uamqp import SendClientAsync from azure.eventhub.common import EventData, _BatchSendEventData -from azure.eventhub.error import EventHubError, ConnectError, \ - AuthenticationError, EventDataError, EventDataSendError, ConnectionLostError, _error_handler +from azure.eventhub.error import _error_handler, OperationTimeoutError +from ..producer import _error, _set_partition_key +from ._consumer_producer_mixin_async import ConsumerProducerMixin + log = logging.getLogger(__name__) -class EventHubProducer(object): +class EventHubProducer(ConsumerProducerMixin): """ A producer responsible for transmitting EventData to a specific Event Hub, - grouped together in batches. Depending on the options specified at creation, the producer may - be created to allow event data to be automatically routed to an available partition or specific - to a partition. + grouped together in batches. Depending on the options specified at creation, the producer may + be created to allow event data to be automatically routed to an available partition or specific + to a partition. """ + _timeout = b'com.microsoft:timeout' def __init__( # pylint: disable=super-init-not-called - self, client, target, partition=None, send_timeout=60, - keep_alive=None, auto_reconnect=True, loop=None): + self, client, target, **kwargs): """ Instantiate an async EventHubProducer. EventHubProducer should be instantiated by calling the `create_producer` method in EventHubClient. @@ -51,6 +54,13 @@ 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 self.client = client @@ -68,6 +78,12 @@ def __init__( # pylint: disable=super-init-not-called if partition: self.target += "/Partitions/" + partition self.name += "-partition{}".format(partition) + 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( self.target, auth=self.client.get_auth(), @@ -76,206 +92,56 @@ def __init__( # pylint: disable=super-init-not-called error_policy=self.retry_policy, keep_alive_interval=self.keep_alive, client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), # pylint: disable=protected-access + link_properties=self._link_properties, + properties=self.client._create_properties( + self.client.config.user_agent), # pylint: disable=protected-access loop=self.loop) - self._outcome = None - self._condition = None - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close(exc_val) - async def _open(self): + async def _open(self, timeout_time=None): """ Open the EventHubProducer using the supplied connection. If the handler has previously been redirected, the redirect context will be used to create a new handler before opening it. """ - if self.redirected: + if not self.running and self.redirected: + self.client._process_redirect_uri(self.redirected) self.target = self.redirected.address - self._handler = SendClientAsync( - self.target, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - msg_timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), # pylint: disable=protected-access - loop=self.loop) - if not self.running: - await self._connect() - self.running = True - - async def _connect(self): - connected = await self._build_connection() - if not connected: - await asyncio.sleep(self.reconnect_backoff) - while not await self._build_connection(is_reconnect=True): - await asyncio.sleep(self.reconnect_backoff) - - async def _build_connection(self, is_reconnect=False): - """ - - :param is_reconnect: True - trying to reconnect after fail to connect or a connection is lost. - False - the 1st time to connect - :return: True - connected. False - not connected - """ - # pylint: disable=protected-access - if is_reconnect: - await self._handler.close_async() - self._handler = SendClientAsync( - self.target, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - msg_timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent), - loop=self.loop) - try: - await self._handler.open_async() - while not await self._handler.client_ready_async(): - await asyncio.sleep(0.05) - return True - except errors.AuthenticationException as shutdown: - if is_reconnect: - log.info("EventHubProducer couldn't authenticate. Shutting down. (%r)", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubProducer couldn't authenticate. Attempting reconnect.") - return False - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubProducer detached. Attempting reconnect.") - return False - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if is_reconnect: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubProducer detached. Attempting reconnect.") - return False - except errors.AMQPConnectionError as shutdown: - if is_reconnect: - log.info("EventHubProducer connection error (%r). Shutting down.", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubProducer couldn't authenticate. Attempting reconnect.") - return False - except compat.TimeoutException as shutdown: - if is_reconnect: - log.info("EventHubProducer authentication timed out. Shutting down.") - error = AuthenticationError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - else: - log.info("EventHubProducer authentication timed out. Attempting reconnect.") - return False - except Exception as e: - log.info("Unexpected error occurred when building connection (%r). Shutting down.", e) - error = EventHubError("Unexpected error occurred when building connection", e) - await self.close(exception=error) - raise error - - async def _reconnect(self): - return await self._build_connection(is_reconnect=True) - - async def _send_event_data(self): - await self._open() + await super(EventHubProducer, self)._open(timeout_time) + + async def _send_event_data(self, timeout=None): + timeout = self.client.config.send_timeout if timeout is None else timeout + if not timeout: + timeout = 100_000 # timeout None or 0 mean no timeout. 100000 seconds is equivalent to no timeout + start_time = time.time() + timeout_time = start_time + timeout max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 + last_exception = None while True: - connecting_count += 1 try: if self.unsent_events: + await self._open(timeout_time) + remaining_time = timeout_time - time.time() + if remaining_time < 0.0: + if last_exception: + error = last_exception + else: + 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: - EventHubProducer._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 (errors.MessageAccepted, - errors.MessageAlreadySettled, - errors.MessageModified, - errors.MessageRejected, - errors.MessageReleased, - errors.MessageContentTooLarge) as msg_error: - raise EventDataError(str(msg_error), msg_error) - except errors.MessageException as failed: - log.error("Send event data error (%r)", failed) - error = EventDataSendError(str(failed), failed) - await self.close(exception=error) - raise error - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubProducer disconnected due to token error. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubProducer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - await self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubProducer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer detached. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer connection lost. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubProducer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - await self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer timed out sending event data. Attempting reconnect.") - await self._reconnect() - else: - log.info("EventHubProducer timed out. Shutting down.") - await self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except Exception as e: - log.info("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Send failed: {}".format(e), e) - await self.close(exception=error) - raise error - - def _check_closed(self): - if self.error: - raise EventHubError("This producer has been closed. Please create a new producer to send event data.", - self.error) + except Exception as exception: + last_exception = await self._handle_exception(exception, retry_count, max_retries, timeout_time) + retry_count += 1 def _on_outcome(self, outcome, condition): """ @@ -289,20 +155,8 @@ def _on_outcome(self, outcome, condition): self._outcome = outcome self._condition = condition - @staticmethod - def _error(outcome, condition): - if outcome != constants.MessageSendResult.Ok: - raise condition - - @staticmethod - def _set_partition_key(event_datas, partition_key): - ed_iter = iter(event_datas) - for ed in ed_iter: - ed._set_partition_key(partition_key) - yield ed - - async def send(self, event_data, partition_key=None): - # type:(Union[EventData, Union[List[EventData], Iterator[EventData], Generator[EventData]]], Union[str, bytes]) -> 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 received or operation times out. @@ -312,6 +166,9 @@ async def send(self, event_data, partition_key=None): :param partition_key: With the given partition_key, event data will land to a particular partition of the Event Hub decided by the service. :type partition_key: str + :param 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. + :type timeout:float :raises: ~azure.eventhub.AuthenticationError, ~azure.eventhub.ConnectError, ~azure.eventhub.ConnectionLostError, ~azure.eventhub.EventDataError, ~azure.eventhub.EventDataSendError, ~azure.eventhub.EventHubError :return: None @@ -326,21 +183,24 @@ async def send(self, event_data, partition_key=None): :caption: Sends an event data and blocks until acknowledgement is received or operation times out. """ + partition_key = kwargs.get("partition_key", None) + timeout = kwargs.get("timeout", None) + self._check_closed() if isinstance(event_data, EventData): if partition_key: event_data._set_partition_key(partition_key) wrapper_event_data = event_data else: - event_data_with_pk = self._set_partition_key(event_data, partition_key) + event_data_with_pk = _set_partition_key(event_data, partition_key) wrapper_event_data = _BatchSendEventData( event_data_with_pk, partition_key=partition_key) if partition_key else _BatchSendEventData(event_data) wrapper_event_data.message.on_send_complete = self._on_outcome self.unsent_events = [wrapper_event_data.message] - await self._send_event_data() + await self._send_event_data(timeout) - async def close(self, exception=None): + async def close(self, **kwargs): # type: (Exception) -> None """ Close down the handler. If the handler has already closed, @@ -360,17 +220,5 @@ async def close(self, exception=None): :caption: Close down the handler. """ - self.running = False - if self.error: - return - if isinstance(exception, errors.LinkRedirect): - self.redirected = exception - elif isinstance(exception, EventHubError): - self.error = exception - elif isinstance(exception, (errors.LinkDetach, errors.ConnectionClose)): - self.error = ConnectError(str(exception), exception) - elif exception: - self.error = EventHubError(str(exception)) - else: - self.error = EventHubError("This send handler is now closed.") - await self._handler.close_async() \ No newline at end of file + exception = kwargs.get("exception", None) + await super(EventHubProducer, self).close(exception) diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py index 308aa2000a6d..deda0ddc01fb 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py @@ -24,10 +24,11 @@ from azure.eventhub.producer import EventHubProducer from azure.eventhub.consumer import EventHubConsumer from azure.eventhub.common import parse_sas_token, EventPosition -from azure.eventhub.error import ConnectError +from azure.eventhub.error import ConnectError, EventHubError from .client_abstract import EventHubClientAbstract from .common import EventHubSASTokenCredential, EventHubSharedKeyCredential - +from ._connection_manager import get_connection_manager +from .error import _handle_exception log = logging.getLogger(__name__) @@ -47,6 +48,16 @@ class EventHubClient(EventHubClientAbstract): """ + def __init__(self, host, event_hub_path, credential, **kwargs): + super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs) + self._conn_manager = get_connection_manager(**kwargs) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + def _create_auth(self, username=None, password=None): """ Create an ~uamqp.authentication.SASTokenAuth instance to authenticate @@ -92,17 +103,21 @@ def _create_auth(self, username=None, password=None): get_jwt_token, http_proxy=http_proxy, transport_type=transport_type) + def _handle_exception(self, exception, retry_count, max_retries): + _handle_exception(exception, retry_count, max_retries, self) + + def _close_connection(self): + self._conn_manager.reset_connection_if_broken() + def _management_request(self, mgmt_msg, op_type): - alt_creds = { - "username": self._auth_config.get("iot_username"), - "password": self._auth_config.get("iot_password")} - connect_count = 0 - while True: - connect_count += 1 - mgmt_auth = self._create_auth(**alt_creds) - mgmt_client = uamqp.AMQPClient(self.mgmt_target, auth=mgmt_auth, debug=self.config.network_tracing) + max_retries = self.config.max_retries + retry_count = 0 + while retry_count <= self.config.max_retries: + mgmt_auth = self._create_auth() + mgmt_client = uamqp.AMQPClient(self.mgmt_target) try: - mgmt_client.open() + conn = self._conn_manager.get_connection(self.host, mgmt_auth) + mgmt_client.open(connection=conn) response = mgmt_client.mgmt_request( mgmt_msg, constants.READ_OPERATION, @@ -110,15 +125,9 @@ def _management_request(self, mgmt_msg, op_type): status_code_field=b'status-code', description_fields=b'status-description') return response - except (errors.AMQPConnectionError, errors.TokenAuthFailure, compat.TimeoutException) as failure: - if connect_count >= self.config.max_retries: - err = ConnectError( - "Can not connect to EventHubs or get management info from the service. " - "Please make sure the connection string or token is correct and retry. " - "Besides, this method doesn't work if you use an IoT connection string.", - failure - ) - raise err + except Exception as exception: + self._handle_exception(exception, retry_count, max_retries) + retry_count += 1 finally: mgmt_client.close() @@ -191,8 +200,7 @@ def get_partition_properties(self, partition): return output def create_consumer( - self, consumer_group, partition_id, event_position, - owner_level=None, operation=None, prefetch=None, + self, consumer_group, partition_id, event_position, **kwargs ): # type: (str, str, EventPosition, int, str, int) -> EventHubConsumer """ @@ -224,8 +232,11 @@ def create_consumer( :caption: Add a consumer to the client for a particular consumer group and partition. """ - prefetch = self.config.prefetch if prefetch is None else prefetch + owner_level = kwargs.get("owner_level", None) + operation = kwargs.get("operation", None) + prefetch = kwargs.get("prefetch", None) + prefetch = prefetch or self.config.prefetch path = self.address.path + operation if operation else self.address.path source_url = "amqps://{}{}/ConsumerGroups/{}/Partitions/{}".format( self.address.hostname, path, consumer_group, partition_id) @@ -234,7 +245,7 @@ def create_consumer( prefetch=prefetch) return handler - def create_producer(self, partition_id=None, operation=None, send_timeout=None): + def create_producer(self, **kwargs): # type: (str, str, float) -> EventHubProducer """ Create an producer to send EventData object to an EventHub. @@ -260,6 +271,10 @@ def create_producer(self, partition_id=None, operation=None, send_timeout=None): :caption: Add a producer to the client to send EventData. """ + partition_id = kwargs.get("partition_id", None) + operation = kwargs.get("operation", None) + send_timeout = kwargs.get("send_timeout", None) + target = "amqps://{}{}".format(self.address.hostname, self.address.path) if operation: target = target + operation @@ -268,3 +283,6 @@ def create_producer(self, partition_id=None, operation=None, send_timeout=None): handler = EventHubProducer( self, target, partition=partition_id, send_timeout=send_timeout) return handler + + def close(self): + self._conn_manager.close_connection() diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py index 38e2afde2615..7f6afb51c7fc 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py @@ -219,7 +219,7 @@ def _process_redirect_uri(self, redirect): self.mgmt_target = redirect_uri @classmethod - def from_connection_string(cls, conn_str, event_hub_path=None, **kwargs): + def from_connection_string(cls, conn_str, **kwargs): """Create an EventHubClient from an EventHub/IotHub connection string. :param conn_str: The connection string of an eventhub or IoT hub @@ -266,6 +266,7 @@ def from_connection_string(cls, conn_str, event_hub_path=None, **kwargs): :caption: Create an EventHubClient from a connection string. """ + event_hub_path = kwargs.get("event_hub_path", None) is_iot_conn_str = conn_str.lstrip().lower().startswith("hostname") if not is_iot_conn_str: address, policy, key, entity = _parse_conn_str(conn_str) @@ -281,12 +282,10 @@ def from_connection_string(cls, conn_str, event_hub_path=None, **kwargs): @abstractmethod def create_consumer( - self, consumer_group, partition_id, event_position, owner_level=None, - operation=None, - prefetch=None, + self, consumer_group, partition_id, event_position, **kwargs ): pass @abstractmethod - def create_producer(self, partition_id=None, operation=None, send_timeout=None): + def create_producer(self, **kwargs): pass diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/common.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/common.py index 5a6702a60324..8faded746a74 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/common.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/common.py @@ -8,6 +8,7 @@ import calendar import json import six +from enum import Enum from uamqp import BatchMessage, Message, types from uamqp.message import MessageHeader, MessageProperties @@ -50,7 +51,7 @@ class EventData(object): PROP_TIMESTAMP = b"x-opt-enqueued-time" PROP_DEVICE_ID = b"iothub-connection-device-id" - def __init__(self, body=None, to_device=None, message=None): + def __init__(self, **kwargs): """ Initialize EventData. @@ -63,6 +64,10 @@ def __init__(self, body=None, to_device=None, message=None): :param message: The received message. :type message: ~uamqp.message.Message """ + body = kwargs.get("body", None) + to_device = kwargs.get("to_device", None) + message = kwargs.get("message", None) + self._partition_key = types.AMQPSymbol(EventData.PROP_PARTITION_KEY) self._annotations = {} self._app_properties = {} @@ -205,7 +210,7 @@ def body(self): except TypeError: raise ValueError("Message data empty.") - def body_as_str(self, encoding='UTF-8'): + def body_as_str(self, **kwargs): """ The body of the event data as a string if the data is of a compatible type. @@ -214,6 +219,7 @@ def body_as_str(self, encoding='UTF-8'): Default is 'UTF-8' :rtype: str or unicode """ + encoding = kwargs.get("encoding", 'UTF-8') data = self.body try: return "".join(b.decode(encoding) for b in data) @@ -226,7 +232,7 @@ def body_as_str(self, encoding='UTF-8'): except Exception as e: raise TypeError("Message data is not compatible with string type: {}".format(e)) - def body_as_json(self, encoding='UTF-8'): + def body_as_json(self, **kwargs): """ The body of the event loaded as a JSON object is the data is compatible. @@ -234,6 +240,7 @@ def body_as_json(self, encoding='UTF-8'): Default is 'UTF-8' :rtype: dict """ + encoding = kwargs.get("encoding", 'UTF-8') data_str = self.body_as_str(encoding=encoding) try: return json.loads(data_str) @@ -279,7 +286,7 @@ class EventPosition(object): >>> event_pos = EventPosition(1506968696002) """ - def __init__(self, value, inclusive=False): + def __init__(self, value, **kwargs): """ Initialize EventPosition. @@ -288,6 +295,7 @@ def __init__(self, value, inclusive=False): :param inclusive: Whether to include the supplied value as the start point. :type inclusive: bool """ + inclusive = kwargs.get("inclusive", False) self.value = value if value is not None else "-1" self.inclusive = inclusive diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/configuration.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/configuration.py index 27eb649628ec..58563bdba0e0 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/configuration.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/configuration.py @@ -8,7 +8,11 @@ class _Configuration(object): def __init__(self, **kwargs): self.user_agent = kwargs.get("user_agent") - self.max_retries = kwargs.get("max_retries", 3) + self.retry_total = kwargs.pop('retry_total', 3) + self.max_retries = self.retry_total or kwargs.get("max_retries", 3) + self.backoff_factor = kwargs.pop('retry_backoff_factor', 0.8) + self.backoff_max = kwargs.pop('retry_backoff_max', 120) + self.network_tracing = kwargs.get("network_tracing", False) self.http_proxy = kwargs.get("http_proxy") self.transport_type = TransportType.AmqpOverWebsocket if self.http_proxy \ diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py index 856c77d6fb65..855e6cf97d53 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py @@ -10,35 +10,34 @@ from typing import List from uamqp import types, errors -from uamqp import compat from uamqp import ReceiveClient, Source from azure.eventhub.common import EventData, EventPosition -from azure.eventhub.error import EventHubError, AuthenticationError, ConnectError, ConnectionLostError, _error_handler - +from azure.eventhub.error import _error_handler, EventHubError +from ._consumer_producer_mixin import ConsumerProducerMixin log = logging.getLogger(__name__) -class EventHubConsumer(object): +class EventHubConsumer(ConsumerProducerMixin): """ A consumer responsible for reading EventData from a specific Event Hub - partition and as a member of a specific consumer group. + partition and as a member of a specific consumer group. A consumer may be exclusive, which asserts ownership over the partition for the consumer - group to ensure that only one consumer from that group is reading the from the partition. - These exclusive consumers are sometimes referred to as "Epoch Consumers." + group to ensure that only one consumer from that group is reading the from the partition. + These exclusive consumers are sometimes referred to as "Epoch Consumers." A consumer may also be non-exclusive, allowing multiple consumers from the same consumer - group to be actively reading events from the partition. These non-exclusive consumers are - sometimes referred to as "Non-Epoch Consumers." + group to be actively reading events from the partition. These non-exclusive consumers are + sometimes referred to as "Non-Epoch Consumers." """ timeout = 0 _epoch = b'com.microsoft:epoch' + _timeout = b'com.microsoft:timeout' - def __init__(self, client, source, event_position=None, prefetch=300, owner_level=None, - keep_alive=None, auto_reconnect=True): + def __init__(self, client, source, **kwargs): """ Instantiate a consumer. EventHubConsumer should be instantiated by calling the `create_consumer` method in EventHubClient. @@ -54,6 +53,13 @@ def __init__(self, client, source, event_position=None, prefetch=300, owner_leve 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 self.source = source @@ -65,118 +71,62 @@ def __init__(self, client, source, event_position=None, prefetch=300, owner_leve self.auto_reconnect = auto_reconnect self.retry_policy = errors.ErrorPolicy(max_retries=self.client.config.max_retries, on_error=_error_handler) self.reconnect_backoff = 1 - self.properties = None + self._link_properties = {} self.redirected = None self.error = None partition = self.source.split('/')[-1] - self.name = "EHReceiver-{}-partition{}".format(uuid.uuid4(), partition) - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) # pylint: disable=protected-access + self.name = "EHConsumer-{}-partition{}".format(uuid.uuid4(), partition) if owner_level: - self.properties = {types.AMQPSymbol(self._epoch): types.AMQPLong(int(owner_level))} - self._handler = ReceiveClient( - source, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent)) # pylint: disable=protected-access - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.close(exc_val) + 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): return self def __next__(self): - self._open() max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 while True: - connecting_count += 1 try: + self._open() if not self.messages_iter: self.messages_iter = self._handler.receive_messages_iter() message = next(self.messages_iter) event_data = EventData(message=message) self.offset = EventPosition(event_data.offset, inclusive=False) return event_data - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubConsumer disconnected due to token error. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry and self.auto_reconnect: - log.info("EventHubConsumer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer connection lost. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer timed out receiving event data. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer timed out. Shutting down.") - self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except StopIteration: - raise - except KeyboardInterrupt: - log.info("EventHubConsumer stops due to keyboard interrupt") - self.close() - raise - except Exception as e: - log.error("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Receive failed: {}".format(e), e) - self.close(exception=error) - raise error - - def _check_closed(self): - if self.error: - raise EventHubError("This consumer has been closed. Please create a new consumer to receive event data.", - self.error) + except Exception as exception: + self._handle_exception(exception, retry_count, max_retries) + retry_count += 1 + + def _create_handler(self): + alt_creds = { + "username": self.client._auth_config.get("iot_username"), + "password": self.client._auth_config.get("iot_password")} + source = Source(self.source) + if self.offset is not None: + source.set_filter(self.offset._selector()) + self._handler = ReceiveClient( + source, + auth=self.client.get_auth(**alt_creds), + debug=self.client.config.network_tracing, + prefetch=self.prefetch, + link_properties=self._link_properties, + timeout=self.timeout, + error_policy=self.retry_policy, + keep_alive_interval=self.keep_alive, + client_name=self.name, + properties=self.client._create_properties( + self.client.config.user_agent)) # pylint: disable=protected-access + self.messages_iter = None def _redirect(self, redirect): - self.redirected = redirect - self.running = False self.messages_iter = None - self._open() + super(EventHubConsumer, self)._redirect(redirect) - def _open(self): + def _open(self, timeout_time=None): """ Open the EventHubConsumer using the supplied connection. If the handler has previously been redirected, the redirect @@ -184,129 +134,10 @@ def _open(self): """ # pylint: disable=protected-access - self._check_closed() - if self.redirected: + if not self.running and self.redirected: self.client._process_redirect_uri(self.redirected) self.source = self.redirected.address - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) - - alt_creds = { - "username": self.client._auth_config.get("iot_username"), - "password":self.client._auth_config.get("iot_password")} - self._handler = ReceiveClient( - source, - auth=self.client.get_auth(**alt_creds), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent)) # pylint: disable=protected-access - if not self.running: - self._connect() - self.running = True - - def _connect(self): - connected = self._build_connection() - if not connected: - time.sleep(self.reconnect_backoff) - while not self._build_connection(is_reconnect=True): - time.sleep(self.reconnect_backoff) - - def _build_connection(self, is_reconnect=False): - """ - - :param is_reconnect: True - trying to reconnect after fail to connect or a connection is lost. - False - the 1st time to connect - :return: True - connected. False - not connected - """ - # pylint: disable=protected-access - if is_reconnect: - alt_creds = { - "username": self.client._auth_config.get("iot_username"), - "password": self.client._auth_config.get("iot_password")} - self._handler.close() - source = Source(self.source) - if self.offset is not None: - source.set_filter(self.offset._selector()) - self._handler = ReceiveClient( - source, - auth=self.client.get_auth(**alt_creds), - debug=self.client.config.network_tracing, - prefetch=self.prefetch, - link_properties=self.properties, - timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties( - self.client.config.user_agent)) # pylint: disable=protected-access - self.messages_iter = None - try: - self._handler.open() - while not self._handler.client_ready(): - time.sleep(0.05) - return True - except errors.AuthenticationException as shutdown: - if is_reconnect: - log.info("EventHubConsumer couldn't authenticate. Shutting down. (%r)", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubConsumer couldn't authenticate. Attempting reconnect.") - return False - except errors.LinkRedirect as redirect: - self._redirect(redirect) - return True - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubConsumer detached. Attempting reconnect.") - return False - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if is_reconnect: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubConsumer detached. Attempting reconnect.") - return False - except errors.AMQPConnectionError as shutdown: - if is_reconnect: - log.info("EventHubConsumer connection error (%r). Shutting down.", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubConsumer couldn't authenticate. Attempting reconnect.") - return False - except compat.TimeoutException as shutdown: - if is_reconnect: - log.info("EventHubConsumer authentication timed out. Shutting down.") - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubConsumer authentication timed out. Attempting reconnect.") - return False - except Exception as e: - log.error("Unexpected error occurred when building connection (%r). Shutting down.", e) - error = EventHubError("Unexpected error occurred when building connection", e) - self.close(exception=error) - raise error - - def _reconnect(self): - return self._build_connection(is_reconnect=True) + super(EventHubConsumer, self)._open(timeout_time) @property def queue_size(self): @@ -321,7 +152,7 @@ def queue_size(self): return self._handler._received_messages.qsize() return 0 - def receive(self, max_batch_size=None, timeout=None): + def receive(self, **kwargs): # type:(int, float) -> List[EventData] """ Receive events from the EventHub. @@ -347,82 +178,46 @@ def receive(self, max_batch_size=None, timeout=None): :caption: Receive events from the EventHub. """ - self._check_closed() - self._open() + 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 + if not timeout: + timeout = 100_000 # timeout None or 0 mean no timeout. 100000 seconds is equivalent to no timeout data_batch = [] # type: List[EventData] + start_time = time.time() + timeout_time = start_time + timeout max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 + last_exception = None while True: - connecting_count += 1 try: - timeout_ms = 1000 * timeout if timeout else 0 + self._open(timeout_time) + remaining_time = timeout_time - time.time() + if remaining_time <= 0.0: + if last_exception: + log.info("%r receive operation timed out. (%r)", self.name, last_exception) + raise last_exception + return data_batch + remaining_time_ms = 1000 * remaining_time message_batch = self._handler.receive_message_batch( max_batch_size=max_batch_size - (len(data_batch) if data_batch else 0), - timeout=timeout_ms) + timeout=remaining_time_ms) for message in message_batch: event_data = EventData(message=message) self.offset = EventPosition(event_data.offset) data_batch.append(event_data) return data_batch - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubConsumer disconnected due to token error. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry and self.auto_reconnect: - log.info("EventHubConsumer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer connection lost. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubConsumer timed out receiving event data. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubConsumer timed out. Shutting down.") - self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except KeyboardInterrupt: - log.info("EventHubConsumer stops due to keyboard interrupt") - self.close() + except EventHubError: raise - except Exception as e: - log.error("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Receive failed: {}".format(e), e) - self.close(exception=error) - raise error + except Exception as exception: + last_exception = self._handle_exception(exception, retry_count, max_retries, timeout_time) + retry_count += 1 - def close(self, exception=None): + def close(self, **kwargs): # type:(Exception) -> None """ Close down the handler. If the handler has already closed, @@ -442,20 +237,10 @@ def close(self, exception=None): :caption: Close down the handler. """ + exception = kwargs.get("exception", None) if self.messages_iter: self.messages_iter.close() self.messages_iter = None - self.running = False - if self.error: - return - if isinstance(exception, errors.LinkRedirect): - self.redirected = exception - elif isinstance(exception, EventHubError): - self.error = exception - elif exception: - self.error = EventHubError(str(exception)) - else: - self.error = EventHubError("This receive handler is now closed.") - self._handler.close() + super(EventHubConsumer, self).close(exception) next = __next__ # for python2.7 diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/error.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/error.py index 6932daa7cc0f..0fb6933e3015 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/error.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/error.py @@ -3,8 +3,10 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- import six +import time +import logging -from uamqp import constants, errors +from uamqp import constants, errors, compat _NO_RETRY_ERRORS = ( @@ -15,6 +17,8 @@ b"com.microsoft:argument-error" ) +log = logging.getLogger(__name__) + def _error_handler(error): """ @@ -53,7 +57,8 @@ class EventHubError(Exception): :vartype details: dict[str, str] """ - def __init__(self, message, details=None): + def __init__(self, message, **kwargs): + details = kwargs.get("details", None) self.error = None self.message = message self.details = details @@ -128,3 +133,108 @@ class EventDataSendError(EventHubError): """ pass + + +class OperationTimeoutError(EventHubError): + """Operation times out + + """ + pass + + +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 + + +def _handle_exception(exception, retry_count, max_retries, closable, timeout_time=None): + try: + name = closable.name + except AttributeError: + name = closable.container_id + if isinstance(exception, KeyboardInterrupt): + log.info("%r stops due to keyboard interrupt", name) + closable.close() + raise + elif isinstance(exception, EventHubError): + closable.close() + raise + elif isinstance(exception, ( + errors.MessageAccepted, + errors.MessageAlreadySettled, + errors.MessageModified, + errors.MessageRejected, + errors.MessageReleased, + errors.MessageContentTooLarge) + ): + log.info("%r Event data error (%r)", name, exception) + error = EventDataError(str(exception), exception) + raise error + elif isinstance(exception, errors.MessageException): + log.info("%r Event data send error (%r)", name, exception) + error = EventDataSendError(str(exception), exception) + raise error + elif retry_count >= max_retries: + error = _create_eventhub_exception(exception) + log.info("%r has exhausted retry. Exception still occurs (%r)", name, exception) + raise error + else: + if isinstance(exception, errors.AuthenticationException): + if hasattr(closable, "_close_connection"): + closable._close_connection() + elif isinstance(exception, errors.LinkRedirect): + log.info("%r link redirect received. Redirecting...", name) + redirect = exception + if hasattr(closable, "_redirect"): + closable._redirect(redirect) + elif isinstance(exception, errors.LinkDetach): + if hasattr(closable, "_close_handler"): + closable._close_handler() + elif isinstance(exception, errors.ConnectionClose): + if hasattr(closable, "_close_connection"): + closable._close_connection() + elif isinstance(exception, errors.MessageHandlerError): + if hasattr(closable, "_close_handler"): + closable._close_handler() + elif isinstance(exception, errors.AMQPConnectionError): + if hasattr(closable, "_close_connection"): + closable._close_connection() + elif isinstance(exception, compat.TimeoutException): + pass # Timeout doesn't need to recreate link or connection to retry + else: + if hasattr(closable, "_close_connection"): + closable._close_connection() + # start processing retry delay + try: + backoff_factor = closable.client.config.backoff_factor + backoff_max = closable.client.config.backoff_max + except AttributeError: + backoff_factor = closable.config.backoff_factor + backoff_max = closable.config.backoff_max + backoff = backoff_factor * 2 ** retry_count + if backoff <= backoff_max and (timeout_time is None or time.time() + backoff <= timeout_time): + time.sleep(backoff) + log.info("%r has an exception (%r). Retrying...", format(name), exception) + return _create_eventhub_exception(exception) + else: + error = _create_eventhub_exception(exception) + log.info("%r operation has timed out. Last exception before timeout is (%r)", name, error) + raise error + # end of processing retry delay diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/producer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/producer.py index 3f95b7be08c3..465fbf45d9ef 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/producer.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/producer.py @@ -7,29 +7,42 @@ import uuid import logging import time -from typing import Iterator, Generator, List, Union +from typing import Iterable, Union -from uamqp import constants, errors +from uamqp import types, constants, errors from uamqp import compat from uamqp import SendClient from azure.eventhub.common import EventData, _BatchSendEventData -from azure.eventhub.error import EventHubError, ConnectError, \ - AuthenticationError, EventDataError, EventDataSendError, ConnectionLostError, _error_handler +from azure.eventhub.error import OperationTimeoutError, _error_handler +from ._consumer_producer_mixin import ConsumerProducerMixin log = logging.getLogger(__name__) -class EventHubProducer(object): +def _error(outcome, condition): + if outcome != constants.MessageSendResult.Ok: + raise condition + + +def _set_partition_key(event_datas, partition_key): + ed_iter = iter(event_datas) + for ed in ed_iter: + ed._set_partition_key(partition_key) + yield ed + + +class EventHubProducer(ConsumerProducerMixin): """ A producer responsible for transmitting EventData to a specific Event Hub, - grouped together in batches. Depending on the options specified at creation, the producer may - be created to allow event data to be automatically routed to an available partition or specific - to a partition. + grouped together in batches. Depending on the options specified at creation, the producer may + be created to allow event data to be automatically routed to an available partition or specific + to a partition. """ + _timeout = b'com.microsoft:timeout' - def __init__(self, client, target, partition=None, send_timeout=60, keep_alive=None, auto_reconnect=True): + def __init__(self, client, target, **kwargs): """ Instantiate an EventHubProducer. EventHubProducer should be instantiated by calling the `create_producer` method in EventHubClient. @@ -51,6 +64,12 @@ def __init__(self, client, target, partition=None, send_timeout=60, keep_alive=N 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 self.target = target @@ -67,6 +86,12 @@ def __init__(self, client, target, partition=None, send_timeout=60, keep_alive=N if partition: self.target += "/Partitions/" + partition self.name += "-partition{}".format(partition) + 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( self.target, auth=self.client.get_auth(), @@ -75,17 +100,10 @@ def __init__(self, client, target, partition=None, send_timeout=60, keep_alive=N 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 - self._outcome = None - self._condition = None - - def __enter__(self): - return self - def __exit__(self, exc_type, exc_val, exc_tb): - self.close(exc_val) - - def _open(self): + def _open(self, timeout_time=None): """ Open the EventHubProducer using the supplied connection. If the handler has previously been redirected, the redirect @@ -93,193 +111,44 @@ def _open(self): """ # pylint: disable=protected-access - self._check_closed() - if self.redirected: + if not self.running and self.redirected: + self.client._process_redirect_uri(self.redirected) self.target = self.redirected.address - self._handler = SendClient( - self.target, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - msg_timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent)) - if not self.running: - self._connect() - self.running = True - - def _connect(self): - connected = self._build_connection() - if not connected: - time.sleep(self.reconnect_backoff) - while not self._build_connection(is_reconnect=True): - time.sleep(self.reconnect_backoff) - - def _build_connection(self, is_reconnect=False): - """ - - :param is_reconnect: True - trying to reconnect after fail to connect or a connection is lost. - False - the 1st time to connect - :return: True - connected. False - not connected - """ - # pylint: disable=protected-access - if is_reconnect: - self._handler.close() - self._handler = SendClient( - self.target, - auth=self.client.get_auth(), - debug=self.client.config.network_tracing, - msg_timeout=self.timeout, - error_policy=self.retry_policy, - keep_alive_interval=self.keep_alive, - client_name=self.name, - properties=self.client._create_properties(self.client.config.user_agent)) - try: - self._handler.open() - while not self._handler.client_ready(): - time.sleep(0.05) - return True - except errors.AuthenticationException as shutdown: - if is_reconnect: - log.info("EventHubProducer couldn't authenticate. Shutting down. (%r)", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubProducer couldn't authenticate. Attempting reconnect.") - return False - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubProducer detached. Attempting reconnect.") - return False - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if is_reconnect: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubProducer detached. Attempting reconnect.") - return False - except errors.AMQPConnectionError as shutdown: - if is_reconnect: - log.info("EventHubProducer connection error (%r). Shutting down.", shutdown) - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubProducer couldn't authenticate. Attempting reconnect.") - return False - except compat.TimeoutException as shutdown: - if is_reconnect: - log.info("EventHubProducer authentication timed out. Shutting down.") - error = AuthenticationError(str(shutdown), shutdown) - self.close(exception=error) - raise error - else: - log.info("EventHubProducer authentication timed out. Attempting reconnect.") - return False - except Exception as e: - log.info("Unexpected error occurred when building connection (%r). Shutting down.", e) - error = EventHubError("Unexpected error occurred when building connection", e) - self.close(exception=error) - raise error - - def _reconnect(self): - return self._build_connection(is_reconnect=True) - - def _send_event_data(self): - self._open() + super(EventHubProducer, self)._open(timeout_time) + + def _send_event_data(self, timeout=None): + timeout = self.client.config.send_timeout if timeout is None else timeout + if not timeout: + timeout = 100_000 # timeout None or 0 mean no timeout. 100000 seconds is equivalent to no timeout + start_time = time.time() + timeout_time = start_time + timeout max_retries = self.client.config.max_retries - connecting_count = 0 + retry_count = 0 + last_exception = None while True: - connecting_count += 1 try: if self.unsent_events: + self._open(timeout_time) + remaining_time = timeout_time - time.time() + if remaining_time <= 0.0: + if last_exception: + error = last_exception + else: + 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: - EventHubProducer._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 (errors.MessageAccepted, - errors.MessageAlreadySettled, - errors.MessageModified, - errors.MessageRejected, - errors.MessageReleased, - errors.MessageContentTooLarge) as msg_error: - raise EventDataError(str(msg_error), msg_error) - except errors.MessageException as failed: - log.error("Send event data error (%r)", failed) - error = EventDataSendError(str(failed), failed) - self.close(exception=error) - raise error - except errors.AuthenticationException as auth_error: - if connecting_count < max_retries: - log.info("EventHubProducer disconnected due to token error. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubProducer authentication failed. Shutting down.") - error = AuthenticationError(str(auth_error), auth_error) - self.close(auth_error) - raise error - except (errors.LinkDetach, errors.ConnectionClose) as shutdown: - if shutdown.action.retry: - log.info("EventHubProducer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(exception=error) - raise error - except errors.MessageHandlerError as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer detached. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubProducer detached. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except errors.AMQPConnectionError as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer connection lost. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubProducer connection lost. Shutting down.") - error = ConnectionLostError(str(shutdown), shutdown) - self.close(error) - raise error - except compat.TimeoutException as shutdown: - if connecting_count < max_retries: - log.info("EventHubProducer timed out sending event data. Attempting reconnect.") - self._reconnect() - else: - log.info("EventHubProducer timed out. Shutting down.") - self.close(shutdown) - raise ConnectionLostError(str(shutdown), shutdown) - except Exception as e: - log.info("Unexpected error occurred (%r). Shutting down.", e) - error = EventHubError("Send failed: {}".format(e), e) - self.close(exception=error) - raise error - - def _check_closed(self): - if self.error: - raise EventHubError("This producer has been closed. Please create a new producer to send event data.", self.error) - - @staticmethod - def _set_partition_key(event_datas, partition_key): - ed_iter = iter(event_datas) - for ed in ed_iter: - ed._set_partition_key(partition_key) - yield ed + except Exception as exception: + last_exception = self._handle_exception(exception, retry_count, max_retries, timeout_time) + retry_count += 1 def _on_outcome(self, outcome, condition): """ @@ -293,13 +162,8 @@ def _on_outcome(self, outcome, condition): self._outcome = outcome self._condition = condition - @staticmethod - def _error(outcome, condition): - if outcome != constants.MessageSendResult.Ok: - raise condition - - def send(self, event_data, partition_key=None): - # type:(Union[EventData, Union[List[EventData], Iterator[EventData], Generator[EventData]]], Union[str, bytes]) -> 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 received or operation times out. @@ -309,6 +173,9 @@ def send(self, event_data, partition_key=None): :param partition_key: With the given partition_key, event data will land to a particular partition of the Event Hub decided by the service. :type partition_key: str + :param 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. + :type timeout:float :raises: ~azure.eventhub.AuthenticationError, ~azure.eventhub.ConnectError, ~azure.eventhub.ConnectionLostError, ~azure.eventhub.EventDataError, ~azure.eventhub.EventDataSendError, ~azure.eventhub.EventHubError @@ -324,21 +191,24 @@ def send(self, event_data, partition_key=None): :caption: Sends an event data and blocks until acknowledgement is received or operation times out. """ + partition_key = kwargs.get("partition_key", None) + timeout = kwargs.get("timeout", None) + self._check_closed() if isinstance(event_data, EventData): if partition_key: event_data._set_partition_key(partition_key) wrapper_event_data = event_data else: - event_data_with_pk = self._set_partition_key(event_data, partition_key) + event_data_with_pk = _set_partition_key(event_data, partition_key) wrapper_event_data = _BatchSendEventData( event_data_with_pk, partition_key=partition_key) if partition_key else _BatchSendEventData(event_data) wrapper_event_data.message.on_send_complete = self._on_outcome self.unsent_events = [wrapper_event_data.message] - self._send_event_data() + self._send_event_data(timeout=timeout) - def close(self, exception=None): + def close(self, **kwargs): # type:(Exception) -> None """ Close down the handler. If the handler has already closed, @@ -358,15 +228,5 @@ def close(self, exception=None): :caption: Close down the handler. """ - self.running = False - if self.error: - return - if isinstance(exception, errors.LinkRedirect): - self.redirected = exception - elif isinstance(exception, EventHubError): - self.error = exception - elif exception: - self.error = EventHubError(str(exception)) - else: - self.error = EventHubError("This send handler is now closed.") - self._handler.close() + exception = kwargs.get("exception", None) + super(EventHubProducer, self).close(exception) diff --git a/sdk/eventhub/azure-eventhubs/tests/test_negative.py b/sdk/eventhub/azure-eventhubs/tests/test_negative.py index 1bf9855c80eb..ac19a01f76c9 100644 --- a/sdk/eventhub/azure-eventhubs/tests/test_negative.py +++ b/sdk/eventhub/azure-eventhubs/tests/test_negative.py @@ -26,7 +26,7 @@ def test_send_with_invalid_hostname(invalid_hostname, connstr_receivers): client = EventHubClient.from_connection_string(invalid_hostname, network_tracing=False) sender = client.create_producer() with pytest.raises(AuthenticationError): - sender._open() + sender.send(EventData("test data")) @pytest.mark.liveTest @@ -34,7 +34,7 @@ def test_receive_with_invalid_hostname_sync(invalid_hostname): client = EventHubClient.from_connection_string(invalid_hostname, network_tracing=False) receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1")) with pytest.raises(AuthenticationError): - receiver._open() + receiver.receive(timeout=3) @pytest.mark.liveTest @@ -43,7 +43,7 @@ def test_send_with_invalid_key(invalid_key, connstr_receivers): client = EventHubClient.from_connection_string(invalid_key, network_tracing=False) sender = client.create_producer() with pytest.raises(AuthenticationError): - sender._open() + sender.send(EventData("test data")) @pytest.mark.liveTest @@ -51,7 +51,7 @@ def test_receive_with_invalid_key_sync(invalid_key): client = EventHubClient.from_connection_string(invalid_key, network_tracing=False) receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1")) with pytest.raises(AuthenticationError): - receiver._open() + receiver.receive(timeout=3) @pytest.mark.liveTest @@ -60,7 +60,7 @@ def test_send_with_invalid_policy(invalid_policy, connstr_receivers): client = EventHubClient.from_connection_string(invalid_policy, network_tracing=False) sender = client.create_producer() with pytest.raises(AuthenticationError): - sender._open() + sender.send(EventData("test data")) @pytest.mark.liveTest @@ -68,7 +68,7 @@ def test_receive_with_invalid_policy_sync(invalid_policy): client = EventHubClient.from_connection_string(invalid_policy, network_tracing=False) receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1")) with pytest.raises(AuthenticationError): - receiver._open() + receiver.receive(timeout=3) @pytest.mark.liveTest @@ -90,7 +90,7 @@ def test_non_existing_entity_sender(connection_str): client = EventHubClient.from_connection_string(connection_str, event_hub_path="nemo", network_tracing=False) sender = client.create_producer(partition_id="1") with pytest.raises(AuthenticationError): - sender._open() + sender.send(EventData("test data")) @pytest.mark.liveTest @@ -98,7 +98,7 @@ def test_non_existing_entity_receiver(connection_str): client = EventHubClient.from_connection_string(connection_str, event_hub_path="nemo", network_tracing=False) receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1")) with pytest.raises(AuthenticationError): - receiver._open() + receiver.receive(timeout=3) @pytest.mark.liveTest @@ -122,7 +122,7 @@ def test_send_to_invalid_partitions(connection_str): sender = client.create_producer(partition_id=p) try: with pytest.raises(ConnectError): - sender._open() + sender.send(EventData("test data")) finally: sender.close()