apache-airflow-providers-microsoft-azure 8.1.0__tar.gz → 8.2.0rc1__tar.gz

{apache-airflow-providers-microsoft-azure-8.1.0/apache_airflow_providers_microsoft_azure.egg-info → apache-airflow-providers-microsoft-azure-8.2.0rc1}/PKG-INFO

@@ -1,14 +1,14 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-microsoft-azure
-Version: 8.1.0
+Version: 8.2.0rc1
 Summary: Provider for Apache Airflow. Implements apache-airflow-providers-microsoft-azure package
 Home-page: https://airflow.apache.org/
 Download-URL: https://archive.apache.org/dist/airflow/providers
 Author: Apache Software Foundation
 Author-email: dev@airflow.apache.org
 License: Apache License 2.0
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/changelog.html
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
@@ -52,8 +52,7 @@ License-File: NOTICE
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.
-
- .. Licensed to the Apache Software Foundation (ASF) under one
+.. Licensed to the Apache Software Foundation (ASF) under one
     or more contributor license agreements.  See the NOTICE file
     distributed with this work for additional information
     regarding copyright ownership.  The ASF licenses this file
@@ -73,7 +72,7 @@ License-File: NOTICE
 
 Package ``apache-airflow-providers-microsoft-azure``
 
-Release: ``8.1.0``
+Release: ``8.2.0rc1``
 
 
 `Microsoft Azure <https://azure.microsoft.com/>`__
@@ -86,8 +85,7 @@ This is a provider package for ``microsoft.azure`` provider. All classes for thi
 are in ``airflow.providers.microsoft.azure`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/>`_.
-
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/>`_.
 
 Installation
 ------------
@@ -149,4 +147,4 @@ Dependent package
 ====================================================================================================  ==========
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/changelog.html>`_.

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/README.rst

@@ -15,8 +15,7 @@
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.
-
- .. Licensed to the Apache Software Foundation (ASF) under one
+.. Licensed to the Apache Software Foundation (ASF) under one
     or more contributor license agreements.  See the NOTICE file
     distributed with this work for additional information
     regarding copyright ownership.  The ASF licenses this file
@@ -36,7 +35,7 @@
 
 Package ``apache-airflow-providers-microsoft-azure``
 
-Release: ``8.1.0``
+Release: ``8.2.0rc1``
 
 
 `Microsoft Azure <https://azure.microsoft.com/>`__
@@ -49,8 +48,7 @@ This is a provider package for ``microsoft.azure`` provider. All classes for thi
 are in ``airflow.providers.microsoft.azure`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/>`_.
-
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/>`_.
 
 Installation
 ------------
@@ -112,4 +110,4 @@ Dependent package
 ====================================================================================================  ==========
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.1.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/8.2.0/changelog.html>`_.

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/__init__.py

@@ -28,7 +28,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.1.0"
+__version__ = "8.2.0"
 
 try:
     from airflow import __version__ as airflow_version

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/get_provider_info.py

@@ -29,6 +29,7 @@ def get_provider_info():
         "description": "`Microsoft Azure <https://azure.microsoft.com/>`__\n",
         "suspended": False,
         "versions": [
+            "8.2.0",
             "8.1.0",
             "8.0.0",
             "7.0.0",

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/adx.py

@@ -29,12 +29,15 @@ import warnings
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from azure.identity import DefaultAzureCredential
 from azure.kusto.data import ClientRequestProperties, KustoClient, KustoConnectionStringBuilder
 from azure.kusto.data.exceptions import KustoServiceError
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.hooks.base import BaseHook
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_sync_default_azure_credential,
+)
 
 if TYPE_CHECKING:
     from azure.kusto.data.response import KustoResponseDataSetV2
@@ -80,6 +83,7 @@ class AzureDataExplorerHook(BaseHook):
     hook_name = "Azure Data Explorer"
 
     @classmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets(cls) -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
@@ -192,9 +196,15 @@ class AzureDataExplorerHook(BaseHook):
         elif auth_method == "AAD_DEVICE":
             kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(cluster)
         elif auth_method == "AZURE_TOKEN_CRED":
+            managed_identity_client_id = conn.extra_dejson.get("managed_identity_client_id")
+            workload_identity_tenant_id = conn.extra_dejson.get("workload_identity_tenant_id")
+            credential = get_sync_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
             kcsb = KustoConnectionStringBuilder.with_azure_token_credential(
                 connection_string=cluster,
-                credential=DefaultAzureCredential(),
+                credential=credential,
             )
         else:
             raise AirflowException(f"Unknown authentication method: {auth_method}")

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/asb.py

@@ -16,14 +16,20 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from azure.identity import DefaultAzureCredential
 from azure.servicebus import ServiceBusClient, ServiceBusMessage, ServiceBusSender
 from azure.servicebus.management import QueueProperties, ServiceBusAdministrationClient
 
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_field,
+    get_sync_default_azure_credential,
+)
+
+if TYPE_CHECKING:
+    from azure.identity import DefaultAzureCredential
 
 
 class BaseAzureServiceBusHook(BaseHook):
@@ -40,6 +46,7 @@ class BaseAzureServiceBusHook(BaseHook):
     hook_name = "Azure Service Bus"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -64,7 +71,7 @@ class BaseAzureServiceBusHook(BaseHook):
                     "<Resource group>.servicebus.windows.net (for Azure AD authenticaltion)"
                 ),
                 "credential": "credential",
-                "schema": "Endpoint=sb://<Resource group>.servicebus.windows.net/;SharedAccessKeyName=<AccessKeyName>;SharedAccessKey=<SharedAccessKey>",  # noqa
+                "schema": "Endpoint=sb://<Resource group>.servicebus.windows.net/;SharedAccessKeyName=<AccessKeyName>;SharedAccessKey=<SharedAccessKey>",
             },
         }
 
@@ -106,7 +113,16 @@ class AdminClientHook(BaseAzureServiceBusHook):
             credential: str | DefaultAzureCredential = self._get_field(extras=extras, field_name="credential")
             fully_qualified_namespace = self._get_field(extras=extras, field_name="fully_qualified_namespace")
             if not credential:
-                credential = DefaultAzureCredential()
+                managed_identity_client_id = self._get_field(
+                    extras=extras, field_name="managed_identity_client_id"
+                )
+                workload_identity_tenant_id = self._get_field(
+                    extras=extras, field_name="workload_identity_tenant_id"
+                )
+                credential = get_sync_default_azure_credential(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
             client = ServiceBusAdministrationClient(
                 fully_qualified_namespace=fully_qualified_namespace,
                 credential=credential,  # type: ignore[arg-type]
@@ -190,7 +206,16 @@ class MessageHook(BaseAzureServiceBusHook):
             credential: str | DefaultAzureCredential = self._get_field(extras=extras, field_name="credential")
             fully_qualified_namespace = self._get_field(extras=extras, field_name="fully_qualified_namespace")
             if not credential:
-                credential = DefaultAzureCredential()
+                managed_identity_client_id = self._get_field(
+                    extras=extras, field_name="managed_identity_client_id"
+                )
+                workload_identity_tenant_id = self._get_field(
+                    extras=extras, field_name="workload_identity_tenant_id"
+                )
+                credential = get_sync_default_azure_credential(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
             client = ServiceBusClient(
                 fully_qualified_namespace=fully_qualified_namespace,
                 credential=credential,  # type: ignore[arg-type]

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/base_azure.py

@@ -24,7 +24,10 @@ from azure.common.credentials import ServicePrincipalCredentials
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import AzureIdentityCredentialAdapter
+from airflow.providers.microsoft.azure.utils import (
+    AzureIdentityCredentialAdapter,
+    add_managed_identity_connection_widgets,
+)
 
 
 class AzureBaseHook(BaseHook):
@@ -45,6 +48,7 @@ class AzureBaseHook(BaseHook):
     hook_name = "Azure"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -133,7 +137,12 @@ class AzureBaseHook(BaseHook):
             )
         else:
             self.log.info("Using DefaultAzureCredential as credential")
-            credentials = AzureIdentityCredentialAdapter()
+            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(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
 
         return self.sdk_client(
             credentials=credentials,

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/batch.py

@@ -26,7 +26,11 @@ from azure.batch import BatchServiceClient, batch_auth, models as batch_models
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import AzureIdentityCredentialAdapter, get_field
+from airflow.providers.microsoft.azure.utils import (
+    AzureIdentityCredentialAdapter,
+    add_managed_identity_connection_widgets,
+    get_field,
+)
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
@@ -46,15 +50,8 @@ class AzureBatchHook(BaseHook):
     conn_type = "azure_batch"
     hook_name = "Azure Batch Service"
 
-    def _get_field(self, extras, name):
-        return get_field(
-            conn_id=self.conn_id,
-            conn_type=self.conn_type,
-            extras=extras,
-            field_name=name,
-        )
-
     @classmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets(cls) -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -80,6 +77,14 @@ class AzureBatchHook(BaseHook):
         super().__init__()
         self.conn_id = azure_batch_conn_id
 
+    def _get_field(self, extras, name):
+        return get_field(
+            conn_id=self.conn_id,
+            conn_type=self.conn_type,
+            extras=extras,
+            field_name=name,
+        )
+
     @cached_property
     def connection(self) -> BatchServiceClient:
         """Get the Batch client connection (cached)."""
@@ -101,8 +106,13 @@ class AzureBatchHook(BaseHook):
         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)

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/container_instance.py

@@ -27,6 +27,7 @@ from azure.mgmt.containerinstance import ContainerInstanceManagementClient
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook
+from airflow.providers.microsoft.azure.utils import get_sync_default_azure_credential
 
 if TYPE_CHECKING:
     from azure.mgmt.containerinstance.models import (
@@ -92,7 +93,12 @@ class AzureContainerInstanceHook(AzureBaseHook):
             )
         else:
             self.log.info("Using DefaultAzureCredential as credential")
-            credential = DefaultAzureCredential()
+            managed_identity_client_id = conn.extra_dejson.get("managed_identity_client_id")
+            workload_identity_tenant_id = conn.extra_dejson.get("workload_identity_tenant_id")
+            credential = get_sync_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
 
         subscription_id = cast(str, conn.extra_dejson.get("subscriptionId"))
         return ContainerInstanceManagementClient(

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/container_registry.py

@@ -21,12 +21,15 @@ from __future__ import annotations
 from functools import cached_property
 from typing import Any
 
-from azure.identity import DefaultAzureCredential
 from azure.mgmt.containerinstance.models import ImageRegistryCredential
 from azure.mgmt.containerregistry import ContainerRegistryManagementClient
 
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_field,
+    get_sync_default_azure_credential,
+)
 
 
 class AzureContainerRegistryHook(BaseHook):
@@ -44,6 +47,7 @@ class AzureContainerRegistryHook(BaseHook):
     hook_name = "Azure Container Registry"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -103,8 +107,15 @@ class AzureContainerRegistryHook(BaseHook):
             extras = conn.extra_dejson
             subscription_id = self._get_field(extras, "subscription_id")
             resource_group = self._get_field(extras, "resource_group")
+            managed_identity_client_id = self._get_field(extras, "managed_identity_client_id")
+            workload_identity_tenant_id = self._get_field(extras, "workload_identity_tenant_id")
+            credential = get_sync_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
             client = ContainerRegistryManagementClient(
-                credential=DefaultAzureCredential(), subscription_id=subscription_id
+                credential=credential,
+                subscription_id=subscription_id,
             )
             credentials = client.registries.list_credentials(resource_group, conn.login).as_dict()
             password = credentials["passwords"][0]["value"]

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/container_volume.py

@@ -1,4 +1,3 @@
-#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
@@ -19,12 +18,15 @@ from __future__ import annotations
 
 from typing import Any
 
-from azure.identity import DefaultAzureCredential
 from azure.mgmt.containerinstance.models import AzureFileVolume, Volume
 from azure.mgmt.storage import StorageManagementClient
 
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_field,
+    get_sync_default_azure_credential,
+)
 
 
 class AzureContainerVolumeHook(BaseHook):
@@ -41,19 +43,8 @@ class AzureContainerVolumeHook(BaseHook):
     conn_type = "azure_container_volume"
     hook_name = "Azure Container Volume"
 
-    def __init__(self, azure_container_volume_conn_id: str = "azure_container_volume_default") -> None:
-        super().__init__()
-        self.conn_id = azure_container_volume_conn_id
-
-    def _get_field(self, extras, name):
-        return get_field(
-            conn_id=self.conn_id,
-            conn_type=self.conn_type,
-            extras=extras,
-            field_name=name,
-        )
-
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
@@ -92,6 +83,18 @@ class AzureContainerVolumeHook(BaseHook):
             },
         }
 
+    def __init__(self, azure_container_volume_conn_id: str = "azure_container_volume_default") -> None:
+        super().__init__()
+        self.conn_id = azure_container_volume_conn_id
+
+    def _get_field(self, extras, name):
+        return get_field(
+            conn_id=self.conn_id,
+            conn_type=self.conn_type,
+            extras=extras,
+            field_name=name,
+        )
+
     def get_storagekey(self, *, storage_account_name: str | None = None) -> str:
         """Get Azure File Volume storage key."""
         conn = self.get_connection(self.conn_id)
@@ -106,8 +109,13 @@ class AzureContainerVolumeHook(BaseHook):
         subscription_id = self._get_field(extras, "subscription_id")
         resource_group = self._get_field(extras, "resource_group")
         if subscription_id and storage_account_name and resource_group:
-            credentials = DefaultAzureCredential()
-            storage_client = StorageManagementClient(credentials, subscription_id)
+            managed_identity_client_id = self._get_field(extras, "managed_identity_client_id")
+            workload_identity_tenant_id = self._get_field(extras, "workload_identity_tenant_id")
+            credential = get_sync_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
+            storage_client = StorageManagementClient(credential, subscription_id)
             storage_account_list_keys_result = storage_client.storage_accounts.list_keys(
                 resource_group, storage_account_name
             )

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/cosmos.py

@@ -31,12 +31,15 @@ from urllib.parse import urlparse
 
 from azure.cosmos.cosmos_client import CosmosClient
 from azure.cosmos.exceptions import CosmosHttpResponseError
-from azure.identity import DefaultAzureCredential
 from azure.mgmt.cosmosdb import CosmosDBManagementClient
 
 from airflow.exceptions import AirflowBadRequest, AirflowException
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_field,
+    get_sync_default_azure_credential,
+)
 
 
 class AzureCosmosDBHook(BaseHook):
@@ -57,6 +60,7 @@ class AzureCosmosDBHook(BaseHook):
     hook_name = "Azure CosmosDB"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -126,9 +130,16 @@ class AzureCosmosDBHook(BaseHook):
             if conn.password:
                 master_key = conn.password
             elif resource_group_name:
+                managed_identity_client_id = self._get_field(extras, "managed_identity_client_id")
+                workload_identity_tenant_id = self._get_field(extras, "workload_identity_tenant_id")
+                subscritption_id = self._get_field(extras, "subscription_id")
+                credential = get_sync_default_azure_credential(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
                 management_client = CosmosDBManagementClient(
-                    credential=DefaultAzureCredential(),
-                    subscription_id=self._get_field(extras, "subscription_id"),
+                    credential=credential,
+                    subscription_id=subscritption_id,
                 )
 
                 database_account = urlparse(conn.login).netloc.split(".")[0]

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/data_factory.py

@@ -48,6 +48,11 @@ from azure.mgmt.datafactory.aio import DataFactoryManagementClient as AsyncDataF
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.hooks.base import BaseHook
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_async_default_azure_credential,
+    get_sync_default_azure_credential,
+)
 
 if TYPE_CHECKING:
     from azure.core.polling import LROPoller
@@ -152,6 +157,7 @@ class AzureDataFactoryHook(BaseHook):
     hook_name: str = "Azure Data Factory"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -205,7 +211,12 @@ class AzureDataFactoryHook(BaseHook):
                 client_id=conn.login, client_secret=conn.password, tenant_id=tenant
             )
         else:
-            credential = DefaultAzureCredential()
+            managed_identity_client_id = get_field(extras, "managed_identity_client_id")
+            workload_identity_tenant_id = get_field(extras, "workload_identity_tenant_id")
+            credential = get_sync_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
         self._conn = self._create_client(credential, subscription_id)
 
         return self._conn
@@ -1140,7 +1151,12 @@ class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
                 client_id=conn.login, client_secret=conn.password, tenant_id=tenant
             )
         else:
-            credential = AsyncDefaultAzureCredential()
+            managed_identity_client_id = get_field(extras, "managed_identity_client_id")
+            workload_identity_tenant_id = get_field(extras, "workload_identity_tenant_id")
+            credential = get_async_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
 
         self._async_conn = AsyncDataFactoryManagementClient(
             credential=credential,

{apache-airflow-providers-microsoft-azure-8.1.0 → apache-airflow-providers-microsoft-azure-8.2.0rc1}/airflow/providers/microsoft/azure/hooks/data_lake.py

@@ -34,7 +34,11 @@ from azure.storage.filedatalake import (
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import AzureIdentityCredentialAdapter, get_field
+from airflow.providers.microsoft.azure.utils import (
+    AzureIdentityCredentialAdapter,
+    add_managed_identity_connection_widgets,
+    get_field,
+)
 
 Credentials = Union[ClientSecretCredential, AzureIdentityCredentialAdapter]
 
@@ -62,6 +66,7 @@ class AzureDataLakeHook(BaseHook):
     hook_name = "Azure Data Lake"
 
     @staticmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets() -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
@@ -118,7 +123,12 @@ class AzureDataLakeHook(BaseHook):
             if tenant:
                 credential = lib.auth(tenant_id=tenant, client_secret=conn.password, client_id=conn.login)
             else:
-                credential = AzureIdentityCredentialAdapter()
+                managed_identity_client_id = self._get_field(extras, "managed_identity_client_id")
+                workload_identity_tenant_id = self._get_field(extras, "workload_identity_tenant_id")
+                credential = AzureIdentityCredentialAdapter(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
             self._conn = core.AzureDLFileSystem(credential, store_name=self.account_name)
             self._conn.connect()
         return self._conn
@@ -265,6 +275,7 @@ class AzureDataLakeStorageV2Hook(BaseHook):
     hook_name = "Azure Date Lake Storage V2"
 
     @classmethod
+    @add_managed_identity_connection_widgets
     def get_connection_form_widgets(cls) -> dict[str, Any]:
         """Returns connection widgets to add to connection form."""
         from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
@@ -305,6 +316,17 @@ class AzureDataLakeStorageV2Hook(BaseHook):
         self.conn_id = adls_conn_id
         self.public_read = public_read
 
+    def _get_field(self, extra_dict, field_name):
+        prefix = "extra__adls__"
+        if field_name.startswith("extra__"):
+            raise ValueError(
+                f"Got prefixed name {field_name}; please remove the '{prefix}' prefix "
+                f"when using this method."
+            )
+        if field_name in extra_dict:
+            return extra_dict[field_name] or None
+        return extra_dict.get(f"{prefix}{field_name}") or None
+
     @cached_property
     def service_client(self) -> DataLakeServiceClient:
         """Return the DataLakeServiceClient object (cached)."""
@@ -330,7 +352,12 @@ class AzureDataLakeStorageV2Hook(BaseHook):
         elif conn.password:
             credential = conn.password
         else:
-            credential = AzureIdentityCredentialAdapter()
+            managed_identity_client_id = self._get_field(extra, "managed_identity_client_id")
+            workload_identity_tenant_id = self._get_field(extra, "workload_identity_tenant_id")
+            credential = AzureIdentityCredentialAdapter(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
 
         return DataLakeServiceClient(
             account_url=f"https://{conn.host}.dfs.core.windows.net",
@@ -338,17 +365,6 @@ class AzureDataLakeStorageV2Hook(BaseHook):
             **extra,
         )
 
-    def _get_field(self, extra_dict, field_name):
-        prefix = "extra__adls__"
-        if field_name.startswith("extra__"):
-            raise ValueError(
-                f"Got prefixed name {field_name}; please remove the '{prefix}' prefix "
-                f"when using this method."
-            )
-        if field_name in extra_dict:
-            return extra_dict[field_name] or None
-        return extra_dict.get(f"{prefix}{field_name}") or None
-
     def create_file_system(self, file_system_name: str) -> None:
         """Create a new file system under the specified account.