This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3e48e8d852a Surface databricks API errors as a separate exception
(#69019)
3e48e8d852a is described below
commit 3e48e8d852a8762bad5e76106a557067e05b3b0f
Author: Amogh Desai <[email protected]>
AuthorDate: Fri Jun 26 15:16:57 2026 +0530
Surface databricks API errors as a separate exception (#69019)
---
.../src/airflow/providers/databricks/exceptions.py | 8 ++++++++
.../airflow/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 59c8f3fb606..0831810f70f 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 7ab51fc7dfa..f846ab23a54 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 tenacity import (
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
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 558e39a2129..1a03934ec83 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 requests.auth import HTTPBasicAuth
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 @@ class TestDatabricksHook:
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 cf9bb1d8647..1fb46185490 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
