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 579aa9f0915 Reconnect to blocked Databricks runs in SubmitRunOperator
durable retries (#69195)
579aa9f0915 is described below
commit 579aa9f0915d219667ddb2d34004b6fa7097d37c
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Jul 1 14:17:29 2026 +0530
Reconnect to blocked Databricks runs in SubmitRunOperator durable retries
(#69195)
---
.../providers/databricks/operators/databricks.py | 9 ++++++++-
.../unit/databricks/operators/test_databricks.py | 19 +++++++++++++++++++
2 files changed, 27 insertions(+), 1 deletion(-)
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index 225bdf7595a..e9f2a016063 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -937,7 +937,14 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
BaseOperator):
def is_job_active(self, status: str) -> bool:
life_cycle_state = status.split(":", 1)[0]
- return life_cycle_state in ("PENDING", "QUEUED", "RUNNING",
"TERMINATING")
+ return life_cycle_state in (
+ "PENDING",
+ "QUEUED",
+ "RUNNING",
+ "TERMINATING",
+ "BLOCKED",
+ "WAITING_FOR_RETRY",
+ )
def is_job_succeeded(self, status: str) -> bool:
return status == "TERMINATED:SUCCESS"
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index a0885ac48f4..ccf742338c0 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -1718,6 +1718,23 @@ class TestDatabricksSubmitRunOperatorDurable:
task_store.set.assert_not_called()
assert op.run_id == RUN_ID
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_reconnects_to_blocked_run_without_resubmitting(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
+ # A gated run sits in BLOCKED; the durable retry must reconnect and
wait, not resubmit.
+ db_mock.get_run.side_effect = [self._state("BLOCKED"),
self._state("TERMINATED", "SUCCESS")]
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = RUN_ID
+
+ op.execute(self._context(task_store))
+
+ db_mock.submit_run.assert_not_called()
+ task_store.set.assert_not_called()
+ assert op.run_id == RUN_ID
+
@mock.patch("airflow.providers.databricks.operators.databricks._handle_databricks_operator_execution")
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_already_succeeded_pushes_xcoms_without_polling(self,
db_mock_class, mock_poll):
@@ -1846,6 +1863,8 @@ class TestDatabricksSubmitRunOperatorDurable:
("PENDING:", True),
("QUEUED:", True),
("TERMINATING:", True),
+ ("BLOCKED:", True),
+ ("WAITING_FOR_RETRY:", True),
("TERMINATED:SUCCESS", False),
("TERMINATED:FAILED", False),
("SKIPPED:", False),
