From b314b31762d119531ea68a0a50fccbe3d4d0db0f Mon Sep 17 00:00:00 2001
From: Amogh Desai <amoghrajesh1999@gmail.com>
Date: Fri, 26 Jun 2026 13:42:34 +0530
Subject: [PATCH] Surface databricks API errors as a separate exception

---
 .../src/airflow/providers/databricks/exceptions.py |  8 ++++++++
 .../providers/databricks/hooks/databricks_base.py  |  7 +++++--
 .../tests/unit/databricks/hooks/test_databricks.py | 14 ++++++++++++++
 scripts/ci/prek/known_airflow_exceptions.txt       |  2 +-
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py
index 59c8f3fb60649..0831810f70f35 100644
--- a/providers/databricks/src/airflow/providers/databricks/exceptions.py
+++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py
@@ -34,3 +34,11 @@ class DatabricksSqlExecutionTimeout(DatabricksSqlExecutionError):
 
 class DatabricksOperatorPayloadError(AirflowException):
     """Raised when a Databricks operator payload is invalid."""
+
+
+class DatabricksApiError(AirflowException):
+    """Raised when a Databricks REST API call returns an error response."""
+
+    def __init__(self, message: str, *, http_status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.http_status_code = http_status_code
diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
index 7ab51fc7dfa12..f846ab23a5437 100644
--- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
+++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -51,6 +51,7 @@
 
 from airflow import __version__
 from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.exceptions import DatabricksApiError
 from airflow.providers_manager import ProvidersManager
 
 try:
@@ -1150,7 +1151,7 @@ def _do_api_call(
         except requests_exceptions.HTTPError as e:
             if wrap_http_errors:
                 msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
-                raise AirflowException(msg)
+                raise DatabricksApiError(msg, http_status_code=e.response.status_code) from e
             raise
 
     async def _a_do_api_call(self, endpoint_info: tuple[str, str], json: dict[str, Any] | None = None):
@@ -1211,7 +1212,9 @@ async def _a_do_api_call(self, endpoint_info: tuple[str, str], json: dict[str, A
         except RetryError:
             raise AirflowException(f"API requests to Databricks failed {self.retry_limit} times. Giving up.")
         except aiohttp.ClientResponseError as err:
-            raise AirflowException(f"Response: {err.message}, Status Code: {err.status}")
+            raise DatabricksApiError(
+                f"Response: {err.message}, Status Code: {err.status}", http_status_code=err.status
+            ) from err
 
     @staticmethod
     def _get_error_code(exception: BaseException) -> str:
diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
index 558e39a212987..1a03934ec8378 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
@@ -35,6 +35,7 @@
 
 from airflow.models import Connection
 from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.databricks.exceptions import DatabricksApiError
 from airflow.providers.databricks.hooks.databricks import (
     GET_RUN_ENDPOINT,
     SUBMIT_RUN_ENDPOINT,
@@ -404,6 +405,19 @@ def test_do_api_call_does_not_retry_with_non_retryable_error(self, mock_requests
 
             mock_errors.assert_not_called()
 
+    @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
+    def test_failing_do_api_call_raises_exception(self, mock_requests):
+        hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+
+        setup_mock_requests(mock_requests, requests_exceptions.HTTPError, status_code=404)
+
+        with pytest.raises(DatabricksApiError) as exc_info:
+            hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+
+        assert exc_info.value.http_status_code == 404
+        # simple backcompat check, so that we do not break existing code that expects AirflowException
+        assert isinstance(exc_info.value, AirflowException)
+
     def test_do_api_call_succeeds_after_retrying(self):
         hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
 
diff --git a/scripts/ci/prek/known_airflow_exceptions.txt b/scripts/ci/prek/known_airflow_exceptions.txt
index cf9bb1d8647e5..1fb461854905d 100644
--- a/scripts/ci/prek/known_airflow_exceptions.txt
+++ b/scripts/ci/prek/known_airflow_exceptions.txt
@@ -174,7 +174,7 @@ providers/common/sql/src/airflow/providers/common/sql/operators/sql.py::7
 providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py::6
 providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py::1
 providers/databricks/src/airflow/providers/databricks/hooks/databricks.py::8
-providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py::46
+providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py::44
 providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py::2
 providers/databricks/src/airflow/providers/databricks/operators/databricks.py::6
 providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::12