Implement on_kill() trigger hook for Databricks triggers (#65672)

Author: amoghrajesh

providers/databricks/src/airflow/providers/databricks/triggers/databricks.py

@@ -84,6 +84,14 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             },
         )
 
+    async def on_kill(self) -> None:
+        """Cancel the Databricks run when the trigger is cancelled by a user action."""
+        if self.run_id:
+            from asgiref.sync import sync_to_async
+
+            self.log.info("Cancelling Databricks run %s.", self.run_id)
+            await sync_to_async(self.hook.cancel_run)(self.run_id)
+
     async def run(self):
         async with self.hook:
             while True:
@@ -167,6 +175,14 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             },
         )
 
+    async def on_kill(self) -> None:
+        """Cancel the Databricks SQL statement when the trigger is cancelled by a user action."""
+        if self.statement_id:
+            from asgiref.sync import sync_to_async
+
+            self.log.info("Cancelling Databricks SQL statement %s.", self.statement_id)
+            await sync_to_async(self.hook.cancel_sql_statement)(self.statement_id)
+
     async def run(self):
         async with self.hook:
             while self.end_time > time.time():

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

@@ -259,6 +259,12 @@ async def test_sleep_between_retries(
         mock_sleep.assert_called_once()
         mock_sleep.assert_called_with(POLLING_INTERVAL_SECONDS)
 
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.cancel_run")
+    async def test_on_kill_cancels_run(self, mock_cancel_run):
+        await self.trigger.on_kill()
+        mock_cancel_run.assert_called_once_with(RUN_ID)
+
 
 class TestDatabricksSQLStatementExecutionTrigger:
     @pytest.fixture(autouse=True)
@@ -361,3 +367,9 @@ async def test_sleep_between_retries(self, mock_a_get_sql_statement_state, mock_
             )
         mock_sleep.assert_called_once()
         mock_sleep.assert_called_with(POLLING_INTERVAL_SECONDS)
+
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.cancel_sql_statement")
+    async def test_on_kill_cancels_statement(self, mock_cancel_sql_statement):
+        await self.trigger.on_kill()
+        mock_cancel_sql_statement.assert_called_once_with(STATEMENT_ID)