[EventHub] Raise error when primary key is rotated (#36245)

* sync side fix for rotate primary key issue in hub

* sync uamqp transport

* add to async

* remove print/cleanup

* mypy/pylint

* changelog

* fix logging formatting

* no retry on all auth exceptions

* remove import error in uamqp recv, b/c no external dependency

* mypy

Author: swathipil

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed a bug where the consumer waited indefinitely when the primary key was rotated while receiving, rather than raising an authentication error. ([#33926](https://github.com/Azure/azure-sdk-for-python/issues/33926))
+
 ### Other Changes
 
 ## 5.12.1 (2024-06-11)

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

@@ -236,6 +236,11 @@ def receive(self, batch=False, max_batch_size=300, max_wait_time=None):
                     # If optional dependency is not installed, do not retry.
                     if isinstance(exception, ImportError):
                         raise exception
+
+                    # If authentication exception, do not retry.
+                    if isinstance(exception, self._amqp_transport.AUTHENTICATION_EXCEPTION):
+                        raise self._handle_exception(exception, is_consumer=True)
+
                     self._amqp_transport.check_link_stolen(self, exception)
                     # TODO: below block hangs when retry_total > 0
                     # need to remove/refactor, issue #27137

sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_base.py

@@ -30,6 +30,7 @@
             ReceiveClient as uamqp_ReceiveClient,
         )
         from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth
+        from uamqp.errors import AuthenticationException as uamqp_AuthenticationException
 
     except ImportError:
         pass
@@ -48,6 +49,7 @@
     from .._pyamqp.constants import (
         ConnectionState as pyamqp_ConnectionState
     )
+    from .._pyamqp.error import AuthenticationException as pyamqp_AuthenticationException
 
 class AmqpTransport(ABC):   # pylint: disable=too-many-public-methods
     """
@@ -76,6 +78,9 @@ class AmqpTransport(ABC):   # pylint: disable=too-many-public-methods
     USER_AGENT_SYMBOL: Union[uamqp_Types_AMQPSymbol, str, bytes]
     PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_Types_AMQPSymbol, str, bytes]
 
+    # exceptions
+    AUTHENTICATION_EXCEPTION: Union["uamqp_AuthenticationException", "pyamqp_AuthenticationException"]
+
     @staticmethod
     @abstractmethod
     def build_message(**kwargs: Any) -> Union["uamqp_Message", "pyamqp_Message"]:

sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_pyamqp_transport.py

@@ -72,6 +72,9 @@ class PyamqpTransport(AmqpTransport):   # pylint: disable=too-many-public-method
 
     ERROR_CONDITIONS = [condition.value for condition in errors.ErrorCondition]
 
+    # define exceptions
+    AUTHENTICATION_EXCEPTION = errors.AuthenticationException
+
     @staticmethod
     def build_message(**kwargs):
         """

sdk/eventhub/azure-eventhub/azure/eventhub/_transport/_uamqp_transport.py

@@ -122,6 +122,9 @@ class UamqpTransport(AmqpTransport):    # pylint: disable=too-many-public-method
         USER_AGENT_SYMBOL = types.AMQPSymbol("user-agent")
         PROP_PARTITION_KEY_AMQP_SYMBOL = types.AMQPSymbol(PROP_PARTITION_KEY)
 
+        # define exceptions
+        AUTHENTICATION_EXCEPTION = errors.AuthenticationException
+
         @staticmethod
         def build_message(**kwargs):
             """

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_connection_manager_async.py

@@ -62,7 +62,7 @@ def __init__(
             **kwargs: Any
     ) -> None:
         self._loop = kwargs.get("loop")
-        self._lock = Lock(loop=self._loop)
+        self._lock = Lock(loop=self._loop) # pylint: disable=unexpected-keyword-arg
         self._conn: Optional[Union[uamqp_ConnectionAsync, ConnectionAsync]] = None
 
         self._container_id = container_id

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_base_async.py

@@ -79,7 +79,6 @@ class AmqpTransportAsync(ABC):  # pylint: disable=too-many-public-methods
     USER_AGENT_SYMBOL: Union[uamqp_Types_AMQPSymbol, Literal["user-agent"]]
     PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_Types_AMQPSymbol, Literal[b'x-opt-partition-key']]
 
-
     @staticmethod
     @abstractmethod
     def build_message(**kwargs: Any) -> Union["uamqp_Message", "pyamqp_Message"]:

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py

@@ -269,6 +269,9 @@ async def _receive_task(consumer):
                     # If optional dependency is not installed, do not retry.
                     if isinstance(exception, ImportError):
                         raise exception
+                    # If authentication exception, do not retry.
+                    if isinstance(exception, errors.AuthenticationException):
+                        raise await consumer._handle_exception(exception)
                     if (
                         isinstance(exception, errors.AMQPLinkError)
                         and exception.condition == errors.ErrorCondition.LinkStolen  # pylint: disable=no-member

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_uamqp_transport_async.py

@@ -283,6 +283,9 @@ async def receive_messages_async(
                     except asyncio.CancelledError:  # pylint: disable=try-except-raise
                         raise
                     except Exception as exception:  # pylint: disable=broad-except
+                        # If authentication exception, do not retry.
+                        if isinstance(exception, errors.AuthenticationException):
+                            raise await consumer._handle_exception(exception)
                         if (
                             isinstance(exception, errors.LinkDetach)
                             and exception.condition == constants.ErrorCodes.LinkStolen  # pylint: disable=no-member