Add durable execution to DatabricksRunNowOperator

[amoghrajesh]

Why we are doing this

DatabricksRunNowOperator triggers a run of an existing Databricks job, gets back a run_id, and then polls synchronously on the worker until the run finishes. That run_id lives only in the worker process. If the worker crashes or is preempted mid-poll (eviction, OOM, a deploy, a spot reclaim), Airflow retries the task in a fresh process with no memory of the run_id, so it calls run-now again, triggering a second run of the same job. The original run keeps executing on Databricks, orphaned, while the retry runs a duplicate.

This is the same crash-safety gap closed for DatabricksSubmitRunOperator in #68974, applied to the run-now path. For long-running jobs it means paying twice for the same work, and the only existing mitigation (cancel_previous_runs=True) cancels rather than reconnects.

Benefits this will bring in

• A worker crash and retry reconnects to the run already executing on Databricks instead of triggering a duplicate, so you do not pay twice for the same job.
• If the prior run already finished successfully, the retry returns immediately without triggering a new run.
• If the stored run no longer exists because its history expired, the retry degrades to a fresh run instead of failing the task.
• Works on the existing synchronous operator with no Triggerer required, and is enabled by default on Airflow 3.3+ with no Dag changes.

Approach

The operator now builds on the AIP-103 task state store. On the first run it persists the Databricks run id to the task state store before polling begins. On a retry it reads that id back and inspects the run's current state: still running means reconnect and keep polling, already succeeded means return immediately, and terminally failed (or a run whose history has expired) means trigger a fresh run. Reconnect keys off the persisted run_id, so it works the same whether the job was identified by job_id or job_name. Deferrable mode is unchanged and takes precedence when enabled.

Changes of note

• The run-state decoding helpers (get_job_status / is_job_active / is_job_succeeded) are intentionally duplicated from DatabricksSubmitRunOperator rather than extracted into a shared base, to keep each operator self-contained and this PR scoped to RunNow. They are small and each operator has its own tests covering them.
• The expired-run degradation reuses the typed DatabricksApiError (carrying http_status_code) introduced alongside the SubmitRun work, so it is a structured 404 check rather than message-string matching.
• The run-now payload build is split into a side-effect-free _build_run_now_payload() (merge, validate, resolve job_name to job_id, inject params) and the cancel_previous_runs step. On reconnect the payload is rebuilt via the side-effect-free path, so repair_run has the resolved job_id it needs and works across the reconnect boundary for both job_id and job_name, without re-firing cancel_previous_runs against the run being reconnected to.

Backcompat

• Fully backward compatible. On a clean run the observable trigger-and-poll behavior is unchanged; the only addition is that the run id is also written to the task state store before polling.
• The task state store is an Airflow 3.3+ capability. On Airflow 2.x / pre-3.3 the operator degrades to exactly the old behavior (always triggers a fresh run on retry) through a compatibility shim, so the provider keeps working on older Airflow.
• If the task state store is unavailable at runtime, the operator logs that crash recovery is disabled and behaves as before.
• One new optional flag is added; no existing parameters change and no migration is needed.

How to opt out

Set durable=False on the operator:

DatabricksRunNowOperator(..., durable=False)

This restores the previous behavior: always trigger a fresh run on retry and never touch the task state store. It can also be set through default_args to opt out across a whole Dag or deployment.

Testing

DAG used:

from datetime import timedelta

import pendulum

from airflow.providers.databricks.operators.databricks import DatabricksRunNowOperator
from airflow.sdk import DAG

JOB_ID = 480676095338270

with DAG(
    dag_id="databricks_runnow_repro",
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1, tz="UTC"),
    catchup=False,
):
    DatabricksRunNowOperator(
        task_id="run_now_sleepy",
        databricks_conn_id="databricks_default",
        job_id=JOB_ID,
        deferrable=False,
        do_xcom_push=True,
        retries=1,
        retry_delay=timedelta(seconds=10),
    )

Before my changes

Kill mid run:

Worker comes back up:

Duplicate submissions:

After my changes

Killed mid run:

Job on dbx:

Restored worker:

No duplicate job:

Killed the job from airflow:

---

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

• Yes (please specify the tool below)

---

• 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 review @kaxil. Merging this one, since i've handled your comments, let me know if there are more comments and I can handle in follow ups.