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 a5c28069af8 Cancel the Dataflow job when a user kills the deferred
task (#69586)
a5c28069af8 is described below
commit a5c28069af8ca623294dd02890d8ae59f7f951b9
Author: Steve Ahn <[email protected]>
AuthorDate: Sat Oct 17 11:57:06 2026 -0700
Cancel the Dataflow job when a user kills the deferred task (#69586)
---
.../providers/google/cloud/operators/dataflow.py | 37 +++-
.../providers/google/cloud/triggers/dataflow.py | 101 +++++++++-
.../unit/google/cloud/operators/test_dataflow.py | 116 +++++++++++
.../unit/google/cloud/triggers/test_dataflow.py | 224 ++++++++++++++++++++-
4 files changed, 472 insertions(+), 6 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py
b/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py
index c5642b07072..9be06bde42e 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py
@@ -298,6 +298,9 @@ class
DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
https://cloud.google.com/dataflow/docs/templates/executing-templates
:param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the
task is killed,
+ both while the operator is running and, for a deferred task, while it
waits in the
+ triggerer (Airflow 3.3+).
"""
template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ class
DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
append_job_name: bool = True,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.template = template
+ self.cancel_on_kill = cancel_on_kill
self.job_name = job_name
self.options = options or {}
self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ class
DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
poll_sleep=self.poll_sleep,
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
)
@@ -453,7 +459,10 @@ class
DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
return job_id
def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through
the trigger instead."""
self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
if self.job is not None:
self.log.info("Cancelling job %s", self.job_name)
self.hook.cancel_job(
@@ -525,6 +534,9 @@ class
DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account
(templated).
:param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the
task is killed,
+ both while the operator is running and, for a deferred task, while it
waits in the
+ triggerer (Airflow 3.3+).
:param expected_terminal_state: The expected final status of the operator
on which the corresponding
Airflow task succeeds. When not specified, it will be determined by
the hook.
:param append_job_name: True if unique suffix has to be appended to job
name.
@@ -550,11 +562,13 @@ class
DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
append_job_name: bool = True,
expected_terminal_state: str | None = None,
poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
*args,
**kwargs,
) -> None:
super().__init__(*args, **kwargs)
self.body = body
+ self.cancel_on_kill = cancel_on_kill
self.location = location
self.project_id = project_id
self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ class
DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
poll_sleep=self.poll_sleep,
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
)
@@ -658,7 +674,10 @@ class
DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
return job
def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through
the trigger instead."""
self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
if self.job is not None:
self.hook.cancel_job(
job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class
DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
:param append_job_name: Optional. Set to True if a unique suffix has to be
appended to the `job_name`.
Defaults to True.
:param drain_pipeline: Optional. Set to True if you want to stop a
streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this
does not work for batch pipeline jobs
- or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this
does not work for batch pipeline jobs.
+ Defaults to False.
For more info see:
https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
:param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the
task is killed,
+ both while the operator is running and, for a deferred task, while it
waits in the
+ triggerer (Airflow 3.3+).
:param expected_terminal_state: Optional. The expected terminal state of
the Dataflow job at which the
operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the
batch jobs and 'JOB_STATE_RUNNING'
for the streaming jobs.
@@ -748,10 +770,12 @@ class
DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
jinja_variables: dict[str, str] | None = None,
options: dict[str, Any] | None = None,
impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
self.yaml_pipeline_file = yaml_pipeline_file
self.region = region
self.project_id = project_id
@@ -793,6 +817,8 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
cancel_timeout=self.cancel_timeout,
expected_terminal_state=self.expected_terminal_state,
impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
)
@@ -818,10 +844,13 @@ class
DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
"""
Cancel the dataflow job if a task instance gets killed.
- This method will not be called if a task instance is killed in a
deferred
- state.
+ This method is not called for a task instance killed in a deferred
state;
+ on Airflow 3.3+, the trigger requests cancellation instead, honoring
cancel_on_kill
+ and drain_pipeline.
"""
self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
if self.job_id:
self.hook.cancel_job(
job_id=self.job_id,
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py
b/providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py
index 4b6889d7325..004fc7d2a84 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py
@@ -22,6 +22,7 @@ from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from asgiref.sync import sync_to_async
from google.api_core.exceptions import ServiceUnavailable
from google.cloud.dataflow_v1beta3 import JobState
from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@ from google.cloud.dataflow_v1beta3.types import (
MetricUpdate,
)
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook,
DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
account from the list granting this role to the originating account
(templated).
:param cancel_timeout: Optional. How long (in seconds) operator should
wait for the pipeline to be
successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the
user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job
to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour
before deferral.
"""
def __init__(
@@ -73,6 +82,8 @@ class TemplateJobStartTrigger(BaseTrigger):
poll_sleep: int = 10,
impersonation_chain: str | Sequence[str] | None = None,
cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
):
super().__init__()
self.project_id = project_id
@@ -82,6 +93,8 @@ class TemplateJobStartTrigger(BaseTrigger):
self.poll_sleep = poll_sleep
self.impersonation_chain = impersonation_chain
self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ class TemplateJobStartTrigger(BaseTrigger):
"poll_sleep": self.poll_sleep,
"impersonation_chain": self.impersonation_chain,
"cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
},
)
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s,
drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread:
the hook resolves the
+ # connection eagerly during construction, which must not run in
the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Cancellation or drain requested for Dataflow job
%s.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be
running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Request cancellation or drain without waiting for the job to reach
a terminal state."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=None,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
async def run(self):
"""
Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
If set as a sequence, the identities from the list must grant
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account
(templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the
user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job
to be stopped by
+ draining it instead of cancelling when the task is killed.
"""
def __init__(
@@ -312,6 +368,8 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
cancel_timeout: int | None = 5 * 60,
expected_terminal_state: str | None = None,
impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
):
super().__init__()
self.project_id = project_id
@@ -322,6 +380,8 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
self.cancel_timeout = cancel_timeout
self.expected_terminal_state = expected_terminal_state
self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
"expected_terminal_state": self.expected_terminal_state,
"impersonation_chain": self.impersonation_chain,
"cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
},
)
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s,
drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread:
the hook resolves the
+ # connection eagerly during construction, which must not run in
the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Cancellation or drain requested for Dataflow job
%s.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be
running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Request cancellation or drain without waiting for the job to reach
a terminal state."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=None,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
async def run(self):
"""
Fetch job and yield events depending on the job's type and state.
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_dataflow.py
b/providers/google/tests/unit/google/cloud/operators/test_dataflow.py
index 90b444dd6c4..9db80b84b1a 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_dataflow.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_dataflow.py
@@ -210,6 +210,32 @@ class TestDataflowTemplatedJobStartOperator:
)
mock_defer_method.assert_called_once()
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowTemplatedJobStartOperator.defer",
autospec=True)
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_deferrable_threads_cancel_on_kill_false_into_trigger(self,
mock_hook, mock_defer_method):
+ """cancel_on_kill=False reaches the trigger, so a deferred kill leaves
the job running."""
+ operator = DataflowTemplatedJobStartOperator(
+ project_id=TEST_PROJECT,
+ task_id=TASK_ID,
+ template=TEMPLATE,
+ job_name=JOB_NAME,
+ location=TEST_LOCATION,
+ deferrable=True,
+ cancel_on_kill=False,
+ )
+ operator.execute(mock.MagicMock(spec=dict))
+
+ assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is
False
+
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_on_kill_respects_cancel_on_kill_false(self, mock_hook,
sync_operator):
+ sync_operator.cancel_on_kill = False
+ sync_operator.job = {"id": "test-job", "projectId": TEST_PROJECT,
"location": TEST_LOCATION}
+
+ sync_operator.on_kill()
+
+ mock_hook.return_value.cancel_job.assert_not_called()
+
def test_validation_deferrable_params_raises_error(self):
init_kwargs = {
"project_id": TEST_PROJECT,
@@ -347,6 +373,46 @@ class TestDataflowStartFlexTemplateOperator:
)
mock_defer_method.assert_called_once()
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowStartFlexTemplateOperator.defer",
autospec=True)
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_deferrable_threads_cancel_on_kill_false_into_trigger(self,
mock_hook, mock_defer_method):
+ """cancel_on_kill=False reaches the trigger, so a deferred kill leaves
the job running."""
+ operator = DataflowStartFlexTemplateOperator(
+ task_id="start_flex_template_streaming_beam_sql",
+ body={"launchParameter": TEST_FLEX_PARAMETERS},
+ project_id=TEST_PROJECT,
+ location=TEST_LOCATION,
+ deferrable=True,
+ cancel_on_kill=False,
+ )
+ operator.execute(mock.MagicMock(spec=dict))
+
+ assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is
False
+
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_on_kill_respects_cancel_on_kill_false(self, mock_hook,
sync_operator):
+ sync_operator.cancel_on_kill = False
+ sync_operator.job = {"id": "test-job", "projectId": TEST_PROJECT,
"location": TEST_LOCATION}
+
+ sync_operator.on_kill()
+
+ mock_hook.return_value.cancel_job.assert_not_called()
+
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowStartFlexTemplateOperator.defer",
autospec=True)
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_deferrable_threads_drain_pipeline_into_trigger(self, mock_hook,
mock_defer_method):
+ operator = DataflowStartFlexTemplateOperator(
+ task_id="start_flex_template_streaming_beam_sql",
+ body={"launchParameter": TEST_FLEX_PARAMETERS},
+ project_id=TEST_PROJECT,
+ location=TEST_LOCATION,
+ deferrable=True,
+ drain_pipeline=True,
+ )
+ operator.execute(mock.MagicMock(spec=dict))
+
+ assert mock_defer_method.call_args.kwargs["trigger"].drain_pipeline is
True
+
class TestDataflowStartYamlJobOperator:
@pytest.fixture
@@ -411,6 +477,56 @@ class TestDataflowStartYamlJobOperator:
)
mock_defer_method.assert_called_once()
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowStartYamlJobOperator.defer",
autospec=True)
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_deferrable_threads_cancel_on_kill_false_into_trigger(self,
mock_hook, mock_defer_method):
+ """cancel_on_kill=False reaches the trigger, so a deferred kill leaves
the job running."""
+ operator = DataflowStartYamlJobOperator(
+ task_id="start_dataflow_yaml_job_cancel_false",
+ job_name="dataflow_yaml_job",
+ yaml_pipeline_file="test_file_path",
+ append_job_name=False,
+ project_id=TEST_PROJECT,
+ region=TEST_LOCATION,
+ deferrable=True,
+ cancel_on_kill=False,
+ expected_terminal_state=DataflowJobStatus.JOB_STATE_RUNNING,
+ )
+ operator.execute(mock.MagicMock(spec=dict))
+
+ assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is
False
+
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_on_kill_respects_cancel_on_kill_false(self, mock_hook,
sync_operator):
+ sync_operator.cancel_on_kill = False
+ sync_operator.job_id = "test-job-id"
+
+ sync_operator.on_kill()
+
+ mock_hook.return_value.cancel_job.assert_not_called()
+
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowStartYamlJobOperator.defer",
autospec=True)
+ @mock.patch(f"{DATAFLOW_PATH}.DataflowHook", autospec=True)
+ def test_deferrable_threads_drain_pipeline_into_trigger(self, mock_hook,
mock_defer_method):
+ """drain_pipeline reaches the trigger so a killed deferred task drains
instead of cancels."""
+ operator = DataflowStartYamlJobOperator(
+ task_id="start_dataflow_yaml_job_drain",
+ job_name="dataflow_yaml_job",
+ yaml_pipeline_file="test_file_path",
+ append_job_name=False,
+ project_id=TEST_PROJECT,
+ region=TEST_LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ deferrable=True,
+ drain_pipeline=True,
+ expected_terminal_state=DataflowJobStatus.JOB_STATE_RUNNING,
+ )
+ operator.execute(mock.MagicMock(spec=dict))
+
+ trigger = mock_defer_method.call_args.kwargs["trigger"]
+ assert trigger.drain_pipeline is True
+ assert trigger.cancel_on_kill is True
+
@mock.patch(f"{DATAFLOW_PATH}.DataflowHook")
def test_execute_complete_success(self, mock_hook, deferrable_operator):
expected_result = {"id": JOB_ID}
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py
b/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py
index b179f40fa3a..5d21878c08c 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py
@@ -25,7 +25,7 @@ import pytest
from google.api_core.exceptions import ServiceUnavailable
from google.cloud.dataflow_v1beta3 import Job, JobState, JobType
-from airflow.providers.google.cloud.hooks.dataflow import DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import DataflowJobStatus,
DataflowJobType
from airflow.providers.google.cloud.triggers.dataflow import (
DataflowJobAutoScalingEventTrigger,
DataflowJobMessagesTrigger,
@@ -143,6 +143,62 @@ def test_dataflow_batch_job():
return Job(id=JOB_ID, current_state=JobState.JOB_STATE_DONE,
type_=JobType.JOB_TYPE_BATCH)
[email protected]
[email protected]("trigger_class", [TemplateJobStartTrigger,
DataflowStartYamlJobTrigger])
[email protected]("project_id", [PROJECT_ID, None],
ids=["explicit-project", "connection-project"])
[email protected](
+ ("drain_pipeline", "requested_state"),
+ [(False, DataflowJobStatus.JOB_STATE_CANCELLED), (True,
DataflowJobStatus.JOB_STATE_DRAINED)],
+)
[email protected](
+
"airflow.providers.google.cloud.hooks.dataflow._DataflowJobsController._wait_for_states",
autospec=True
+)
[email protected]("airflow.providers.google.cloud.hooks.dataflow.DataflowHook.get_conn",
autospec=True)
[email protected]("airflow.providers.google.common.hooks.base_google.BaseHook.get_connection",
autospec=True)
[email protected](
+
"airflow.providers.google.common.hooks.base_google.GoogleBaseHook.get_credentials_and_project_id",
+ autospec=True,
+ return_value=(None, PROJECT_ID),
+)
+async def test_on_kill_requests_stop_without_waiting(
+ mock_credentials,
+ mock_connection,
+ mock_get_conn,
+ mock_wait_for_states,
+ trigger_class,
+ project_id,
+ drain_pipeline,
+ requested_state,
+):
+ mock_connection.return_value.extra_dejson = {}
+ jobs =
mock_get_conn.return_value.projects.return_value.locations.return_value.jobs.return_value
+ jobs.get.return_value.execute.return_value = {
+ "id": JOB_ID,
+ "name": "test-job",
+ "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
+ "type": DataflowJobType.JOB_TYPE_STREAMING,
+ }
+ trigger = trigger_class(
+ project_id=project_id,
+ job_id=JOB_ID,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ cancel_timeout=CANCEL_TIMEOUT,
+ drain_pipeline=drain_pipeline,
+ )
+
+ await trigger.on_kill()
+
+ jobs.update.assert_called_once_with(
+ projectId=PROJECT_ID,
+ location=LOCATION,
+ jobId=JOB_ID,
+ body={"requestedState": requested_state},
+ )
+ jobs.update.return_value.execute.assert_called_once()
+ mock_wait_for_states.assert_not_called()
+
+
class TestTemplateJobStartTrigger:
def test_serialize(self, template_job_start_trigger):
actual_data = template_job_start_trigger.serialize()
@@ -156,6 +212,8 @@ class TestTemplateJobStartTrigger:
"poll_sleep": POLL_SLEEP,
"impersonation_chain": IMPERSONATION_CHAIN,
"cancel_timeout": CANCEL_TIMEOUT,
+ "cancel_on_kill": True,
+ "drain_pipeline": False,
},
)
assert actual_data == expected_data
@@ -175,6 +233,87 @@ class TestTemplateJobStartTrigger:
assert actual is not None
assert actual == expected
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("project_id", [PROJECT_ID, None],
ids=["explicit-project", "connection-project"])
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_stops_the_job(self, mock_hook_class,
template_job_start_trigger, project_id):
+ template_job_start_trigger.project_id = project_id
+ await template_job_start_trigger.on_kill()
+
+ mock_hook_class.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ drain_pipeline=False,
+ cancel_timeout=None,
+ poll_sleep=POLL_SLEEP,
+ )
+ mock_hook_class.return_value.cancel_job.assert_called_once_with(
+ job_id=JOB_ID,
+ project_id=project_id,
+ location=LOCATION,
+ )
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_passes_drain_pipeline_to_the_hook(self,
mock_hook_class):
+ trigger = TemplateJobStartTrigger(
+ project_id=PROJECT_ID,
+ job_id=JOB_ID,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ poll_sleep=POLL_SLEEP,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ cancel_timeout=CANCEL_TIMEOUT,
+ drain_pipeline=True,
+ )
+
+ await trigger.on_kill()
+
+ assert mock_hook_class.call_args.kwargs["drain_pipeline"] is True
+
+ @pytest.mark.parametrize(
+ ("cancel_on_kill", "job_id", "project_id"),
+ [
+ pytest.param(False, JOB_ID, PROJECT_ID,
id="cancel-on-kill-disabled"),
+ pytest.param(True, None, PROJECT_ID, id="no-job-id"),
+ ],
+ )
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_does_not_stop_the_job(self, mock_hook_class,
cancel_on_kill, job_id, project_id):
+ trigger = TemplateJobStartTrigger(
+ project_id=project_id,
+ job_id=job_id,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ cancel_on_kill=cancel_on_kill,
+ )
+
+ await trigger.on_kill()
+
+ mock_hook_class.assert_not_called()
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_swallows_hook_errors(self, mock_hook_class,
template_job_start_trigger):
+ mock_hook_class.return_value.cancel_job.side_effect =
RuntimeError("api down")
+
+ await template_job_start_trigger.on_kill()
+
+ mock_hook_class.return_value.cancel_job.assert_called_once()
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_swallows_hook_construction_errors(
+ self, mock_hook_class, template_job_start_trigger
+ ):
+ # The hook resolves its connection during construction; that must not
escape on_kill.
+ mock_hook_class.side_effect = RuntimeError("no connection")
+
+ await template_job_start_trigger.on_kill()
+
+ mock_hook_class.assert_called_once()
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.hooks.dataflow.AsyncDataflowHook.get_job_status")
async def test_run_loop_return_success_event(self, mock_job_status,
template_job_start_trigger):
@@ -819,10 +958,93 @@ class TestDataflowStartYamlJobTrigger:
"expected_terminal_state": None,
"impersonation_chain": IMPERSONATION_CHAIN,
"cancel_timeout": CANCEL_TIMEOUT,
+ "cancel_on_kill": True,
+ "drain_pipeline": False,
},
)
assert actual_data == expected_data
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("project_id", [PROJECT_ID, None],
ids=["explicit-project", "connection-project"])
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_stops_the_job(self, mock_hook_class,
dataflow_start_yaml_job_trigger, project_id):
+ dataflow_start_yaml_job_trigger.project_id = project_id
+ await dataflow_start_yaml_job_trigger.on_kill()
+
+ mock_hook_class.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ drain_pipeline=False,
+ cancel_timeout=None,
+ poll_sleep=POLL_SLEEP,
+ )
+ mock_hook_class.return_value.cancel_job.assert_called_once_with(
+ job_id=JOB_ID,
+ project_id=project_id,
+ location=LOCATION,
+ )
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_passes_drain_pipeline_to_the_hook(self,
mock_hook_class):
+ trigger = DataflowStartYamlJobTrigger(
+ project_id=PROJECT_ID,
+ job_id=JOB_ID,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ poll_sleep=POLL_SLEEP,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ cancel_timeout=CANCEL_TIMEOUT,
+ drain_pipeline=True,
+ )
+
+ await trigger.on_kill()
+
+ assert mock_hook_class.call_args.kwargs["drain_pipeline"] is True
+
+ @pytest.mark.parametrize(
+ ("cancel_on_kill", "job_id", "project_id"),
+ [
+ pytest.param(False, JOB_ID, PROJECT_ID,
id="cancel-on-kill-disabled"),
+ pytest.param(True, None, PROJECT_ID, id="no-job-id"),
+ ],
+ )
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_does_not_stop_the_job(self, mock_hook_class,
cancel_on_kill, job_id, project_id):
+ trigger = DataflowStartYamlJobTrigger(
+ project_id=project_id,
+ job_id=job_id,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ cancel_on_kill=cancel_on_kill,
+ )
+
+ await trigger.on_kill()
+
+ mock_hook_class.assert_not_called()
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_swallows_hook_errors(self, mock_hook_class,
dataflow_start_yaml_job_trigger):
+ mock_hook_class.return_value.cancel_job.side_effect =
RuntimeError("api down")
+
+ await dataflow_start_yaml_job_trigger.on_kill()
+
+ mock_hook_class.return_value.cancel_job.assert_called_once()
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook",
autospec=True)
+ async def test_on_kill_swallows_hook_construction_errors(
+ self, mock_hook_class, dataflow_start_yaml_job_trigger
+ ):
+ # The hook resolves its connection during construction; that must not
escape on_kill.
+ mock_hook_class.side_effect = RuntimeError("no connection")
+
+ await dataflow_start_yaml_job_trigger.on_kill()
+
+ mock_hook_class.assert_called_once()
+
@pytest.mark.parametrize(
("attr", "expected"),
[
