apache-airflow-providers-databricks 7.7.2rc1__tar.gz → 7.7.3rc1__tar.gz

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-databricks
-Version: 7.7.2rc1
+Version: 7.7.3rc1
 Summary: Provider package apache-airflow-providers-databricks for Apache Airflow
 Keywords: airflow-provider,databricks,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -38,8 +38,8 @@ Requires-Dist: apache-airflow-providers-openlineage>=2.3.0rc1 ; extra == "openli
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -75,7 +75,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.7.2``
+Release: ``7.7.3``
 
 
 `Databricks <https://databricks.com/>`__
@@ -88,7 +88,7 @@ This is a provider package for ``databricks`` provider. All classes for this pro
 are in ``airflow.providers.databricks`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/>`_.
 
 Installation
 ------------
@@ -141,6 +141,19 @@ Dependent package
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
+Optional dependencies
+----------------------
+
+==================  ================================================================
+Extra               Dependencies
+==================  ================================================================
+``sdk``             ``databricks-sdk==0.10.0``
+``azure-identity``  ``azure-identity>=1.3.1``
+``fab``             ``apache-airflow-providers-fab>=2.2.0; python_version < '3.13'``
+``standard``        ``apache-airflow-providers-standard``
+``openlineage``     ``apache-airflow-providers-openlineage>=2.3.0``
+==================  ================================================================
+
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html>`_.
 

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.7.2``
+Release: ``7.7.3``
 
 
 `Databricks <https://databricks.com/>`__
@@ -36,7 +36,7 @@ This is a provider package for ``databricks`` provider. All classes for this pro
 are in ``airflow.providers.databricks`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/>`_.
 
 Installation
 ------------
@@ -89,5 +89,18 @@ Dependent package
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
+Optional dependencies
+----------------------
+
+==================  ================================================================
+Extra               Dependencies
+==================  ================================================================
+``sdk``             ``databricks-sdk==0.10.0``
+``azure-identity``  ``azure-identity>=1.3.1``
+``fab``             ``apache-airflow-providers-fab>=2.2.0; python_version < '3.13'``
+``standard``        ``apache-airflow-providers-standard``
+``openlineage``     ``apache-airflow-providers-openlineage>=2.3.0``
+==================  ================================================================
+
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html>`_.

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/docs/changelog.rst

@@ -26,6 +26,25 @@
 Changelog
 ---------
 
+7.7.3
+.....
+
+
+Release Date: ``|PypiReleaseDate|``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``Fix metadata service check handle 429 (#55462)``
+
+Misc
+~~~~
+
+* ``Switch all airflow logging to structlog (#52651)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+
 7.7.2
 .....
 

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/docs/index.rst

@@ -78,7 +78,7 @@ apache-airflow-providers-databricks package
 `Databricks <https://databricks.com/>`__
 
 
-Release: 7.7.2
+Release: 7.7.3
 
 Provider package
 ----------------
@@ -143,5 +143,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-databricks 7.7.2 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2.tar.gz.sha512>`__)
-* `The apache-airflow-providers-databricks 7.7.2 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.2-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-databricks 7.7.3 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3.tar.gz.sha512>`__)
+* `The apache-airflow-providers-databricks 7.7.3 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.3-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/provider.yaml

@@ -22,12 +22,13 @@ description: |
     `Databricks <https://databricks.com/>`__
 
 state: ready
-source-date-epoch: 1756876778
+source-date-epoch: 1757950169
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 7.7.3
   - 7.7.2
   - 7.7.1
   - 7.7.0

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "7.7.2rc1"
+version = "7.7.3rc1"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -133,8 +133,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2"
-"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html"
+"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3"
+"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/src/airflow/providers/databricks/__init__.py

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

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/src/airflow/providers/databricks/hooks/databricks_base.py

@@ -121,6 +121,9 @@ class BaseDatabricksHook(BaseHook):
         self.oauth_tokens: dict[str, dict] = {}
         self.token_timeout_seconds = 10
         self.caller = caller
+        self._metadata_cache: dict[str, Any] = {}
+        self._metadata_expiry: float = 0
+        self._metadata_ttl: int = 300
 
         def my_after_func(retry_state):
             self._log_request_error(retry_state.attempt_number, retry_state.outcome)
@@ -515,43 +518,64 @@ class BaseDatabricksHook(BaseHook):
 
         return int(token[time_key]) > (int(time.time()) + TOKEN_REFRESH_LEAD_TIME)
 
-    @staticmethod
-    def _check_azure_metadata_service() -> None:
+    def _check_azure_metadata_service(self) -> None:
         """
-        Check for Azure Metadata Service.
+        Check for Azure Metadata Service (with caching).
 
         https://docs.microsoft.com/en-us/azure/virtual-machines/linux/instance-metadata-service
         """
+        if self._metadata_cache and time.time() < self._metadata_expiry:
+            return
         try:
-            jsn = requests.get(
-                AZURE_METADATA_SERVICE_INSTANCE_URL,
-                params={"api-version": "2021-02-01"},
-                headers={"Metadata": "true"},
-                timeout=2,
-            ).json()
-            if "compute" not in jsn or "azEnvironment" not in jsn["compute"]:
-                raise AirflowException(
-                    f"Was able to fetch some metadata, but it doesn't look like Azure Metadata: {jsn}"
-                )
+            for attempt in self._get_retry_object():
+                with attempt:
+                    response = requests.get(
+                        AZURE_METADATA_SERVICE_INSTANCE_URL,
+                        params={"api-version": "2021-02-01"},
+                        headers={"Metadata": "true"},
+                        timeout=2,
+                    )
+                    response.raise_for_status()
+                    response_json = response.json()
+
+                    self._validate_azure_metadata_service(response_json)
+                    self._metadata_cache = response_json
+                    self._metadata_expiry = time.time() + self._metadata_ttl
+                    break
+        except RetryError:
+            raise ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.")
         except (requests_exceptions.RequestException, ValueError) as e:
-            raise AirflowException(f"Can't reach Azure Metadata Service: {e}")
+            raise ConnectionError(f"Can't reach Azure Metadata Service: {e}")
 
     async def _a_check_azure_metadata_service(self):
         """Async version of `_check_azure_metadata_service()`."""
+        if self._metadata_cache and time.time() < self._metadata_expiry:
+            return
         try:
-            async with self._session.get(
-                url=AZURE_METADATA_SERVICE_INSTANCE_URL,
-                params={"api-version": "2021-02-01"},
-                headers={"Metadata": "true"},
-                timeout=2,
-            ) as resp:
-                jsn = await resp.json()
-            if "compute" not in jsn or "azEnvironment" not in jsn["compute"]:
-                raise AirflowException(
-                    f"Was able to fetch some metadata, but it doesn't look like Azure Metadata: {jsn}"
-                )
-        except (requests_exceptions.RequestException, ValueError) as e:
-            raise AirflowException(f"Can't reach Azure Metadata Service: {e}")
+            async for attempt in self._a_get_retry_object():
+                with attempt:
+                    async with self._session.get(
+                        url=AZURE_METADATA_SERVICE_INSTANCE_URL,
+                        params={"api-version": "2021-02-01"},
+                        headers={"Metadata": "true"},
+                        timeout=2,
+                    ) as resp:
+                        resp.raise_for_status()
+                        response_json = await resp.json()
+                    self._validate_azure_metadata_service(response_json)
+                    self._metadata_cache = response_json
+                    self._metadata_expiry = time.time() + self._metadata_ttl
+                    break
+        except RetryError:
+            raise ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.")
+        except (aiohttp.ClientError, ValueError) as e:
+            raise ConnectionError(f"Can't reach Azure Metadata Service: {e}")
+
+    def _validate_azure_metadata_service(self, response_json: dict) -> None:
+        if "compute" not in response_json or "azEnvironment" not in response_json["compute"]:
+            raise ValueError(
+                f"Was able to fetch some metadata, but it doesn't look like Azure Metadata: {response_json}"
+            )
 
     def _get_token(self, raise_error: bool = False) -> str | None:
         if "token" in self.databricks_conn.extra_dejson:

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/src/airflow/providers/databricks/operators/databricks.py

@@ -24,7 +24,6 @@ import time
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from functools import cached_property
-from logging import Logger
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
@@ -60,12 +59,8 @@ if TYPE_CHECKING:
         DatabricksWorkflowTaskGroup,
     )
     from airflow.providers.openlineage.extractors import OperatorLineage
-    from airflow.utils.context import Context
-
-    try:
-        from airflow.sdk import TaskGroup
-    except ImportError:
-        from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
+    from airflow.sdk import TaskGroup
+    from airflow.sdk.types import Context, Logger
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk import BaseOperatorLink

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/src/airflow/providers/databricks/plugins/databricks_workflow.py

@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import logging
 import os
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote
@@ -45,6 +44,7 @@ if TYPE_CHECKING:
 
     from airflow.models import BaseOperator
     from airflow.providers.databricks.operators.databricks import DatabricksTaskBaseOperator
+    from airflow.sdk.types import Logger
     from airflow.utils.context import Context
 
 
@@ -62,7 +62,7 @@ def get_auth_decorator():
 
 
 def get_databricks_task_ids(
-    group_id: str, task_map: dict[str, DatabricksTaskBaseOperator], log: logging.Logger
+    group_id: str, task_map: dict[str, DatabricksTaskBaseOperator], log: Logger
 ) -> list[str]:
     """
     Return a list of all Databricks task IDs for a dictionary of Airflow tasks.
@@ -112,7 +112,7 @@ if not AIRFLOW_V_3_0_PLUS:
 
     @provide_session
     def _clear_task_instances(
-        dag_id: str, run_id: str, task_ids: list[str], log: logging.Logger, session: Session = NEW_SESSION
+        dag_id: str, run_id: str, task_ids: list[str], log: Logger, session: Session = NEW_SESSION
     ) -> None:
         dag = _get_dag(dag_id, session=session)
         log.debug("task_ids %s to clear", str(task_ids))
@@ -145,7 +145,7 @@ def _repair_task(
     databricks_conn_id: str,
     databricks_run_id: int,
     tasks_to_repair: list[str],
-    logger: logging.Logger,
+    logger: Logger,
 ) -> int:
     """
     Repair a Databricks task using the Databricks API.
@@ -294,7 +294,7 @@ class WorkflowJobRunLink(BaseOperatorLink, LoggingMixin):
 def store_databricks_job_run_link(
     context: Context,
     metadata: Any,
-    logger: logging.Logger,
+    logger: Logger,
 ) -> None:
     """
     Store the Databricks job run link in XCom during task execution.
@@ -368,7 +368,7 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
                 children[child_id] = child
         return children
 
-    def get_tasks_to_run(self, ti_key: TaskInstanceKey, operator: BaseOperator, log: logging.Logger) -> str:
+    def get_tasks_to_run(self, ti_key: TaskInstanceKey, operator: BaseOperator, log: Logger) -> str:
         task_group = operator.task_group
         if not task_group:
             raise AirflowException("Task group is required for generating repair link.")

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/tests/unit/databricks/hooks/test_databricks.py

@@ -1449,7 +1449,6 @@ class TestDatabricksHookConnSettings(TestDatabricksHookToken):
     @pytest.mark.asyncio
     @mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
     async def test_async_do_api_call_only_existing_response_properties_are_read(self, mock_get):
-        self.hook.log.setLevel("DEBUG")
         response = mock_get.return_value.__aenter__.return_value
         response.mock_add_spec(aiohttp.ClientResponse, spec_set=True)
         response.json = AsyncMock(return_value={"bar": "baz"})

{apache_airflow_providers_databricks-7.7.2rc1 → apache_airflow_providers_databricks-7.7.3rc1}/tests/unit/databricks/hooks/test_databricks_base.py

@@ -26,7 +26,7 @@ import time_machine
 from aiohttp.client_exceptions import ClientConnectorError
 from requests import exceptions as requests_exceptions
 from requests.auth import HTTPBasicAuth
-from tenacity import Future, RetryError
+from tenacity import AsyncRetrying, Future, RetryError, retry_if_exception, stop_after_attempt, wait_fixed
 
 from airflow.exceptions import AirflowException
 from airflow.models import Connection
@@ -768,3 +768,145 @@ class TestBaseDatabricksHook:
         exception.response = mock_response
         hook = BaseDatabricksHook()
         assert hook._get_error_code(exception) == "INVALID_REQUEST"
+
+    @mock.patch("requests.get")
+    @time_machine.travel("2025-07-12 12:00:00")
+    def test_check_azure_metadata_service_normal(self, mock_get):
+        travel_time = int(datetime(2025, 7, 12, 12, 0, 0).timestamp())
+        hook = BaseDatabricksHook()
+        mock_response = {"compute": {"azEnvironment": "AzurePublicCloud"}}
+        mock_get.return_value.json.return_value = mock_response
+
+        hook._check_azure_metadata_service()
+
+        assert hook._metadata_cache == mock_response
+        assert int(hook._metadata_expiry) == travel_time + hook._metadata_ttl
+
+    @mock.patch("requests.get")
+    @time_machine.travel("2025-07-12 12:00:00")
+    def test_check_azure_metadata_service_cached(self, mock_get):
+        travel_time = int(datetime(2025, 7, 12, 12, 0, 0).timestamp())
+        hook = BaseDatabricksHook()
+        mock_response = {"compute": {"azEnvironment": "AzurePublicCloud"}}
+        hook._metadata_cache = mock_response
+        hook._metadata_expiry = travel_time + 1000
+
+        hook._check_azure_metadata_service()
+        mock_get.assert_not_called()
+
+    @mock.patch("requests.get")
+    def test_check_azure_metadata_service_http_error(self, mock_get):
+        hook = BaseDatabricksHook()
+        mock_get.side_effect = requests_exceptions.RequestException("Fail")
+
+        with pytest.raises(ConnectionError, match="Can't reach Azure Metadata Service"):
+            hook._check_azure_metadata_service()
+        assert hook._metadata_cache == {}
+        assert hook._metadata_expiry == 0
+
+    @mock.patch("requests.get")
+    def test_check_azure_metadata_service_retry_error(self, mock_get):
+        hook = BaseDatabricksHook()
+
+        resp_429 = mock.Mock()
+        resp_429.status_code = 429
+        resp_429.content = b"Too many requests"
+        http_error = requests_exceptions.HTTPError(response=resp_429)
+        mock_get.side_effect = http_error
+
+        with pytest.raises(ConnectionError, match="Failed to reach Azure Metadata Service after 3 retries."):
+            hook._check_azure_metadata_service()
+        assert mock_get.call_count == 3
+
+    @pytest.mark.asyncio
+    @mock.patch("aiohttp.ClientSession.get")
+    async def test_a_check_azure_metadata_service_normal(self, mock_get):
+        hook = BaseDatabricksHook()
+
+        async_mock = mock.AsyncMock()
+        async_mock.__aenter__.return_value = async_mock
+        async_mock.__aexit__.return_value = None
+        async_mock.json.return_value = {"compute": {"azEnvironment": "AzurePublicCloud"}}
+
+        mock_get.return_value = async_mock
+
+        async with aiohttp.ClientSession() as session:
+            hook._session = session
+            mock_attempt = mock.Mock()
+            mock_attempt.__enter__ = mock.Mock(return_value=None)
+            mock_attempt.__exit__ = mock.Mock(return_value=None)
+
+            async def mock_retry_generator():
+                yield mock_attempt
+
+            hook._a_get_retry_object = mock.Mock(return_value=mock_retry_generator())
+            await hook._a_check_azure_metadata_service()
+
+            assert hook._metadata_cache["compute"]["azEnvironment"] == "AzurePublicCloud"
+            assert hook._metadata_expiry > 0
+
+    @pytest.mark.asyncio
+    @mock.patch("aiohttp.ClientSession.get")
+    @time_machine.travel("2025-07-12 12:00:00")
+    async def test_a_check_azure_metadata_service_cached(self, mock_get):
+        travel_time = int(datetime(2025, 7, 12, 12, 0, 0).timestamp())
+        hook = BaseDatabricksHook()
+        hook._metadata_cache = {"compute": {"azEnvironment": "AzurePublicCloud"}}
+        hook._metadata_expiry = travel_time + 1000
+
+        async with aiohttp.ClientSession() as session:
+            hook._session = session
+            await hook._a_check_azure_metadata_service()
+            mock_get.assert_not_called()
+
+    @pytest.mark.asyncio
+    @mock.patch("aiohttp.ClientSession.get")
+    async def test_a_check_azure_metadata_service_http_error(self, mock_get):
+        hook = BaseDatabricksHook()
+
+        async_mock = mock.AsyncMock()
+        async_mock.__aenter__.side_effect = aiohttp.ClientError("Fail")
+        async_mock.__aexit__.return_value = None
+        mock_get.return_value = async_mock
+
+        async with aiohttp.ClientSession() as session:
+            hook._session = session
+            mock_attempt = mock.Mock()
+            mock_attempt.__enter__ = mock.Mock(return_value=None)
+            mock_attempt.__exit__ = mock.Mock(return_value=None)
+
+            async def mock_retry_generator():
+                yield mock_attempt
+
+            hook._a_get_retry_object = mock.Mock(return_value=mock_retry_generator())
+
+            with pytest.raises(ConnectionError, match="Can't reach Azure Metadata Service"):
+                await hook._a_check_azure_metadata_service()
+            assert hook._metadata_cache == {}
+            assert hook._metadata_expiry == 0
+
+    @pytest.mark.asyncio
+    @mock.patch("aiohttp.ClientSession.get")
+    async def test_a_check_azure_metadata_service_retry_error(self, mock_get):
+        hook = BaseDatabricksHook()
+
+        mock_get.side_effect = aiohttp.ClientResponseError(
+            request_info=mock.Mock(), history=(), status=429, message="429 Too Many Requests"
+        )
+
+        async with aiohttp.ClientSession() as session:
+            hook._session = session
+
+            hook._a_get_retry_object = lambda: AsyncRetrying(
+                stop=stop_after_attempt(hook.retry_limit),
+                wait=wait_fixed(0),
+                retry=retry_if_exception(hook._retryable_error),
+            )
+
+            hook._validate_azure_metadata_service = mock.Mock()
+
+            with pytest.raises(
+                ConnectionError, match="Failed to reach Azure Metadata Service after 3 retries."
+            ):
+                await hook._a_check_azure_metadata_service()
+            assert mock_get.call_count == 3