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 a1e66b1f1d3 Check hooks before cancelling BigQuery jobs in on_kill
(#74946)
a1e66b1f1d3 is described below
commit a1e66b1f1d3906b161889052555c2735f98466b4
Author: Yuseok Jo <[email protected]>
AuthorDate: Sun Oct 18 16:50:47 2026 +0900
Check hooks before cancelling BigQuery jobs in on_kill (#74946)
---
.../providers/google/cloud/operators/bigquery.py | 6 ++----
.../google/cloud/transfers/bigquery_to_gcs.py | 6 ++----
.../unit/google/cloud/operators/test_bigquery.py | 12 +++++++++++
.../google/cloud/transfers/test_bigquery_to_gcs.py | 25 ++++++++++++++++------
4 files changed, 34 insertions(+), 15 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
index 8b60e461ead..76c2389e3c5 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
@@ -2708,10 +2708,8 @@ class BigQueryInsertJobOperator(
return self.job_id
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
- )
+ 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)
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 a15a6eb08b7..db5b40b8cc6 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
@@ -304,10 +304,8 @@ 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
- )
+ 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)
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
index c44947bf1f6..86172a145f0 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
@@ -1200,6 +1200,18 @@ class TestBigQueryInsertJobOperator:
project_id=TEST_GCP_PROJECT_ID,
)
+ def test_on_kill_before_execute(self):
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration={"query": {"query": "SELECT * FROM any",
"useLegacySql": False}},
+ location=TEST_DATASET_LOCATION,
+ job_id="123456",
+ project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
+ )
+
+ op.on_kill()
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryJob")
def test_on_kill_after_execution_timeout(self, mock_job, mock_hook):
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 1231375beb8..c9c28f3366f 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
@@ -209,8 +209,9 @@ class TestBigQueryToGCSOperator:
assert exc.value.trigger.cancel_on_kill is False
+ @pytest.mark.parametrize("cancel_on_kill", [True, False])
@mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_gcs.BigQueryHook")
- def test_on_kill(self, mock_hook):
+ def test_on_kill(self, mock_hook, cancel_on_kill):
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)
@@ -220,19 +221,29 @@ class TestBigQueryToGCSOperator:
destination_cloud_storage_uris=["gs://some-bucket/some-file.txt"],
project_id=JOB_PROJECT_ID,
location="US",
- cancel_on_kill=False,
+ cancel_on_kill=cancel_on_kill,
)
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"
+ if cancel_on_kill:
+ mock_hook.return_value.cancel_job.assert_called_once_with(
+ job_id="real_job_id", project_id=JOB_PROJECT_ID, location="US"
+ )
+ else:
+ mock_hook.return_value.cancel_job.assert_not_called()
+
+ def test_on_kill_before_execute(self):
+ 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"],
+ job_id="given_job_id",
)
+ operator.on_kill()
+
def test_execute_complete_reassigns_job_id(self):
"""Assert that we use job_id from event after deferral."""
