Add on_kill() to DatabricksTaskBaseOperator to cancel runs on task kill

providers/databricks/src/airflow/providers/databricks/operators/databricks.py

@@ -1941,6 +1941,31 @@ def execute_complete(self, context: dict | None, event: dict) -> None:
         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):
     """

providers/databricks/tests/unit/databricks/operators/test_databricks.py

@@ -4034,3 +4034,67 @@ def test_user_databricks_task_key(self):
         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()