Enable repair of failed Databricks Workflow tasks from Airflow 3 UI - #69998
Open
PrakshiGoyal10 wants to merge 2 commits into
Open
Enable repair of failed Databricks Workflow tasks from Airflow 3 UI#69998PrakshiGoyal10 wants to merge 2 commits into
PrakshiGoyal10 wants to merge 2 commits into
Conversation
|
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| if not task_keys: | ||
| log.info("No failed Databricks tasks to repair for run %s", databricks_run_id) | ||
| return RedirectResponse(return_url, status_code=303) |
| raise HTTPException(status_code=502, detail=f"Databricks repair request failed: {e}") from e | ||
|
|
||
| _clear_repaired_and_downstream(dag_id, run_id, task_keys, log) | ||
| return RedirectResponse(return_url, status_code=303) |
Contributor
|
I rechecked the current head and think the repair path needs a security/compatibility redesign before it is safe to validate against a real workspace:
The PR body should also include Drafted-by: Codex (GPT-5) |
Contributor
|
@PrakshiGoyal10 , Some CI signals are failing. PTAL and fix those. Thanks! |
On Airflow 3 the "Repair a single task" and "Repair All Failed Tasks" links on DatabricksWorkflowTaskGroup tasks were gated off, so on-call users had to leave Airflow and repair failed runs in the Databricks UI. This restores repair-from-Airflow on Airflow 3: - Register a FastAPI app on the API server that performs the repair, authorized with Dag-run edit access and authenticated via either the UI's bearer token (XHR) or its _token cookie (plain navigation), so a repair link clicked in the browser is authorized. - Resolve the set of failed tasks from the live Databricks run state rather than Airflow's metadata DB (new hook helper get_run_failed_task_keys), and call repair_run with rerun_dependent_tasks so downstream tasks resume too. - Clear the repaired task instances and their downstream instances, mapping Databricks task keys to Airflow task_ids via the same md5(dag_id__task_id) scheme used by the operators, which works on the serialized Dag. - Remove the Airflow 3 gates on the plugin registration and the per-operator extra links so the repair buttons render again. Verified end-to-end against a real Databricks workflow on serverless and job clusters: a failed task plus its upstream_failed downstream both return to success after clicking repair.
…irect The repair action mutated state on a cookie-authenticated GET and trusted the Databricks connection, run id, and task keys from the request, so a crafted repair link could repair an arbitrary run and CodeQL flagged the redirect as an open redirect. The endpoint also required an auth helper unavailable on 3.0.6. A GET now renders a read-only confirmation page and only a same-site POST repairs and clears tasks, which SameSite=Lax on the auth cookie protects from cross-site forgery. The connection and Databricks run id are derived server-side from the launch task's trusted WorkflowRunMetadata XCom, the redirect target is a same-site relative path, Databricks failures return a generic 502, and the backend is gated to Airflow 3.1+ where its authentication resolver exists.
PrakshiGoyal10
force-pushed
the
databricks-repair-from-airflow-3
branch
from
August 5, 2026 13:44
5f287ab to
5e536ea
Compare
|
|
||
| if not task_keys: | ||
| log.info("No failed Databricks tasks to repair for run %s", metadata.run_id) | ||
| return RedirectResponse(return_url, status_code=303) |
| _clear_repaired_and_downstream(dag, run_id, repaired_task_ids, session, log) | ||
| session.commit() | ||
|
|
||
| return RedirectResponse(return_url, status_code=303) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On Airflow 3 the "Repair a single task" and "Repair All Failed Tasks" links on DatabricksWorkflowTaskGroup tasks were gated off, so on-call users had to leave Airflow and repair failed runs in the Databricks UI.
This restores repair-from-Airflow on Airflow 3 (3.1+), served by a FastAPI sub-application on the API server since Flask-AppBuilder was dropped:
GETrenders a read-only confirmation page; only a same-sitePOSTperforms the repair and clears tasks. Because the API server's auth cookie isSameSite=Lax, it is not sent on a cross-sitePOST, so moving the mutation offGETis what protects the action from cross-site request forgery — matching how Airflow 3's own UI mutations are protected.task_id, and for a single-task repair the targettask_id). The Databricks connection and run id are derived server-side from the launch task's trustedWorkflowRunMetadataXCom, never from the request, so a crafted link cannot point the repair at an arbitrary connection or Databricks run.get_run_failed_task_keys), andrepair_runis called withrerun_dependent_tasksso downstream tasks resume too.task_ids via the samemd5(dag_id__task_id)scheme the operators use, resolved from the serialized Dag.502without reflecting the upstream exception text, and confirmation-page identifiers are HTML-escaped.Verified end-to-end against a real Databricks workflow on serverless and job clusters: a failed task plus its
upstream_faileddownstream both return to success after confirming the repair. Unit tests cover the GET-does-not-mutate behaviour, confirmation-page escaping, server-side identifier resolution on the POST, and the generic 502.related: #52280
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 4.8) following the guidelines
Important
🛠️ Maintainer triage note for @PrakshiGoyal10 · by
@potiuk· 2026-07-28 16:20 UTCHelpful heads-up from the maintainers — please address before this PR can be reviewed:
The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.
See the Pull Request quality criteria for how to fix each item. There is no rush.
Note: your branch is 294 commits behind
main— please rebase and push again to get up-to-date CI results.Automated triage — may be imperfect; a maintainer takes the next look. We use this two-stage triage process so maintainers' limited time goes to the conversation with you.