From 781832b708337a52ba1adbe1c1e1087441fb3271 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Sun, 6 Sep 2026 19:48:22 +0800 Subject: [PATCH] Fix Dataproc job not cancelled when killing a deferred task On Airflow 3.3+ the triggerer invokes a trigger's on_kill() only when the trigger lets asyncio.CancelledError propagate out of run(). Both Dataproc submit triggers caught it and returned instead, so killing a deferred task left the Dataproc job running and the triggerer additionally reported the trigger as having exited without sending an event. --- .../google/cloud/triggers/dataproc.py | 6 ++++ .../google/cloud/triggers/test_dataproc.py | 28 +++++++++++++++++++ 2 files changed, 34 insertions(+) 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 127eaea8921cf..47d47462136f9 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,9 @@ async def run(self): except Exception as e: self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e)) raise e + else: + # The triggerer only calls on_kill() when CancelledError propagates out of run(). + raise class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger): @@ -437,6 +440,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: except Exception as e: self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e)) raise e + else: + # The triggerer only calls on_kill() when CancelledError propagates out of run(). + raise class DataprocClusterTrigger(DataprocBaseTrigger): 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 781bcf079a05a..aee3dd325528e 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py @@ -799,6 +799,18 @@ async def test_submit_trigger_run_cancelled( await async_gen.aclose() + @pytest.mark.asyncio + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Cancellation is handled inline before Airflow 3.3.0") + @mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_async_hook") + async def test_submit_trigger_run_reraises_cancelled_error(self, mock_get_async_hook, submit_trigger): + mock_async_hook = mock_get_async_hook.return_value + mock_async_hook.get_job_client = mock.AsyncMock() + mock_async_hook.get_job.side_effect = asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + async for _ in submit_trigger.run(): + pass + class TestDataprocSubmitJobDirectTrigger: def test_serialization(self, submit_job_direct_trigger): @@ -1183,3 +1195,19 @@ async def test_run_cancelled_tolerates_a_job_dataproc_never_received( await async_gen.asend(None) mock_get_sync_hook.return_value.cancel_job.assert_called_once() + + @pytest.mark.asyncio + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Cancellation is handled inline before Airflow 3.3.0") + @mock.patch( + "airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitJobDirectTrigger.get_async_hook" + ) + async def test_run_reraises_cancelled_error(self, mock_get_async_hook, submit_job_direct_trigger): + mock_hook = mock_get_async_hook.return_value + submit_future = asyncio.Future() + submit_future.set_result(Job(reference={"job_id": TEST_JOB_ID})) + mock_hook.submit_job.return_value = submit_future + mock_hook.get_job.side_effect = asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + async for _ in submit_job_direct_trigger.run(): + pass