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 3a37fb1c33e Fix BigQuery cancellation and links with inferred job
locations (#75717)
3a37fb1c33e is described below
commit 3a37fb1c33ed2aabd49e24471dac4d98b7ed8e6b
Author: Django.Jiang <[email protected]>
AuthorDate: Tue Oct 27 00:00:41 2026 -0700
Fix BigQuery cancellation and links with inferred job locations (#75717)
---
.../providers/google/cloud/operators/bigquery.py | 14 +-
.../unit/google/cloud/operators/test_bigquery.py | 183 ++++++++++++++++++++-
2 files changed, 191 insertions(+), 6 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 800d6fa711e..85ebd84fd42 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
@@ -2447,6 +2447,7 @@ class BigQueryInsertJobOperator(
raise AirflowException(f"Job failed with state: {job.state}")
def _persist_job_links(self, job: BigQueryJob | UnknownJob, context: Any)
-> None:
+ location = job.location or self.location
job_types = {
LoadJob._JOB_TYPE: ["sourceTable", "destinationTable"],
CopyJob._JOB_TYPE: ["sourceTable", "destinationTable"],
@@ -2477,7 +2478,7 @@ class BigQueryInsertJobOperator(
job_id_path = convert_job_id(
job_id=self.job_id,
project_id=self.project_id,
- location=self.location,
+ location=location,
)
warnings.warn(
"BigQueryInsertJobOperator's `job_id_path` XCom is deprecated
and will be removed in a "
@@ -2490,7 +2491,7 @@ class BigQueryInsertJobOperator(
persist_kwargs = {
"context": context,
"project_id": self.project_id,
- "location": self.location,
+ "location": location,
"job_id": self.job_id,
}
BigQueryJobDetailLink.persist(**persist_kwargs)
@@ -2530,7 +2531,7 @@ class BigQueryInsertJobOperator(
conn_id=self.gcp_conn_id,
job_id=self.job_id,
project_id=self.project_id,
- location=self.location or self.hook.location, # type:
ignore[union-attr]
+ location=job.location or self.location or
self.hook.location, # type: ignore[union-attr]
poll_interval=self.poll_interval,
impersonation_chain=self.impersonation_chain,
cancel_on_kill=self.cancel_on_kill,
@@ -2698,7 +2699,12 @@ class BigQueryInsertJobOperator(
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)
+ location = self.location
+ job: BigQueryJob | UnknownJob | None = getattr(self, "_job", None)
+ # A replacement may be submitting while _job still refers to the
previous job.
+ if job is not None and job.job_id == self.job_id:
+ location = job.location or location
+ self.hook.cancel_job(job_id=self.job_id,
project_id=self.project_id, location=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 ecc35874454..e7993051a2b 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
@@ -21,6 +21,7 @@ import json
import logging
import os
from contextlib import suppress
+from typing import Any
from unittest import mock
from unittest.mock import ANY, MagicMock
@@ -1035,6 +1036,121 @@ class TestBigQueryUpsertTableOperator:
class TestBigQueryInsertJobOperator:
+ @pytest.mark.parametrize("configured_location", [None, "EU"])
+ @pytest.mark.parametrize("durable", [False, True])
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook",
autospec=True)
+ def test_submitted_job_uses_actual_location_for_links_and_cancellation(
+ self, mock_hook: MagicMock, durable: bool, configured_location: str |
None
+ ) -> None:
+ hook = mock_hook.return_value
+ hook.labels = {}
+ hook.generate_job_id.return_value = "submitted-job"
+ job = MagicMock(
+ spec=QueryJob,
+ job_id="submitted-job",
+ location="asia-northeast3",
+ state="DONE",
+ error_result=None,
+ )
+ job.to_api_repr.return_value = {"configuration": {"query": {}}}
+ hook.insert_job.return_value = job
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration={"query": {"query": "SELECT 1", "useLegacySql":
False}},
+ project_id=TEST_GCP_PROJECT_ID,
+ location=configured_location,
+ deferrable=False,
+ durable=durable,
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+ ti = MagicMock(spec_set=["try_number", "stats_tags", "xcom_push"],
try_number=1, stats_tags={})
+ context = {"ti": ti, "task": op, "task_state_store": task_store}
+
+ assert op.execute(context) == "submitted-job"
+ op.on_kill()
+
+ ti.xcom_push.assert_any_call(
+ key="bigquery_job_detail",
+ value={
+ "project_id": TEST_GCP_PROJECT_ID,
+ "job_id": "submitted-job",
+ "location": "asia-northeast3",
+ },
+ )
+ ti.xcom_push.assert_any_call(
+ key="job_id_path",
value=f"{TEST_GCP_PROJECT_ID}:asia-northeast3:submitted-job"
+ )
+ hook.cancel_job.assert_called_once_with(
+ job_id="submitted-job", project_id=TEST_GCP_PROJECT_ID,
location="asia-northeast3"
+ )
+ assert hook.insert_job.call_args.kwargs["location"] ==
configured_location
+ assert op.location == configured_location
+
+ job.location = None
+ hook.cancel_job.reset_mock()
+ ti.xcom_push.reset_mock()
+ op._persist_job_links(job, context)
+ op.on_kill()
+
+ ti.xcom_push.assert_any_call(
+ key="bigquery_job_detail",
+ value={
+ "project_id": TEST_GCP_PROJECT_ID,
+ "job_id": "submitted-job",
+ "location": configured_location,
+ },
+ )
+ ti.xcom_push.assert_any_call(
+ key="job_id_path",
value=f"{TEST_GCP_PROJECT_ID}:{configured_location or 'US'}:submitted-job"
+ )
+ hook.cancel_job.assert_called_once_with(
+ job_id="submitted-job", project_id=TEST_GCP_PROJECT_ID,
location=configured_location
+ )
+
+ @pytest.mark.parametrize("configured_location", [None, "EU"])
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook",
autospec=True)
+ def test_defers_with_actual_job_location(
+ self, mock_hook: MagicMock, configured_location: str | None
+ ) -> None:
+ hook = mock_hook.return_value
+ hook.labels = {}
+ hook.location = "US"
+ hook.generate_job_id.return_value = "submitted-job"
+ for actual_location in ("asia-northeast3", None):
+ job = MagicMock(spec=QueryJob, job_id="submitted-job",
location=actual_location)
+ job.to_api_repr.return_value = {"configuration": {"query": {}}}
+ job.running.return_value = True
+ hook.insert_job.return_value = job
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration={"query": {"query": "SELECT 1", "useLegacySql":
False}},
+ project_id=TEST_GCP_PROJECT_ID,
+ location=configured_location,
+ deferrable=True,
+ durable=False,
+ )
+ ti = MagicMock(spec_set=["try_number", "xcom_push"], try_number=1)
+ context = {"ti": ti, "task": op}
+
+ with pytest.raises(TaskDeferred) as deferred:
+ op.execute(context)
+
+ assert isinstance(deferred.value.trigger, BigQueryInsertJobTrigger)
+ assert deferred.value.trigger.location == (
+ actual_location or configured_location or hook.location
+ )
+ assert deferred.value.trigger.job_id == "submitted-job"
+ ti.xcom_push.assert_any_call(
+ key="bigquery_job_detail",
+ value={
+ "project_id": TEST_GCP_PROJECT_ID,
+ "job_id": "submitted-job",
+ "location": actual_location or configured_location,
+ },
+ )
+ assert op.location == configured_location
+
def test_init_does_not_snapshot_templated_job_id(self):
op = BigQueryInsertJobOperator(
task_id="insert_query_job",
@@ -1171,7 +1287,7 @@ class TestBigQueryInsertJobOperator:
}
}
mock_hook.return_value.insert_job.return_value = MagicMock(
- state="DONE", job_id=real_job_id, error_result=False
+ state="DONE", job_id=real_job_id, error_result=False, location=None
)
mock_hook.return_value.generate_job_id.return_value = real_job_id
@@ -1224,6 +1340,7 @@ class TestBigQueryInsertJobOperator:
}
mock_job.job_id = real_job_id
+ mock_job.location = None
mock_job.error_result = False
mock_job.state = "DONE"
mock_job.result.side_effect = AirflowTaskTimeout()
@@ -2306,6 +2423,48 @@ class TestBigQueryInsertJobOperator:
not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
)
class TestBigQueryInsertJobOperatorDurable:
+ @pytest.mark.parametrize("keep_previous_job", [False, True])
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook",
autospec=True)
+ def test_cancel_during_replacement_does_not_use_previous_job_location(
+ self, mock_hook: MagicMock, keep_previous_job: bool
+ ) -> None:
+ hook = mock_hook.return_value
+ hook.labels = {}
+ hook.generate_job_id.return_value = "replacement-job"
+ hook.get_job.return_value = MagicMock(
+ spec=QueryJob,
+ job_id="previous-job",
+ location="asia-northeast3",
+ state="DONE",
+ error_result="failed",
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = {"job_id": "previous-job", "location":
"asia-northeast3"}
+ context = self._context(task_store)
+ op = self._make_operator(deferrable=False)
+
+ assert op.get_job_status("previous-job", context) == "error"
+ op.on_kill()
+ hook.cancel_job.assert_called_once_with(
+ job_id="previous-job", project_id=TEST_GCP_PROJECT_ID,
location="asia-northeast3"
+ )
+ hook.cancel_job.reset_mock()
+ if not keep_previous_job:
+ del op._job
+
+ def interrupt_submission(**kwargs: Any) -> None:
+ op.on_kill()
+ raise AirflowTaskTimeout()
+
+ hook.insert_job.side_effect = interrupt_submission
+ with pytest.raises(AirflowTaskTimeout):
+ op.submit_job(context)
+
+ hook.cancel_job.assert_called_once_with(
+ job_id="replacement-job", project_id=TEST_GCP_PROJECT_ID,
location=TEST_DATASET_LOCATION
+ )
+ assert op.location == TEST_DATASET_LOCATION
+
@pytest.mark.parametrize("status", ["RUNNING", "success", "error",
"not_found"])
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
def test_recovers_returned_location_after_worker_failure(self, mock_hook:
MagicMock, status: str) -> None:
@@ -2334,6 +2493,11 @@ class TestBigQueryInsertJobOperatorDurable:
"bigquery_job_id": "original-job",
"bigquery_job_location": {"job_id": "original-job", "location":
"asia-northeast3"},
}
+ op.on_kill()
+ hook.cancel_job.assert_called_once_with(
+ job_id="original-job", project_id=TEST_GCP_PROJECT_ID,
location="asia-northeast3"
+ )
+ hook.cancel_job.reset_mock()
hook.insert_job.reset_mock()
job.result.side_effect = lambda **_: setattr(job, "state", "DONE")
job.state = "RUNNING" if status == "RUNNING" else "DONE"
@@ -2348,7 +2512,8 @@ class TestBigQueryInsertJobOperatorDurable:
op = self._make_operator()
op.location = "EU" if status in ("error", "not_found") else None
- result = op.execute(self._context(task_store))
+ context = self._context(task_store)
+ result = op.execute(context)
hook.get_job.assert_called_once_with(
project_id=TEST_GCP_PROJECT_ID, location="asia-northeast3",
job_id="original-job"
@@ -2361,6 +2526,20 @@ class TestBigQueryInsertJobOperatorDurable:
assert result == "original-job"
hook.insert_job.assert_not_called()
+ expected_location = "EU" if status in ("error", "not_found") else
"asia-northeast3"
+ context["ti"].xcom_push.assert_any_call(
+ key="bigquery_job_detail",
+ value={"project_id": TEST_GCP_PROJECT_ID, "job_id": result,
"location": expected_location},
+ )
+ context["ti"].xcom_push.assert_any_call(
+ key="job_id_path",
value=f"{TEST_GCP_PROJECT_ID}:{expected_location}:{result}"
+ )
+ op.on_kill()
+ hook.cancel_job.assert_called_once_with(
+ job_id=result, project_id=TEST_GCP_PROJECT_ID,
location=expected_location
+ )
+ assert op.location == ("EU" if status in ("error", "not_found") else
None)
+
@pytest.mark.parametrize(
"saved_location",
[
