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 794cbd8b51c Add durable execution to `DatabricksRunNowOperator`
(#69174)
794cbd8b51c is described below
commit 794cbd8b51cd8c992e5d1212657d47ea851dd0ca
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Jul 1 14:18:28 2026 +0530
Add durable execution to `DatabricksRunNowOperator` (#69174)
---
providers/databricks/docs/operators/run_now.rst | 39 +++
.../providers/databricks/operators/databricks.py | 98 +++++++-
.../unit/databricks/operators/test_databricks.py | 270 +++++++++++++++++++--
3 files changed, 381 insertions(+), 26 deletions(-)
diff --git a/providers/databricks/docs/operators/run_now.rst
b/providers/databricks/docs/operators/run_now.rst
index f39d872a0c1..30f64139b9c 100644
--- a/providers/databricks/docs/operators/run_now.rst
+++ b/providers/databricks/docs/operators/run_now.rst
@@ -74,6 +74,45 @@ contains ``job_parameters``, it is left untouched.
# i.e. the same dict, passed straight through to the run-now request body.
+Durable execution
+^^^^^^^^^^^^^^^^^
+
+``DatabricksRunNowOperator`` triggers a run of an existing job and then polls
it to completion on
+the worker. By default the operator runs in a *durable* mode that makes this
crash-safe: the
+Databricks run id is persisted to Airflow's task state store before polling
begins, so if the
+worker crashes or is preempted and the task is retried, the operator
reconnects to the run that is
+already executing on Databricks instead of triggering a duplicate run of the
same job.
+
+On retry the operator checks the prior run's state:
+
+* if it is still running, the operator reconnects and continues polling
+* if it already succeeded, the operator returns immediately without triggering
a new run
+* if it failed terminally, or the run no longer exists because its history
expired, the operator
+ triggers a fresh run
+
+This avoids triggering the same Databricks job twice when a worker is lost
mid-run, which is common
+for long-running jobs.
+
+Durable execution requires Airflow 3.3 or newer, since it relies on the task
state store. On earlier
+Airflow versions the flag is a no-op and the operator always triggers a fresh
run on retry,
+exactly as before. If the task state store is unavailable at runtime, the
operator logs that crash
+recovery is disabled and behaves the same way.
+
+To opt out and always trigger a fresh run on retry, set ``durable=False``:
+
+.. code-block:: python
+
+ run_now = DatabricksRunNowOperator(
+ task_id="run_now",
+ job_id=123,
+ durable=False,
+ )
+
+Durable execution applies to the synchronous path. When ``deferrable=True`` is
set, the Triggerer
+already tracks the run across the wait, so deferrable mode takes precedence
and ``durable`` has no
+effect.
+
+
DatabricksRunNowDeferrableOperator
==================================
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index e9f2a016063..16827406466 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -980,7 +980,7 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
BaseOperator):
_handle_deferrable_databricks_operator_completion(event, self.log)
-class DatabricksRunNowOperator(BaseOperator):
+class DatabricksRunNowOperator(ResumableJobMixin, BaseOperator):
"""
Runs an existing Spark job run to Databricks using the
api/2.2/jobs/run-now API endpoint.
@@ -1163,6 +1163,10 @@ class DatabricksRunNowOperator(BaseOperator):
(https://docs.databricks.com/api/workspace/jobs/update). If
nothing is matched, then repair
will not get triggered.
:param cancel_previous_runs: Cancel all existing running jobs before
submitting new one.
+ :param durable: When ``True`` (the default), the Databricks run id is
persisted to task state
+ before polling begins so that a worker crash and retry reconnects to
the existing run
+ instead of triggering a duplicate run of the same job. Set to
``False`` to always trigger a
+ fresh run on retry. Requires Airflow 3.3+; on earlier versions it is
silently ignored.
.. note::
If ``job_parameters`` is not set in ``json`` and the operator's
``params`` dict is
@@ -1171,6 +1175,8 @@ class DatabricksRunNowOperator(BaseOperator):
hardcoding them in ``json``.
"""
+ external_id_key = "databricks_run_now_id"
+
# Used in airflow.models.BaseOperator
template_fields: Sequence[str] = (
"json",
@@ -1283,10 +1289,21 @@ class DatabricksRunNowOperator(BaseOperator):
)
def execute(self, context: Context):
+ if self.deferrable:
+ json = self._prepare_run_now_json()
+ self.run_id = self._hook.run_now(json)
+ _handle_deferrable_databricks_operator_execution(self, self._hook,
self.log, context)
+ else:
+ return self.execute_resumable(context)
+
+ def _build_run_now_payload(self) -> dict[str, Any]:
+ # Utility to build the run payload: merge, validate, resolve job_name
-> job_id, inject params.
+ # Kept separate from cancel_previous_runs so the reconnect path can
rebuild the payload
+ # (for repair_run) without re-cancelling the run it is reconnecting to.
json = self._get_merged_json()
self._validate_merged_json(json)
# Validate payload types before touching the hook so an invalid
payload fails fast,
- # before find_job_id_by_name / cancel_all_runs hit the Databricks API.
+ # before find_job_id_by_name hits the Databricks API.
json = cast("dict[str, Any]", normalise_json_content(json))
hook = self._hook
if "job_name" in json:
@@ -1298,23 +1315,82 @@ class DatabricksRunNowOperator(BaseOperator):
json["job_id"] = job_id
del json["job_name"]
+ if not json.get("job_parameters") and self.params:
+ json["job_parameters"] = dict(self.params)
+
+ return json
+
+ def _prepare_run_now_json(self) -> dict[str, Any]:
+ json = self._build_run_now_payload()
if self.cancel_previous_runs:
if (job_id := json.get("job_id")) is None:
raise ValueError(
"cancel_previous_runs=True requires either job_id or
job_name to be provided."
)
+ self._hook.cancel_all_runs(job_id)
- hook.cancel_all_runs(job_id)
+ self._merged_json = json
+ return json
- if not json.get("job_parameters") and self.params:
- json["job_parameters"] = dict(self.params)
+ def submit_job(self, context: Context) -> int:
+ json = self._prepare_run_now_json()
+ # Set run_id the instant the run exists so on_kill can cancel it even
if the worker dies
+ # before polling begins.
+ self.run_id = self._hook.run_now(json)
+ self.log.info("Run submitted with run_id: %s", self.run_id)
+ return self.run_id
- self._merged_json = json
- self.run_id = hook.run_now(json)
- if self.deferrable:
- _handle_deferrable_databricks_operator_execution(self, hook,
self.log, context)
- else:
- _handle_databricks_operator_execution(self, hook, self.log,
context)
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ # Databricks splits run state across life_cycle_state and
result_state; the mixin's status
+ # interface is a single string, so encode both as "LIFE_CYCLE:RESULT"
and decode them in
+ # is_job_active / is_job_succeeded.
+ 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:
+ return "NOT_FOUND:"
+ raise
+ state = RunState(**run_info["state"])
+ return f"{state.life_cycle_state}:{state.result_state}"
+
+ def is_job_active(self, status: str) -> bool:
+ life_cycle_state = status.split(":", 1)[0]
+ 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"
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
+ # submit_job and the stored external id are always Databricks run ids
(int).
+ self.run_id = cast("int", external_id)
+ # On reconnect submit_job did not run, so rebuild the merged payload
(read by the repair path
+ # in the poll helper). _build_run_now_payload resolves job_name ->
job_id but omits
+ # cancel_previous_runs, which would otherwise cancel the run we are
reconnecting to.
+ if not getattr(self, "_merged_json", None):
+ self._merged_json = self._build_run_now_payload()
+ # The run already exists here (fresh submit logged in submit_job, or
reconnect logged by the
+ # mixin), so the poll helper must not announce a submission.
+ _handle_databricks_operator_execution(self, self._hook, self.log,
context, announce_submission=False)
+ # The helper pushed run_id/run_page_url xcoms; record that so
get_job_result does not push again.
+ self._run_xcoms_pushed = True
+
+ def get_job_result(self, external_id: JsonValue, context: Context) -> None:
+ self.run_id = cast("int", external_id)
+ # The already-succeeded retry path skips polling, so push the run
xcoms here for parity with the
+ # normal success path. When polling ran, poll_until_complete already
pushed them.
+ if not getattr(self, "_run_xcoms_pushed", False) and self.do_xcom_push
and context is not None:
+ context["ti"].xcom_push(key=XCOM_RUN_ID_KEY, value=self.run_id)
+ context["ti"].xcom_push(key=XCOM_RUN_PAGE_URL_KEY,
value=self._hook.get_run_page_url(self.run_id))
+ return None
def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
if event:
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index ccf742338c0..9c83c19ebe9 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -1981,6 +1981,7 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_exec_with_json_string_and_templated_named_parameters(self,
db_mock_class):
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
json='{"job_id": "1", "notebook_params": {"source": "json"},
"jar_params": ["json"]}',
job_id="{{ params.job_id }}",
@@ -2010,7 +2011,7 @@ class TestDatabricksRunNowOperator:
r"Type \<(type|class) \'datetime.datetime\'\> used "
r"for parameter json\[test\] is not a number or a string"
)
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=json)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=json, durable=False)
with pytest.raises(AirflowException, match=exception_message):
op.execute(None)
@@ -2024,6 +2025,7 @@ class TestDatabricksRunNowOperator:
into the ``json`` template field, so a retry / deferral-resume
re-renders from the original
template instead of a clobbered dict."""
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
job_id=JOB_ID,
json={"notebook_params": {"a": "b"}},
@@ -2055,7 +2057,7 @@ class TestDatabricksRunNowOperator:
def test_exec_invalid_payload_fails_before_api_call(self, db_mock_class,
kwargs):
"""An invalid payload type must fail before ``find_job_id_by_name`` /
``cancel_all_runs`` /
``run_now`` touch the Databricks API."""
- op = DatabricksRunNowOperator(task_id=TASK_ID, json={"bad":
datetime.now()}, **kwargs)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, json={"bad":
datetime.now()}, **kwargs, durable=False)
db_mock = db_mock_class.return_value
with pytest.raises(AirflowException, match="is not a number or a
string"):
@@ -2108,6 +2110,7 @@ class TestDatabricksRunNowOperator:
``_handle_databricks_operator_execution`` -- the only path that reads
``operator._merged_json`` --
so a regression there (e.g. the attribute being unset) fails loudly
instead of passing CI."""
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
job_id=JOB_ID,
json={"job_parameters": {"k": "v"}},
@@ -2133,7 +2136,7 @@ class TestDatabricksRunNowOperator:
Test the execute function in case where the run is successful.
"""
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
@@ -2167,7 +2170,7 @@ class TestDatabricksRunNowOperator:
Test the execute function in case where the run failed.
"""
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
db_mock.get_run = make_run_with_state_mock("TERMINATED", "FAILED")
@@ -2201,7 +2204,7 @@ class TestDatabricksRunNowOperator:
Test the execute function in case where the run failed.
"""
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
db_mock.get_run = mock_dict(
@@ -2257,7 +2260,7 @@ class TestDatabricksRunNowOperator:
Test the execute function in case where the run failed.
"""
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
db_mock.get_run = mock_dict(
@@ -2334,7 +2337,7 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_wait_for_termination(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
@@ -2366,7 +2369,9 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_no_wait_for_termination(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
wait_for_termination=False, json=run)
+ op = DatabricksRunNowOperator(
+ task_id=TASK_ID, job_id=JOB_ID, wait_for_termination=False,
json=run, durable=False
+ )
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
@@ -2398,17 +2403,17 @@ class TestDatabricksRunNowOperator:
def test_init_exception_with_job_name_and_job_id(self, db_mock_class):
exception_message = "Argument 'job_name' is not allowed with argument
'job_id'"
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
job_name=JOB_NAME)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
job_name=JOB_NAME, durable=False)
with pytest.raises(AirflowException, match=exception_message):
op.execute(None)
run = {"job_id": JOB_ID, "job_name": JOB_NAME}
- op = DatabricksRunNowOperator(task_id=TASK_ID, json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, json=run, durable=False)
with pytest.raises(AirflowException, match=exception_message):
op.execute(None)
run = {"job_id": JOB_ID}
- op = DatabricksRunNowOperator(task_id=TASK_ID, json=run,
job_name=JOB_NAME)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, json=run,
job_name=JOB_NAME, durable=False)
with pytest.raises(AirflowException, match=exception_message):
op.execute(None)
@@ -2417,6 +2422,7 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_exec_exception_with_rendered_job_name_and_job_id(self,
db_mock_class):
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
json='{"job_id": "42", "job_name": "job-name"}',
)
@@ -2431,7 +2437,7 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_exec_with_job_name(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_name=JOB_NAME,
json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_name=JOB_NAME,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.find_job_id_by_name.return_value = JOB_ID
db_mock.run_now.return_value = RUN_ID
@@ -2464,7 +2470,7 @@ class TestDatabricksRunNowOperator:
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
def test_exec_failure_if_job_id_not_found(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
- op = DatabricksRunNowOperator(task_id=TASK_ID, job_name=JOB_NAME,
json=run)
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_name=JOB_NAME,
json=run, durable=False)
db_mock = db_mock_class.return_value
db_mock.find_job_id_by_name.return_value = None
@@ -2478,7 +2484,12 @@ class TestDatabricksRunNowOperator:
def test_cancel_previous_runs(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
op = DatabricksRunNowOperator(
- task_id=TASK_ID, job_id=JOB_ID, cancel_previous_runs=True,
wait_for_termination=False, json=run
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ cancel_previous_runs=True,
+ wait_for_termination=False,
+ json=run,
)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
@@ -2512,7 +2523,12 @@ class TestDatabricksRunNowOperator:
def test_no_cancel_previous_runs(self, db_mock_class):
run = {"notebook_params": NOTEBOOK_PARAMS, "notebook_task":
NOTEBOOK_TASK, "jar_params": JAR_PARAMS}
op = DatabricksRunNowOperator(
- task_id=TASK_ID, job_id=JOB_ID, cancel_previous_runs=False,
wait_for_termination=False, json=run
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ cancel_previous_runs=False,
+ wait_for_termination=False,
+ json=run,
)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = RUN_ID
@@ -2551,6 +2567,7 @@ class TestDatabricksRunNowOperator:
}
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
json=run,
cancel_previous_runs=True,
@@ -2872,6 +2889,7 @@ class TestDatabricksRunNowOperator:
(regression test for GH-39002).
"""
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
job_id=JOB_ID,
params={"env": "prod", "batch_size": 100},
@@ -2892,6 +2910,7 @@ class TestDatabricksRunNowOperator:
not override it.
"""
op = DatabricksRunNowOperator(
+ durable=False,
task_id=TASK_ID,
job_id=JOB_ID,
json={"job_parameters": {"explicit": "value"}},
@@ -2907,6 +2926,227 @@ class TestDatabricksRunNowOperator:
assert actual["job_parameters"] == {"explicit": "value"}
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
+class TestDatabricksRunNowOperatorDurable:
+ @staticmethod
+ def _context(task_store=None):
+ ctx: dict = {"ti": MagicMock(stats_tags={})}
+ if task_store is not None:
+ ctx["task_state_store"] = task_store
+ return ctx
+
+ @staticmethod
+ def _state(lifecycle: str, result: str = ""):
+ return {"state": {"life_cycle_state": lifecycle, "result_state":
result, "state_message": ""}}
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_persists_run_id_to_task_state_store_on_fresh_submit(self,
db_mock_class):
+ 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
+ db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+
+ op.execute(self._context(task_store))
+
+ db_mock.run_now.assert_called_once()
+ 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_reconnects_to_active_run_without_resubmitting(self,
db_mock_class):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ db_mock = db_mock_class.return_value
+ # status check sees the stored run still RUNNING, the poll then sees
it finish.
+ db_mock.get_run.side_effect = [self._state("RUNNING"),
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.run_now.assert_not_called()
+ 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 = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ 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.run_now.assert_not_called()
+ task_store.set.assert_not_called()
+ assert op.run_id == RUN_ID
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_reconnect_with_job_name_resolves_job_id_for_repair(self,
db_mock_class):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_name=JOB_NAME)
+ db_mock = db_mock_class.return_value
+ db_mock.find_job_id_by_name.return_value = JOB_ID
+ db_mock.get_run.side_effect = [self._state("RUNNING"),
self._state("TERMINATED", "SUCCESS")]
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = RUN_ID
+
+ op.execute(self._context(task_store))
+
+ # The reconnect rebuilds the payload and resolves job_name -> job_id
so repair_run has it,
+ # without resubmitting the run.
+ db_mock.run_now.assert_not_called()
+ db_mock.find_job_id_by_name.assert_called_once_with(JOB_NAME)
+ assert op._merged_json["job_id"] == JOB_ID
+ assert "job_name" not in op._merged_json
+
+
@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):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ db_mock = db_mock_class.return_value
+ db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
+ db_mock.get_run_page_url.return_value = RUN_PAGE_URL
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = RUN_ID
+ ti = MagicMock(stats_tags={})
+
+ op.execute({"ti": ti, "task_state_store": task_store})
+
+ # No duplicate run and no poll loop, but the run xcoms are still
pushed for parity.
+ db_mock.run_now.assert_not_called()
+ mock_poll.assert_not_called()
+ db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
+ ti.xcom_push.assert_any_call(key="run_id", value=RUN_ID)
+ ti.xcom_push.assert_any_call(key="run_page_url", value=RUN_PAGE_URL)
+ assert op.run_id == RUN_ID
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_fresh_submit_pushes_run_xcoms_exactly_once(self, db_mock_class):
+ 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
+ db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
+ db_mock.get_run_page_url.return_value = RUN_PAGE_URL
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+ ti = MagicMock(stats_tags={})
+
+ op.execute({"ti": ti, "task_state_store": task_store})
+
+ # Poll helper pushes run_id + run_page_url; get_job_result must not
push them again.
+ assert ti.xcom_push.call_count == 2
+ db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_resubmits_when_stored_run_in_terminal_failure(self,
db_mock_class):
+ 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
+ db_mock.get_run.side_effect = [
+ self._state("TERMINATED", "FAILED"),
+ self._state("TERMINATED", "SUCCESS"),
+ ]
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = 999
+
+ op.execute(self._context(task_store))
+
+ db_mock.run_now.assert_called_once()
+ 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_resubmits_when_stored_run_no_longer_exists(self, db_mock_class):
+ 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.
+ db_mock.get_run.side_effect = [
+ DatabricksApiError("Response: RESOURCE_DOES_NOT_EXIST, Status
Code: 404", http_status_code=404),
+ self._state("TERMINATED", "SUCCESS"),
+ ]
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = 999
+
+ op.execute(self._context(task_store))
+
+ db_mock.run_now.assert_called_once()
+ 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_durable_false_never_touches_task_state_store(self, db_mock_class):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID,
durable=False)
+ db_mock = db_mock_class.return_value
+ db_mock.run_now.return_value = RUN_ID
+ db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
+ task_store = MagicMock(spec_set=["get", "set"])
+
+ op.execute(self._context(task_store))
+
+ db_mock.run_now.assert_called_once()
+ task_store.get.assert_not_called()
+ task_store.set.assert_not_called()
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_durable_true_submits_fresh_missing_task_state_store(self,
db_mock_class):
+ 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
+ db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS")
+
+ op.execute(self._context()) # no task_state_store in context
+
+ db_mock.run_now.assert_called_once()
+ assert op.run_id == RUN_ID
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_get_job_status_encodes_life_cycle_and_result(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.return_value = self._state("TERMINATED", "SUCCESS")
+
+ assert op.get_job_status(RUN_ID, {}) == "TERMINATED:SUCCESS"
+
+ @pytest.mark.parametrize(
+ ("status", "expected"),
+ [
+ ("RUNNING:", True),
+ ("PENDING:", True),
+ ("QUEUED:", True),
+ ("TERMINATING:", True),
+ ("BLOCKED:", True),
+ ("WAITING_FOR_RETRY:", True),
+ ("TERMINATED:SUCCESS", False),
+ ("TERMINATED:FAILED", False),
+ ("SKIPPED:", False),
+ ("INTERNAL_ERROR:", False),
+ ("NOT_FOUND:", False),
+ ],
+ )
+ def test_is_job_active(self, status, expected):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ assert op.is_job_active(status) is expected
+
+ @pytest.mark.parametrize(
+ ("status", "expected"),
+ [
+ ("TERMINATED:SUCCESS", True),
+ ("TERMINATED:FAILED", False),
+ ("RUNNING:", False),
+ ("SKIPPED:", False),
+ ],
+ )
+ def test_is_job_succeeded(self, status, expected):
+ op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID)
+ assert op.is_job_succeeded(status) is expected
+
+
class TestDatabricksSQLStatementsOperator:
def test_init(self):
"""
