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 f3784d2e9e1 Cancel the BigQuery job when BigQueryToGCSOperator is
killed (#74564)
f3784d2e9e1 is described below
commit f3784d2e9e10c175ba4421e2afbb14b61feb3519
Author: Yuseok Jo <[email protected]>
AuthorDate: Sat Oct 17 22:44:18 2026 +0900
Cancel the BigQuery job when BigQueryToGCSOperator is killed (#74564)
---
.../google/cloud/transfers/bigquery_to_gcs.py | 15 ++++++++
.../google/cloud/transfers/test_bigquery_to_gcs.py | 42 ++++++++++++++++++++++
2 files changed, 57 insertions(+)
diff --git
a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py
b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py
index c9ec3df8023..a15a6eb08b7 100644
---
a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py
+++
b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py
@@ -86,6 +86,10 @@ class BigQueryToGCSOperator(BaseOperator):
:param reattach_states: Set of BigQuery job's states in case of which we
should reattach
to the job. Should be other than final states.
:param deferrable: Run operator in the deferrable mode
+ :param cancel_on_kill: Flag which indicates whether cancel the hook's job
or not, when on_kill is called.
+ With ``force_rerun=False``, a retry reuses the same job ID, and a
cancelled job cannot be
+ reattached through ``reattach_states``. Set ``force_rerun=True`` for
retries to start a new
+ job, or ``cancel_on_kill=False`` to keep reattaching to the running
one.
:return: URIs for the objects created in Google Cloud Storage
"""
@@ -122,6 +126,7 @@ class BigQueryToGCSOperator(BaseOperator):
force_rerun: bool = False,
reattach_states: set[str] | None = None,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ cancel_on_kill: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -143,6 +148,7 @@ class BigQueryToGCSOperator(BaseOperator):
self.reattach_states: set[str] = reattach_states or set()
self.hook: BigQueryHook | None = None
self.deferrable = deferrable
+ self.cancel_on_kill = cancel_on_kill
@staticmethod
def _handle_job_error(job: BigQueryJob | UnknownJob) -> None:
@@ -270,6 +276,7 @@ class BigQueryToGCSOperator(BaseOperator):
project_id=self.project_id or self.hook.project_id,
location=self.location or self.hook.location,
impersonation_chain=self.impersonation_chain,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name="execute_complete",
)
@@ -296,6 +303,14 @@ class BigQueryToGCSOperator(BaseOperator):
return self.destination_cloud_storage_uris
+ def on_kill(self) -> None:
+ if self.job_id and self.cancel_on_kill:
+ self.hook.cancel_job( # type: ignore[union-attr]
+ job_id=self.job_id, project_id=self.project_id,
location=self.location
+ )
+ else:
+ self.log.info("Skipping to cancel job: %s:%s.%s", self.project_id,
self.location, self.job_id)
+
def get_openlineage_facets_on_complete(self, task_instance):
"""Implement on_complete as we will include final BQ job id."""
from airflow.providers.common.compat.openlineage.facet import (
diff --git
a/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_gcs.py
b/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_gcs.py
index ba1beed291d..1231375beb8 100644
--- a/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_gcs.py
+++ b/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_gcs.py
@@ -191,6 +191,48 @@ class TestBigQueryToGCSOperator:
nowait=True,
)
+
@mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_gcs.BigQueryHook")
+ def test_execute_deferrable_mode_passes_cancel_on_kill_to_trigger(self,
mock_hook):
+ mock_hook.return_value.split_tablename.return_value = (PROJECT_ID,
TEST_DATASET, TEST_TABLE_ID)
+ mock_hook.return_value.insert_job.return_value =
MagicMock(job_id="real_job_id", error_result=False)
+
+ operator = BigQueryToGCSOperator(
+ task_id=TASK_ID,
+
source_project_dataset_table=f"{PROJECT_ID}:{TEST_DATASET}.{TEST_TABLE_ID}",
+ destination_cloud_storage_uris=["gs://some-bucket/some-file.txt"],
+ project_id=JOB_PROJECT_ID,
+ deferrable=True,
+ cancel_on_kill=False,
+ )
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context=mock.MagicMock())
+
+ assert exc.value.trigger.cancel_on_kill is False
+
+
@mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_gcs.BigQueryHook")
+ def test_on_kill(self, mock_hook):
+ mock_hook.return_value.split_tablename.return_value = (PROJECT_ID,
TEST_DATASET, TEST_TABLE_ID)
+ mock_hook.return_value.insert_job.return_value =
MagicMock(job_id="real_job_id", error_result=False)
+
+ operator = BigQueryToGCSOperator(
+ task_id=TASK_ID,
+
source_project_dataset_table=f"{PROJECT_ID}:{TEST_DATASET}.{TEST_TABLE_ID}",
+ destination_cloud_storage_uris=["gs://some-bucket/some-file.txt"],
+ project_id=JOB_PROJECT_ID,
+ location="US",
+ cancel_on_kill=False,
+ )
+ operator.execute(context=mock.MagicMock())
+
+ operator.on_kill()
+ mock_hook.return_value.cancel_job.assert_not_called()
+
+ operator.cancel_on_kill = True
+ operator.on_kill()
+ mock_hook.return_value.cancel_job.assert_called_once_with(
+ job_id="real_job_id", project_id=JOB_PROJECT_ID, location="US"
+ )
+
def test_execute_complete_reassigns_job_id(self):
"""Assert that we use job_id from event after deferral."""
