amoghrajesh opened a new pull request, #76918:
URL: https://github.com/apache/airflow/pull/76918
<!-- 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:
https://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".
-->
- [x] Yes - claude sonnet
---
### What
`DbtCloudRunJobOperator` triggers a dbt Cloud job run and waits for it. If
the worker dies while waiting (OOM, pod eviction, spot instance loss), the run
keeps going in dbt Cloud. The next Airflow retry has no record of that run, so
it triggers a new one. You pay for two runs, and both write to the same tables
at the same time.
### Current behaviour
Each try calls `trigger_job_run`. A retry after a worker crash starts a
second dbt Cloud run while the first one is still `RUNNING`.
### Proposed change
`DbtCloudRunJobOperator` now uses `ResumableJobMixin`. Before waiting, it
stores the dbt Cloud run id in task state under `dbt_cloud_job_run_id`. On a
retry, it reads that run's status:
- `QUEUED`, `STARTING` or `RUNNING`: waits on the same run.
- `SUCCESS`: returns the run id without triggering anything.
- `ERROR`, `CANCELLED`, or the run is gone (404): triggers a new run as
before.
This works in both modes. In deferrable mode, a retry defers on the stored
run instead of a new one.
`durable=False` turns it off for a task. It applies only when
`wait_for_termination=True`. With `wait_for_termination=False` the operator
does not wait, so there is nothing to reconnect to.
### Changes of Note
- **dbt's `retry_from_failure` still works:** when the stored run failed,
the new run goes through `_get_or_trigger_run`. So `retry_from_failure=True`
still calls dbt Cloud's rerun endpoint, and `reuse_existing_run` still applies.
- **The waiting moved:** both waiting branches of `execute` (sync and
deferrable) moved into `poll_until_complete`. The deferral code moved into
`_defer_until_complete` unchanged.
- **XComs:** `job_run_url` and `job_run_id` are still pushed before waiting.
On a reconnect or an already finished run, they point to the stored run.
- **Run not found:** `HttpHook.check_response` raises again without the
status code. So `_is_not_found` reads the 404 from the chained `HTTPError`.
- **`durable` is a named parameter** so that `default_args` reaches the
`common.compat` stand-in below Airflow 3.3. This matches the Databricks
operators.
- **common.compat:** the floor stays as is, with the `# use next version`
marker, because `ResumableJobMixin` comes from `common.compat`.
### User implications / backcompat
- **Airflow 3.3 and later:** durable is on by default. A retry after a crash
now reconnects to the running dbt Cloud run instead of starting a new one. Set
`durable=False` to keep the old behaviour.
- **Below Airflow 3.3:** no change. The `common.compat` stand-in always
triggers a new run.
- **Clearing the task:** this discards the stored run id, so a cleared task
always triggers a new run.
### Testing
**Manual test against a real dbt Cloud job**
Setup:
- A dbt Cloud trial account, with Databricks as the warehouse and Unity
Catalog.
- A job running a model that takes a few minutes.
- Airflow 3.5 (`main`) in Breeze, with a `dbt_cloud_default` connection.
Dag:
```python
from datetime import datetime, timedelta
from airflow.providers.dbt.cloud.operators.dbt import DbtCloudRunJobOperator
from airflow.sdk import DAG
with DAG(
dag_id="dbt_durable_demo",
start_date=datetime(2026, 1, 1),
schedule=None,
catchup=False,
default_args={"retries": 1, "retry_delay": timedelta(seconds=0)},
):
DbtCloudRunJobOperator(
task_id="run_slow_model",
dbt_cloud_conn_id="dbt_cloud_default",
account_id=<account_id>,
job_id=<job_id>,
check_interval=10,
timeout=3600,
)
```
Steps:
1. Trigger the Dag, and wait until try 1 logs that it is waiting on the dbt
Cloud run.
2. Kill the task process (`airflow worker -- <ti-uuid>`) with `kill -9`.
This simulates a worker death: `on_kill` does not run, so the dbt Cloud run
keeps going.
3. Let try 2 run, then compare which dbt Cloud run each try waited on.
**Before (`main`):** try 2 triggered a second run (`70506184179958`) while
try 1's run (`70506184179955`) was still `RUNNING`.
<img width="3200" height="2000" alt="before_try1"
src="https://github.com/user-attachments/assets/ceb532d7-a589-4915-b620-639398b30207";
/>
<img width="3200" height="2000" alt="before_try2"
src="https://github.com/user-attachments/assets/9dfba0e3-74c6-46ae-a91d-167403df5b94";
/>
**After, run still going:** try 2 logged `Reconnecting to existing job ...
70506184179975 status=RUNNING`, waited on the same run, and succeeded. No
second run.
<img width="3200" height="2000" alt="after_reconnect_try1"
src="https://github.com/user-attachments/assets/22f87669-f93c-4346-9d6d-ba41116fbbc7";
/>
<img width="3200" height="2000" alt="after_reconnect_try2"
src="https://github.com/user-attachments/assets/e5470d1f-03a5-4896-a5b3-c30b297949e0";
/>
**After, run already finished:** with `retry_delay=10s`, the run finished
before try 2 started. Try 2 logged `Job already completed successfully,
skipping resubmission` and triggered nothing.
<img width="3200" height="2000" alt="after_finished_try1"
src="https://github.com/user-attachments/assets/0b74a11d-b681-414e-80e6-b88ee4e286b5";
/>
<img width="3200" height="2000" alt="after_finished_try2"
src="https://github.com/user-attachments/assets/24702603-8bc1-4dc7-aba9-95b58bb39e21";
/>
---
* 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]
