diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index 1682740646674..90def077023d4 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/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): """ diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks.py b/providers/databricks/tests/unit/databricks/operators/test_databricks.py index 9c83c19ebe915..a4222908aebda 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py +++ b/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()