From 2ebda593655283a3f045bb8543af66df0ef34c71 Mon Sep 17 00:00:00 2001 From: Jake VanCampen Date: Sun, 28 Jul 2024 11:36:34 -0700 Subject: [PATCH] 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) --- docs/further.md | 2 +- .../__init__.py | 14 +++++-------- snakemake_executor_plugin_azure_batch/util.py | 20 ++++++++++++++++--- 3 files changed, 23 insertions(+), 13 deletions(-) diff --git a/docs/further.md b/docs/further.md index 8910b82..2494234 100644 --- a/docs/further.md +++ b/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/): diff --git a/snakemake_executor_plugin_azure_batch/__init__.py b/snakemake_executor_plugin_azure_batch/__init__.py index 40a4586..ba28868 100644 --- a/snakemake_executor_plugin_azure_batch/__init__.py +++ b/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,8 +278,7 @@ 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 @@ -287,12 +286,9 @@ def init_batch_client(self): """ 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 @@ -300,7 +296,7 @@ def init_batch_client(self): # initialize BatchManagementClient self.batch_mgmt_client = BatchManagementClient( - credential=default_credential, + credential=CustomAzureCredential(), subscription_id=self.settings.subscription_id, ) diff --git a/snakemake_executor_plugin_azure_batch/util.py b/snakemake_executor_plugin_azure_batch/util.py index e735531..479970e 100644 --- a/snakemake_executor_plugin_azure_batch/util.py +++ b/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):