apache-airflow-providers-microsoft-azure 12.4.1rc1__py3-none-any.whl → 12.5.0rc1__py3-none-any.whl

airflow/providers/microsoft/azure/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "12.4.1"
+__version__ = "12.5.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/microsoft/azure/fs/adls.py

@@ -20,9 +20,13 @@ from typing import TYPE_CHECKING, Any
 
 from azure.identity import ClientSecretCredential
 
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import get_field, parse_blob_account_url
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     from fsspec import AbstractFileSystem
 

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

@@ -28,18 +28,22 @@ from __future__ import annotations
 
 import warnings
 from functools import cached_property
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from azure.kusto.data import ClientRequestProperties, KustoClient, KustoConnectionStringBuilder
 from azure.kusto.data.exceptions import KustoServiceError
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     from azure.kusto.data.response import KustoResponseDataSet
 
@@ -170,7 +174,7 @@ class AzureDataExplorerHook(BaseHook):
         if auth_method == "AAD_APP":
             tenant = get_required_param("tenant")
             kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
-                cluster, conn.login, conn.password, tenant
+                cluster, cast("str", conn.login), cast("str", conn.password), tenant
             )
         elif auth_method == "AAD_APP_CERT":
             certificate = get_required_param("certificate")
@@ -178,7 +182,7 @@ class AzureDataExplorerHook(BaseHook):
             tenant = get_required_param("tenant")
             kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication(
                 cluster,
-                conn.login,
+                cast("str", conn.login),
                 certificate,
                 thumbprint,
                 tenant,
@@ -186,7 +190,7 @@ class AzureDataExplorerHook(BaseHook):
         elif auth_method == "AAD_CREDS":
             tenant = get_required_param("tenant")
             kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication(
-                cluster, conn.login, conn.password, tenant
+                cluster, cast("str", conn.login), cast("str", conn.password), tenant
             )
         elif auth_method == "AAD_DEVICE":
             kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(cluster)

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 from uuid import UUID, uuid4
 
 from azure.core.exceptions import ResourceNotFoundError
@@ -36,13 +37,17 @@ from azure.servicebus.management import (
     SubscriptionProperties,
 )
 
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     import datetime
 

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

@@ -22,12 +22,16 @@ from azure.common.client_factory import get_client_from_auth_file, get_client_fr
 from azure.common.credentials import ServicePrincipalCredentials
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     AzureIdentityCredentialAdapter,
     add_managed_identity_connection_widgets,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 
 class AzureBaseHook(BaseHook):
     """
@@ -41,7 +45,7 @@ class AzureBaseHook(BaseHook):
         which refers to the information to connect to the service.
     """
 
-    conn_name_attr = "azure_conn_id"
+    conn_name_attr = "conn_id"
     default_conn_name = "azure_default"
     conn_type = "azure"
     hook_name = "Azure"

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

@@ -25,12 +25,16 @@ from typing import TYPE_CHECKING, Any
 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,
     add_managed_identity_connection_widgets,
     get_field,
 )
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 from airflow.utils import timezone
 
 if TYPE_CHECKING:

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

@@ -85,7 +85,9 @@ class AzureContainerInstanceHook(AzureBaseHook):
         if all([conn.login, conn.password, tenant]):
             self.log.info("Getting connection using specific credentials and subscription_id.")
             credential = ClientSecretCredential(
-                client_id=conn.login, client_secret=conn.password, tenant_id=cast("str", tenant)
+                client_id=cast("str", conn.login),
+                client_secret=cast("str", conn.password),
+                tenant_id=cast("str", tenant),
             )
         else:
             self.log.info("Using DefaultAzureCredential as credential")

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

@@ -20,18 +20,22 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import Any
+from typing import Any, cast
 
 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 (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 
 class AzureContainerRegistryHook(BaseHook):
     """
@@ -121,4 +125,6 @@ class AzureContainerRegistryHook(BaseHook):
             credentials = client.registries.list_credentials(resource_group, conn.login).as_dict()
             password = credentials["passwords"][0]["value"]
 
-        return ImageRegistryCredential(server=conn.host, username=conn.login, password=password)
+        return ImageRegistryCredential(
+            server=cast("str", conn.host), username=cast("str", conn.login), password=password
+        )

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

@@ -16,18 +16,22 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, cast
 
 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 (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 
 class AzureContainerVolumeHook(BaseHook):
     """
@@ -121,7 +125,7 @@ class AzureContainerVolumeHook(BaseHook):
             )
             return storage_account_list_keys_result.as_dict()["keys"][0]["value"]
 
-        return conn.password
+        return cast("str", conn.password)
 
     def get_file_volume(
         self, mount_name: str, share_name: str, storage_account_name: str, read_only: bool = False

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

@@ -27,7 +27,7 @@ the default database and collection to use (see connection `azure_cosmos_default
 from __future__ import annotations
 
 import uuid
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import urlparse
 
 from azure.cosmos import PartitionKey
@@ -36,15 +36,19 @@ from azure.cosmos.exceptions import CosmosHttpResponseError
 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 (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
-    PartitionKeyType = Union[str, list[str]]
+    PartitionKeyType = str | list[str]
 
 
 class AzureCosmosDBHook(BaseHook):
@@ -131,6 +135,7 @@ class AzureCosmosDBHook(BaseHook):
             conn = self.get_connection(self.conn_id)
             extras = conn.extra_dejson
             endpoint_uri = conn.login
+            endpoint_uri = cast("str", endpoint_uri)
             resource_group_name = self._get_field(extras, "resource_group_name")
 
             if conn.password:
@@ -147,12 +152,12 @@ class AzureCosmosDBHook(BaseHook):
                     credential=credential,
                     subscription_id=subscritption_id,
                 )
-
+                conn.login = cast("str", conn.login)
                 database_account = urlparse(conn.login).netloc.split(".")[0]
                 database_account_keys = management_client.database_accounts.list_keys(
                     resource_group_name, database_account
                 )
-                master_key = database_account_keys.primary_master_key
+                master_key = cast("str", database_account_keys.primary_master_key)
             else:
                 raise AirflowException("Either password or resource_group_name is required")
 

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

@@ -35,8 +35,9 @@ from __future__ import annotations
 
 import inspect
 import time
+from collections.abc import Callable
 from functools import wraps
-from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, Union, cast
+from typing import IO, TYPE_CHECKING, Any, TypeVar, cast
 
 from asgiref.sync import sync_to_async
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
@@ -48,13 +49,17 @@ from azure.mgmt.datafactory import DataFactoryManagementClient
 from azure.mgmt.datafactory.aio import DataFactoryManagementClient as AsyncDataFactoryManagementClient
 
 from airflow.exceptions import AirflowException
-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,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     from azure.core.polling import LROPoller
     from azure.mgmt.datafactory.models import (
@@ -68,8 +73,8 @@ if TYPE_CHECKING:
         TriggerResource,
     )
 
-Credentials = Union[ClientSecretCredential, DefaultAzureCredential]
-AsyncCredentials = Union[AsyncClientSecretCredential, AsyncDefaultAzureCredential]
+Credentials = ClientSecretCredential | DefaultAzureCredential
+AsyncCredentials = AsyncClientSecretCredential | AsyncDefaultAzureCredential
 
 T = TypeVar("T", bound=Any)
 

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

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import Any, Union
+from typing import Any, cast
 
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
 from azure.datalake.store import core, lib, multithread
@@ -33,7 +33,6 @@ from azure.storage.filedatalake import (
 )
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     AzureIdentityCredentialAdapter,
     add_managed_identity_connection_widgets,
@@ -41,7 +40,12 @@ from airflow.providers.microsoft.azure.utils import (
     get_sync_default_azure_credential,
 )
 
-Credentials = Union[ClientSecretCredential, AzureIdentityCredentialAdapter, DefaultAzureCredential]
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
+Credentials = ClientSecretCredential | AzureIdentityCredentialAdapter | DefaultAzureCredential
 
 
 class AzureDataLakeHook(BaseHook):
@@ -355,12 +359,13 @@ class AzureDataLakeStorageV2Hook(BaseHook):
             app_id = conn.login
             app_secret = conn.password
             proxies = extra.get("proxies", {})
-
+            app_id = cast("str", app_id)
+            app_secret = cast("str", app_secret)
             credential = ClientSecretCredential(
                 tenant_id=tenant, client_id=app_id, client_secret=app_secret, proxies=proxies
             )
         elif conn.password:
-            credential = conn.password
+            credential = conn.password  # type: ignore[assignment]
         else:
             managed_identity_client_id = self._get_field(extra, "managed_identity_client_id")
             workload_identity_tenant_id = self._get_field(extra, "workload_identity_tenant_id")
@@ -424,7 +429,7 @@ class AzureDataLakeStorageV2Hook(BaseHook):
         :param file_system_name: Name of the file system or instance of FileSystemProperties.
         :param directory_name: Name of the directory which needs to be created in the file system.
         """
-        result = self.get_file_system(file_system_name).create_directory(directory_name, kwargs)
+        result = self.get_file_system(file_system_name).create_directory(directory_name, **kwargs)
         return result
 
     def get_directory_client(
@@ -479,7 +484,7 @@ class AzureDataLakeStorageV2Hook(BaseHook):
         """
         file_client = self.create_file(file_system_name, file_name)
         with open(file_path, "rb") as data:
-            file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs)
+            file_client.upload_data(data, overwrite=overwrite, **kwargs)
 
     def upload_file_to_directory(
         self,
@@ -500,9 +505,9 @@ class AzureDataLakeStorageV2Hook(BaseHook):
         :param overwrite: Boolean flag to overwrite an existing file or not.
         """
         directory_client = self.get_directory_client(file_system_name, directory_name=directory_name)
-        file_client = directory_client.create_file(file_name, kwargs=kwargs)
+        file_client = directory_client.create_file(file_name, **kwargs)
         with open(file_path, "rb") as data:
-            file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs)
+            file_client.upload_data(data, overwrite=overwrite, **kwargs)
 
     def list_files_directory(
         self, file_system_name: FileSystemProperties | str, directory_name: str

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

@@ -21,12 +21,16 @@ from typing import IO, Any
 
 from azure.storage.fileshare import FileProperties, ShareDirectoryClient, ShareFileClient, ShareServiceClient
 
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 
 class AzureFileShareHook(BaseHook):
     """

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

@@ -50,7 +50,11 @@ from airflow.exceptions import (
     AirflowException,
     AirflowNotFoundException,
 )
-from airflow.hooks.base import BaseHook
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 
 if TYPE_CHECKING:
     from azure.identity._internal.client_credential_base import ClientCredentialBase
@@ -254,7 +258,7 @@ class KiotaRequestAdapterHook(BaseHook):
             client_secret = connection.password
             config = connection.extra_dejson if connection.extra else {}
             api_version = self.get_api_version(config)
-            host = self.get_host(connection)
+            host = self.get_host(connection)  # type: ignore[arg-type]
             base_url = config.get("base_url", urljoin(host, api_version))
             authority = config.get("authority")
             proxies = self.get_proxies(config)
@@ -299,6 +303,7 @@ class KiotaRequestAdapterHook(BaseHook):
                     timeout=Timeout(timeout=self.timeout),
                     verify=verify,
                     trust_env=trust_env,
+                    base_url=base_url,
                 ),
                 host=host,  # type: ignore
             )
@@ -446,9 +451,9 @@ class KiotaRequestAdapterHook(BaseHook):
             request_information.url = url
         elif request_information.query_parameters.keys():
             query = ",".join(request_information.query_parameters.keys())
-            request_information.url_template = f"{{+baseurl}}/{self.normalize_url(url)}{{?{query}}}"
+            request_information.url_template = f"{{+baseurl}}{self.normalize_url(url)}{{?{query}}}"
         else:
-            request_information.url_template = f"{{+baseurl}}/{self.normalize_url(url)}"
+            request_information.url_template = f"{{+baseurl}}{self.normalize_url(url)}"
         if not response_type:
             request_information.request_options[ResponseHandlerOption.get_key()] = ResponseHandlerOption(
                 response_handler=DefaultResponseHandler()
@@ -467,6 +472,7 @@ class KiotaRequestAdapterHook(BaseHook):
                 header_name=RequestInformation.CONTENT_TYPE_HEADER, header_value="application/json"
             )
             request_information.content = json.dumps(data).encode("utf-8")
+        print("Request Information:", request_information.url)
         return request_information
 
     @staticmethod

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

@@ -189,12 +189,15 @@ class PowerBIHook(KiotaRequestAdapterHook):
 
         return refresh_details
 
-    async def trigger_dataset_refresh(self, *, dataset_id: str, group_id: str) -> str:
+    async def trigger_dataset_refresh(
+        self, *, dataset_id: str, group_id: str, request_body: dict[str, Any] | None = None
+    ) -> str:
         """
         Triggers a refresh for the specified dataset from the given group id.
 
         :param dataset_id: The dataset id.
         :param group_id: The workspace id.
+        :param request_body: Additional arguments to pass to the request body, as described in https://learn.microsoft.com/en-us/rest/api/power-bi/datasets/refresh-dataset-in-group#request-body.
 
         :return: Request id of the dataset refresh request.
         """
@@ -207,6 +210,7 @@ class PowerBIHook(KiotaRequestAdapterHook):
                     "group_id": group_id,
                     "dataset_id": dataset_id,
                 },
+                data=request_body,
             )
 
             request_id = response.get("requestid")

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

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 from azure.core.exceptions import ServiceRequestError
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
@@ -25,18 +25,22 @@ from azure.synapse.artifacts import ArtifactsClient
 from azure.synapse.spark import SparkClient
 
 from airflow.exceptions import AirflowException, AirflowTaskTimeout
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     from azure.synapse.artifacts.models import CreateRunResponse, PipelineRun
     from azure.synapse.spark.models import SparkBatchJobOptions
 
-Credentials = Union[ClientSecretCredential, DefaultAzureCredential]
+Credentials = ClientSecretCredential | DefaultAzureCredential
 
 
 class AzureSynapseSparkBatchRunStatus:

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

@@ -29,7 +29,7 @@ from __future__ import annotations
 import logging
 import os
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, cast
 
 from asgiref.sync import sync_to_async
 from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError
@@ -46,7 +46,6 @@ from azure.storage.blob.aio import (
 )
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_async_default_azure_credential,
@@ -54,10 +53,16 @@ from airflow.providers.microsoft.azure.utils import (
     parse_blob_account_url,
 )
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
     from azure.storage.blob._models import BlobProperties
 
-AsyncCredentials = Union[AsyncClientSecretCredential, AsyncDefaultAzureCredential]
+AsyncCredentials = AsyncClientSecretCredential | AsyncDefaultAzureCredential
 
 
 class WasbHook(BaseHook):
@@ -67,8 +72,9 @@ class WasbHook(BaseHook):
     These parameters have to be passed in Airflow Data Base: account_name and account_key.
 
     Additional options passed in the 'extra' field of the connection will be
-    passed to the `BlockBlockService()` constructor. For example, authenticate
-    using a SAS token by adding {"sas_token": "YOUR_TOKEN"}.
+    passed to the `BlobServiceClient()` constructor. For example, authenticate
+    using a SAS token by adding {"sas_token": "YOUR_TOKEN"} or using an account key
+    by adding {"account_key": "YOUR_ACCOUNT_KEY"}.
 
     If no authentication configuration is provided, DefaultAzureCredential will be used (applicable
     when using Azure compute infrastructure).
@@ -121,7 +127,7 @@ class WasbHook(BaseHook):
                 "tenant_id": "tenant",
                 "shared_access_key": "shared access key",
                 "sas_token": "account url or token",
-                "extra": "additional options for use with ClientSecretCredential or DefaultAzureCredential",
+                "extra": "additional options for use with ClientSecretCredential, DefaultAzureCredential, or account_key authentication",
             },
         }
 
@@ -171,8 +177,8 @@ class WasbHook(BaseHook):
         tenant = self._get_field(extra, "tenant_id")
         if tenant:
             # use Active Directory auth
-            app_id = conn.login
-            app_secret = conn.password
+            app_id = cast("str", conn.login)
+            app_secret = cast("str", conn.password)
             token_credential = ClientSecretCredential(
                 tenant_id=tenant, client_id=app_id, client_secret=app_secret, **client_secret_auth_config
             )
@@ -196,15 +202,20 @@ class WasbHook(BaseHook):
             return BlobServiceClient(account_url=f"{account_url.rstrip('/')}/{sas_token}", **extra)
 
         # Fall back to old auth (password) or use managed identity if not provided.
-        credential = conn.password
+        credential: str | TokenCredential | None = conn.password
         if not credential:
-            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 = get_sync_default_azure_credential(
-                managed_identity_client_id=managed_identity_client_id,
-                workload_identity_tenant_id=workload_identity_tenant_id,
-            )
-            self.log.info("Using DefaultAzureCredential as credential")
+            # Check for account_key in extra fields before falling back to DefaultAzureCredential
+            account_key = self._get_field(extra, "account_key")
+            if account_key:
+                credential = account_key
+            else:
+                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 = get_sync_default_azure_credential(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
+                self.log.info("Using DefaultAzureCredential as credential")
         return BlobServiceClient(
             account_url=account_url,
             credential=credential,
@@ -646,13 +657,18 @@ class WasbAsyncHook(WasbHook):
         # Fall back to old auth (password) or use managed identity if not provided.
         credential = conn.password
         if not credential:
-            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 = get_async_default_azure_credential(
-                managed_identity_client_id=managed_identity_client_id,
-                workload_identity_tenant_id=workload_identity_tenant_id,
-            )
-            self.log.info("Using DefaultAzureCredential as credential")
+            # Check for account_key in extra fields before falling back to DefaultAzureCredential
+            account_key = self._get_field(extra, "account_key")
+            if account_key:
+                credential = account_key
+            else:
+                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 = get_async_default_azure_credential(
+                    managed_identity_client_id=managed_identity_client_id,
+                    workload_identity_tenant_id=workload_identity_tenant_id,
+                )
+                self.log.info("Using DefaultAzureCredential as credential")
         self.blob_service_client = AsyncBlobServiceClient(
             account_url=account_url,
             credential=credential,

airflow/providers/microsoft/azure/operators/adls.py

@@ -19,8 +19,8 @@ from __future__ import annotations
 from collections.abc import Iterable, Sequence
 from typing import IO, TYPE_CHECKING, Any, AnyStr
 
-from airflow.models import BaseOperator
 from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook, AzureDataLakeStorageV2Hook
+from airflow.providers.microsoft.azure.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context