diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py index fd12b439d3d4..96479a6ef31b 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py @@ -73,7 +73,7 @@ def _open(self, timeout_time=None): # pylint:disable=unused-argument # TODO: to else: alt_creds = {} self._create_handler() - self._handler.open(connection=self.client._conn_manager.get_connection( + self._handler.open(connection=self.client._conn_manager.get_connection( # pylint: disable=protected-access self.client.address.hostname, self.client.get_auth(**alt_creds) )) 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 84853ce1534b..00b136b8bed5 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py @@ -6,6 +6,7 @@ import datetime import functools import asyncio + from typing import Any, List, Dict, Union, TYPE_CHECKING from uamqp import authentication, constants # type: ignore @@ -47,6 +48,7 @@ class EventHubClient(EventHubClientAbstract): def __init__(self, host, event_hub_path, credential, **kwargs): # type:(str, str, Union[EventHubSharedKeyCredential, EventHubSASTokenCredential, TokenCredential], Any) -> None super(EventHubClient, self).__init__(host=host, event_hub_path=event_hub_path, credential=credential, **kwargs) + self._lock = asyncio.Lock() self._conn_manager = get_connection_manager(**kwargs) async def __aenter__(self): @@ -105,10 +107,17 @@ async def _close_connection(self): await self._conn_manager.reset_connection_if_broken() async def _management_request(self, mgmt_msg, op_type): + if self._is_iothub and not self._iothub_redirect_info: + await self._iothub_redirect() + + alt_creds = { + "username": self._auth_config.get("iot_username"), + "password": self._auth_config.get("iot_password") + } max_retries = self.config.max_retries retry_count = 0 while True: - mgmt_auth = self._create_auth() + mgmt_auth = self._create_auth(**alt_creds) mgmt_client = AMQPClientAsync(self.mgmt_target, auth=mgmt_auth, debug=self.config.network_tracing) try: conn = await self._conn_manager.get_connection(self.host, mgmt_auth) @@ -126,6 +135,18 @@ async def _management_request(self, mgmt_msg, op_type): finally: await mgmt_client.close_async() + async def _iothub_redirect(self): + async with self._lock: + if self._is_iothub and not self._iothub_redirect_info: + if not self._redirect_consumer: + self._redirect_consumer = self.create_consumer(consumer_group='$default', + partition_id='0', + event_position=EventPosition('-1'), + operation='/messages/events') + async with self._redirect_consumer: + await self._redirect_consumer._open_with_retry(timeout=self.config.receive_timeout) # pylint: disable=protected-access + self._redirect_consumer = None + async def get_properties(self): # type:() -> Dict[str, Any] """ @@ -139,6 +160,8 @@ async def get_properties(self): :rtype: dict :raises: ~azure.eventhub.ConnectError """ + if self._is_iothub and not self._iothub_redirect_info: + await self._iothub_redirect() mgmt_msg = Message(application_properties={'name': self.eh_name}) response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub') output = {} @@ -178,6 +201,8 @@ async def get_partition_properties(self, partition): :rtype: dict :raises: ~azure.eventhub.ConnectError """ + if self._is_iothub and not self._iothub_redirect_info: + await self._iothub_redirect() mgmt_msg = Message(application_properties={'name': self.eh_name, 'partition': partition}) response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:partition') 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 3747d1af4d9d..49e85fe2e811 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py @@ -110,8 +110,10 @@ async def __anext__(self): def _create_handler(self): alt_creds = { - "username": self.client._auth_config.get("iot_username"), # pylint:disable=protected-access - "password": self.client._auth_config.get("iot_password")} # pylint:disable=protected-access + "username": self.client._auth_config.get("iot_username") if self.redirected else None, # pylint:disable=protected-access + "password": self.client._auth_config.get("iot_password") if self.redirected else None # pylint:disable=protected-access + } + source = Source(self.source) if self.offset is not None: source.set_filter(self.offset._selector()) # pylint:disable=protected-access @@ -134,7 +136,7 @@ async def _redirect(self, redirect): self.messages_iter = None await super(EventHubConsumer, self)._redirect(redirect) - async def _open(self, timeout_time=None): + async def _open(self, timeout_time=None, **kwargs): """ Open the EventHubConsumer using the supplied connection. If the handler has previously been redirected, the redirect @@ -142,11 +144,17 @@ async def _open(self, timeout_time=None): """ # pylint: disable=protected-access + self.redirected = self.redirected or self.client._iothub_redirect_info + if not self.running and self.redirected: self.client._process_redirect_uri(self.redirected) self.source = self.redirected.address await super(EventHubConsumer, self)._open(timeout_time) + @_retry_decorator + async def _open_with_retry(self, timeout_time=None, **kwargs): + return await self._open(timeout_time=timeout_time, **kwargs) + async def _receive(self, timeout_time=None, max_batch_size=None, **kwargs): last_exception = kwargs.get("last_exception") data_batch = kwargs.get("data_batch") @@ -254,4 +262,5 @@ 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() + if self._handler: + await self._handler.close_async() diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py index 600faaf31041..f0e4eca2ac47 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py @@ -7,6 +7,8 @@ import logging import datetime import functools +import threading + from typing import Any, List, Dict, Union, TYPE_CHECKING import uamqp # type: ignore @@ -46,6 +48,7 @@ class EventHubClient(EventHubClientAbstract): def __init__(self, host, event_hub_path, credential, **kwargs): # type:(str, str, Union[EventHubSharedKeyCredential, EventHubSASTokenCredential, TokenCredential], Any) -> None super(EventHubClient, self).__init__(host=host, event_hub_path=event_hub_path, credential=credential, **kwargs) + self._lock = threading.RLock() self._conn_manager = get_connection_manager(**kwargs) def __enter__(self): @@ -106,10 +109,15 @@ 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") + } + max_retries = self.config.max_retries retry_count = 0 while retry_count <= self.config.max_retries: - mgmt_auth = self._create_auth() + mgmt_auth = self._create_auth(**alt_creds) mgmt_client = uamqp.AMQPClient(self.mgmt_target) try: conn = self._conn_manager.get_connection(self.host, mgmt_auth) #pylint:disable=assignment-from-none @@ -127,6 +135,18 @@ def _management_request(self, mgmt_msg, op_type): finally: mgmt_client.close() + def _iothub_redirect(self): + with self._lock: + if self._is_iothub and not self._iothub_redirect_info: + if not self._redirect_consumer: + self._redirect_consumer = self.create_consumer(consumer_group='$default', + partition_id='0', + event_position=EventPosition('-1'), + operation='/messages/events') + with self._redirect_consumer: + self._redirect_consumer._open_with_retry(timeout=self.config.receive_timeout) # pylint: disable=protected-access + self._redirect_consumer = None + def get_properties(self): # type:() -> Dict[str, Any] """ @@ -140,6 +160,8 @@ def get_properties(self): :rtype: dict :raises: ~azure.eventhub.ConnectError """ + if self._is_iothub and not self._iothub_redirect_info: + self._iothub_redirect() mgmt_msg = Message(application_properties={'name': self.eh_name}) response = self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub') output = {} @@ -179,6 +201,8 @@ def get_partition_properties(self, partition): :rtype: dict :raises: ~azure.eventhub.ConnectError """ + if self._is_iothub and not self._iothub_redirect_info: + self._iothub_redirect() mgmt_msg = Message(application_properties={'name': self.eh_name, 'partition': partition}) response = self._management_request(mgmt_msg, op_type=b'com.microsoft:partition') diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py index 1d33091755d0..8b6712ce207e 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py @@ -11,7 +11,8 @@ import functools from abc import abstractmethod from typing import Dict, Union, Any, TYPE_CHECKING -from azure.eventhub import __version__ + +from azure.eventhub import __version__, EventPosition from azure.eventhub.configuration import _Configuration from .common import EventHubSharedKeyCredential, EventHubSASTokenCredential, _Address @@ -153,6 +154,8 @@ def __init__(self, host, event_hub_path, credential, **kwargs): self.get_auth = functools.partial(self._create_auth) self.config = _Configuration(**kwargs) self.debug = self.config.network_tracing + self._is_iothub = False + self._iothub_redirect_info = None log.info("%r: Created the Event Hub client", self.container_id) @@ -173,6 +176,11 @@ def _from_iothub_connection_string(cls, conn_str, **kwargs): 'iot_password': key, 'username': username, 'password': password} + client._is_iothub = True + client._redirect_consumer = client.create_consumer(consumer_group='$default', + partition_id='0', + event_position=EventPosition('-1'), + operation='/messages/events') return client @abstractmethod @@ -213,6 +221,8 @@ def _process_redirect_uri(self, redirect): self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path) self.eh_name = self.address.path.lstrip('/') self.mgmt_target = redirect_uri + if self._is_iothub: + self._iothub_redirect_info = redirect @classmethod def from_connection_string(cls, conn_str, **kwargs): diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py index 44be38386b73..6306c750260a 100644 --- a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py +++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py @@ -106,8 +106,10 @@ def __next__(self): def _create_handler(self): alt_creds = { - "username": self.client._auth_config.get("iot_username"), # pylint:disable=protected-access - "password": self.client._auth_config.get("iot_password")} # pylint:disable=protected-access + "username": self.client._auth_config.get("iot_username") if self.redirected else None, # pylint:disable=protected-access + "password": self.client._auth_config.get("iot_password") if self.redirected else None # pylint:disable=protected-access + } + source = Source(self.source) if self.offset is not None: source.set_filter(self.offset._selector()) # pylint:disable=protected-access @@ -129,7 +131,7 @@ def _redirect(self, redirect): self.messages_iter = None super(EventHubConsumer, self)._redirect(redirect) - def _open(self, timeout_time=None): + def _open(self, timeout_time=None, **kwargs): """ Open the EventHubConsumer using the supplied connection. If the handler has previously been redirected, the redirect @@ -137,11 +139,17 @@ def _open(self, timeout_time=None): """ # pylint: disable=protected-access + self.redirected = self.redirected or self.client._iothub_redirect_info + if not self.running and self.redirected: self.client._process_redirect_uri(self.redirected) self.source = self.redirected.address super(EventHubConsumer, self)._open(timeout_time) + @_retry_decorator + def _open_with_retry(self, timeout_time=None, **kwargs): + return self._open(timeout_time=timeout_time, **kwargs) + def _receive(self, timeout_time=None, max_batch_size=None, **kwargs): last_exception = kwargs.get("last_exception") data_batch = kwargs.get("data_batch") diff --git a/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py b/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py index f581f64584ab..9226a5d62a93 100644 --- a/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py +++ b/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py @@ -4,13 +4,11 @@ # license information. #-------------------------------------------------------------------------- -import os import asyncio import pytest -import time from azure.eventhub.aio import EventHubClient -from azure.eventhub import EventData, EventPosition, EventHubError +from azure.eventhub import EventPosition async def pump(receiver, sleep=None): @@ -18,25 +16,16 @@ async def pump(receiver, sleep=None): if sleep: await asyncio.sleep(sleep) async with receiver: - batch = await receiver.receive(timeout=1) + batch = await receiver.receive(timeout=3) messages += len(batch) return messages -async def get_partitions(iot_connection_str): - client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) - receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), prefetch=1000, operation='/messages/events') - async with receiver: - partitions = await client.get_properties() - return partitions["partition_ids"] - - @pytest.mark.liveTest @pytest.mark.asyncio async def test_iothub_receive_multiple_async(iot_connection_str): - pytest.skip("This will get AuthenticationError. We're investigating...") - partitions = await get_partitions(iot_connection_str) client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partitions = await client.get_partition_ids() receivers = [] for p in partitions: receivers.append(client.create_consumer(consumer_group="$default", partition_id=p, event_position=EventPosition("-1"), prefetch=10, operation='/messages/events')) @@ -44,3 +33,53 @@ async def test_iothub_receive_multiple_async(iot_connection_str): assert isinstance(outputs[0], int) and outputs[0] <= 10 assert isinstance(outputs[1], int) and outputs[1] <= 10 + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_iothub_get_properties_async(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + properties = await client.get_properties() + assert properties["partition_ids"] == ["0", "1", "2", "3"] + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_iothub_get_partition_ids_async(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partitions = await client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"] + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_iothub_get_partition_properties_async(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partition_properties = await client.get_partition_properties("0") + assert partition_properties["id"] == "0" + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_iothub_receive_after_mgmt_ops_async(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partitions = await client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"] + receiver = client.create_consumer(consumer_group="$default", partition_id=partitions[0], event_position=EventPosition("-1"), operation='/messages/events') + async with receiver: + received = await receiver.receive(timeout=5) + assert len(received) == 0 + + +@pytest.mark.liveTest +@pytest.mark.asyncio +async def test_iothub_mgmt_ops_after_receive_async(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events') + async with receiver: + received = await receiver.receive(timeout=5) + assert len(received) == 0 + + partitions = await client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"] + diff --git a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py index ac5787b6b12e..ac7e211bd736 100644 --- a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py +++ b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py @@ -4,23 +4,61 @@ # license information. #-------------------------------------------------------------------------- -import os import pytest -import time -from azure.eventhub import EventData, EventPosition, EventHubClient +from azure.eventhub import EventPosition, EventHubClient @pytest.mark.liveTest def test_iothub_receive_sync(iot_connection_str, device_id): - pytest.skip("current code will cause ErrorCodes.LinkRedirect") client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events') - receiver._open() try: - partitions = client.get_properties() - assert partitions["partition_ids"] == ["0", "1", "2", "3"] received = receiver.receive(timeout=5) assert len(received) == 0 finally: receiver.close() + + +@pytest.mark.liveTest +def test_iothub_get_properties_sync(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + properties = client.get_properties() + assert properties["partition_ids"] == ["0", "1", "2", "3"] + + +@pytest.mark.liveTest +def test_iothub_get_partition_ids_sync(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partitions = client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"] + + +@pytest.mark.liveTest +def test_iothub_get_partition_properties_sync(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partition_properties = client.get_partition_properties("0") + assert partition_properties["id"] == "0" + + +@pytest.mark.liveTest +def test_iothub_receive_after_mgmt_ops_sync(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + partitions = client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"] + receiver = client.create_consumer(consumer_group="$default", partition_id=partitions[0], event_position=EventPosition("-1"), operation='/messages/events') + with receiver: + received = receiver.receive(timeout=5) + assert len(received) == 0 + + +@pytest.mark.liveTest +def test_iothub_mgmt_ops_after_receive_sync(iot_connection_str, device_id): + client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False) + receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events') + with receiver: + received = receiver.receive(timeout=5) + assert len(received) == 0 + + partitions = client.get_partition_ids() + assert partitions == ["0", "1", "2", "3"]