Add azure-identity as an additional-extras dependency

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:

airflow/providers/databricks/provider.yaml

@@ -76,14 +76,18 @@ 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]
   - name: sdk
     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

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\"",