This is an automated email from the ASF dual-hosted git repository.
amoghrajesh 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 3217e297f99 Adding durable execution to `BigQueryInsertJobOperator`
(#69542)
3217e297f99 is described below
commit 3217e297f994fbfa49079e11d945b5af9369795f
Author: Amogh Desai <[email protected]>
AuthorDate: Fri Jul 10 10:38:10 2026 +0530
Adding durable execution to `BigQueryInsertJobOperator` (#69542)
---
providers/google/docs/operators/cloud/bigquery.rst | 57 ++++
.../providers/google/cloud/operators/bigquery.py | 218 ++++++++++-----
.../unit/google/cloud/operators/test_bigquery.py | 292 ++++++++++++++++++++-
3 files changed, 496 insertions(+), 71 deletions(-)
diff --git a/providers/google/docs/operators/cloud/bigquery.rst
b/providers/google/docs/operators/cloud/bigquery.rst
index d42cda58656..1513a958670 100644
--- a/providers/google/docs/operators/cloud/bigquery.rst
+++ b/providers/google/docs/operators/cloud/bigquery.rst
@@ -369,6 +369,63 @@ Also for all this action you can use operator in the
deferrable mode:
:start-after: [START howto_operator_bigquery_insert_job_async]
:end-before: [END howto_operator_bigquery_insert_job_async]
+Durable execution
+^^^^^^^^^^^^^^^^^
+
+``BigQueryInsertJobOperator`` submits a job and then polls it to completion on
the worker. By
+default the operator runs in a *durable* mode that makes this crash-safe: the
submitted BigQuery
+job id is persisted to task state before polling begins, so if the worker
crashes or is preempted
+and the task is retried, the operator reconnects to the job that is already
running in BigQuery
+instead of submitting a duplicate.
+
+On retry the operator checks the prior job's state:
+
+* if it is still running, the operator reconnects and continues polling
+* if it already succeeded, the operator returns immediately without
resubmitting
+* if it failed, or its id is no longer found, the operator submits the job
fresh
+
+This makes reattachment on retry work regardless of ``force_rerun``, which
defaults to ``True``.
+Without durable execution, ``force_rerun=True`` computes a new random job id
on every attempt, so
+the existing ``reattach_states``/``Conflict`` mechanism (recomputing an
identical id and relying
+on BigQuery rejecting the duplicate) almost never has a matching id to find on
retry -- setting
+``force_rerun=False`` was the only way to make that mechanism reliable.
Durable execution replaces
+that specific use of ``force_rerun=False``: the operator now reconnects using
the id it actually
+persisted, so ``force_rerun`` no longer needs to change purely to make
*retries of the same task
+instance* reattach.
+
+``force_rerun`` still matters for a different, narrower purpose it always had:
idempotency across
+*separate* invocations, such as two different DAG runs or a manual re-trigger
that reuses the same
+explicit ``job_id``. Durable execution's persisted state is scoped to a single
task instance
+(``dag_id``/``task_id``/``run_id``/``map_index``); a new DAG run has no
durable state to reconnect
+to at all. If you rely on ``force_rerun=False`` with an explicit ``job_id`` so
that re-triggering
+the same logical run doesn't submit a second job, that guarantee is unrelated
to durable execution
+and keeping ``force_rerun=False`` is still the right choice for it.
+
+Durable execution requires Airflow 3.3 or newer, since it relies on the task
state store. On
+earlier Airflow versions the flag is a no-op and the operator always submits a
fresh job on retry,
+exactly as before -- including the pre-existing
``reattach_states``/``Conflict`` behavior, which is
+unchanged. If the task state store is unavailable at runtime, the operator
logs that crash
+recovery is disabled and behaves the same way.
+
+To opt out and always submit a fresh job on retry, set ``durable=False``:
+
+.. code-block:: python
+
+ insert_query_job = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration={
+ "query": {
+ "query": "SELECT 1",
+ "useLegacySql": False,
+ }
+ },
+ durable=False,
+ )
+
+Durable execution applies to the synchronous path. When ``deferrable=True`` is
set, the Triggerer
+already tracks the job across the wait, so deferrable mode takes precedence
and ``durable`` has no
+effect.
+
Validate data
^^^^^^^^^^^^^
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 06b9e905555..7800a5ad44d 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
@@ -27,7 +27,7 @@ from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any, SupportsAbs
-from google.api_core.exceptions import Conflict
+from google.api_core.exceptions import Conflict, NotFound
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
from google.cloud.bigquery import DEFAULT_RETRY, CopyJob, ExtractJob, LoadJob,
QueryJob, Row
from google.cloud.bigquery.routine import Routine
@@ -73,9 +73,29 @@ except ImportError:
return value is not NOTSET
+try:
+ from airflow.sdk import ResumableJobMixin
+except ImportError:
+
+ class ResumableJobMixin: # type: ignore[no-redef]
+ """Airflow <3.3 stub, task_state_store unavailable, always submits
fresh."""
+
+ external_id_key: str = "bigquery_job_id"
+
+ def __init__(self, *, durable: bool = True, **kwargs: Any) -> None:
+ super().__init__(**kwargs)
+ self.durable = durable
+
+ def execute_resumable(self, context):
+ external_id = self.submit_job(context)
+ self.poll_until_complete(external_id, context)
+ return self.get_job_result(external_id, context)
+
+
if TYPE_CHECKING:
from google.api_core.retry import Retry
from google.cloud.bigquery import UnknownJob
+ from pydantic import JsonValue
from airflow.providers.common.compat.sdk import Context
@@ -2260,7 +2280,9 @@ class
BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
return OperatorLineage(outputs=[output_dataset])
-class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOperatorOpenLineageMixin):
+class BigQueryInsertJobOperator(
+ ResumableJobMixin, GoogleCloudBaseOperator,
_BigQueryInsertJobOperatorOpenLineageMixin
+):
"""
Execute a BigQuery job.
@@ -2313,6 +2335,11 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
:param deferrable: Run operator in the deferrable mode
:param poll_interval: (Deferrable mode only) polling period in seconds to
check for the status of job.
Defaults to 4 seconds.
+ :param durable: When ``True`` (the default), the submitted BigQuery job id
is persisted to task
+ state before polling begins. A worker crash on retry reconnects to the
existing job instead of
+ submitting a duplicate, this works regardless of ``force_rerun``,
since the persisted id is
+ read back directly rather than recomputed. Set to ``False`` to always
submit fresh on retry.
+ Requires Airflow 3.3+; no-op on earlier versions.
"""
template_fields: Sequence[str] = (
@@ -2329,6 +2356,7 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
template_fields_renderers = {"configuration": "json",
"configuration.query.query": "sql"}
ui_color = BigQueryUIColors.QUERY.value
operator_extra_links = (BigQueryTableLink(), BigQueryJobDetailLink())
+ external_id_key = "bigquery_job_id"
def __init__(
self,
@@ -2351,6 +2379,7 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
self.configuration = configuration
self.location = location
self.job_id = job_id
+ self._configured_job_id = job_id
self.project_id = project_id
self.gcp_conn_id = gcp_conn_id
self.force_rerun = force_rerun
@@ -2422,6 +2451,55 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
if job.state != "DONE":
raise AirflowException(f"Job failed with state: {job.state}")
+ def _persist_job_links(self, job: BigQueryJob | UnknownJob, context: Any)
-> None:
+ job_types = {
+ LoadJob._JOB_TYPE: ["sourceTable", "destinationTable"],
+ CopyJob._JOB_TYPE: ["sourceTable", "destinationTable"],
+ ExtractJob._JOB_TYPE: ["sourceTable"],
+ QueryJob._JOB_TYPE: ["destinationTable"],
+ }
+
+ if self.project_id:
+ for job_type, tables_prop in job_types.items():
+ job_configuration = job.to_api_repr()["configuration"]
+ if job_type in job_configuration:
+ for table_prop in tables_prop:
+ if table_prop in job_configuration[job_type]:
+ table = job_configuration[job_type][table_prop]
+ persist_kwargs = {
+ "context": context,
+ "project_id": self.project_id,
+ "table_id": table,
+ }
+ if not isinstance(table, str):
+ persist_kwargs["table_id"] = table["tableId"]
+ persist_kwargs["dataset_id"] =
table["datasetId"]
+ persist_kwargs["project_id"] =
table["projectId"]
+ BigQueryTableLink.persist(**persist_kwargs)
+
+ self.job_id = job.job_id
+ if self.project_id:
+ job_id_path = convert_job_id(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ warnings.warn(
+ "BigQueryInsertJobOperator's `job_id_path` XCom is deprecated
and will be removed in a "
+ "future provider release. Use the operator return value or
BigQuery job extra link instead.",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ context["ti"].xcom_push(key="job_id_path", value=job_id_path)
+
+ persist_kwargs = {
+ "context": context,
+ "project_id": self.project_id,
+ "location": self.location,
+ "job_id": self.job_id,
+ }
+ BigQueryJobDetailLink.persist(**persist_kwargs)
+
def _submit_new_job_on_retry(
self,
context: Any,
@@ -2442,6 +2520,34 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
return job
def execute(self, context: Any):
+ self._configured_job_id = self.job_id
+ if not self.deferrable:
+ self.execute_resumable(context)
+ return self.job_id
+
+ self.job_id = self.submit_job(context)
+ job = self._job
+
+ if job.running():
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=BigQueryInsertJobTrigger(
+ 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]
+ poll_interval=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ cancel_on_kill=self.cancel_on_kill,
+ ),
+ method_name="execute_complete",
+ )
+ self.log.info("Current state of job %s is %s", job.job_id, job.state)
+ self._handle_job_error(job)
+ return self.job_id
+
+ def submit_job(self, context: Any) -> str:
+ """Submit the job (or reattach per today's
Conflict/reattach_states/429 rules) and return its id."""
hook = BigQueryHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
@@ -2459,11 +2565,11 @@ class
BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOpera
# To maintain backward compatibility, the try_number is appended to
the job name
# only starting from the 2nd attempt.
ti_try_number = None
- if self.job_id is None and context["ti"].try_number > 2:
+ if self._configured_job_id is None and context["ti"].try_number > 2:
ti_try_number = context["ti"].try_number - 1
self.job_id = hook.generate_job_id(
- job_id=self.job_id,
+ job_id=self._configured_job_id,
dag_id=self.dag_id,
task_id=self.task_id,
logical_date=None,
@@ -2516,77 +2622,49 @@ class
BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOpera
# We are reattaching to a job
self.log.info("Reattaching to existing Job in state %s",
job.state)
- job_types = {
- LoadJob._JOB_TYPE: ["sourceTable", "destinationTable"],
- CopyJob._JOB_TYPE: ["sourceTable", "destinationTable"],
- ExtractJob._JOB_TYPE: ["sourceTable"],
- QueryJob._JOB_TYPE: ["destinationTable"],
- }
-
- if self.project_id:
- for job_type, tables_prop in job_types.items():
- job_configuration = job.to_api_repr()["configuration"]
- if job_type in job_configuration:
- for table_prop in tables_prop:
- if table_prop in job_configuration[job_type]:
- table = job_configuration[job_type][table_prop]
- persist_kwargs = {
- "context": context,
- "project_id": self.project_id,
- "table_id": table,
- }
- if not isinstance(table, str):
- persist_kwargs["table_id"] = table["tableId"]
- persist_kwargs["dataset_id"] =
table["datasetId"]
- persist_kwargs["project_id"] =
table["projectId"]
- BigQueryTableLink.persist(**persist_kwargs)
-
- self.job_id = job.job_id
- if self.project_id:
- job_id_path = convert_job_id(
- job_id=self.job_id,
- project_id=self.project_id,
- location=self.location,
+ self._job = job
+ self._persist_job_links(job, context)
+ return job.job_id
+
+ def get_job_status(self, external_id: JsonValue, context: Any) -> str:
+ """Query the raw job status; a missing job degrades to a not_found
sentinel."""
+ job_id = str(external_id)
+ # On reconnect/already-succeeded, submit_job never runs, so
hook/project_id normally
+ # resolved there must be resolved here too.
+ if self.hook is None:
+ self.hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain
)
- warnings.warn(
- "BigQueryInsertJobOperator's `job_id_path` XCom is deprecated
and will be removed in a "
- "future provider release. Use the operator return value or
BigQuery job extra link instead.",
- AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- context["ti"].xcom_push(key="job_id_path", value=job_id_path)
+ if self.project_id is None:
+ self.project_id = self.hook.project_id
- persist_kwargs = {
- "context": context,
- "project_id": self.project_id,
- "location": self.location,
- "job_id": self.job_id,
- }
- BigQueryJobDetailLink.persist(**persist_kwargs)
+ try:
+ job = self.hook.get_job(project_id=self.project_id,
location=self.location, job_id=job_id)
+ except NotFound:
+ return "not_found"
+ # Reuse the same link bookkeeping as submit_job: this is the only
place a job object is
+ # obtained when reconnecting to (or finding already-succeeded) a
previously stored id.
+ self._job = job
+ self._persist_job_links(job, context)
+ if job.state != "DONE":
+ return job.state
+ return "error" if job.error_result else "success"
- # Wait for the job to complete
- if not self.deferrable:
- job.result(timeout=self.result_timeout, retry=self.result_retry)
- self._handle_job_error(job)
- return self.job_id
+ def is_job_active(self, status: str) -> bool:
+ return status not in ("success", "error", "not_found")
- if job.running():
- self.defer(
- timeout=self.execution_timeout,
- trigger=BigQueryInsertJobTrigger(
- conn_id=self.gcp_conn_id,
- job_id=self.job_id,
- project_id=self.project_id,
- location=self.location or hook.location,
- poll_interval=self.poll_interval,
- impersonation_chain=self.impersonation_chain,
- cancel_on_kill=self.cancel_on_kill,
- ),
- method_name="execute_complete",
- )
- self.log.info("Current state of job %s is %s", job.job_id, job.state)
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == "success"
+
+ def poll_until_complete(self, external_id: JsonValue, context: Any) ->
None:
+ # self._job is set by whichever of submit_job / get_job_status last
obtained it,
+ # never fetched again here, since neither of those calls skip setting
it.
+ job = self._job
+ job.result(timeout=self.result_timeout, retry=self.result_retry)
self._handle_job_error(job)
- return self.job_id
+
+ def get_job_result(self, external_id: JsonValue, context: Any) -> None:
+ return None
def execute_complete(self, context: Context, event: dict[str, Any]) -> str
| None:
"""
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 284342b7f70..28e46235dce 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
@@ -28,7 +28,7 @@ import pandas as pd
import pytest
from google.cloud.bigquery import DEFAULT_RETRY, ScalarQueryParameter, Table
from google.cloud.bigquery.routine import Routine
-from google.cloud.exceptions import Conflict
+from google.cloud.exceptions import Conflict, NotFound
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.openlineage.facet import (
@@ -84,6 +84,8 @@ from airflow.providers.google.cloud.triggers.bigquery import (
from airflow.utils.task_group import TaskGroup
from airflow.utils.timezone import datetime
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+
TASK_ID = "test-bq-generic-operator"
TEST_DATASET = "test-dataset"
TEST_DATASET_LOCATION = "EU"
@@ -1056,6 +1058,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
result = op.execute(context=MagicMock())
assert configuration["labels"] == {"airflow-dag": "adhoc_airflow",
"airflow-task": "insert_query_job"}
@@ -1094,6 +1097,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
with pytest.warns(AirflowProviderDeprecationWarning,
match="`job_id_path` XCom is deprecated"):
@@ -1127,6 +1131,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
result = op.execute(context=MagicMock())
assert configuration["labels"] == {"airflow-dag": "adhoc_airflow",
"airflow-task": "copy_query_job"}
@@ -1167,6 +1172,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
cancel_on_kill=False,
+ durable=False,
)
op.execute(context=MagicMock())
@@ -1210,6 +1216,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
cancel_on_kill=True,
+ durable=False,
)
with pytest.raises(AirflowTaskTimeout):
op.execute(context=MagicMock())
@@ -1244,6 +1251,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
with pytest.raises(AirflowException):
op.execute(context=MagicMock())
@@ -1280,6 +1288,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
reattach_states={"PENDING", "RUNNING"},
+ durable=False,
)
result = op.execute(context=MagicMock())
@@ -1326,6 +1335,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
reattach_states={"PENDING"},
+ durable=False,
)
with pytest.raises(AirflowException):
# Not possible to reattach to any state if job is already DONE
@@ -1359,6 +1369,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
force_rerun=True,
+ durable=False,
)
result = op.execute(context=MagicMock())
@@ -1404,6 +1415,7 @@ class TestBigQueryInsertJobOperator:
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
reattach_states={"PENDING"},
+ durable=False,
)
# No force rerun
with pytest.raises(AirflowException):
@@ -1696,6 +1708,7 @@ class TestBigQueryInsertJobOperator:
job_id=None, # job_id must be None to trigger the try_number logic
project_id=TEST_GCP_PROJECT_ID,
force_rerun=False,
+ durable=False,
)
op.execute(context=context)
@@ -1754,6 +1767,7 @@ class TestBigQueryInsertJobOperator:
job_id=None,
project_id=TEST_GCP_PROJECT_ID,
force_rerun=False,
+ durable=False,
)
result = op.execute(context=context)
@@ -1811,6 +1825,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
result = op.execute(context=MagicMock())
@@ -1858,6 +1873,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
mock_hook.return_value.generate_job_id.return_value = "1234"
mock_hook.return_value.get_client.return_value.get_job.side_effect =
RuntimeError()
@@ -1944,6 +1960,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
op.execute(context=MagicMock())
assert configuration["labels"] == {
@@ -1977,6 +1994,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
op.execute(context=MagicMock())
assert configuration["labels"] is None
@@ -2005,6 +2023,7 @@ class TestBigQueryInsertJobOperator:
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
+ durable=False,
)
op.execute(context=MagicMock())
@@ -2261,6 +2280,277 @@ class TestBigQueryInsertJobOperator:
op._handle_job_error(job_empty_error)
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
+class TestBigQueryInsertJobOperatorDurable:
+ @staticmethod
+ def _context(task_store=None):
+ ctx: dict = {"ti": MagicMock(stats_tags={}, try_number=1)}
+ if task_store is not None:
+ ctx["task_state_store"] = task_store
+ return ctx
+
+ @staticmethod
+ def _make_operator(**kwargs):
+ configuration = {"query": {"query": "SELECT 1", "useLegacySql": False}}
+ kwargs.setdefault("project_id", TEST_GCP_PROJECT_ID)
+ return BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ **kwargs,
+ )
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_persists_job_id_to_task_state_store_on_fresh_submit(self,
mock_hook):
+ op = self._make_operator()
+ mock_hook.return_value.generate_job_id.return_value =
"generated-job-id"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="generated-job-id", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_called_once()
+ task_store.set.assert_called_once_with("bigquery_job_id",
"generated-job-id")
+ assert result == "generated-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_reconnects_to_running_job_without_resubmitting(self, mock_hook):
+ op = self._make_operator()
+ job = MagicMock(state="RUNNING", job_id="stored-job-id",
error_result=False)
+ job.result.side_effect = lambda **_: setattr(job, "state", "DONE")
+ mock_hook.return_value.get_job.return_value = job
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "stored-job-id"
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_not_called()
+ mock_hook.return_value.generate_job_id.assert_not_called()
+ task_store.set.assert_not_called()
+ job.result.assert_called_once()
+ assert result == "stored-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_reconnect_wins_over_force_rerun_default(self, mock_hook):
+ """Headline behavior: force_rerun=True (the operator default) no
longer defeats
+ reattach, since a retry reconnects via the persisted id instead of
recomputing a
+ fresh uuid-based id that would never match the original."""
+ op = self._make_operator() # force_rerun defaults to True
+ job = MagicMock(state="RUNNING", job_id="stored-job-id",
error_result=False)
+ job.result.side_effect = lambda **_: setattr(job, "state", "DONE")
+ mock_hook.return_value.get_job.return_value = job
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "stored-job-id"
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.generate_job_id.assert_not_called()
+ mock_hook.return_value.insert_job.assert_not_called()
+ assert result == "stored-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_already_succeeded_returns_result_without_polling(self, mock_hook):
+ op = self._make_operator()
+ job = MagicMock(state="DONE", job_id="stored-job-id",
error_result=False)
+ mock_hook.return_value.get_job.return_value = job
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "stored-job-id"
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_not_called()
+ job.result.assert_not_called()
+ assert result == "stored-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_resubmits_when_stored_job_in_terminal_failure(self, mock_hook):
+ op = self._make_operator()
+ failed_job = MagicMock(state="DONE", job_id="stored-job-id",
error_result="boom")
+ mock_hook.return_value.get_job.return_value = failed_job
+ mock_hook.return_value.generate_job_id.return_value = "new-job-id"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="new-job-id", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "stored-job-id"
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_called_once()
+ task_store.set.assert_called_once_with("bigquery_job_id", "new-job-id")
+ assert result == "new-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def
test_resubmit_computes_id_from_configured_job_id_not_stale_stored_id(self,
mock_hook):
+ """Regression test: get_job_status persists the stored (failed) job's
id onto
+ self.job_id before terminal_resubmit falls through to submit_job.
submit_job must
+ still compute the new id from the user's originally configured job_id,
not from
+ that stale value -- otherwise ids grow an extra suffix on every failed
retry."""
+ op = self._make_operator(job_id="my_job", force_rerun=False)
+ failed_job = MagicMock(state="DONE", job_id="my_job_oldsuffix",
error_result="boom")
+ mock_hook.return_value.get_job.return_value = failed_job
+ mock_hook.return_value.generate_job_id.return_value =
"my_job_newsuffix"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="my_job_newsuffix", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "my_job_oldsuffix"
+
+ op.execute(self._context(task_store))
+
+ mock_hook.return_value.generate_job_id.assert_called_once_with(
+ job_id="my_job",
+ dag_id=op.dag_id,
+ task_id=op.task_id,
+ logical_date=None,
+ configuration=op.configuration,
+
run_after=mock_hook.return_value.get_run_after_or_logical_date.return_value,
+ force_rerun=False,
+ try_number=None,
+ )
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_configured_job_id_uses_rendered_template_not_raw_jinja(self,
mock_hook):
+ """
+ Test to validate that _configured_job_id must be snapshotted after
template
+ rendering. job_id is a template field, rendering happens between
__init__ and
+ execute(), so capturing it in __init__ would feed the raw, unrendered
Jinja
+ string into generate_job_id instead of the rendered value.
+ """
+ op = self._make_operator(job_id="report_{{ ds_nodash }}",
force_rerun=False)
+ op.job_id = "report_20260101" # simulate template rendering having
already run
+ mock_hook.return_value.generate_job_id.return_value = "report_20260101"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="report_20260101", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+
+ op.execute(self._context(task_store))
+
+ mock_hook.return_value.generate_job_id.assert_called_once_with(
+ job_id="report_20260101",
+ dag_id=op.dag_id,
+ task_id=op.task_id,
+ logical_date=None,
+ configuration=op.configuration,
+
run_after=mock_hook.return_value.get_run_after_or_logical_date.return_value,
+ force_rerun=False,
+ try_number=None,
+ )
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_resubmits_when_stored_job_not_found(self, mock_hook):
+ op = self._make_operator()
+ mock_hook.return_value.get_job.side_effect = NotFound("gone")
+ mock_hook.return_value.generate_job_id.return_value = "new-job-id"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="new-job-id", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = "stored-job-id"
+
+ result = op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_called_once()
+ task_store.set.assert_called_once_with("bigquery_job_id", "new-job-id")
+ assert result == "new-job-id"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_durable_false_never_touches_task_state_store(self, mock_hook):
+ op = self._make_operator(durable=False)
+ mock_hook.return_value.generate_job_id.return_value =
"generated-job-id"
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", job_id="generated-job-id", error_result=False
+ )
+ task_store = MagicMock(spec_set=["get", "set"])
+
+ op.execute(self._context(task_store))
+
+ mock_hook.return_value.insert_job.assert_called_once()
+ task_store.get.assert_not_called()
+ task_store.set.assert_not_called()
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_deferrable_unaffected_by_durable(self, mock_hook):
+ op = self._make_operator(deferrable=True)
+ mock_hook.return_value.generate_job_id.return_value =
"generated-job-id"
+ running_job = MagicMock(state="RUNNING", job_id="generated-job-id",
error_result=False)
+ running_job.running.return_value = True
+ mock_hook.return_value.insert_job.return_value = running_job
+ task_store = MagicMock(spec_set=["get", "set"])
+
+ with pytest.raises(TaskDeferred) as exc:
+ op.execute(self._context(task_store))
+
+ assert isinstance(exc.value.trigger, BigQueryInsertJobTrigger)
+ task_store.get.assert_not_called()
+ task_store.set.assert_not_called()
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_get_job_status_success_when_done_without_error(self, mock_hook):
+ op = self._make_operator()
+ op.hook = mock_hook.return_value
+ mock_hook.return_value.get_job.return_value = MagicMock(state="DONE",
error_result=False)
+
+ assert op.get_job_status("some-id", context=self._context()) ==
"success"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_get_job_status_error_when_done_with_error(self, mock_hook):
+ op = self._make_operator()
+ op.hook = mock_hook.return_value
+ mock_hook.return_value.get_job.return_value = MagicMock(state="DONE",
error_result="boom")
+
+ assert op.get_job_status("some-id", context=self._context()) == "error"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_get_job_status_returns_raw_state_when_not_done(self, mock_hook):
+ op = self._make_operator()
+ op.hook = mock_hook.return_value
+ mock_hook.return_value.get_job.return_value =
MagicMock(state="RUNNING", error_result=False)
+
+ assert op.get_job_status("some-id", context=self._context()) ==
"RUNNING"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_get_job_status_not_found(self, mock_hook):
+ op = self._make_operator()
+ op.hook = mock_hook.return_value
+ mock_hook.return_value.get_job.side_effect = NotFound("gone")
+
+ assert op.get_job_status("some-id", context=self._context()) ==
"not_found"
+
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_get_job_status_resolves_hook_and_project_id_when_unset(self,
mock_hook):
+ """Regression test: on reconnect, submit_job never runs, so
get_job_status must
+ resolve self.hook / self.project_id itself rather than assuming
they're set."""
+ op = self._make_operator(project_id=None)
+ assert op.hook is None
+ mock_hook.return_value.project_id = "resolved-project"
+ mock_hook.return_value.get_job.return_value = MagicMock(state="DONE",
error_result=False)
+
+ status = op.get_job_status("some-id", context=self._context())
+
+ assert status == "success"
+ assert op.hook is not None
+ assert op.project_id == "resolved-project"
+
+ def test_is_job_active_and_is_job_succeeded_predicates(self):
+ op = self._make_operator()
+
+ for status in ("RUNNING", "PENDING"):
+ assert op.is_job_active(status) is True
+ for status in ("success", "error", "not_found"):
+ assert op.is_job_active(status) is False
+
+ assert op.is_job_succeeded("success") is True
+ assert op.is_job_succeeded("RUNNING") is False
+
+
class TestBigQueryIntervalCheckOperator:
def test_bigquery_interval_check_operator_execute_complete(self):
"""Asserts that logging occurs as expected"""
