From 1d08e57fee817dda2d084d7c0940d0bc5aab08c6 Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Mon, 27 Apr 2026 15:34:46 -0700 Subject: [PATCH 1/5] Add system tests for Dataproc trigger on_kill cancel behavior Follow-up to #65742 (requested by @shahar1). Adds a system test DAG that exercises the cancel_on_kill plumbing introduced when DataprocSubmitTrigger and DataprocSubmitJobDirectTrigger migrated to BaseTrigger.on_kill(). Test A: deferrable submit with cancel_on_kill=True completes normally. Test B: submits a long-running job, cancels it via DataprocHook.cancel_job() (the same call on_kill delegates to), and asserts CANCELLED state. --- .../example_dataproc_cancel_on_kill.py | 194 ++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py diff --git a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py new file mode 100644 index 0000000000000..793181b74d5f5 --- /dev/null +++ b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py @@ -0,0 +1,194 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Example Airflow DAG that tests cancel_on_kill behavior for Dataproc triggers. + +Test A (happy path): Submits a Spark job in deferrable mode with cancel_on_kill=True +and verifies it completes successfully. + +Test B (cancel path): Submits a long-running Spark job asynchronously, cancels it +via DataprocHook.cancel_job() — the same call that trigger.on_kill() delegates to — +and verifies the job reaches CANCELLED state. +""" + +from __future__ import annotations + +import os +import time +from datetime import datetime + +from google.api_core.retry import Retry +from google.cloud.dataproc_v1 import JobStatus + +from airflow.decorators import task +from airflow.models.dag import DAG +from airflow.providers.google.cloud.hooks.dataproc import DataprocHook +from airflow.providers.google.cloud.operators.dataproc import ( + DataprocCreateClusterOperator, + DataprocDeleteClusterOperator, + DataprocSubmitJobOperator, +) + +try: + from airflow.sdk import TriggerRule +except ImportError: + from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] + +from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default") +DAG_ID = "dataproc_cancel_on_kill" +PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID + +CLUSTER_NAME_BASE = f"cluster-{DAG_ID}".replace("_", "-") +CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-") +CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL + +REGION = "europe-west1" + +CLUSTER_CONFIG = { + "master_config": { + "num_instances": 1, + "machine_type_uri": "n1-standard-4", + "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 32}, + }, + "worker_config": { + "num_instances": 2, + "machine_type_uri": "n1-standard-4", + "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 32}, + }, +} + +SPARK_JOB = { + "reference": {"project_id": PROJECT_ID}, + "placement": {"cluster_name": CLUSTER_NAME}, + "spark_job": { + "jar_file_uris": ["file:///usr/lib/spark/examples/jars/spark-examples.jar"], + "main_class": "org.apache.spark.examples.SparkPi", + }, +} + +# [START how_to_cloud_dataproc_cancel_on_kill_config] +LONG_RUNNING_SPARK_JOB = { + "reference": {"project_id": PROJECT_ID}, + "placement": {"cluster_name": CLUSTER_NAME}, + "spark_job": { + "jar_file_uris": ["file:///usr/lib/spark/examples/jars/spark-examples.jar"], + "main_class": "org.apache.spark.examples.SparkPi", + "args": ["1000000"], + }, +} +# [END how_to_cloud_dataproc_cancel_on_kill_config] + + +with DAG( + DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "dataproc", "cancel_on_kill", "deferrable"], +) as dag: + create_cluster = DataprocCreateClusterOperator( + task_id="create_cluster", + project_id=PROJECT_ID, + cluster_config=CLUSTER_CONFIG, + region=REGION, + cluster_name=CLUSTER_NAME, + retry=Retry(maximum=100.0, initial=10.0, multiplier=1.0), + num_retries_if_resource_is_not_ready=3, + ) + + # Test A: deferrable submit with cancel_on_kill=True completes normally + # [START how_to_cloud_dataproc_deferrable_cancel_on_kill] + spark_task_deferrable = DataprocSubmitJobOperator( + task_id="spark_task_deferrable", + job=SPARK_JOB, + region=REGION, + project_id=PROJECT_ID, + deferrable=True, + cancel_on_kill=True, + ) + # [END how_to_cloud_dataproc_deferrable_cancel_on_kill] + + # Test B: submit a long-running job, cancel it, verify CANCELLED state + submit_long_job = DataprocSubmitJobOperator( + task_id="submit_long_job", + job=LONG_RUNNING_SPARK_JOB, + region=REGION, + project_id=PROJECT_ID, + asynchronous=True, + ) + + @task(task_id="cancel_and_verify") + def cancel_and_verify_job(job_id: str, project_id: str, region: str): + """Cancel a running Dataproc job and verify it reaches CANCELLED state. + + Exercises the same DataprocHook.cancel_job() call that + DataprocSubmitTrigger.on_kill() and DataprocSubmitJobDirectTrigger.on_kill() + delegate to. + """ + hook = DataprocHook(gcp_conn_id="google_cloud_default") + + hook.cancel_job(job_id=job_id, project_id=project_id, region=region) + + for _ in range(30): + job = hook.get_job(job_id=job_id, project_id=project_id, region=region) + state = job.status.state + if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR): + break + time.sleep(5) + else: + raise RuntimeError(f"Job {job_id} did not reach terminal state within 150s") + + assert job.status.state == JobStatus.State.CANCELLED, ( + f"Expected CANCELLED, got {JobStatus.State(job.status.state).name}" + ) + + cancel_task = cancel_and_verify_job( + job_id=submit_long_job.output, + project_id=PROJECT_ID, + region=REGION, + ) + + delete_cluster = DataprocDeleteClusterOperator( + task_id="delete_cluster", + project_id=PROJECT_ID, + cluster_name=CLUSTER_NAME, + region=REGION, + trigger_rule=TriggerRule.ALL_DONE, + ) + + ( + # TEST SETUP + create_cluster + # TEST BODY + >> spark_task_deferrable + >> submit_long_job + >> cancel_task + # TEST TEARDOWN + >> delete_cluster + ) + + from tests_common.test_utils.watcher import watcher + + list(dag.tasks) >> watcher() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +test_run = get_test_run(dag) From 070267d22382311edf3312ad39d2a85d3b0eec4a Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Mon, 27 Apr 2026 17:51:02 -0700 Subject: [PATCH 2/5] Fix MyPy failure: use version-compat import pattern for task decorator The bare `from airflow.decorators import task` import fails MyPy in the providers check because `airflow.decorators` does not export `task` in Airflow 3.x. Switch to the standard `AIRFLOW_V_3_0_PLUS` conditional import pattern used by all other system tests. --- .../cloud/dataproc/example_dataproc_cancel_on_kill.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py index 793181b74d5f5..226d8e69d1922 100644 --- a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py +++ b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py @@ -35,7 +35,6 @@ from google.api_core.retry import Retry from google.cloud.dataproc_v1 import JobStatus -from airflow.decorators import task from airflow.models.dag import DAG from airflow.providers.google.cloud.hooks.dataproc import DataprocHook from airflow.providers.google.cloud.operators.dataproc import ( @@ -44,9 +43,12 @@ DataprocSubmitJobOperator, ) -try: - from airflow.sdk import TriggerRule -except ImportError: +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import TriggerRule, task +else: + from airflow.decorators import task # type: ignore[attr-defined,no-redef] from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID From dc691af338ae804ec67da7d24f71b0997d171774 Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Wed, 29 Apr 2026 07:06:12 -0700 Subject: [PATCH 3/5] Gate cancel_on_kill system test behind opt-in env var Per reviewer feedback (#65982), exclude the cancel_on_kill system test from automated Google CI runs. Test now skips at collection time unless RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST=1 is set, allowing manual end-to-end runs without burdening automated infrastructure. --- .../cloud/dataproc/example_dataproc_cancel_on_kill.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py index 226d8e69d1922..9a391e4a40ed5 100644 --- a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py +++ b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py @@ -32,6 +32,7 @@ import time from datetime import datetime +import pytest from google.api_core.retry import Retry from google.cloud.dataproc_v1 import JobStatus @@ -51,6 +52,12 @@ from airflow.decorators import task # type: ignore[attr-defined,no-redef] from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] +if not os.environ.get("RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST"): + pytest.skip( + "Manual-only system test: set RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST=1 to run.", + allow_module_level=True, + ) + from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default") From 649d39779cae6408e18bfc27e134a00e243368ca Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Wed, 29 Apr 2026 14:15:25 -0700 Subject: [PATCH 4/5] Replace module-level pytest.skip with pytestmark skipif pytest.skip(allow_module_level=True) raises a Skipped exception at import time that DagBag records as an import error, breaking the test_should_be_importable parametrized test in DB-core CI jobs. pytestmark = pytest.mark.skipif(...) lets the file import cleanly (DagBag happy, importability test passes) while still skipping test_run when RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST is unset. --- .../dataproc/example_dataproc_cancel_on_kill.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py index 9a391e4a40ed5..fd5f06e12459d 100644 --- a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py +++ b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py @@ -44,6 +44,7 @@ DataprocSubmitJobOperator, ) +from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS if AIRFLOW_V_3_0_PLUS: @@ -52,13 +53,10 @@ from airflow.decorators import task # type: ignore[attr-defined,no-redef] from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] -if not os.environ.get("RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST"): - pytest.skip( - "Manual-only system test: set RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST=1 to run.", - allow_module_level=True, - ) - -from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID +pytestmark = pytest.mark.skipif( + not os.environ.get("RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST"), + reason="Manual-only system test: set RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST=1 to run.", +) ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default") DAG_ID = "dataproc_cancel_on_kill" From d9f467cd88dc133ef3d033427802bdb2fa705172 Mon Sep 17 00:00:00 2001 From: Sadha Chilukoori Date: Tue, 5 May 2026 06:43:41 -0700 Subject: [PATCH 5/5] Rename env flag to RUN_MANUAL_GOOGLE_SYSTEM_TESTS for generalization --- .../google/cloud/dataproc/example_dataproc_cancel_on_kill.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py index fd5f06e12459d..457ad240798d4 100644 --- a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py +++ b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_cancel_on_kill.py @@ -54,8 +54,8 @@ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] pytestmark = pytest.mark.skipif( - not os.environ.get("RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST"), - reason="Manual-only system test: set RUN_MANUAL_DATAPROC_CANCEL_ON_KILL_TEST=1 to run.", + not os.environ.get("RUN_MANUAL_GOOGLE_SYSTEM_TESTS"), + reason="Manual-only system test: set RUN_MANUAL_GOOGLE_SYSTEM_TESTS=1 to run.", ) ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")