diff --git a/airflow/providers/databricks/hooks/databricks_base.py b/airflow/providers/databricks/hooks/databricks_base.py index 94c7c10223d6e..4dd95d755c72d 100644 --- a/airflow/providers/databricks/hooks/databricks_base.py +++ b/airflow/providers/databricks/hooks/databricks_base.py @@ -34,11 +34,6 @@ import aiohttp import requests -from azure.identity import ClientSecretCredential, ManagedIdentityCredential -from azure.identity.aio import ( - ClientSecretCredential as AsyncClientSecretCredential, - ManagedIdentityCredential as AsyncManagedIdentityCredential, -) from requests import PreparedRequest, exceptions as requests_exceptions from requests.auth import AuthBase, HTTPBasicAuth from requests.exceptions import JSONDecodeError @@ -52,7 +47,7 @@ ) from airflow import __version__ -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, AirflowOptionalProviderFeatureException from airflow.hooks.base import BaseHook from airflow.providers_manager import ProvidersManager @@ -302,6 +297,8 @@ def _get_aad_token(self, resource: str) -> str: self.log.info("Existing AAD token is expired, or going to expire soon. Refreshing...") try: + from azure.identity import ClientSecretCredential, ManagedIdentityCredential + for attempt in self._get_retry_object(): with attempt: if self.databricks_conn.extra_dejson.get("use_azure_managed_identity", False): @@ -321,6 +318,8 @@ def _get_aad_token(self, resource: str) -> str: self._is_oauth_token_valid(jsn) self.oauth_tokens[resource] = jsn break + except ImportError as e: + raise AirflowOptionalProviderFeatureException(e) except RetryError: raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.") except requests_exceptions.HTTPError as e: @@ -342,6 +341,11 @@ async def _a_get_aad_token(self, resource: str) -> str: self.log.info("Existing AAD token is expired, or going to expire soon. Refreshing...") try: + from azure.identity.aio import ( + ClientSecretCredential as AsyncClientSecretCredential, + ManagedIdentityCredential as AsyncManagedIdentityCredential, + ) + async for attempt in self._a_get_retry_object(): with attempt: if self.databricks_conn.extra_dejson.get("use_azure_managed_identity", False): @@ -361,6 +365,8 @@ async def _a_get_aad_token(self, resource: str) -> str: self._is_oauth_token_valid(jsn) self.oauth_tokens[resource] = jsn break + except ImportError as e: + raise AirflowOptionalProviderFeatureException(e) except RetryError: raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.") except aiohttp.ClientResponseError as err: diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml index 5bc1d901f7ea9..d790d5707209f 100644 --- a/airflow/providers/databricks/provider.yaml +++ b/airflow/providers/databricks/provider.yaml @@ -76,7 +76,7 @@ dependencies: - pandas>=2.1.2,<2.2;python_version>="3.9" - pandas>=1.5.3,<2.2;python_version<"3.9" - pyarrow>=14.0.1 - - azure-identity>=1.3.1 + additional-extras: # pip install apache-airflow-providers-databricks[sdk] @@ -84,6 +84,10 @@ additional-extras: description: Install Databricks SDK dependencies: - databricks-sdk==0.10.0 + - name: azure-identity + description: Install Azure Identity client library + dependencies: + - azure-identity>=1.3.1 devel-dependencies: - deltalake>=0.12.0 diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 69cb5bda217f6..1663a9ed2a1d3 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -415,7 +415,6 @@ "aiohttp>=3.9.2, <4", "apache-airflow-providers-common-sql>=1.10.0", "apache-airflow>=2.7.0", - "azure-identity>=1.3.1", "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0", "mergedeep>=1.3.4", "pandas>=1.5.3,<2.2;python_version<\"3.9\"",