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 a8031341e0e Resubmit lost Databricks runs instead of failing every
retry (#76718)
a8031341e0e is described below
commit a8031341e0e52baf3fb0a88b79767cbf7f6b1eba
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Nov 10 10:08:08 2026 +0530
Resubmit lost Databricks runs instead of failing every retry (#76718)
---
.../src/airflow/providers/databricks/exceptions.py | 5 +-
.../providers/databricks/hooks/databricks_base.py | 6 +-
.../providers/databricks/operators/databricks.py | 38 +++++++++--
.../tests/unit/databricks/hooks/test_databricks.py | 25 ++++++++
.../unit/databricks/operators/test_databricks.py | 73 +++++++++++++++++++++-
5 files changed, 136 insertions(+), 11 deletions(-)
diff --git
a/providers/databricks/src/airflow/providers/databricks/exceptions.py
b/providers/databricks/src/airflow/providers/databricks/exceptions.py
index 01d47523929..88e44011a0b 100644
--- a/providers/databricks/src/airflow/providers/databricks/exceptions.py
+++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py
@@ -51,9 +51,12 @@ class
DatabricksAgentInvocationTimeout(DatabricksAgentInvocationError):
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:
+ def __init__(
+ self, message: str, *, http_status_code: int | None = None,
error_code: str | None = None
+ ) -> None:
super().__init__(message)
self.http_status_code = http_status_code
+ self.error_code = error_code
class DatabricksUnityMCPError(DatabricksApiError):
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 d5dcedaa451..f9a637c4f12 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -1388,7 +1388,11 @@ 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 DatabricksApiError(msg,
http_status_code=e.response.status_code) from e
+ raise DatabricksApiError(
+ msg,
+ http_status_code=e.response.status_code,
+ error_code=self._get_error_code(e) or None,
+ ) from e
raise
async def _a_do_api_call(self, endpoint_info: tuple[str, str], json:
dict[str, Any] | None = None):
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index 0f01fa2ae55..b17259b6775 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -23,6 +23,7 @@ import ast
import copy
import hashlib
import json as json_utils
+import re
import time
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
@@ -82,6 +83,19 @@ XCOM_JOB_ID_KEY = "job_id"
XCOM_RUN_PAGE_URL_KEY = "run_page_url"
XCOM_STATEMENT_ID_KEY = "statement_id"
+# jobs/runs/get answers 400 INVALID_PARAMETER_VALUE for a run it no longer
knows about, never 404.
+# These are the two message forms the Databricks SDK overrides to
ResourceDoesNotExist; the code
+# alone also covers malformed input, which must stay an error rather than
trigger a resubmit.
+_RUN_MISSING_MESSAGE_PATTERN = re.compile(r"Run .* does not exist|Run: .* in
job: .* doesn't exist")
+
+
+def _is_run_missing_error(exc: DatabricksApiError) -> bool:
+ """Whether ``exc`` means Databricks has no record of the run (expired,
deleted, or foreign)."""
+ if exc.error_code == "RESOURCE_DOES_NOT_EXIST":
+ return True
+ return exc.error_code == "INVALID_PARAMETER_VALUE" and
bool(_RUN_MISSING_MESSAGE_PATTERN.search(str(exc)))
+
+
# Mapping from Airflow TriggerRule to Databricks RunIf condition.
# Only non-default rules are included; ALL_SUCCESS is the Databricks default
# and is omitted to keep the task JSON minimal.
@@ -1024,9 +1038,15 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
BaseOperator):
try:
run_info = self._hook.get_run(external_id)
except DatabricksApiError as e:
- # A stored run whose history expired (or was deleted) returns 404.
Report a sentinel that
- # is neither active nor succeeded so the mixin resubmits fresh
instead of failing the task.
- if e.http_status_code == 404:
+ # Neither active nor succeeded, so the mixin resubmits instead of
failing every retry.
+ if _is_run_missing_error(e):
+ self.log.warning(
+ "Databricks has no record of run %s, so a fresh run will
be submitted. "
+ "If the workspace connection changed, the original run may
still be going there. "
+ "Databricks said: %s",
+ external_id,
+ e,
+ )
return "NOT_FOUND:"
raise
state = RunState(**run_info["state"])
@@ -1520,9 +1540,15 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
try:
run_info = self._hook.get_run(external_id)
except DatabricksApiError as e:
- # A stored run whose history expired (or was deleted) returns 404.
Report a sentinel that
- # is neither active nor succeeded so the mixin resubmits fresh
instead of failing the task.
- if e.http_status_code == 404:
+ # Neither active nor succeeded, so the mixin resubmits instead of
failing every retry.
+ if _is_run_missing_error(e):
+ self.log.warning(
+ "Databricks has no record of run %s, so a fresh run will
be submitted. "
+ "If the workspace connection changed, the original run may
still be going there. "
+ "Databricks said: %s",
+ external_id,
+ e,
+ )
return "NOT_FOUND:"
raise
state = RunState(**run_info["state"])
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
index ac6d0705777..b73a6990de9 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
@@ -402,6 +402,31 @@ class TestDatabricksHook:
# simple backcompat check, so that we do not break existing code that
expects AirflowException
assert isinstance(exc_info.value, AirflowException)
+ @pytest.mark.parametrize(
+ ("body", "expected_error_code"),
+ [
+ (
+ {"error_code": "INVALID_PARAMETER_VALUE", "message": "Run 1
does not exist."},
+ "INVALID_PARAMETER_VALUE",
+ ),
+ ({"message": "no code in this body"}, None),
+ ],
+ )
+ @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
+ def test_failing_do_api_call_carries_error_code(self, mock_requests, body,
expected_error_code):
+ hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+ response = mock.MagicMock()
+ response.status_code = 400
+ response.json.return_value = body
+ response.content = json.dumps(body).encode()
+ response.raise_for_status.side_effect =
requests_exceptions.HTTPError(response=response)
+ mock_requests.post.side_effect = itertools.repeat(response)
+
+ with pytest.raises(DatabricksApiError) as exc_info:
+ hook._do_api_call(SUBMIT_RUN_ENDPOINT, {})
+
+ assert exc_info.value.error_code == expected_error_code
+
def test_do_api_call_succeeds_after_retrying(self):
hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index 40672c42f2e..1f284e530f9 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import copy
import hashlib
+import json
from datetime import datetime, timedelta
from typing import Any
from unittest import mock
@@ -47,6 +48,7 @@ from airflow.providers.databricks.operators.databricks import
(
DatabricksSubmitRunOperator,
DatabricksTaskBaseOperator,
DatabricksTaskOperator,
+ _is_run_missing_error,
)
from airflow.providers.databricks.triggers.databricks import (
DatabricksExecutionTrigger,
@@ -345,6 +347,39 @@ def make_run_with_state_mock(
)
+def make_api_error(status_code: int, error_code: str, message: str) ->
DatabricksApiError:
+ body = json.dumps({"error_code": error_code, "message": message})
+ return DatabricksApiError(
+ f"Response: {body}, Status Code: {status_code}",
+ http_status_code=status_code,
+ error_code=error_code,
+ )
+
+
+class TestIsRunMissingError:
+ @pytest.mark.parametrize(
+ ("status_code", "error_code", "message", "expected"),
+ [
+ (400, "INVALID_PARAMETER_VALUE", "Run 999 does not exist.", True),
+ (400, "INVALID_PARAMETER_VALUE", "Run: 999 in job: 42 doesn't
exist", True),
+ (400, "RESOURCE_DOES_NOT_EXIST", "Run 999 does not exist.", True),
+ (404, "RESOURCE_DOES_NOT_EXIST", "Run 999 does not exist.", True),
+ # A status code on its own proves nothing: an unauthenticated
proxy can answer 404.
+ (404, "", "Run 999 does not exist.", False),
+ # Same code, but bad input rather than a lost run: must stay an
error.
+ (400, "INVALID_PARAMETER_VALUE", "Invalid value for run_id:
'abc'", False),
+ (403, "PERMISSION_DENIED", "User not authorized", False),
+ (500, "INTERNAL_ERROR", "Run 999 does not exist.", False),
+ ],
+ )
+ def test_detects_only_a_lost_run(self, status_code, error_code, message,
expected):
+ assert _is_run_missing_error(make_api_error(status_code, error_code,
message)) is expected
+
+ def test_unparsed_error_code_is_not_a_lost_run(self):
+ exc = DatabricksApiError("Response: Bad Gateway, Status Code: 502",
http_status_code=502)
+ assert _is_run_missing_error(exc) is False
+
+
class TestDatabricksCreateJobsOperator:
def test_init_with_named_parameters(self):
"""
@@ -2229,7 +2264,7 @@ class TestDatabricksSubmitRunOperatorDurable:
db_mock = db_mock_class.return_value
db_mock.submit_run.return_value = RUN_ID
db_mock.get_run.side_effect = [
- DatabricksApiError("Response: RESOURCE_DOES_NOT_EXIST, Status
Code: 404", http_status_code=404),
+ make_api_error(400, "INVALID_PARAMETER_VALUE", "Run 999 does not
exist."),
self._state("TERMINATED", "SUCCESS"),
]
task_store = MagicMock(spec_set=["get", "set"])
@@ -2240,6 +2275,23 @@ class TestDatabricksSubmitRunOperatorDurable:
task_store.set.assert_called_once_with("databricks_run_id", RUN_ID)
assert op.run_id == RUN_ID
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def
test_does_not_resubmit_when_the_stored_run_id_is_rejected_as_malformed(self,
db_mock_class):
+ op = DatabricksSubmitRunOperator(
+ task_id=TASK_ID, json={"new_cluster": NEW_CLUSTER,
"notebook_task": NOTEBOOK_TASK}
+ )
+ db_mock = db_mock_class.return_value
+ db_mock.get_run.side_effect = make_api_error(
+ 400, "INVALID_PARAMETER_VALUE", "Invalid value for run_id: 'abc'"
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = 999
+
+ with pytest.raises(DatabricksApiError):
+ op.execute(self._context(task_store))
+
+ db_mock.submit_run.assert_not_called()
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_durable_false_never_touches_task_state_store(self, db_mock_class):
op = DatabricksSubmitRunOperator(
@@ -3584,9 +3636,9 @@ class TestDatabricksRunNowOperatorDurable:
op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
- # the stored run's history has expired -> get_run 404s; the fresh run
then polls to success.
+ # the stored run's history has expired -> get_run rejects it; the
fresh run polls to success.
db_mock.get_run.side_effect = [
- DatabricksApiError("Response: RESOURCE_DOES_NOT_EXIST, Status
Code: 404", http_status_code=404),
+ make_api_error(400, "INVALID_PARAMETER_VALUE", "Run 999 does not
exist."),
self._state("TERMINATED", "SUCCESS"),
]
task_store = MagicMock(spec_set=["get", "set"])
@@ -3598,6 +3650,21 @@ class TestDatabricksRunNowOperatorDurable:
task_store.set.assert_called_once_with("databricks_run_now_id", RUN_ID)
assert op.run_id == RUN_ID
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def
test_does_not_resubmit_when_the_stored_run_id_is_rejected_as_malformed(self,
db_mock_class):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ db_mock = db_mock_class.return_value
+ db_mock.get_run.side_effect = make_api_error(
+ 400, "INVALID_PARAMETER_VALUE", "Invalid value for run_id: 'abc'"
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = 999
+
+ with pytest.raises(DatabricksApiError):
+ op.execute(self._context(task_store))
+
+ db_mock.run_now.assert_not_called()
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_durable_false_never_touches_task_state_store(self, db_mock_class):
op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
durable=False)
