Adding durable execution to BigQueryInsertJobOperator

[amoghrajesh]

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

What

BigQueryInsertJobOperator submits a BigQuery job and polls it on the worker. If the worker crashes mid-poll and the task is retried, there was no reliable way to reconnect to the still running job & the retry submits a brand new one, orphaning the original (which keeps running and costing money) instead of reusing it.

Current behaviour

The operator already has a partial, fragile reattach mechanism: it optimistically resubmits with a recomputed job id and relies on BigQuery rejecting the duplicate (Conflict/HTTP 409) to detect and reattach to an existing job. This only works when force_rerun=False (or a fixed job_id) and reattach_states is populated but force_rerun=True is the operator's own default, which makes generate_job_id return a fresh random id on every attempt, so by default the retry's recomputed id never matches and reattach never triggers. A crash-and-retry under default settings always submits a duplicate job and orphans the first one.

Proposed change

Adopt ResumableJobMixin (from airflow.sdk, AIP-103) so the operator persists the actual submitted job id to task state and reads it back on retry, rather than trying to recompute an identical id. This makes reattach work regardless of force_rerun:

• submit_job: extracted from the old execute(): generates/submits the job (still honoring the existing Conflict/reattach_states/429-retry rules on the fresh-submit path) and persists the job id.
• get_job_status: fetches the job by id (get_job), mapping DONE to "success"/"error" based on error_result, with NotFound degrading to a "not_found" sentinel so an expired/unknown id triggers a fresh submit instead of crashing.
• is_job_active / is_job_succeeded: simple predicates over that status.
• poll_until_complete / get_job_result: wait on the job and return its id, reusing the existing link/XCom bookkeeping (_persist_job_links) so behavior is identical whether this attempt submitted the job or reconnected to it.
• execute() now checks deferrable first (unchanged) and routes the synchronous path through execute_resumable; submit_job is shared as-is by the deferrable branch since its logic doesn't depend on deferrable at all.

Changes of Note

• self._configured_job_id: get_job_status (called during a durable resubmit's initial status check) mutates self.job_id to the old failed job's id via _persist_job_links, before submit_job runs again. If submit_job used self.job_id as input to generate_job_id, it would compute a corrupted id from the stale value. Fixed by capturing the user's original job_id input once in __init__ as self._configured_job_id (never mutated) and using that instead.
• poll_until_complete must return the job id: on the mixin's reconnect path, execute_resumable returns poll_until_complete's result directly, skipping get_job_result entirely, an easy bug to reintroduce (also hit independently in the Redshift port).
• get_job_status self-resolves self.hook/self.project_id: these are normally set by submit_job, which is skipped entirely on the reconnect path, so get_job_status resolves them itself when unset.
• The existing reattach_states/Conflict/429/DONE-can't-reattach special cases are preserved unchanged on the fresh-submit path (durable=False, or pre-Airflow-3.3).

User implications / backcompat

New durable parameter, defaulting to True on Airflow 3.3+. On older Airflow the ResumableJobMixin import is stubbed and the operator falls back to today's regenerate-id + Conflict-reattach behavior — no behavior change for those versions. force_rerun's default is unchanged; durable execution makes its non-determinism irrelevant to crash-recovery specifically, but it's still relevant to idempotency across separate DAG runs (task state is scoped per task instance), so force_rerun is not being deprecated.

Testing

Create a connection with your GCP credentials, like so:

airflow connections add google_cloud_default   --conn-type google_cloud_platform   --conn-extra '{"key_path": "/opt/airflow/dev/sigma-night-269811-6107ef06a57f.json", "project": "sigma-night-269811", "scope": "https://www.googleapis.com/auth/cloud-platform"}'

DAG:

from __future__ import annotations

from datetime import timedelta

import pendulum

from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from airflow.sdk import DAG

with DAG(
    dag_id="bigquery_insert_job_demo",
    start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
    schedule=None,
    catchup=False,
):
    # A slow-ish query so there's a real window to kill the worker mid-poll.
    # Plain aggregation over a cross join finishes too fast -- BigQuery parallelizes it
    # across many slots. A JS UDF forces genuinely slow, hard-to-parallelize per-row work,
    # so wall-clock time is controlled by ROWS * LOOP_ITERS, not just row count.
    # Tune ROWS (currently 3000 -> 9M rows) or the JS loop bound (1000) up/down to hit
    # whatever poll window you want.
    run_job = BigQueryInsertJobOperator(
        task_id="run_job",
        configuration={
            "query": {
                "query": (
                    "CREATE TEMP FUNCTION slow_calc(x INT64) RETURNS FLOAT64 LANGUAGE js AS '''\n"
                    "  var acc = 0;\n"
                    "  for (var i = 0; i < 1000; i++) { acc += Math.sqrt(x + i); }\n"
                    "  return acc;\n"
                    "''';\n"
                    "SELECT COUNT(*) FROM UNNEST(GENERATE_ARRAY(1, 3000)) AS a "
                    "CROSS JOIN UNNEST(GENERATE_ARRAY(1, 3000)) AS b "
                    "WHERE MOD(CAST(slow_calc(a * b) AS INT64), 7) = 0"
                ),
                "useLegacySql": False,
            }
        },
        location="US",
        # durable=True,
        retries=2,
        retry_delay=timedelta(seconds=5),
    )

Before changes

First try where worker was killed:

This job started running:

When worker comes back up, a new bigquery job gets submitted:

Hence lot of compute wastage

Before changes but with force_rerun=False

Try 1

Worker killed and brought back up

Try 2:

Try 3:

You can see that:

• Try 1's job is still running in BigQuery, nobody ever read its result, and Try 3 submitted a second, independent job for the same task instance - one orphaned job, one "real" one. This is precisely the force_rerun=False-without-reattach_states fragility, and it's why the mixin's persisted-id-and-reconnect approach (on bigquery-durable-execution) is needed, it reconnects on Try 2 directly instead of depending on reattach_states being populated correctly.

After Changes:

Try 1:

Worker came back up, try 2:

What's next

Identified several other BigQuery operators with the same crash-vulnerability shape (BigQueryCheckOperator, BigQueryValueCheckOperator, BigQueryIntervalCheckOperator, BigQueryColumnCheckOperator, BigQueryTableCheckOperator), will pick them and work one by one.

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[amoghrajesh]

Thanks for the reviews, @kaxil. Merging this one in!

[shahar1]

Thanks for the reviews, @kaxil. Merging this one in!

Didn't manage to review it, but looks great :)