This is an automated email from the ASF dual-hosted git repository.
shahar1 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 d59d32a9595 Implement on_kill() trigger hook for Databricks triggers
(#65672)
d59d32a9595 is described below
commit d59d32a95951fa40be3f4e223926065a09b98bd9
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Apr 22 20:35:37 2026 +0530
Implement on_kill() trigger hook for Databricks triggers (#65672)
---
.../airflow/providers/databricks/triggers/databricks.py | 16 ++++++++++++++++
.../tests/unit/databricks/triggers/test_databricks.py | 12 ++++++++++++
2 files changed, 28 insertions(+)
diff --git
a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py
b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py
index 8f1ad0d422b..67f8a392a0c 100644
---
a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py
@@ -84,6 +84,14 @@ class DatabricksExecutionTrigger(BaseTrigger):
},
)
+ 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 @@ class DatabricksSQLStatementExecutionTrigger(BaseTrigger):
},
)
+ 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():
diff --git
a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py
b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py
index a34e4723388..d2465077534 100644
--- a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py
@@ -259,6 +259,12 @@ class TestDatabricksExecutionTrigger:
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 @@ class TestDatabricksSQLStatementExecutionTrigger:
)
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)
