Eh named key (#18292)

Author: Rakshith Bhyravabhotla

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -1,6 +1,13 @@
 # Release History
 
-## 5.4.1 (Unreleased)
+## 5.5.0 (Unreleased)
+
+**New Features**
+- Added support for using `azure.core.credentials.AzureNamedKeyCredential` as credential for authenticating producer and consumer clients.
+
+**Notes**
+
+- Updated azure-core dependency to 1.14.0.
 
 **Bug Fixes**
 

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

@@ -20,8 +20,9 @@
 
 from uamqp import AMQPClient, Message, authentication, constants, errors, compat, utils
 import six
+from azure.core.credentials import AccessToken, AzureSasCredential, AzureNamedKeyCredential
 from azure.core.utils import parse_connection_string as core_parse_connection_string
-from azure.core.credentials import AccessToken, AzureSasCredential
+
 
 from .exceptions import _handle_exception, ClientClosedError, ConnectError
 from ._configuration import Configuration
@@ -173,6 +174,25 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
             raise ValueError("No token scope provided.")
         return _generate_sas_token(scopes[0], self.policy, self.key)
 
+class EventhubAzureNamedKeyTokenCredential(object):
+    """The named key credential used for authentication.
+
+    :param credential: The AzureNamedKeyCredential that should be used.
+    :type credential: ~azure.core.credentials.AzureNamedKeyCredential
+    """
+
+    def __init__(self, azure_named_key_credential):
+        # type: (AzureNamedKeyCredential) -> None
+        self._credential = azure_named_key_credential
+        self.token_type = b"servicebus.windows.net:sastoken"
+
+    def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (str, Any) -> _AccessToken
+        if not scopes:
+            raise ValueError("No token scope provided.")
+        name, key = self._credential.named_key
+        return _generate_sas_token(scopes[0], name, key)
+
 
 class EventHubSASTokenCredential(object):
     """The shared access token credential used for authentication.
@@ -197,7 +217,7 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         """
         return AccessToken(self.token, self.expiry)
 
-class AzureSasTokenCredential(object):
+class EventhubAzureSasTokenCredential(object):
     """The shared access token credential used for authentication
     when AzureSasCredential is provided.
 
@@ -226,15 +246,17 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
 
 class ClientBase(object):  # pylint:disable=too-many-instance-attributes
     def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwargs):
-        # type: (str, str, Union[AzureSasCredential, TokenCredential], Any) -> None
+        # type: (str, str, Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential], Any) -> None
         self.eventhub_name = eventhub_name
         if not eventhub_name:
             raise ValueError("The eventhub name can not be None or empty.")
         path = "/" + eventhub_name if eventhub_name else ""
         self._address = _Address(hostname=fully_qualified_namespace, path=path)
         self._container_id = CONTAINER_PREFIX + str(uuid.uuid4())[:8]
         if isinstance(credential, AzureSasCredential):
-            self._credential = AzureSasTokenCredential(credential)
+            self._credential = EventhubAzureSasTokenCredential(credential)
+        elif isinstance(credential, AzureNamedKeyCredential):
+            self._credential = EventhubAzureNamedKeyTokenCredential(credential) # type: ignore
         else:
             self._credential = credential #type: ignore
         self._keep_alive = kwargs.get("keep_alive", 30)

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py

@@ -15,7 +15,7 @@
 
 if TYPE_CHECKING:
     import datetime
-    from azure.core.credentials import TokenCredential, AzureSasCredential
+    from azure.core.credentials import TokenCredential, AzureSasCredential, AzureNamedKeyCredential
     from typing import (  # pylint: disable=ungrouped-imports
         Any,
         Union,
@@ -59,6 +59,7 @@ class EventHubConsumerClient(ClientBase):
      :class:`EventHubSharedKeyCredential<azure.eventhub.EventHubSharedKeyCredential>`, or credential objects generated
      by the azure-identity library and objects that implement the `get_token(self, *scopes)` method.
     :type credential: ~azure.core.credentials.TokenCredential or ~azure.core.credentials.AzureSasCredential
+     or ~azure.core.credentials.AzureNamedKeyCredential
     :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
     :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service.
      The default value is 60 seconds. If set to 0, no timeout will be enforced from the client.
@@ -129,7 +130,7 @@ def __init__(
         fully_qualified_namespace,  # type: str
         eventhub_name,  # type: str
         consumer_group,  # type: str
-        credential,  # type: Union[AzureSasCredential, TokenCredential]
+        credential,  # type: Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential]
         **kwargs  # type: Any
     ):
         # type: (...) -> None

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

@@ -16,7 +16,7 @@
 from ._common import EventDataBatch, EventData
 
 if TYPE_CHECKING:
-    from azure.core.credentials import TokenCredential, AzureSasCredential
+    from azure.core.credentials import TokenCredential, AzureSasCredential, AzureNamedKeyCredential
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,6 +33,7 @@ class EventHubProducerClient(ClientBase):
      :class:`EventHubSharedKeyCredential<azure.eventhub.EventHubSharedKeyCredential>`, or credential objects generated
      by the azure-identity library and objects that implement the `get_token(self, *scopes)` method.
     :type credential: ~azure.core.credentials.TokenCredential or ~azure.core.credentials.AzureSasCredential
+     or ~azure.core.credentials.AzureNamedKeyCredential
     :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
     :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service.
      The default value is 60 seconds. If set to 0, no timeout will be enforced from the client.
@@ -74,7 +75,7 @@ def __init__(
         self,
         fully_qualified_namespace,  # type: str
         eventhub_name,  # type: str
-        credential,  # type: Union[AzureSasCredential, TokenCredential]
+        credential,  # type: Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential]
         **kwargs  # type: Any
     ):
         # type:(...) -> None

sdk/eventhub/azure-eventhub/azure/eventhub/_version.py

@@ -3,4 +3,4 @@
 # Licensed under the MIT License.
 # ------------------------------------
 
-VERSION = "5.4.1"
+VERSION = "5.5.0"

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

@@ -19,7 +19,7 @@
     Message,
     AMQPClientAsync,
 )
-from azure.core.credentials import AccessToken, AzureSasCredential
+from azure.core.credentials import AccessToken, AzureSasCredential, AzureNamedKeyCredential
 
 from .._client_base import ClientBase, _generate_sas_token, _parse_conn_str
 from .._utils import utc_from_timestamp, parse_sas_credential
@@ -83,7 +83,26 @@ async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:  # pylint
         """
         return AccessToken(self.token, self.expiry)
 
-class AzureSasTokenCredentialAsync(object):
+class EventhubAzureNamedKeyTokenCredentialAsync(object):
+    """The named key credential used for authentication.
+
+    :param credential: The AzureNamedKeyCredential that should be used.
+    :type credential: ~azure.core.credentials.AzureNamedKeyCredential
+    """
+
+    def __init__(self, azure_named_key_credential):
+        # type: (AzureNamedKeyCredential) -> None
+        self._credential = azure_named_key_credential
+        self.token_type = b"servicebus.windows.net:sastoken"
+
+    async def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
+        if not scopes:
+            raise ValueError("No token scope provided.")
+        name, key = self._credential.named_key
+        return _generate_sas_token(scopes[0], name, key)
+
+
+class EventhubAzureSasTokenCredentialAsync(object):
     """The shared access token credential used for authentication
     when AzureSasCredential is provided.
 
@@ -107,12 +126,14 @@ def __init__(
         self,
         fully_qualified_namespace: str,
         eventhub_name: str,
-        credential: Union["AsyncTokenCredential", AzureSasCredential],
+        credential: Union["AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential],
         **kwargs: Any
     ) -> None:
         self._loop = kwargs.pop("loop", None)
         if isinstance(credential, AzureSasCredential):
-            self._credential = AzureSasTokenCredentialAsync(credential) # type: ignore
+            self._credential = EventhubAzureSasTokenCredentialAsync(credential) # type: ignore
+        elif isinstance(credential, AzureNamedKeyCredential):
+            self._credential = EventhubAzureNamedKeyTokenCredentialAsync(credential) # type: ignore
         else:
             self._credential = credential # type: ignore
         super(ClientBaseAsync, self).__init__(

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

@@ -18,7 +18,7 @@
     Awaitable,
 )
 
-from azure.core.credentials import AzureSasCredential
+from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential
 
 from ._eventprocessor.event_processor import EventProcessor
 from ._consumer_async import EventHubConsumer
@@ -66,6 +66,7 @@ class EventHubConsumerClient(ClientBaseAsync):
      :class:`EventHubSharedKeyCredential<azure.eventhub.aio.EventHubSharedKeyCredential>`, or credential objects
      generated by the azure-identity library and objects that implement the `get_token(self, *scopes)` method.
     :type credential: ~azure.core.credentials_async.AsyncTokenCredential or ~azure.core.credentials.AzureSasCredential
+     or ~azure.core.credentials.AzureNamedKeyCredential
     :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
     :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service.
      The default value is 60 seconds. If set to 0, no timeout will be enforced from the client.
@@ -136,7 +137,7 @@ def __init__(
         fully_qualified_namespace: str,
         eventhub_name: str,
         consumer_group: str,
-        credential: Union["AsyncTokenCredential", AzureSasCredential],
+        credential: Union["AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential],
         **kwargs
     ) -> None:
         self._checkpoint_store = kwargs.pop("checkpoint_store", None)

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

@@ -8,7 +8,7 @@
 from typing import Any, Union, TYPE_CHECKING, List, Optional, Dict, cast
 from uamqp import constants
 
-from azure.core.credentials import AzureSasCredential
+from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential
 
 from ..exceptions import ConnectError, EventHubError
 from ._client_base_async import ClientBaseAsync
@@ -36,6 +36,7 @@ class EventHubProducerClient(ClientBaseAsync):
      :class:`EventHubSharedKeyCredential<azure.eventhub.aio.EventHubSharedKeyCredential>`, or credential objects
      generated by the azure-identity library and objects that implement the `get_token(self, *scopes)` method.
     :type credential: ~azure.core.credentials_async.AsyncTokenCredential or ~azure.core.credentials.AzureSasCredential
+     or ~azure.core.credentials.AzureNamedKeyCredential
     :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
     :keyword float auth_timeout: The time in seconds to wait for a token to be authorized by the service.
      The default value is 60 seconds. If set to 0, no timeout will be enforced from the client.
@@ -76,7 +77,7 @@ def __init__(
         self,
         fully_qualified_namespace: str,
         eventhub_name: str,
-        credential: Union["AsyncTokenCredential", AzureSasCredential],
+        credential: Union["AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential],
         **kwargs
     ) -> None:
         super(EventHubProducerClient, self).__init__(

sdk/eventhub/azure-eventhub/samples/async_samples/authenticate_with_named_key_credential_async.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+"""
+Example to demonstrate utilizing AzureNamedKeyCredential to authenticate with Event Hubs asynchronously.
+"""
+
+# pylint: disable=C0111
+
+import os
+import asyncio
+import time
+from azure.core.credentials import AzureNamedKeyCredential
+from azure.eventhub.aio import EventHubProducerClient
+from azure.eventhub import EventData
+
+# Target namespace and hub must also be specified.
+FULLY_QUALIFIED_NAMESPACE = os.environ['EVENT_HUB_HOSTNAME']
+EVENTHUB_NAME = os.environ['EVENT_HUB_NAME']
+
+EVENTHUB_POLICY_NAME = os.environ['EVENT_HUB_SAS_POLICY']
+EVENTHUB_KEY = os.environ['EVENT_HUB_SAS_KEY']
+
+credential = AzureNamedKeyCredential(EVENTHUB_POLICY_NAME, EVENTHUB_KEY)
+
+producer_client = EventHubProducerClient(
+    fully_qualified_namespace=FULLY_QUALIFIED_NAMESPACE,
+    eventhub_name=EVENTHUB_NAME,
+    credential=credential,
+    logging_enable=True
+)
+
+start_time = time.time()
+async def authenticate_with_named_key():
+    async with producer_client:
+        event_data_batch = await producer_client.create_batch()
+        event_data_batch.add(EventData('Single message'))
+        await producer_client.send_batch(event_data_batch)
+
+loop = asyncio.get_event_loop()
+start_time = time.time()
+loop.run_until_complete(authenticate_with_named_key())
+print("Send messages in {} seconds.".format(time.time() - start_time))

sdk/eventhub/azure-eventhub/samples/sync_samples/authenticate_with_named_key_credential.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+"""
+Example to demonstrate utilizing AzureNamedKeyCredential to authenticate with Event Hubs.
+"""
+
+# pylint: disable=C0111
+
+import os
+import time
+from azure.core.credentials import AzureNamedKeyCredential
+from azure.eventhub import EventHubProducerClient, EventData
+
+# Target namespace and hub must also be specified.
+FULLY_QUALIFIED_NAMESPACE = os.environ['EVENT_HUB_HOSTNAME']
+EVENTHUB_NAME = os.environ['EVENT_HUB_NAME']
+
+EVENTHUB_POLICY_NAME = os.environ['EVENT_HUB_SAS_POLICY']
+EVENTHUB_KEY = os.environ['EVENT_HUB_SAS_KEY']
+
+credential = AzureNamedKeyCredential(EVENTHUB_POLICY_NAME, EVENTHUB_KEY)
+
+producer_client = EventHubProducerClient(
+    fully_qualified_namespace=FULLY_QUALIFIED_NAMESPACE,
+    eventhub_name=EVENTHUB_NAME,
+    credential=credential,
+    logging_enable=True
+)
+
+start_time = time.time()
+with producer_client:
+    event_data_batch = producer_client.create_batch()
+    event_data_batch.add(EventData('Single message'))
+    producer_client.send_batch(event_data_batch)
+
+print("Send messages in {} seconds.".format(time.time() - start_time))