feat: custom azure credential (#29)

This PR replaces DefaultAzureCredential with a CustomAzureCredential
chain. This custom credential chain reorders the credential preference,
to prefer AzureCliCredential, then ManagedIdentityCredential, then
EnvironmentCredential (service principal)

Author: jakevc

docs/further.md

@@ -1,6 +1,6 @@
 # Azure Batch Authentication
 
-The plugin uses [DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) to create and destroy Azure Batch resources. The caller must have Contributor permissions on the Azure Batch account for the plugin to work properly. If you are using the Azure Storage plugin you should also have the Storage Blob Data Contributor role for the storage account(s) you use. 
+The plugin uses a CustomAzureCredential chain that prefers the use of AzureCliCredential, then falls back to a ManagedIdentityCredential, and finally, an EnvironmentCredential (service principal) to create and destroy Azure Batch resources. The caller must have Contributor permissions on the Azure Batch account for the plugin to work properly. If you are using the Azure Storage plugin you should also have the Storage Blob Data Contributor role for the storage account(s) you use. 
 
 To run a Snakemake workflow using your azure identity you need to ensure you are logged in using the [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/):
 

snakemake_executor_plugin_azure_batch/__init__.py

@@ -15,7 +15,6 @@
 import azure.batch.models as bm
 from azure.batch import BatchServiceClient
 from azure.core.exceptions import HttpResponseError
-from azure.identity import DefaultAzureCredential
 from azure.mgmt.batch import BatchManagementClient
 from snakemake_interface_common.exceptions import WorkflowError
 from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo
@@ -30,6 +29,7 @@
 from snakemake_executor_plugin_azure_batch.constant import AZURE_BATCH_RESOURCE_ENDPOINT
 from snakemake_executor_plugin_azure_batch.util import (
     AzureIdentityCredentialAdapter,
+    CustomAzureCredential,
     read_stream_as_string,
     unpack_compute_node_errors,
     unpack_task_failure_information,
@@ -278,29 +278,25 @@ def __post_init__(self):
 
     def init_batch_client(self):
         """
-        Initialize the BatchServiceClient and BatchManagementClient using
-        DefaultAzureCredential.
+        Initialize the BatchServiceClient and BatchManagementClient
 
         Sets:
             - self.batch_client
             - self.batch_mgmt_client
         """
         try:
 
-            # initialize BatchServiceClient
-            default_credential = DefaultAzureCredential(
-                exclude_managed_identity_credential=True
-            )
             adapted_credential = AzureIdentityCredentialAdapter(
-                credential=default_credential, resource_id=AZURE_BATCH_RESOURCE_ENDPOINT
+                credential=CustomAzureCredential(),
+                resource_id=AZURE_BATCH_RESOURCE_ENDPOINT,
             )
             self.batch_client = BatchServiceClient(
                 adapted_credential, self.settings.account_url
             )
 
             # initialize BatchManagementClient
             self.batch_mgmt_client = BatchManagementClient(
-                credential=default_credential,
+                credential=CustomAzureCredential(),
                 subscription_id=self.settings.subscription_id,
             )
 

snakemake_executor_plugin_azure_batch/util.py

@@ -6,7 +6,21 @@
 from azure.core.pipeline import PipelineContext, PipelineRequest
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.core.pipeline.transport import HttpRequest
-from azure.identity import DefaultAzureCredential
+from azure.identity import (
+    AzureCliCredential,
+    ChainedTokenCredential,
+    EnvironmentCredential,
+    ManagedIdentityCredential,
+)
+
+
+def CustomAzureCredential() -> ChainedTokenCredential:
+    credential_chain = (
+        AzureCliCredential(),
+        ManagedIdentityCredential(),
+        EnvironmentCredential(),
+    )
+    return ChainedTokenCredential(*credential_chain)
 
 
 # The usage of this credential helper is required to authenticate batch with managed
@@ -26,13 +40,13 @@ def __init__(
         azure.common.credentials or msrestazure.
 
         Args:
-            credential: Any azure-identity credential (DefaultAzureCredential by
+            credential: Any azure-identity credential (CustomAzureCredential by
                    default)
             resource_id (str): The scope to use to get the token (default ARM)
         """
         super(AzureIdentityCredentialAdapter, self).__init__(None)
         if credential is None:
-            credential = DefaultAzureCredential()
+            credential = CustomAzureCredential()
         self._policy = BearerTokenCredentialPolicy(credential, resource_id, **kwargs)
 
     def _make_request(self):