fix azure servicebus using managed identity support (#1801)

* azure servicebus: use DefaultAzureCredential in documentation

* azure servicebus: only use connection string when using sas key

* azure servicebus: add two small tests for paring of connection string

* azure servicebus: fix lint issues

kombu/transport/azureservicebus.py

@@ -35,7 +35,7 @@
 .. code-block::
 
     azureservicebus://SAS_POLICY_NAME:SAS_KEY@SERVICE_BUSNAMESPACE
-    azureservicebus://DefaultAzureIdentity@SERVICE_BUSNAMESPACE
+    azureservicebus://DefaultAzureCredential@SERVICE_BUSNAMESPACE
     azureservicebus://ManagedIdentityCredential@SERVICE_BUSNAMESPACE
 
 Transport Options
@@ -140,6 +140,11 @@ def __init__(self, *args, **kwargs):
     def _try_parse_connection_string(self) -> None:
         self._namespace, self._credential = Transport.parse_uri(
             self.conninfo.hostname)
+
+        if (isinstance(self._credential, DefaultAzureCredential) or
+                isinstance(self._credential, ManagedIdentityCredential)):
+            return None
+
         if ":" in self._credential:
             self._policy, self._sas_key = self._credential.split(':', 1)
 
@@ -434,7 +439,8 @@ class Transport(virtual.Transport):
     can_parse_url = True
 
     @staticmethod
-    def parse_uri(uri: str) -> tuple[str, str, str]:
+    def parse_uri(uri: str) -> tuple[str, str | DefaultAzureCredential |
+                                     ManagedIdentityCredential]:
         # URL like:
         #  azureservicebus://{SAS policy name}:{SAS key}@{ServiceBus Namespace}
         # urllib parse does not work as the sas key could contain a slash

t/unit/transport/test_azureservicebus.py

@@ -16,6 +16,13 @@
 import azure.servicebus.exceptions  # noqa
 from azure.servicebus import ServiceBusMessage, ServiceBusReceiveMode  # noqa
 
+try:
+    from azure.identity import (DefaultAzureCredential,
+                                ManagedIdentityCredential)
+except ImportError:
+    DefaultAzureCredential = None
+    ManagedIdentityCredential = None
+
 from kombu.transport import azureservicebus  # noqa
 
 
@@ -95,7 +102,9 @@ def get_queue_runtime_properties(self, queue_name):
 
 
 URL_NOCREDS = 'azureservicebus://'
-URL_CREDS = 'azureservicebus://policyname:ke/y@hostname'
+URL_CREDS_SAS = 'azureservicebus://policyname:ke/y@hostname'
+URL_CREDS_DA = 'azureservicebus://DefaultAzureCredential@hostname'
+URL_CREDS_MI = 'azureservicebus://ManagedIdentityCredential@hostname'
 
 
 def test_queue_service_nocredentials():
@@ -105,9 +114,9 @@ def test_queue_service_nocredentials():
         assert exc == 'Need an URI like azureservicebus://{SAS policy name}:{SAS key}@{ServiceBus Namespace}'   # noqa
 
 
-def test_queue_service():
+def test_queue_service_sas():
     # Test gettings queue service without credentials
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     with patch('kombu.transport.azureservicebus.ServiceBusClient') as m:
         channel = conn.channel()
 
@@ -126,20 +135,38 @@ def test_queue_service():
         assert m.from_connection_string.call_count == 1
 
 
+def test_queue_service_da():
+    conn = Connection(URL_CREDS_DA, transport=azureservicebus.Transport)
+    channel = conn.channel()
+
+    # Check the DefaultAzureCredential has been parsed from the url correctly
+    # and the credential is a ManagedIdentityCredential
+    assert isinstance(channel._credential, DefaultAzureCredential)
+
+
+def test_queue_service_mi():
+    conn = Connection(URL_CREDS_MI, transport=azureservicebus.Transport)
+    channel = conn.channel()
+
+    # Check the ManagedIdentityCredential has been parsed from the url
+    # correctly and the credential is a ManagedIdentityCredential
+    assert isinstance(channel._credential, ManagedIdentityCredential)
+
+
 def test_conninfo():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
     assert channel.conninfo is conn
 
 
 def test_transport_type():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
     assert not channel.transport_options
 
 
 def test_default_wait_timeout_seconds():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
 
     assert channel.wait_time_seconds == \
@@ -148,7 +175,7 @@ def test_default_wait_timeout_seconds():
 
 def test_custom_wait_timeout_seconds():
     conn = Connection(
-        URL_CREDS,
+        URL_CREDS_SAS,
         transport=azureservicebus.Transport,
         transport_options={'wait_time_seconds': 10}
     )
@@ -158,15 +185,15 @@ def test_custom_wait_timeout_seconds():
 
 
 def test_default_peek_lock_seconds():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
 
     assert channel.peek_lock_seconds == \
         azureservicebus.Channel.default_peek_lock_seconds
 
 
 def test_custom_peek_lock_seconds():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport,
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport,
                       transport_options={'peek_lock_seconds': 65})
     channel = conn.channel()
 
@@ -175,7 +202,7 @@ def test_custom_peek_lock_seconds():
 
 def test_invalid_peek_lock_seconds():
     # Max is 300
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport,
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport,
                       transport_options={'peek_lock_seconds': 900})
     channel = conn.channel()
 
@@ -230,7 +257,7 @@ def mock_clients(
 def mock_queue(mock_asb, mock_asb_management, random_queue) -> MockQueue:
     exchange = Exchange('test_servicebus', type='direct')
     queue = Queue(random_queue, exchange, random_queue)
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
 
     queue(channel).declare()
@@ -312,7 +339,7 @@ def test_purge(mock_queue: MockQueue):
 
 def test_custom_queue_name_prefix():
     conn = Connection(
-        URL_CREDS,
+        URL_CREDS_SAS,
         transport=azureservicebus.Transport,
         transport_options={'queue_name_prefix': 'test-queue'}
     )
@@ -322,7 +349,7 @@ def test_custom_queue_name_prefix():
 
 
 def test_custom_entity_name():
-    conn = Connection(URL_CREDS, transport=azureservicebus.Transport)
+    conn = Connection(URL_CREDS_SAS, transport=azureservicebus.Transport)
     channel = conn.channel()
 
     # dashes allowed and dots replaced by dashes