This is an automated email from the ASF dual-hosted git repository.
eladkal 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 dc0e0bb084d Add on_kill() to DatabricksTaskBaseOperator to cancel runs
on task kill (#69442)
dc0e0bb084d is described below
commit dc0e0bb084d685e2cb475e312f729fbb2a262efa
Author: Victory <[email protected]>
AuthorDate: Thu Aug 6 13:46:27 2026 +0800
Add on_kill() to DatabricksTaskBaseOperator to cancel runs on task kill
(#69442)
DatabricksSubmitRunOperator and DatabricksRunNowOperator both implement
on_kill() to cancel the Databricks run when an Airflow task is killed
(SIGTERM or execution_timeout). DatabricksTaskBaseOperator — the base for
DatabricksTaskOperator and DatabricksNotebookOperator — was missing the
same implementation, leaving Databricks jobs running after the Airflow task
was killed and orphaning compute resources.
DatabricksWorkflowTaskGroup received on_kill() in #42115; this PR closes
the remaining gap for standalone task operators.
For workflow members self.databricks_run_id is the shared parent run ID;
cancelling it would stop all sibling tasks. on_kill() therefore calls
_get_current_databricks_task()["run_id"] to target only the current task's
own child run, mirroring monitor_databricks_job. Standalone operators
continue to cancel via self.databricks_run_id directly.
If resolving the child run_id fails (API error, task_key mismatch), on_kill
logs the exception and returns without cancelling anything — falling back to
the parent run_id would stop sibling tasks, defeating the purpose.
Unit tests cover: cancel called for standalone operator, no-op when
databricks_run_id is None, workflow-member cancels child run (parent=1,
child=999, asserts cancel_run(999)), and workflow-member where
_get_current_databricks_task raises asserts cancel_run not called.
---
.../providers/databricks/operators/databricks.py | 25 +++++++++
.../unit/databricks/operators/test_databricks.py | 64 ++++++++++++++++++++++
2 files changed, 89 insertions(+)
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index f6bf1cf6aca..a64acb78598 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -1973,6 +1973,31 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
errors = event.get("errors", [])
self._handle_terminal_run_state(run_state, errors)
+ def on_kill(self) -> None:
+ if self.databricks_run_id is None:
+ return
+ if self._databricks_workflow_task_group:
+ # Workflow member: cancel only this task's child run, not the
shared parent workflow run.
+ # Cancelling the parent would also stop all sibling tasks.
+ # If the child run_id cannot be resolved, log and bail out — do
NOT fall back to the
+ # parent run_id as that would cancel sibling tasks.
+ try:
+ run_id_to_cancel =
self._get_current_databricks_task()["run_id"]
+ except Exception:
+ self.log.exception(
+ "Task: %s could not resolve child run_id; skipping cancel
to avoid stopping sibling tasks.",
+ self.task_id,
+ )
+ return
+ else:
+ run_id_to_cancel = self.databricks_run_id
+ self._hook.cancel_run(run_id_to_cancel)
+ self.log.info(
+ "Task: %s with run_id: %s was requested to be cancelled.",
+ self.task_id,
+ run_id_to_cancel,
+ )
+
class DatabricksNotebookOperator(DatabricksTaskBaseOperator):
"""
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index ad9d718e34f..65c990a5926 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -4090,3 +4090,67 @@ class TestDatabricksTaskOperator:
expected_task_key = "test_task_key"
assert expected_task_key == operator.databricks_task_key
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_on_kill_cancels_run(self, db_mock_class):
+ operator = DatabricksTaskOperator(
+ task_id="task",
+ task_config={"sql_task": {"query": {"query_id": "abc"}}},
+ )
+ db_mock = db_mock_class.return_value
+ operator.databricks_run_id = 1
+ operator.on_kill()
+ db_mock.cancel_run.assert_called_once_with(1)
+
+ def test_on_kill_does_nothing_when_run_id_is_none(self):
+ operator = DatabricksTaskOperator(
+ task_id="task",
+ task_config={"sql_task": {"query": {"query_id": "abc"}}},
+ )
+ with
mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
as db_mock_class:
+ operator.on_kill()
+ db_mock_class.return_value.cancel_run.assert_not_called()
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def test_on_kill_workflow_member_cancels_child_run(self, db_mock_class):
+ operator = DatabricksTaskOperator(
+ task_id="task",
+ task_config={"sql_task": {"query": {"query_id": "abc"}}},
+ )
+ db_mock = db_mock_class.return_value
+ operator.databricks_run_id = 1
+ with mock.patch.object(
+ operator,
+ "_get_current_databricks_task",
+ return_value={"run_id": 999, "task_key": "task"},
+ ):
+ with mock.patch(
+ "airflow.providers.databricks.operators.databricks"
+ ".DatabricksTaskBaseOperator._databricks_workflow_task_group",
+ new_callable=mock.PropertyMock,
+ return_value=object(),
+ ):
+ operator.on_kill()
+ db_mock.cancel_run.assert_called_once_with(999)
+
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+ def
test_on_kill_workflow_member_get_task_raises_does_not_cancel_parent(self,
db_mock_class):
+ operator = DatabricksTaskOperator(
+ task_id="task",
+ task_config={"sql_task": {"query": {"query_id": "abc"}}},
+ )
+ db_mock = db_mock_class.return_value
+ operator.databricks_run_id = 1
+ with mock.patch.object(
+ operator,
+ "_get_current_databricks_task",
+ side_effect=Exception("API error"),
+ ):
+ with mock.patch(
+ "airflow.providers.databricks.operators.databricks"
+ ".DatabricksTaskBaseOperator._databricks_workflow_task_group",
+ new_callable=mock.PropertyMock,
+ return_value=object(),
+ ):
+ operator.on_kill()
+ db_mock.cancel_run.assert_not_called()
