amoghrajesh opened a new pull request, #69542:
URL: https://github.com/apache/airflow/pull/69542
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
### 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:
```shell
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:
```python
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:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/31e5bb51-9a24-465d-8a39-9135fc0aefbf";
/>
This job started running:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/e2c162b0-9184-41b6-b558-97947b9ead48";
/>
When worker comes back up, a new bigquery job gets submitted:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/6ae58edf-0c0b-44d6-bc2f-613aedca3232";
/>
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/e6879318-6cfd-42dc-9b44-b233e241b010";
/>
Hence lot of compute wastage
### Before changes but with `force_rerun=False`
Try 1
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/15798fb3-4428-46f1-95ba-eb0e27fe4779";
/>
Worker killed and brought back up
Try 2:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/efc121ae-0fee-4afc-aba0-f1f6c7ed1f8d";
/>
Try 3:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/22269680-5542-4895-a199-f4401a283f79";
/>
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:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/762843b2-570b-40a9-ab35-e8d2aaf7ba75";
/>
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/bf9567f0-5005-4bc9-b537-82654f1397a8";
/>
Worker came back up, try 2:
<img width="1725" height="977" alt="image"
src="https://github.com/user-attachments/assets/ac2758e9-6f7e-430c-a8aa-aee4e7f92e99";
/>
### 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](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
