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 ac71404000e45..d1134afee473f 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. """ template_fields: Sequence[str] = ( @@ -334,11 +337,13 @@ def __init__( 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 @@ def set_current_job(current_job): 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 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str: 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. :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 @@ def __init__( 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 @@ def set_current_job(current_job): 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 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]: 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. :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 @@ def __init__( 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 @@ def execute(self, context: Context) -> dict[str, Any]: 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 @@ def on_kill(self): """ 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; + in that case the trigger cancels the job 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 4b6889d73258c..024bae14124c9 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 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 @@ 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 @@ def __init__( 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 @@ def __init__( 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 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "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 or not self.project_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("Dataflow job %s stopped.", 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: + """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop).""" + hook = DataflowHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + drain_pipeline=self.drain_pipeline, + cancel_timeout=self.cancel_timeout, + 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 @@ def __init__( 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 @@ def __init__( 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 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "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 or not self.project_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("Dataflow job %s stopped.", 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: + """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop).""" + hook = DataflowHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + drain_pipeline=self.drain_pipeline, + cancel_timeout=self.cancel_timeout, + 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 d4442a1c52079..5cfd9f5154c5b 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_dataflow.py +++ b/providers/google/tests/unit/google/cloud/operators/test_dataflow.py @@ -214,6 +214,32 @@ def test_execute_with_deferrable_mode(self, mock_hook, mock_defer_method, deferr ) mock_defer_method.assert_called_once() + @mock.patch(f"{DATAFLOW_PATH}.DataflowTemplatedJobStartOperator.defer") + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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()) + + assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is False + + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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, @@ -351,6 +377,31 @@ def test_execute_with_deferrable_mode(self, mock_hook, mock_defer_method, deferr ) mock_defer_method.assert_called_once() + @mock.patch(f"{DATAFLOW_PATH}.DataflowStartFlexTemplateOperator.defer") + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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()) + + assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is False + + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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() + class TestDataflowStartYamlJobOperator: @pytest.fixture @@ -415,6 +466,56 @@ def test_execute_with_deferrable_mode(self, mock_hook, mock_defer_method, deferr ) mock_defer_method.assert_called_once() + @mock.patch(f"{DATAFLOW_PATH}.DataflowStartYamlJobOperator.defer") + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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()) + + assert mock_defer_method.call_args.kwargs["trigger"].cancel_on_kill is False + + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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") + @mock.patch(f"{DATAFLOW_PATH}.DataflowHook") + 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()) + + 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 bc2eb646e813b..161ab4c01ebff 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_dataflow.py @@ -200,6 +200,8 @@ def test_serialize(self, template_job_start_trigger): "poll_sleep": POLL_SLEEP, "impersonation_chain": IMPERSONATION_CHAIN, "cancel_timeout": CANCEL_TIMEOUT, + "cancel_on_kill": True, + "drain_pipeline": False, }, ) assert actual_data == expected_data @@ -219,6 +221,79 @@ def test_get_async_hook(self, template_job_start_trigger, attr, expected): assert actual is not None assert actual == expected + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + def test_on_kill_stops_the_job(self, mock_hook_class, template_job_start_trigger): + asyncio.run(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=CANCEL_TIMEOUT, + 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, + ) + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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, + ) + + asyncio.run(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.param(True, JOB_ID, None, id="no-project-id"), + ], + ) + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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, + ) + + asyncio.run(trigger.on_kill()) + + mock_hook_class.assert_not_called() + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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") + + asyncio.run(template_job_start_trigger.on_kill()) + + mock_hook_class.return_value.cancel_job.assert_called_once() + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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") + + asyncio.run(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): @@ -863,10 +938,87 @@ def test_serialize(self, dataflow_start_yaml_job_trigger): "expected_terminal_state": None, "impersonation_chain": IMPERSONATION_CHAIN, "cancel_timeout": CANCEL_TIMEOUT, + "cancel_on_kill": True, + "drain_pipeline": False, }, ) assert actual_data == expected_data + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + def test_on_kill_stops_the_job(self, mock_hook_class, dataflow_start_yaml_job_trigger): + asyncio.run(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=CANCEL_TIMEOUT, + 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, + ) + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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, + ) + + asyncio.run(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.param(True, JOB_ID, None, id="no-project-id"), + ], + ) + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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, + ) + + asyncio.run(trigger.on_kill()) + + mock_hook_class.assert_not_called() + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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") + + asyncio.run(dataflow_start_yaml_job_trigger.on_kill()) + + mock_hook_class.return_value.cancel_job.assert_called_once() + + @mock.patch("airflow.providers.google.cloud.triggers.dataflow.DataflowHook") + 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") + + asyncio.run(dataflow_start_yaml_job_trigger.on_kill()) + + mock_hook_class.assert_called_once() + @pytest.mark.parametrize( ("attr", "expected"), [