This is an automated email from the ASF dual-hosted git repository.
potiuk 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 d339bb2767a Fix Dataproc triggers swallowing task cancellation (#74192)
d339bb2767a is described below
commit d339bb2767afafe77603e357747b90cd0408902e
Author: Atharva Ajmera <[email protected]>
AuthorDate: Mon Oct 5 19:38:52 2026 +0530
Fix Dataproc triggers swallowing task cancellation (#74192)
The except asyncio.CancelledError handlers in DataprocSubmitTrigger.run(),
DataprocSubmitJobDirectTrigger.run() and DataprocClusterTrigger.run() performed
cleanup but never re-raised, so triggerer migration (sibling steals trigger
rows, original cancels its coroutines) looked like a crash: Trigger exited
without sending an event, failing all dependent deferred tasks. Add the missing
bare raise, mirroring the BigQuery fix in #63730, so cancellation propagates
and cleanup_finished_trigg [...]
---
.../providers/google/cloud/triggers/dataproc.py | 3 ++
.../unit/google/cloud/triggers/test_dataproc.py | 63 ++++++++++++++--------
2 files changed, 45 insertions(+), 21 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
index 127eaea8921..665e336c0b0 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
@@ -233,6 +233,7 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
except Exception as e:
self.log.error("Failed to cancel the job: %s with error :
%s", self.job_id, str(e))
raise e
+ raise
class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
@@ -437,6 +438,7 @@ class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
except Exception as e:
self.log.error("Failed to cancel the job: %s with error :
%s", self.job_id, str(e))
raise e
+ raise
class DataprocClusterTrigger(DataprocBaseTrigger):
@@ -584,6 +586,7 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
self.log.error("Error during cancellation handling: %s", e)
msg = f"Error during cancellation handling: {e}"
raise AirflowException(msg)
+ raise
async def fetch_cluster(self) -> Cluster:
"""Fetch the cluster status."""
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
index 781bcf079a0..d469cdf478c 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
@@ -17,7 +17,6 @@
from __future__ import annotations
import asyncio
-import contextlib
import logging
from asyncio import CancelledError, Future, sleep
from types import SimpleNamespace
@@ -378,10 +377,7 @@ class TestDataprocClusterTrigger:
):
"""Test the trigger's cancellation behavior when it is not safe to
cancel."""
mock_safe_to_cancel.return_value = False
- cluster =
Cluster(status=ClusterStatus(state=ClusterStatus.State.RUNNING))
- future_cluster = asyncio.Future()
- future_cluster.set_result(cluster)
- mock_get_async_hook.return_value.get_cluster.return_value =
future_cluster
+ mock_get_async_hook.return_value.get_cluster.return_value = Future()
mock_delete_cluster = mock.MagicMock()
mock_get_sync_hook.return_value.delete_cluster = mock_delete_cluster
@@ -393,12 +389,43 @@ class TestDataprocClusterTrigger:
await sleep(0)
task.cancel()
- with contextlib.suppress(CancelledError):
+ with pytest.raises(CancelledError):
await task
assert mock_delete_cluster.call_count == 0
mock_delete_cluster.assert_not_called()
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocClusterTrigger.get_async_hook")
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocClusterTrigger.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocClusterTrigger.safe_to_cancel")
+ @mock.patch.object(DataprocClusterTrigger, "log")
+ async def test_cluster_trigger_run_cancelled_safe_to_cancel(
+ self, mock_log, mock_safe_to_cancel, mock_get_sync_hook,
mock_get_async_hook, cluster_trigger
+ ):
+ """Cancellation must delete the cluster and still propagate so
migration stays transparent."""
+ mock_safe_to_cancel.return_value = True
+ mock_get_async_hook.return_value.get_cluster.return_value = Future()
+
+ mock_delete_cluster = mock.MagicMock()
+ mock_get_sync_hook.return_value.delete_cluster = mock_delete_cluster
+
+ cluster_trigger.delete_on_error = True
+
+ async_gen = cluster_trigger.run()
+ task = asyncio.create_task(async_gen.__anext__())
+ await sleep(0)
+ task.cancel()
+
+ with pytest.raises(CancelledError):
+ await task
+
+ mock_delete_cluster.assert_called_once_with(
+ region=cluster_trigger.region,
+ cluster_name=cluster_trigger.cluster_name,
+ project_id=cluster_trigger.project_id,
+ )
+
@pytest.mark.asyncio
@pytest.mark.skipif(
not AIRFLOW_V_3_0_PLUS, reason="The task state is read through the
task SDK on Airflow 3 only"
@@ -778,15 +805,11 @@ class TestDataprocSubmitTrigger:
mock_sync_hook.cancel_job = mock.MagicMock()
async_gen = submit_trigger.run()
- try:
- await async_gen.asend(None)
+ if is_safe_to_cancel:
+ event = await async_gen.asend(None)
+ assert event.payload["job_state"] ==
ClusterStatus.State.DELETING.name
+ with pytest.raises(asyncio.CancelledError):
await async_gen.asend(None)
- except asyncio.CancelledError:
- pass
- except StopAsyncIteration:
- pass
- except Exception as e:
- pytest.fail(f"Unexpected exception raised: {e}")
if submit_trigger.cancel_on_kill and is_safe_to_cancel:
mock_sync_hook.cancel_job.assert_called_once_with(
@@ -954,13 +977,11 @@ class TestDataprocSubmitJobDirectTrigger:
mock_sync_hook.cancel_job = mock.MagicMock()
async_gen = submit_job_direct_trigger.run()
- try:
- await async_gen.asend(None)
+ if is_safe_to_cancel:
+ event = await async_gen.asend(None)
+ assert event.payload["job_state"] ==
ClusterStatus.State.DELETING.name
+ with pytest.raises(asyncio.CancelledError):
await async_gen.asend(None)
- except (asyncio.CancelledError, StopAsyncIteration):
- pass
- except Exception as e:
- pytest.fail(f"Unexpected exception raised: {e}")
if submit_job_direct_trigger.cancel_on_kill and is_safe_to_cancel:
mock_sync_hook.cancel_job.assert_called_once_with(
@@ -1179,7 +1200,7 @@ class TestDataprocSubmitJobDirectTrigger:
mock_get_sync_hook.return_value.cancel_job.side_effect = NotFound("no
such job")
async_gen = submit_job_direct_trigger.run()
- with contextlib.suppress(asyncio.CancelledError, StopAsyncIteration):
+ with pytest.raises(asyncio.CancelledError):
await async_gen.asend(None)
mock_get_sync_hook.return_value.cancel_job.assert_called_once()
