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 a1976c57411 Cancel the BigQuery job when BigQueryToBigQueryOperator is
killed (#75101)
a1976c57411 is described below
commit a1976c57411f64d2115b2d8a85fda3057b0c911e
Author: Yuseok Jo <[email protected]>
AuthorDate: Sun Oct 25 15:39:58 2026 +0900
Cancel the BigQuery job when BigQueryToBigQueryOperator is killed (#75101)
---
.../google/cloud/transfers/bigquery_to_bigquery.py | 22 +++++++++++-
.../cloud/transfers/test_bigquery_to_bigquery.py | 42 ++++++++++++++++++++++
2 files changed, 63 insertions(+), 1 deletion(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py
b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py
index a581fc560fa..b259ed7911b 100644
---
a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py
+++
b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py
@@ -75,6 +75,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
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 project_id: Google Cloud Project where the job is running
+ :param cancel_on_kill: Flag which indicates whether cancel the hook's job
or not, when on_kill is called
"""
template_fields: Sequence[str] = (
@@ -101,6 +102,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
encryption_configuration: dict | None = None,
location: str | None = None,
impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -117,6 +119,8 @@ class BigQueryToBigQueryOperator(BaseOperator):
self.hook: BigQueryHook | None = None
self._job_conf: dict = {}
self.project_id = project_id
+ self.cancel_on_kill = cancel_on_kill
+ self.job_id: str | None = None
def _prepare_job_configuration(self):
self.source_project_dataset_tables = (
@@ -177,8 +181,18 @@ class BigQueryToBigQueryOperator(BaseOperator):
self.project_id = self.hook.project_id
configuration = self._prepare_job_configuration()
+ # The ID is set before submitting so on_kill can cancel the job while
insert_job waits for it.
+ self.job_id = self.hook.generate_job_id(
+ job_id=None,
+ dag_id=self.dag_id,
+ task_id=self.task_id,
+ logical_date=None,
+ configuration=configuration,
+ run_after=self.hook.get_run_after_or_logical_date(context),
+ force_rerun=True,
+ )
self._job_conf = self.hook.insert_job(
- configuration=configuration, project_id=self.project_id
+ configuration=configuration, job_id=self.job_id,
project_id=self.project_id
).to_api_repr()
dest_table_info =
self._job_conf["configuration"]["copy"]["destinationTable"]
@@ -189,6 +203,12 @@ class BigQueryToBigQueryOperator(BaseOperator):
table_id=dest_table_info["tableId"],
)
+ def on_kill(self) -> None:
+ if self.hook and self.job_id and self.cancel_on_kill:
+ self.hook.cancel_job(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_bigquery.py
b/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_bigquery.py
index f6da2c4f2bf..2143944b7f1 100644
---
a/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_bigquery.py
+++
b/providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_bigquery.py
@@ -19,6 +19,7 @@ from __future__ import annotations
from unittest import mock
+import pytest
from google.cloud.bigquery import Table
from airflow.providers.common.compat.openlineage.facet import (
@@ -105,6 +106,7 @@ class TestBigQueryToBigQueryOperator:
},
"labels": LABELS,
},
+ job_id=mock_hook.return_value.generate_job_id.return_value,
project_id=mock_hook.return_value.project_id,
)
@@ -146,9 +148,49 @@ class TestBigQueryToBigQueryOperator:
},
"labels": LABELS,
},
+ job_id=mock_hook.return_value.generate_job_id.return_value,
project_id=mock_hook.return_value.project_id,
)
+ @mock.patch(BQ_HOOK_PATH)
+ def test_execute_generates_new_job_id_for_each_attempt(self, mock_hook):
+ operator = BigQueryToBigQueryOperator(
+ task_id=TASK_ID,
+ source_project_dataset_tables=SOURCE_PROJECT_DATASET_TABLE,
+
destination_project_dataset_table=DESTINATION_PROJECT_DATASET_TABLE,
+ )
+
+ mock_hook.return_value.split_tablename.side_effect =
split_tablename_side_effect
+ operator.execute(context=mock.MagicMock())
+
+ assert
mock_hook.return_value.generate_job_id.call_args.kwargs["force_rerun"] is True
+ assert operator.job_id ==
mock_hook.return_value.generate_job_id.return_value
+
+ @pytest.mark.parametrize("cancel_on_kill", [True, False])
+ @mock.patch(BQ_HOOK_PATH)
+ def test_on_kill(self, mock_hook, cancel_on_kill):
+ operator = BigQueryToBigQueryOperator(
+ task_id=TASK_ID,
+ source_project_dataset_tables=SOURCE_PROJECT_DATASET_TABLE,
+
destination_project_dataset_table=DESTINATION_PROJECT_DATASET_TABLE,
+ project_id=TEST_GCP_PROJECT_ID,
+ location="US",
+ cancel_on_kill=cancel_on_kill,
+ )
+
+ mock_hook.return_value.split_tablename.side_effect =
split_tablename_side_effect
+ mock_hook.return_value.generate_job_id.return_value = "real_job_id"
+ operator.execute(context=mock.MagicMock())
+
+ operator.on_kill()
+
+ if cancel_on_kill:
+ mock_hook.return_value.cancel_job.assert_called_once_with(
+ job_id="real_job_id", project_id=TEST_GCP_PROJECT_ID,
location="US"
+ )
+ else:
+ mock_hook.return_value.cancel_job.assert_not_called()
+
@mock.patch(BQ_HOOK_PATH)
def test_get_openlineage_facets_on_complete_single_source_table(self,
mock_hook):
location = "us-central1"
