You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by sh...@apache.org on 2026/04/27 18:03:25 UTC
(airflow) branch main updated: Migrate DataprocSubmitTrigger and DataprocSubmitJobDirectTrigger to on_kill() (#65742)
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 b5af76dfe8e Migrate DataprocSubmitTrigger and DataprocSubmitJobDirectTrigger to on_kill() (#65742)
b5af76dfe8e is described below
commit b5af76dfe8e47d8fab10571399823562ba0b6fe7
Author: Sadha Chilukoori <sr...@gmail.com>
AuthorDate: Mon Apr 27 11:03:17 2026 -0700
Migrate DataprocSubmitTrigger and DataprocSubmitJobDirectTrigger to on_kill() (#65742)
---
.../providers/google/cloud/triggers/dataproc.py | 226 +++++++++++----------
.../src/airflow/providers/google/version_compat.py | 4 +-
.../unit/google/cloud/triggers/test_dataproc.py | 89 ++++++--
3 files changed, 197 insertions(+), 122 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 76aa2f5021f..56449fec0bc 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
@@ -33,14 +33,14 @@ from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook, DataprocHook
from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_3_0_PLUS
from airflow.triggers.base import BaseTrigger, TriggerEvent
from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
from sqlalchemy.orm.session import Session
-if not AIRFLOW_V_3_0_PLUS:
+if not AIRFLOW_V_3_3_0_PLUS:
from sqlalchemy import select
from airflow.models.taskinstance import TaskInstance
@@ -123,7 +123,16 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
},
)
- if not AIRFLOW_V_3_0_PLUS:
+ async def on_kill(self) -> None:
+ """Cancel the Dataproc job when the task is killed by a user action."""
+ if self.job_id and self.cancel_on_kill:
+ self.log.info("Cancelling Dataproc job: %s.", self.job_id)
+ await sync_to_async(self.get_sync_hook().cancel_job)(
+ job_id=self.job_id, project_id=self.project_id, region=self.region
+ )
+ self.log.info("Job: %s is cancelled.", self.job_id)
+
+ if not AIRFLOW_V_3_3_0_PLUS:
@provide_session
def get_task_instance(self, session: Session) -> TaskInstance:
@@ -150,41 +159,41 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
)
return task_instance
- async def get_task_state(self):
- from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+ async def get_task_state(self):
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
- task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
- dag_id=self.task_instance.dag_id,
- task_ids=[self.task_instance.task_id],
- run_ids=[self.task_instance.run_id],
- map_index=self.task_instance.map_index,
- )
- try:
- task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
- except Exception:
- raise AirflowException(
- "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
- self.task_instance.dag_id,
- self.task_instance.task_id,
- self.task_instance.run_id,
- self.task_instance.map_index,
+ task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+ dag_id=self.task_instance.dag_id,
+ task_ids=[self.task_instance.task_id],
+ run_ids=[self.task_instance.run_id],
+ map_index=self.task_instance.map_index,
)
- return task_state
+ try:
+ task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
+ except Exception:
+ raise AirflowException(
+ "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
+ self.task_instance.dag_id,
+ self.task_instance.task_id,
+ self.task_instance.run_id,
+ self.task_instance.map_index,
+ )
+ return task_state
- async def safe_to_cancel(self) -> bool:
- """
- Whether it is safe to cancel the external job which is being executed by this trigger.
+ async def safe_to_cancel(self) -> bool:
+ """
+ Whether it is safe to cancel the external job which is being executed by this trigger.
- This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
- Because in those cases, we should NOT cancel the external job.
- """
- if AIRFLOW_V_3_0_PLUS:
- task_state = await self.get_task_state()
- else:
- # Database query is needed to get the latest state of the task instance.
- task_instance = self.get_task_instance() # type: ignore[call-arg]
- task_state = task_instance.state
- return task_state != TaskInstanceState.DEFERRED
+ This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
+ Because in those cases, we should NOT cancel the external job.
+ """
+ if AIRFLOW_V_3_0_PLUS:
+ task_state = await self.get_task_state()
+ else:
+ # Database query is needed to get the latest state of the task instance.
+ task_instance = self.get_task_instance() # type: ignore[call-arg]
+ task_state = task_instance.state
+ return task_state != TaskInstanceState.DEFERRED
async def run(self):
hook = self.get_async_hook()
@@ -203,31 +212,24 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
{"job_id": self.job_id, "job_state": JobStatus.State(state).name, "job": Job.to_dict(job)}
)
except asyncio.CancelledError:
- self.log.info("Task got cancelled.")
- try:
- if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
- self.log.info(
- "Cancelling the job as it is safe to do so. Note that the airflow TaskInstance is not"
- " in deferred state."
- )
- self.log.info("Cancelling the job: %s", self.job_id)
- # The synchronous hook is utilized to delete the cluster when a task is cancelled. This
- # is because the asynchronous hook deletion is not awaited when the trigger task is
- # cancelled. The call for deleting the cluster or job through the sync hook is not a
- # blocking call, which means it does not wait until the cluster or job is deleted.
- self.get_sync_hook().cancel_job(
- job_id=self.job_id, project_id=self.project_id, region=self.region
- )
- self.log.info("Job: %s is cancelled", self.job_id)
- yield TriggerEvent(
- {
- "job_id": self.job_id,
- "job_state": ClusterStatus.State.DELETING.name,
- }
- )
- except Exception as e:
- self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
- raise e
+ if not AIRFLOW_V_3_3_0_PLUS:
+ self.log.info("Task got cancelled.")
+ try:
+ if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
+ self.log.info("Cancelling the job: %s", self.job_id)
+ self.get_sync_hook().cancel_job(
+ job_id=self.job_id, project_id=self.project_id, region=self.region
+ )
+ self.log.info("Job: %s is cancelled", self.job_id)
+ yield TriggerEvent(
+ {
+ "job_id": self.job_id,
+ "job_state": ClusterStatus.State.DELETING.name,
+ }
+ )
+ except Exception as e:
+ self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
+ raise e
class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
@@ -273,7 +275,16 @@ class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
},
)
- if not AIRFLOW_V_3_0_PLUS:
+ async def on_kill(self) -> None:
+ """Cancel the Dataproc job when the task is killed by a user action."""
+ if self.job_id and self.cancel_on_kill:
+ self.log.info("Cancelling Dataproc job: %s.", self.job_id)
+ await sync_to_async(self.get_sync_hook().cancel_job)(
+ job_id=self.job_id, project_id=self.project_id, region=self.region
+ )
+ self.log.info("Job: %s is cancelled.", self.job_id)
+
+ if not AIRFLOW_V_3_3_0_PLUS:
@provide_session
def get_task_instance(self, session: Session) -> TaskInstance:
@@ -300,40 +311,40 @@ class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
)
return task_instance
- async def get_task_state(self):
- from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+ async def get_task_state(self):
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
- task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
- dag_id=self.task_instance.dag_id,
- task_ids=[self.task_instance.task_id],
- run_ids=[self.task_instance.run_id],
- map_index=self.task_instance.map_index,
- )
- try:
- task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
- except Exception:
- raise RuntimeError(
- "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
- self.task_instance.dag_id,
- self.task_instance.task_id,
- self.task_instance.run_id,
- self.task_instance.map_index,
+ task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+ dag_id=self.task_instance.dag_id,
+ task_ids=[self.task_instance.task_id],
+ run_ids=[self.task_instance.run_id],
+ map_index=self.task_instance.map_index,
)
- return task_state
+ try:
+ task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
+ except Exception:
+ raise RuntimeError(
+ "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
+ self.task_instance.dag_id,
+ self.task_instance.task_id,
+ self.task_instance.run_id,
+ self.task_instance.map_index,
+ )
+ return task_state
- async def safe_to_cancel(self) -> bool:
- """
- Whether it is safe to cancel the external job which is being executed by this trigger.
+ async def safe_to_cancel(self) -> bool:
+ """
+ Whether it is safe to cancel the external job which is being executed by this trigger.
- This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
- Because in those cases, we should NOT cancel the external job.
- """
- if AIRFLOW_V_3_0_PLUS:
- task_state = await self.get_task_state()
- else:
- task_instance = self.get_task_instance() # type: ignore[call-arg]
- task_state = task_instance.state
- return task_state != TaskInstanceState.DEFERRED
+ This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
+ Because in those cases, we should NOT cancel the external job.
+ """
+ if AIRFLOW_V_3_0_PLUS:
+ task_state = await self.get_task_state()
+ else:
+ task_instance = self.get_task_instance() # type: ignore[call-arg]
+ task_state = task_instance.state
+ return task_state != TaskInstanceState.DEFERRED
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
@@ -360,23 +371,24 @@ class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
{"job_id": self.job_id, "job_state": JobStatus.State(state).name, "job": Job.to_dict(job)}
)
except asyncio.CancelledError:
- self.log.info("Task got cancelled.")
- try:
- if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
- self.log.info("Cancelling the job: %s", self.job_id)
- self.get_sync_hook().cancel_job(
- job_id=self.job_id, project_id=self.project_id, region=self.region
- )
- self.log.info("Job: %s is cancelled", self.job_id)
- yield TriggerEvent(
- {
- "job_id": self.job_id,
- "job_state": ClusterStatus.State.DELETING.name, # type: ignore[attr-defined]
- }
- )
- except Exception as e:
- self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
- raise e
+ if not AIRFLOW_V_3_3_0_PLUS:
+ self.log.info("Task got cancelled.")
+ try:
+ if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
+ self.log.info("Cancelling the job: %s", self.job_id)
+ self.get_sync_hook().cancel_job(
+ job_id=self.job_id, project_id=self.project_id, region=self.region
+ )
+ self.log.info("Job: %s is cancelled", self.job_id)
+ yield TriggerEvent(
+ {
+ "job_id": self.job_id,
+ "job_state": ClusterStatus.State.DELETING.name, # type: ignore[attr-defined]
+ }
+ )
+ except Exception as e:
+ self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
+ raise e
class DataprocClusterTrigger(DataprocBaseTrigger):
diff --git a/providers/google/src/airflow/providers/google/version_compat.py b/providers/google/src/airflow/providers/google/version_compat.py
index 23ef9e614c9..47dbdc0aa74 100644
--- a/providers/google/src/airflow/providers/google/version_compat.py
+++ b/providers/google/src/airflow/providers/google/version_compat.py
@@ -34,6 +34,8 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
+# BaseTrigger.on_kill() was introduced in 3.3.0; use this flag to guard trigger migrations
+AIRFLOW_V_3_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
# Version-compatible imports
# BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator 3.0.x
@@ -47,4 +49,4 @@ else:
# Explicitly export these imports to protect them from being removed by linters
-__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseOperator"]
+__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_3_0_PLUS", "BaseOperator"]
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 1e0f9176497..cfae5af0dcf 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
@@ -37,6 +37,8 @@ from airflow.providers.google.cloud.triggers.dataproc import (
from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
from airflow.triggers.base import TriggerEvent
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+
TEST_PROJECT_ID = "project-id"
TEST_REGION = "region"
TEST_BATCH_ID = "batch-id"
@@ -634,6 +636,33 @@ class TestDataprocSubmitTrigger:
assert event.payload == expected_event.payload
@pytest.mark.asyncio
+ @mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
+ async def test_on_kill_cancels_job(self, mock_get_sync_hook, submit_trigger):
+ mock_sync_hook = mock_get_sync_hook.return_value
+ mock_sync_hook.cancel_job = mock.MagicMock()
+ await submit_trigger.on_kill()
+ mock_sync_hook.cancel_job.assert_called_once_with(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
+ async def test_on_kill_respects_cancel_on_kill_false(self, mock_get_sync_hook):
+ trigger = DataprocSubmitTrigger(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ polling_interval_seconds=TEST_POLL_INTERVAL,
+ cancel_on_kill=False,
+ )
+ await trigger.on_kill()
+ mock_get_sync_hook.return_value.cancel_job.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="on_kill() handles cancellation for Airflow 3.3.0+")
@pytest.mark.parametrize("is_safe_to_cancel", [True, False])
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_async_hook")
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
@@ -641,33 +670,25 @@ class TestDataprocSubmitTrigger:
async def test_submit_trigger_run_cancelled(
self, mock_safe_to_cancel, mock_get_sync_hook, mock_get_async_hook, submit_trigger, is_safe_to_cancel
):
- """Test the trigger correctly handles an asyncio.CancelledError."""
+ """Test the trigger correctly handles an asyncio.CancelledError on Airflow < 3.3.0."""
mock_safe_to_cancel.return_value = is_safe_to_cancel
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
-
mock_sync_hook = mock_get_sync_hook.return_value
mock_sync_hook.cancel_job = mock.MagicMock()
async_gen = submit_trigger.run()
-
try:
await async_gen.asend(None)
- # Should raise StopAsyncIteration if no more items to yield
await async_gen.asend(None)
except asyncio.CancelledError:
- # Handle the cancellation as expected
pass
except StopAsyncIteration:
- # The generator should be properly closed after handling the cancellation
pass
except Exception as e:
- # Catch any other exceptions that should not occur
pytest.fail(f"Unexpected exception raised: {e}")
- # Check if cancel_job was correctly called
if submit_trigger.cancel_on_kill and is_safe_to_cancel:
mock_sync_hook.cancel_job.assert_called_once_with(
job_id=submit_trigger.job_id,
@@ -677,7 +698,6 @@ class TestDataprocSubmitTrigger:
else:
mock_sync_hook.cancel_job.assert_not_called()
- # Clean up the generator
await async_gen.aclose()
@@ -760,6 +780,50 @@ class TestDataprocSubmitJobDirectTrigger:
assert event.payload == expected_event.payload
@pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitJobDirectTrigger.get_sync_hook"
+ )
+ async def test_on_kill_cancels_job(self, mock_get_sync_hook, submit_job_direct_trigger):
+ submit_job_direct_trigger.job_id = TEST_JOB_ID
+ mock_sync_hook = mock_get_sync_hook.return_value
+ mock_sync_hook.cancel_job = mock.MagicMock()
+ await submit_job_direct_trigger.on_kill()
+ mock_sync_hook.cancel_job.assert_called_once_with(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitJobDirectTrigger.get_sync_hook"
+ )
+ async def test_on_kill_no_job_id_does_not_cancel(self, mock_get_sync_hook, submit_job_direct_trigger):
+ # job_id is None before submission completes; on_kill must be a no-op
+ assert submit_job_direct_trigger.job_id is None
+ await submit_job_direct_trigger.on_kill()
+ mock_get_sync_hook.return_value.cancel_job.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitJobDirectTrigger.get_sync_hook"
+ )
+ async def test_on_kill_respects_cancel_on_kill_false(self, mock_get_sync_hook):
+ trigger = DataprocSubmitJobDirectTrigger(
+ job=TEST_JOB,
+ request_id=TEST_REQUEST_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ polling_interval_seconds=TEST_POLL_INTERVAL,
+ cancel_on_kill=False,
+ )
+ trigger.job_id = TEST_JOB_ID
+ await trigger.on_kill()
+ mock_get_sync_hook.return_value.cancel_job.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="on_kill() handles cancellation for Airflow 3.3.0+")
@pytest.mark.parametrize("is_safe_to_cancel", [True, False])
@mock.patch(
"airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitJobDirectTrigger.get_async_hook"
@@ -778,22 +842,19 @@ class TestDataprocSubmitJobDirectTrigger:
submit_job_direct_trigger,
is_safe_to_cancel,
):
+ """Test CancelledError handling for Airflow < 3.3.0."""
mock_safe_to_cancel.return_value = is_safe_to_cancel
mock_hook = mock_get_async_hook.return_value
-
mock_submitted_job = mock.MagicMock()
mock_submitted_job.reference.job_id = TEST_JOB_ID
submit_future = asyncio.Future()
submit_future.set_result(mock_submitted_job)
mock_hook.submit_job.return_value = submit_future
-
mock_hook.get_job.side_effect = asyncio.CancelledError
-
mock_sync_hook = mock_get_sync_hook.return_value
mock_sync_hook.cancel_job = mock.MagicMock()
async_gen = submit_job_direct_trigger.run()
-
try:
await async_gen.asend(None)
await async_gen.asend(None)