feat(provider/azure): add managed identity support to batch hook (#35327

)

airflow/providers/microsoft/azure/hooks/batch.py

@@ -63,6 +63,12 @@ def get_connection_form_widgets(cls) -> dict[str, Any]:
 
         return {
             "account_url": StringField(lazy_gettext("Batch Account URL"), widget=BS3TextFieldWidget()),
+            "managed_identity_client_id": StringField(
+                lazy_gettext("Managed Identity Client ID"), widget=BS3TextFieldWidget()
+            ),
+            "workload_identity_tenant_id": StringField(
+                lazy_gettext("Workload Identity Tenant ID"), widget=BS3TextFieldWidget()
+            ),
         }
 
     @classmethod
@@ -73,6 +79,8 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]:
             "relabeling": {
                 "login": "Batch Account Name",
                 "password": "Batch Account Access Key",
+                "managed_identity_client_id": "Managed Identity Client ID",
+                "workload_identity_tenant_id": "Workload Identity Tenant ID",
             },
         }
 
@@ -101,8 +109,13 @@ def get_conn(self) -> BatchServiceClient:
         if all([conn.login, conn.password]):
             credentials = batch_auth.SharedKeyCredentials(conn.login, conn.password)
         else:
+            managed_identity_client_id = conn.extra_dejson.get("managed_identity_client_id")
+            workload_identity_tenant_id = conn.extra_dejson.get("workload_identity_tenant_id")
             credentials = AzureIdentityCredentialAdapter(
-                None, resource_id="https://batch.core.windows.net/.default"
+                None,
+                resource_id="https://batch.core.windows.net/.default",
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
             )
 
         batch_client = BatchServiceClient(credentials, batch_url=batch_account_url)

docs/apache-airflow-providers-microsoft-azure/connections/azure_batch.rst

@@ -32,8 +32,8 @@ There is one way to connect to Azure Batch using Airflow.
 1. Use `Azure Shared Key Credential
    <https://docs.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key>`_
    i.e. add shared key credentials to the Airflow connection.
-2. Fallback on `DefaultAzureCredential
-   <https://docs.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential>`_.
+2. Use managed identity by setting ``managed_identity_client_id``, ``workload_identity_tenant_id`` (under the hook, it uses DefaultAzureCredential_ with these arguments)
+3. Fallback on DefaultAzureCredential_.
    This includes a mechanism to try different options to authenticate: Managed System Identity, environment variables, authentication through Azure CLI and etc.
 
 
@@ -47,15 +47,22 @@ Configuring the Connection
 
 Batch Account Name (optional)
     Specify the Azure Batch Account Name used for the initial connection.
-    It can be left out to fall back on ``DefaultAzureCredential``.
+    It can be left out to fall back on DefaultAzureCredential_.
 
 Batch Account Access Key (optional)
     Specify the access key used for the initial connection.
-    It can be left out to fall back on ``DefaultAzureCredential``.
+    It can be left out to fall back on DefaultAzureCredential_.
 
 Batch Account URL
     Specify the batch account URL you would like to use.
 
+Managed Identity Client ID (optional)
+    The client ID of a user-assigned managed identity. If provided with ``workload_identity_tenant_id``, they'll pass to DefaultAzureCredential_.
+
+Workload Identity Tenant ID (optional)
+    ID of the application's Microsoft Entra tenant. Also called its "directory" ID. If provided with ``managed_identity_client_id``, they'll pass to DefaultAzureCredential_.
+
+
 When specifying the connection in environment variable you should specify it using URI syntax.
 
 Note that all components of the URI should be URL-encoded.
@@ -65,3 +72,9 @@ For example:
 .. code-block:: bash
 
    export AIRFLOW_CONN_AZURE_BATCH_DEFAULT='azure-batch://batch%20acount:batch%20key@?account_url=mybatchaccount.com'
+
+
+.. _DefaultAzureCredential: https://docs.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential
+
+.. spelling:word-list::
+    Entra

tests/providers/microsoft/azure/hooks/test_azure_batch.py

@@ -79,7 +79,10 @@ def test_fallback_to_azure_identity_credential_adppter_when_name_and_key_is_not_
         hook = AzureBatchHook(azure_batch_conn_id=self.test_vm_conn_id)
         assert isinstance(hook.get_conn(), BatchServiceClient)
         mock_azure_identity_credential_adapter.assert_called_with(
-            None, resource_id="https://batch.core.windows.net/.default"
+            None,
+            resource_id="https://batch.core.windows.net/.default",
+            managed_identity_client_id=None,
+            workload_identity_tenant_id=None,
         )
         assert not mock_shared_key_credentials.auth.called