You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org> on 2026/07/17 06:54:26 UTC
[PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
PrakshiGoyal10 opened a new pull request, #69998:
URL: https://github.com/apache/airflow/pull/69998
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_id`s 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.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Opus 4.8)
Generated-by: Claude Code (Opus 4.8) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-6037730498
All review feedback is now addressed in `999f9c054d` (rebased onto latest `main`). Auth uses `requires_access_dag`, the plugin is gated to Airflow 3.1.1+, clearing reuses the same `find_relevant_relatives` + `dag.clear` path as the core route, and both open questions have replies inline (the `:742` clearing thread and the `:793` `DagRun` thread).
@ashb, could you take another look when you have a chance?
One CI note: **Integration: providers redis** was cancelled after timing out in the `apache/infrastructure-actions` download step, before any tests ran. It doesn't touch the Databricks provider, but it marks the Tests (AMD) run as cancelled. Could someone with permissions re-run that job? Everything else passes (143/143).
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4228580941
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +385,40 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
Review Comment:
I think this is fine, but I also don't think the logic holds.
Given the plugin includes a custom app, couldn't the link could be hard-coded to `/databricks/workflow/repair/...` with task_id etc pre-filled"?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5215552250
cc @moomindani for review
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3988918336
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
Review Comment:
Thanks for the precise catch — you're right, it's 3.1.1, not 3.1.0. Gated the backend, its extra links, and the test at `AIRFLOW_V_3_1_1_PLUS` (`>= (3, 1, 1)`), following the patch-level precedent in the amazon/fab providers, so a browser-clicked link is only offered where `get_user` actually honours the `_token` cookie. Also confirmed relaxing to 3.0+ is closed off for the same reason (3.0.6's `get_user` is oauth-only). Pushed in 549f946.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3797613285
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,266 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ from urllib.parse import urlencode
+
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from sqlalchemy import select
+
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ request: Request,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Same-site relative action; SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = f"{request.url.path}?{request.url.query}"
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
Review Comment:
This is being flagged by security scanner as risk. I haven't checked the code yet but please look into this
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "moomindani (via GitHub)" <gi...@apache.org>.
moomindani commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5343140938
Thanks for checking this independently — your reading matches mine on both counts, including the detail that the Airflow 3 branches in `get_link` are unreachable rather than merely wrong.
On who does the work: that is @PrakshiGoyal10's call as the author, not mine. I have no objection either way, and offering a patch on this PR rather than opening a competing one is the right way round.
One technical note before anyone writes code, because there may be a cleaner fix than persisting the URLs. Plugin links are not unconditionally replaced by `XComOperatorLink`: `_deserialize_operator` walks `plugins_manager.get_operator_extra_links()` and matches each link's `operators` against the serialized `task_type` / `_task_module`, and plugin links then take precedence over the xcom-key form (`serialization/serialized_objects.py:1088-1101` and the `op_predefined_extra_links.update(op_extra_links_from_plugin)` line). All three links in this plugin currently declare no `operators`, which is why they get replaced. Declaring them on `WorkflowJobRepairAllFailedLink` and `WorkflowJobRepairSingleTaskLink` would keep the real objects on the API server and let `_build_repair_url` run at request time — no URL duplicated into XCom, no `base_url` frozen at execute time, and the branches this PR already added would actually be exercised.
Persisting into XCom the way the job-run link does also works and has the virtue of matching existing precedent, so I am not arguing hard for one over the other. Worth deciding deliberately rather than by default, and either way the `_get_launch_task_id_v3` fix and the `SerializedDagModel.get_dag` round-trip test are needed — that test is the part that stops this class of break from returning silently, so I would treat it as non-optional.
On scope: the two blockers plus that test are what gate merge from my side. The `start_time` fragility is a genuine crash risk in the same flow and is welcome in the same pass; the unmapped-key warning on legacy repair-all and the `base_url` form-action / redirect notes are lower priority and can equally be follow-ups if this PR is already long in the tooth.
---
Drafted-by: Claude Code (Opus 5); reviewed by @moomindani before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3949124386
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
Review Comment:
There's no `XComModel.get_one()` — it's mentioned in a docstring but not implemented. `get_many(..., limit=1).first()` + `deserialize_value` is the current server-side read pattern (it's what the core XCom route uses). Happy to switch if you'd rather a different call.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
Review Comment:
Agreed — fixed in 924c9ee. Replaced the hand-rolled header/cookie parsing with `Depends(requires_access_dag("PUT", DagAccessEntity.RUN))`, the same pattern `providers/common/ai/plugins/hitl_review.py` uses. The `get_user` behind it already resolves both the bearer header and the `_token` cookie, so a browser-clicked link stays authorized, and nothing here parses the request or assumes the auth-manager mapping anymore.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
Review Comment:
I can replace the hand-rolled downstream walk + `clear_task_instances` with `find_relevant_relatives(direction="downstream")` + `dag.clear(task_ids=..., run_id=...)`, which is what the core clear-task-instances endpoint uses. Before I rework it — are you OK with `dag.clear` being called here from the provider's FastAPI app?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
Review Comment:
You're right — removed (924c9ee). FastAPI decodes path params already, so it was a redundant double-decode.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
Review Comment:
The `dag_run` is used for two things: the clear step, and building the redirect target from the run's own persisted `dag_id`/`run_id` rather than the request path — that was to resolve the CodeQL open-redirect finding (a same-site path from DB-sourced values, not request input). If I move clearing to `dag.clear(run_id=...)`, the object would only be needed for that redirect. Would you prefer I keep it DB-sourced for CodeQL, or drop the query and build the redirect from the (validated) path params?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5393155352
Thanks for the thorough re-review, and for verifying the buttons through the route's own resolution path plus the revert-each-fix tampering — that's a better check than the test alone.
Handled the Airflow 2 case you raised. In the pre-3.0 branch of `WorkflowJobRepairSingleTaskLink.get_link`, the `get_launch_task_id` call is now wrapped so the "no launch task in this group" `AirflowException` returns `""` instead of propagating — a standalone Databricks task renders no repair link rather than surfacing a 500 behind the button, matching the graceful degradation the Airflow 3 branch already has. I reproduced the raise for a top-level task first to confirm the path before changing anything. `WorkflowJobRepairAllFailedLink` is untouched since `_CreateDatabricksWorkflowOperator` only exists inside a workflow group.
Added an Airflow-2 test that asserts the standalone-task case returns `""`. Provider unit tests pass locally, ruff and mypy clean.
@moomindani — I think that closes the last item. PTAL.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3804655605
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,266 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ from urllib.parse import urlencode
+
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from sqlalchemy import select
+
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ request: Request,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Same-site relative action; SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = f"{request.url.path}?{request.url.query}"
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
Review Comment:
Thanks — addressed the scanner flag on the repair redirect.
The redirect and the confirmation-page form action were assembled from the request's own path and query. The values were percent-encoded into a fixed same-site `/dags/{dag_id}/runs/{run_id}` path and could not actually leave the origin, but CodeQL traced request input into a redirect target and kept the open-redirect alert open regardless.
Pushed a change so no request-controlled string reaches either target:
- The redirect is now built from the `DagRun`'s persisted `dag_id` / `run_id`, read from the database, rather than the request path params.
- The confirmation form action is rebuilt from the already-validated identifiers instead of echoing the request URL.
- Fetching the run also lets the endpoint return a `404` for an unknown run and reuse that row when clearing task instances.
That breaks the taint flow the scanner was following, so the finding should clear on the next analysis. All provider unit tests pass locally.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3949111289
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,39 @@ def get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED", "MAXIMUM_CONCURRENT_RUNS_REACHED"}
+
+ # ``get_run_tasks`` returns one entry per attempt and is not ordered by attempt, so a
+ # retried or already-repaired task appears several times under the same ``task_key``. Keep
+ # only the latest attempt per key (sort by ``start_time``, same idiom as
+ # ``DatabricksTaskBaseOperator._get_current_databricks_task``) before judging its state, so
+ # the result never contains duplicate keys — Databricks rejects duplicates in
+ # ``rerun_tasks`` — and a task whose latest attempt succeeded is not reported as failed.
+ # Never-started sub-tasks omit ``start_time`` or send null; treat those as 0.
+ sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda task: task.get("start_time") or 0)
+ latest_by_key = {task["task_key"]: task for task in sorted_tasks}
+
+ failed_task_keys = []
+ for task_key, task in latest_by_key.items():
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
Review Comment:
The `dag_run` is used for two things: the clear step, and building the redirect target from the run's own persisted `dag_id`/`run_id` rather than the request path — that was to resolve the CodeQL open-redirect finding (a same-site path from DB-sourced values, not request input). If I move clearing to `dag.clear(run_id=...)`, the object would only be needed for that redirect. Would you prefer I keep it DB-sourced for CodeQL, or drop the query and build the redirect from the (validated) path params?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "Vamsi-klu (via GitHub)" <gi...@apache.org>.
Vamsi-klu commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5344730413
@PrakshiGoyal10 please feel free to use the patch. cc @moomindani
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4228592329
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +385,40 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
Review Comment:
The possible advantage of doing that is fewer imports needed on the API server to make it work.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-6166401383
Pushed `e33d691113`, which styles the repair confirmation page to match the Airflow UI: a card layout that follows the OS light/dark setting, plus a **Cancel** link back to the run. It is still a single inline stylesheet with no JavaScript and no front-end build. Screenshot 4 in the description is updated.
On why a separate confirmation page is needed at all:
- **Extra links can only navigate.** Core renders every extra link as a plain `<a href>` (`ExtraLinks.tsx`), so a click is always a `GET`, and a plugin has no hook to turn it into a form submit or a `POST`.
- **A `GET` must not repair.** The UI's `_token` cookie is `SameSite=Lax`, so browsers still send it on top-level cross-site `GET` navigations. If the link repaired directly, a link posted elsewhere or a link preview fetching the URL could trigger a repair. So the `GET` only renders the page, and the repair happens on the page's same-site `POST`, which a cross-site request cannot authenticate.
- **The new tab comes from core.** Airflow 3.1.x opens every extra link with `target="_blank"`. Current `main` opens same-origin links in the same tab, so on newer versions the page replaces the task view, and Repair or Cancel lands back on the run in that tab.
To remove the separate page entirely, a follow-up could add a `react_apps` tab on the task instance (`destination: "task_instance"`, as `common.ai`'s HITL Review plugin does) with a real Repair button that posts to the same endpoint from inside the UI. The endpoint and its protections would stay as they are. I've kept that out of this PR because it adds a front-end toolchain to the provider. Happy to open it as a follow-up if that direction sounds right.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub) via commits" <co...@airflow.apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4444592652
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, task_id=task_id)
+ ti.xcom_push(key=key, value=link)
+ logger.info("Stored Databricks repair link in XCom: %s", link)
+ except Exception as e:
+ logger.warning("Failed to store Databricks repair link: %s", e)
+
+
+def _get_link_from_xcom(ti_key: TaskInstanceKey | None, key: str, logger: Logger) -> str:
+ if ti_key is None:
+ return ""
+ try:
+ return XCom.get_value(ti_key=ti_key, key=key) or ""
+ except Exception as e:
+ logger.warning("Failed to retrieve Databricks link from XCom: %s", e)
+ return ""
+
+
class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
"""Constructs a link to send a request to repair all failed tasks in the Databricks workflow."""
name = "Repair All Failed Tasks"
+ @property
+ def xcom_key(self) -> str:
+ """XCom key where the link is stored during task execution."""
+ return "databricks_repair_all_failed_link"
Review Comment:
Is this meant to be settable? If not
```suggestion
xcom_key: ClassVar[str] = "databricks_repair_all_failed_link"
"""XCom key where the link is stored during task execution."""
```
And then we also don't need to instantiate this class to read it
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
Review Comment:
key is the same in both arms.
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
Review Comment:
Nit:
```suggestion
if mapped := task_key_map.get(task_id)
return mapped
```
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +531,21 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def xcom_key(self) -> str:
+ """XCom key where the link is stored during task execution."""
+ return "databricks_repair_single_task_link"
Review Comment:
Same class var here.
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, task_id=task_id)
+ ti.xcom_push(key=key, value=link)
Review Comment:
It honestly feels a bit odd to store this URL in xcom, given that the returned URL is _already_ stored in XCom. What's the reason for storing it twice?
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
Review Comment:
As this is currently written, this assumes that the server is able to read XCom values, which is not always the case if custom XCom backends are in use. For now it might be sufficient to document that the repair doesn't work on those situations.
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks and their downstream task instances for this run.
+
+ Expands to downstream with ``dag.partial_subset`` and clears with ``dag.clear``, as the core
+ clear-task-instances route does, so the upstream-failed dependents resume once the repaired
+ Databricks sub-runs succeed — without clearing the whole Dag. ``find_relevant_relatives``,
+ which newer core uses for this, does not exist on Airflow 3.1.
+ """
+ if not task_ids:
+ return
+ to_clear: set[str] = set(task_ids)
+ downstream = dag.partial_subset(
+ task_ids=task_ids, include_downstream=True, include_upstream=False, exclude_original=True
+ )
+ to_clear.update(downstream.task_dict)
+ logger.info("Clearing %s task(s) and downstream after Databricks repair: %s", len(task_ids), to_clear)
+ dag.clear(task_ids=list(to_clear), run_id=run_id, session=session)
+
+ # Plain-CSS approximation of the Airflow UI, so the page does not look out of place next to it.
+ # No JavaScript: the repair runs only when the user submits the form.
+ _CONFIRMATION_PAGE_STYLE = """
+ :root { color-scheme: light dark; --bg: #f4f6f8; --card: #ffffff; --text: #1a202c;
+ --muted: #4a5568; --border: #e2e8f0; --primary: #017cee; --primary-hover: #0165c3; }
+ @media (prefers-color-scheme: dark) { :root { --bg: #0f1520; --card: #1a2230;
+ --text: #e2e8f0; --muted: #a0aec0; --border: #2d3748; --primary: #3b9cff;
+ --primary-hover: #63b0ff; } }
+ body { margin: 0; min-height: 100vh; display: flex; align-items: center; justify-content: center;
+ background: var(--bg); color: var(--text);
+ font: 15px/1.5 system-ui, -apple-system, "Segoe UI", Roboto, sans-serif; }
+ .card { width: min(560px, calc(100vw - 32px)); background: var(--card);
+ border: 1px solid var(--border); border-radius: 12px; padding: 28px 32px;
+ box-shadow: 0 4px 16px rgba(0, 0, 0, 0.08); }
+ h1 { font-size: 20px; margin: 0 0 20px; }
+ dl { display: grid; grid-template-columns: max-content 1fr; gap: 6px 16px; margin: 0 0 20px; }
+ dt { color: var(--muted); }
+ dd { margin: 0; font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }
+ p { margin: 0 0 24px; }
+ .actions { display: flex; gap: 12px; justify-content: flex-end; align-items: center; }
+ button, .cancel { font: inherit; border-radius: 6px; padding: 8px 18px; cursor: pointer;
+ text-decoration: none; }
+ button { background: var(--primary); color: #fff; border: 1px solid var(--primary); }
+ button:hover { background: var(--primary-hover); }
+ .cancel { color: var(--text); border: 1px solid var(--border); }
+ .cancel:hover { background: var(--bg); }
+ """
+
+ def _repair_confirmation_page(
+ dag_id: str, run_id: str, action: str, summary: str, cancel_url: str
+ ) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
Review Comment:
Nit: `"""` and `textwrap.dedent()` here please.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3949111289
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,39 @@ def get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED", "MAXIMUM_CONCURRENT_RUNS_REACHED"}
+
+ # ``get_run_tasks`` returns one entry per attempt and is not ordered by attempt, so a
+ # retried or already-repaired task appears several times under the same ``task_key``. Keep
+ # only the latest attempt per key (sort by ``start_time``, same idiom as
+ # ``DatabricksTaskBaseOperator._get_current_databricks_task``) before judging its state, so
+ # the result never contains duplicate keys — Databricks rejects duplicates in
+ # ``rerun_tasks`` — and a task whose latest attempt succeeded is not reported as failed.
+ # Never-started sub-tasks omit ``start_time`` or send null; treat those as 0.
+ sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda task: task.get("start_time") or 0)
+ latest_by_key = {task["task_key"]: task for task in sorted_tasks}
+
+ failed_task_keys = []
+ for task_key, task in latest_by_key.items():
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
Review Comment:
The `dag_run` is used for two things: the clear step, and building the redirect target from the run's own persisted `dag_id`/`run_id` rather than the request path — that was to resolve the CodeQL open-redirect finding (a same-site path from DB-sourced values, not request input). If I move clearing to `dag.clear(run_id=...)`, the object would only be needed for that redirect. Would you prefer I keep it DB-sourced for CodeQL, or drop the query and build the redirect from the (validated) path params?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,39 @@ def get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED", "MAXIMUM_CONCURRENT_RUNS_REACHED"}
+
+ # ``get_run_tasks`` returns one entry per attempt and is not ordered by attempt, so a
+ # retried or already-repaired task appears several times under the same ``task_key``. Keep
+ # only the latest attempt per key (sort by ``start_time``, same idiom as
+ # ``DatabricksTaskBaseOperator._get_current_databricks_task``) before judging its state, so
+ # the result never contains duplicate keys — Databricks rejects duplicates in
+ # ``rerun_tasks`` — and a task whose latest attempt succeeded is not reported as failed.
+ # Never-started sub-tasks omit ``start_time`` or send null; treat those as 0.
+ sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda task: task.get("start_time") or 0)
+ latest_by_key = {task["task_key"]: task for task in sorted_tasks}
+
+ failed_task_keys = []
+ for task_key, task in latest_by_key.items():
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
Review Comment:
The Databricks Jobs API `GET 2.2/jobs/runs/get` returns the whole run with every task and its state — there's no server-side filter for task state — so selecting the failed tasks is necessarily done client-side here.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -49,11 +49,15 @@ class WorkflowRunMetadata:
:param run_id: The ID of the Databricks workflow run.
:param job_id: The ID of the Databricks workflow job.
:param conn_id: The connection ID used to connect to Databricks.
+ :param task_key_map: Airflow ``task_id`` → Databricks ``task_key`` for the workflow's tasks.
+ Optional and defaulted for backward compatibility with runs launched before it was added
+ (those XComs carry only conn_id/job_id/run_id).
"""
conn_id: str
job_id: int
run_id: int
+ task_key_map: dict[str, str] = field(default_factory=dict)
Review Comment:
This is Airflow-3-specific. An explicit `databricks_task_key` doesn't survive Dag serialization, so on the API server we can't reconstruct the real key for such a task. The launch task records the `task_id → task_key` map in its `WorkflowRunMetadata` XCom so the repair endpoint can map keys back to tasks. Airflow 2 rendered the links from the live operator (which has the key), so it never needed this. I'll add a comment saying so.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +382,33 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
Review Comment:
Verified true — a plugin extra link that declares no `operators` is replaced at deserialization by `XComOperatorLink`, which just returns a URL stored under `xcom_key`. This link stores no such URL; it builds one at request time in `get_link`, so it has to survive as the real object. Rewrote the comment to say exactly that (924c9ee).
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +382,33 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
+ # XComOperatorLink. Lazy import avoids a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks_workflow import (
+ _CreateDatabricksWorkflowOperator,
+ )
+
+ return [_CreateDatabricksWorkflowOperator]
+
def get_link( # type: ignore[override] # Signature intentionally kept this way for Airflow 2.x compatibility
self,
operator,
dttm=None,
*,
ti_key: TaskInstanceKey | None = None,
) -> str:
+ if AIRFLOW_V_3_0_PLUS:
+ if not AIRFLOW_V_3_1_PLUS or ti_key is None:
Review Comment:
This is tied to the 3.1 gate below — if we relax it (see that thread) this collapses to `if ti_key is None:`. I'll simplify it once the gate question is settled.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
Review Comment:
After the auth change, nothing in this block imports anything 3.1-only anymore. The one remaining 3.1 dependency is behavioural: `get_user` only honours the `_token` cookie from 3.1 onward, so on 3.0.x a repair link clicked in the browser wouldn't authenticate (only an XHR carrying a bearer header would). So I can relax the gate to `AIRFLOW_V_3_0_PLUS` and accept that degradation on 3.0.x, or keep it at 3.1 — which would you prefer?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "moomindani (via GitHub)" <gi...@apache.org>.
moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3974626224
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
Review Comment:
One correction to the premise here, since it moves where the gate can sit rather than whether it should.
`get_user` picks up the `_token` cookie from **3.1.1**, not 3.1.0. In 3.1.0 it takes only the two header schemes:
```python
async def get_user(
oauth_token: str | None = Depends(oauth2_scheme),
bearer_credentials: HTTPAuthorizationCredentials | None = Depends(bearer_scheme),
) -> BaseUser:
```
The `token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)` branch appears in 3.1.1 and is present in every release I checked after it (3.1.2-3.1.5, 3.1.8, 3.2.x, 3.3.x).
`AIRFLOW_V_3_1_PLUS` is `>= (3, 1, 0)`, so on exactly 3.1.0 a repair link clicked in the browser carries only the cookie, `requires_access_dag` resolves no token, and the click 401s — which the hand-rolled resolver did handle. So the options are to gate at 3.1.1 (patch-level gates are already used elsewhere, e.g. `AIRFLOW_V_3_1_9_PLUS` in the celery provider) or to keep 3.1 and say in that comment that browser clicks need 3.1.1+.
Relaxing to `AIRFLOW_V_3_0_PLUS` looks closed off for the same reason: 3.0.6's `get_user` reads only `oauth2_scheme`, so a cookie-only navigation never authenticates there either.
---
Drafted-by: Claude Code (Opus 5); reviewed by @moomindani before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "moomindani (via GitHub)" <gi...@apache.org>.
moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3745948631
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,28 @@ def get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED", "MAXIMUM_CONCURRENT_RUNS_REACHED"}
+ failed_task_keys = []
+ for task in self.get_run_tasks(run_id):
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
+ failed_task_keys.append(task["task_key"])
+ return failed_task_keys
Review Comment:
This assumes one entry per `task_key`, but `get_run_tasks` returns one entry **per attempt**, so a task that retried (or a run that was already repaired) yields the same key several times. Databricks rejects that:
```
Error: rerun_tasks cannot contain duplicate task keys.
```
I confirmed this against a live run — with `max_retries=1` on a failing task this returns `['flaky', 'flaky']` and the repair fails; deduped to `['flaky']` it succeeds. Attempts accumulate across repairs, so the duplicate count grows (I saw four).
There is a second, quieter symptom: because every historical attempt is inspected, a `task_key` whose **latest** attempt succeeded is still reported as failed. After my run finished fully green (`parent: SUCCESS`, both tasks `SUCCESS`), this still returned `['flaky', 'flaky', 'flaky', 'flaky']`. Databricks does *not* reject repairing an already-successful task — it accepts it and returns a `repair_id` — so that path silently re-runs a healthy task and clears its Airflow TI.
Both symptoms have one root cause, and the fix already exists in this repo: `_get_current_databricks_task` in `operators/databricks.py` resolves the latest attempt by sorting on `start_time` before building the map. Same idiom here:
```python
sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda t: t["start_time"])
latest_by_key = {task["task_key"]: task for task in sorted_tasks}
```
then evaluate the failure states over `latest_by_key.values()`. That dedupes and fixes the stale-attempt judgement in one move.
##########
providers/databricks/tests/unit/databricks/hooks/test_databricks.py:
##########
@@ -749,6 +749,20 @@ def test_get_run_tasks_success_multiple_pages(self, mock_requests):
assert len(tasks) == 2
assert tasks == GET_RUN_RESPONSE["tasks"] * 2
+ def test_get_run_failed_task_keys(self):
Review Comment:
This fixture has one entry per `task_key`, which is why CI is green despite the duplicate-key bug — the real API returns one entry per *attempt*.
Worth adding a case shaped like a live run, e.g. two `flaky` entries with `attempt_number` 0/1 both `FAILED` plus a later attempt that `SUCCESS`ed, asserting the key appears once and is judged by its latest attempt. That would fail on the current implementation.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5318283409
cc @moomindani if all checks out please approve :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub) via commits" <co...@airflow.apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4451887559
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
Review Comment:
Agreed. I added a note to `docs/operators/workflow.rst`. On Airflow 3 the repair endpoint reads the launch task's XCom result straight from the metadata database, so repair isn't supported when a custom XCom backend stores values outside it.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks and their downstream task instances for this run.
+
+ Expands to downstream with ``dag.partial_subset`` and clears with ``dag.clear``, as the core
+ clear-task-instances route does, so the upstream-failed dependents resume once the repaired
+ Databricks sub-runs succeed — without clearing the whole Dag. ``find_relevant_relatives``,
+ which newer core uses for this, does not exist on Airflow 3.1.
+ """
+ if not task_ids:
+ return
+ to_clear: set[str] = set(task_ids)
+ downstream = dag.partial_subset(
+ task_ids=task_ids, include_downstream=True, include_upstream=False, exclude_original=True
+ )
+ to_clear.update(downstream.task_dict)
+ logger.info("Clearing %s task(s) and downstream after Databricks repair: %s", len(task_ids), to_clear)
+ dag.clear(task_ids=list(to_clear), run_id=run_id, session=session)
+
+ # Plain-CSS approximation of the Airflow UI, so the page does not look out of place next to it.
+ # No JavaScript: the repair runs only when the user submits the form.
+ _CONFIRMATION_PAGE_STYLE = """
+ :root { color-scheme: light dark; --bg: #f4f6f8; --card: #ffffff; --text: #1a202c;
+ --muted: #4a5568; --border: #e2e8f0; --primary: #017cee; --primary-hover: #0165c3; }
+ @media (prefers-color-scheme: dark) { :root { --bg: #0f1520; --card: #1a2230;
+ --text: #e2e8f0; --muted: #a0aec0; --border: #2d3748; --primary: #3b9cff;
+ --primary-hover: #63b0ff; } }
+ body { margin: 0; min-height: 100vh; display: flex; align-items: center; justify-content: center;
+ background: var(--bg); color: var(--text);
+ font: 15px/1.5 system-ui, -apple-system, "Segoe UI", Roboto, sans-serif; }
+ .card { width: min(560px, calc(100vw - 32px)); background: var(--card);
+ border: 1px solid var(--border); border-radius: 12px; padding: 28px 32px;
+ box-shadow: 0 4px 16px rgba(0, 0, 0, 0.08); }
+ h1 { font-size: 20px; margin: 0 0 20px; }
+ dl { display: grid; grid-template-columns: max-content 1fr; gap: 6px 16px; margin: 0 0 20px; }
+ dt { color: var(--muted); }
+ dd { margin: 0; font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }
+ p { margin: 0 0 24px; }
+ .actions { display: flex; gap: 12px; justify-content: flex-end; align-items: center; }
+ button, .cancel { font: inherit; border-radius: 6px; padding: 8px 18px; cursor: pointer;
+ text-decoration: none; }
+ button { background: var(--primary); color: #fff; border: 1px solid var(--primary); }
+ button:hover { background: var(--primary-hover); }
+ .cancel { color: var(--text); border: 1px solid var(--border); }
+ .cancel:hover { background: var(--bg); }
+ """
+
+ def _repair_confirmation_page(
+ dag_id: str, run_id: str, action: str, summary: str, cancel_url: str
+ ) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
Review Comment:
Done. The page is now a `textwrap.dedent("""...""")` template, filled with `str.format`. The multi-line stylesheet goes in after the dedent so it doesn't affect it, and every request-derived value is still `escape()`d.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4184531854
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
Review Comment:
Reworked. The hand-rolled downstream walk + `clear_task_instances` is replaced with `find_relevant_relatives(direction="downstream")` + `dag.clear(task_ids=..., run_id=...)` — the same pair the core `clearTaskInstances` route uses (`api_fastapi/core_api/routes/public/task_instances.py`). Pushed in 999f9c0.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
Review Comment:
Kept the `DagRun` fetch, but it's now load-bearing rather than just an existence check: the same-site redirect target is built from `dag_run.dag_id` / `dag_run.run_id` (DB-sourced) specifically so the CodeQL open-redirect finding doesn't re-trigger — building it from the request path param reintroduces that alert. With the clearing rework above, the object no longer feeds `_clear_repaired_and_downstream` (that now takes `run_id`), so the query earns its keep via the trusted redirect plus the 404. Happy to drop the explicit 404 and lean on the natural failure path if you'd prefer, but the redirect source needs to stay DB-sourced for CodeQL. Pushed in 999f9c0.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5474619979
Thanks for the approval and for measuring each step rather than taking the description's word for it — the verification on the standalone-task path is much appreciated.
Good catch on the new Airflow-2 test never executing — you're right, the `AIRFLOW_V_3_0_PLUS` and `RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES` guards cover every environment between them, so it skips everywhere, same as the neighbouring `*_airflow2` tests. I left it in that shape for consistency with those, but I take the point that it reads as a live guard for the fix when it isn't.
Noted the pre-existing Airflow-2 repair-all `databricks_task_key` clear-parity item as out of scope for this PR.
One question on merging: the PR still shows `REVIEW_REQUIRED` / blocked despite the approval — is a committer approval (or a second reviewer) still needed to land it, or is there anything else you'd like from me first?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3608706108
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,193 @@
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ *,
+ repair_all: bool = False,
+ tasks_to_repair: list[str] | None = None,
+) -> str:
+ """Build the URL to the Airflow-3 FastAPI repair endpoint for a workflow run."""
+ from urllib.parse import quote, urlencode
+
+ from airflow.configuration import conf
+
+ query: dict[str, Any] = {
+ "databricks_conn_id": databricks_conn_id,
+ "databricks_run_id": databricks_run_id,
+ }
+ if repair_all:
+ query["repair_all"] = "true"
+ if tasks_to_repair:
+ query["tasks_to_repair"] = ",".join(tasks_to_repair)
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return f"{base_url}{REPAIR_URL_PREFIX}/{dag_id}/{quote(run_id, safe='')}?{urlencode(query)}"
+
+
+def _get_launch_metadata_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> Any:
+ """
+ Read the launch task's ``WorkflowRunMetadata`` (conn_id, job_id, run_id) from XCom.
+
+ Uses only the public XCom API keyed by ``ti_key`` — no metadata-DB access — which is the
+ only mechanism available to extra links on Airflow 3.
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" not in ti_key.task_id:
+ launch_task_id = get_launch_task_id(task_group)
+ ti_key = _get_launch_task_key(ti_key, task_id=launch_task_id)
+ result = XCom.get_value(ti_key=ti_key, key="return_value")
+ if not result:
+ return None
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ return WorkflowRunMetadata(**result)
+
+
+if AIRFLOW_V_3_0_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import RedirectResponse
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _clear_repaired_and_downstream(
+ dag_id: str, run_id: str, task_keys: list[str], logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ import hashlib
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.models.taskinstance import clear_task_instances
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail=f"Dag {dag_id} not found.")
+
+ # Serialized tasks don't expose the operator's ``databricks_task_key`` property, so
+ # reproduce its default: an explicit key if present, else md5(dag_id__task_id).
+ key_to_task_id = {}
+ for task in dag.tasks:
+ task_key = (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+ key_to_task_id[task_key] = task.task_id
+
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ target_task_ids: set[str] = set(repaired_task_ids)
+ for task_id in repaired_task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [
+ ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+ session.commit()
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ tasks_to_repair: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+ from airflow.configuration import conf
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return_url = f"{base_url}/dags/{dag_id}/runs/{run_id}"
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a clear
+ # error to the UI instead of a bare 500.
+ try:
+ hook = DatabricksHook(databricks_conn_id=databricks_conn_id)
+ if repair_all:
+ task_keys = hook.get_run_failed_task_keys(databricks_run_id)
+ elif tasks_to_repair:
+ task_keys = tasks_to_repair.split(",")
+ else:
+ task_keys = []
+
+ 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)
+
+ log.info("Repairing Databricks run %s tasks %s", databricks_run_id, task_keys)
+ _repair_task(
+ databricks_conn_id=databricks_conn_id,
+ databricks_run_id=databricks_run_id,
+ tasks_to_repair=task_keys,
+ logger=log,
+ )
+ except HTTPException:
+ raise
+ except Exception as e:
+ log.exception("Databricks repair failed for run %s", databricks_run_id)
+ 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)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/618)
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,193 @@
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ *,
+ repair_all: bool = False,
+ tasks_to_repair: list[str] | None = None,
+) -> str:
+ """Build the URL to the Airflow-3 FastAPI repair endpoint for a workflow run."""
+ from urllib.parse import quote, urlencode
+
+ from airflow.configuration import conf
+
+ query: dict[str, Any] = {
+ "databricks_conn_id": databricks_conn_id,
+ "databricks_run_id": databricks_run_id,
+ }
+ if repair_all:
+ query["repair_all"] = "true"
+ if tasks_to_repair:
+ query["tasks_to_repair"] = ",".join(tasks_to_repair)
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return f"{base_url}{REPAIR_URL_PREFIX}/{dag_id}/{quote(run_id, safe='')}?{urlencode(query)}"
+
+
+def _get_launch_metadata_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> Any:
+ """
+ Read the launch task's ``WorkflowRunMetadata`` (conn_id, job_id, run_id) from XCom.
+
+ Uses only the public XCom API keyed by ``ti_key`` — no metadata-DB access — which is the
+ only mechanism available to extra links on Airflow 3.
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" not in ti_key.task_id:
+ launch_task_id = get_launch_task_id(task_group)
+ ti_key = _get_launch_task_key(ti_key, task_id=launch_task_id)
+ result = XCom.get_value(ti_key=ti_key, key="return_value")
+ if not result:
+ return None
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ return WorkflowRunMetadata(**result)
+
+
+if AIRFLOW_V_3_0_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import RedirectResponse
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _clear_repaired_and_downstream(
+ dag_id: str, run_id: str, task_keys: list[str], logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ import hashlib
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.models.taskinstance import clear_task_instances
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail=f"Dag {dag_id} not found.")
+
+ # Serialized tasks don't expose the operator's ``databricks_task_key`` property, so
+ # reproduce its default: an explicit key if present, else md5(dag_id__task_id).
+ key_to_task_id = {}
+ for task in dag.tasks:
+ task_key = (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+ key_to_task_id[task_key] = task.task_id
+
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ target_task_ids: set[str] = set(repaired_task_ids)
+ for task_id in repaired_task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [
+ ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+ session.commit()
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ tasks_to_repair: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+ from airflow.configuration import conf
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return_url = f"{base_url}/dags/{dag_id}/runs/{run_id}"
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a clear
+ # error to the UI instead of a bare 500.
+ try:
+ hook = DatabricksHook(databricks_conn_id=databricks_conn_id)
+ if repair_all:
+ task_keys = hook.get_run_failed_task_keys(databricks_run_id)
+ elif tasks_to_repair:
+ task_keys = tasks_to_repair.split(",")
+ else:
+ task_keys = []
+
+ 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)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/617)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5294456127
Thanks for the careful review and for validating it against a live workspace — the duplicate-key bug is exactly right, and I confirmed it the same way. Pushed a fix.
**Duplicate task keys.** `get_run_failed_task_keys` now keeps only the latest attempt per `task_key` before judging state, using the sort-by-`start_time` idiom from `_get_current_databricks_task` as you suggested. That removes the duplicates that Databricks rejects and also fixes the stale-attempt symptom (a key whose newest attempt succeeded is no longer reported failed). Added a hook test shaped like a live run — multiple `flaky` attempts (FAILED, FAILED, then SUCCESS) plus a task whose latest attempt failed — asserting each key appears once and is judged by its latest attempt.
**A second, related defect I found while fixing that.** Reconstructing the Databricks `task_key` on the API-server side as `md5(dag_id__task_id)` is only correct for tasks that use the default key. An explicit `databricks_task_key` does not survive Dag serialization (I verified the attribute is absent from the deserialized operator), so for such tasks single-task repair targeted a non-existent key and repair-all silently skipped clearing that task's instance. The launch task now records the `task_id` to `task_key` map in its run-metadata XCom, and the endpoint reads it server-side (so no request-supplied keys, consistent with the trusted-identifier model); runs launched before this fall back to the md5 derivation. Covered by tests for both the explicit-key and legacy-fallback paths.
All provider unit tests pass locally. On the still-open CodeQL redirect finding: the redirect target is a same-site relative path built with `quote(..., safe='')`, so it can't leave the origin — I believe it's a false positive, but happy to refactor if you'd prefer it read differently.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "Vamsi-klu (via GitHub)" <gi...@apache.org>.
Vamsi-klu commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5337966305
@PrakshiGoyal10 @moomindani I caught the latest review and went through the plugin plus the extra-link path. Those two blockers look like the ones that keep this from actually showing up in the Airflow 3 UI, and I wanted to ask if you are both okay with me taking a pass at them on this PR.
The repair buttons never render because Airflow 3 only serializes extra links as a name-to-xcom-key map, then the API server rebuilds them as XComOperatorLink. The job-run link already follows that: it sets xcom_key and store_databricks_job_run_link writes the URL during execute. The two repair links do neither, so /links comes back empty and the UI drops the button. The Airflow 3 branches in get_link never run on that path.
The other issue is separate. Even if the URL were built at request time, _get_launch_task_id_v3 still calls get_child_by_label, and SerializedTaskGroup does not have that method. That is the member-task case the single-task link is registered on.
If that works for both of you, I would persist the repair URLs into XCom at the same point as the job-run link (before wait_for_termination, so a later Databricks failure still leaves a button), resolve the launch task from the group id or the upstream .launch task instead of get_child_by_label, and add the SerializedDagModel.get_dag round-trip test so this class of break cannot come back silently.
I can also take the start_time fragility, the unmapped-key warning on repair-all, and the base_url form-action / redirect notes in the same pass if you want those now. I will not open another PR. I can send a patch you can pull, or I can wait if you already have a fix in flight.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-4999891934
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](https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [prek-hooks]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-prek-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example Dag that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "Vamsi-klu (via GitHub)" <gi...@apache.org>.
Vamsi-klu commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5014062857
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:
- `GET` should render a read-only confirmation page; only a CSRF-protected, Dag-authorized `POST` should repair or clear tasks.
- The URL should carry only validated Dag/run/task identifiers. `databricks_conn_id`, the Databricks run ID, and Databricks task keys should be derived server-side from trusted `WorkflowRunMetadata` XCom instead of accepted from the request.
- Redirect targets should be same-site relative paths built from validated identifiers, resolving the two open CodeQL redirect findings.
- The FastAPI repair app should be registered only on Airflow 3.1+, since `resolve_user_from_token` is unavailable on 3.0.6. Provider config must come from `airflow.providers.common.compat.sdk`, not `airflow.configuration`, to clear the current static check.
- Databricks failures should produce a generic 502 without reflecting raw upstream exception text, and task clearing should happen only after a successful repair call.
- Confirmation-page values must be safely encoded so identifiers cannot become reflected HTML or JavaScript.
The PR body should also include `related: #52280`. I can prepare a patch branch with the route split, trusted identifier resolution, access checks, and focused DB/API tests if that collaboration route would be useful; I will not open a competing Apache PR while this one is active. Until the mutation-on-GET issue is fixed, I do not recommend clicking the current repair link against a live Databricks workspace.
---
Drafted-by: Codex (GPT-5) (no human review before posting)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3721149126
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,260 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ from urllib.parse import urlencode
+
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _serialized_task_key(dag_id: str, task: Any) -> str:
+ """
+ Reproduce an operator's ``databricks_task_key`` from a serialized task.
+
+ Serialized tasks don't expose the operator property, so mirror its default: an explicit key
+ if one was set, else ``md5(dag_id__task_id)``.
+ """
+ import hashlib
+
+ return (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from sqlalchemy import select
+
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ request: Request,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Same-site relative action; SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = f"{request.url.path}?{request.url.query}"
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ # Redirect to a same-site relative path with the identifiers percent-encoded, so the
+ # target can never be steered to another host or scheme (CodeQL open-redirect).
+ return_url = f"/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
+
+ metadata = _read_launch_metadata(dag_id, run_id, launch_task_id, session)
+
+ if repair_all:
+ repaired_task_ids: list[str] = [] # resolved from live Databricks state below
+ else:
+ if task_id is None or not dag.has_task(task_id):
+ raise HTTPException(status_code=404, detail="Task not found in Dag.")
+ repaired_task_ids = [task_id]
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a
+ # generic error to the UI without leaking the upstream exception text.
+ try:
+ if repair_all:
+ hook = DatabricksHook(databricks_conn_id=metadata.conn_id)
+ task_keys = hook.get_run_failed_task_keys(metadata.run_id)
+ key_to_task_id = {_serialized_task_key(dag_id, t): t.task_id for t in dag.tasks}
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ else:
+ task_keys = [_serialized_task_key(dag_id, dag.get_task(repaired_task_ids[0]))]
+
+ 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)
+
+ log.info("Repairing Databricks run %s tasks %s", metadata.run_id, task_keys)
+ _repair_task(
+ databricks_conn_id=metadata.conn_id,
+ databricks_run_id=metadata.run_id,
+ tasks_to_repair=task_keys,
+ logger=log,
+ )
+ except HTTPException:
+ raise
+ except Exception:
+ log.exception("Databricks repair failed for run %s", metadata.run_id)
+ raise HTTPException(status_code=502, detail="Databricks repair request failed.")
+
+ # Clear only after a successful repair call, so a failed repair leaves state untouched.
+ _clear_repaired_and_downstream(dag, run_id, repaired_task_ids, session, log)
+ session.commit()
+
+ return RedirectResponse(return_url, status_code=303)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/3)
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,260 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ from urllib.parse import urlencode
+
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _serialized_task_key(dag_id: str, task: Any) -> str:
+ """
+ Reproduce an operator's ``databricks_task_key`` from a serialized task.
+
+ Serialized tasks don't expose the operator property, so mirror its default: an explicit key
+ if one was set, else ``md5(dag_id__task_id)``.
+ """
+ import hashlib
+
+ return (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from sqlalchemy import select
+
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ request: Request,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Same-site relative action; SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = f"{request.url.path}?{request.url.query}"
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ # Redirect to a same-site relative path with the identifiers percent-encoded, so the
+ # target can never be steered to another host or scheme (CodeQL open-redirect).
+ return_url = f"/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
+
+ metadata = _read_launch_metadata(dag_id, run_id, launch_task_id, session)
+
+ if repair_all:
+ repaired_task_ids: list[str] = [] # resolved from live Databricks state below
+ else:
+ if task_id is None or not dag.has_task(task_id):
+ raise HTTPException(status_code=404, detail="Task not found in Dag.")
+ repaired_task_ids = [task_id]
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a
+ # generic error to the UI without leaking the upstream exception text.
+ try:
+ if repair_all:
+ hook = DatabricksHook(databricks_conn_id=metadata.conn_id)
+ task_keys = hook.get_run_failed_task_keys(metadata.run_id)
+ key_to_task_id = {_serialized_task_key(dag_id, t): t.task_id for t in dag.tasks}
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ else:
+ task_keys = [_serialized_task_key(dag_id, dag.get_task(repaired_task_ids[0]))]
+
+ 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)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/31)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5214770323
Thanks for the detailed review. I've reworked the repair path along these lines. Rebased onto `main` and pushed; CI is green (static checks, provider tests, and the CodeQL analysis all pass now).
Point by point:
- **Mutation on GET → confirmation + POST.** `GET` now renders a read-only confirmation page; only a same-site `POST` repairs and clears tasks. Rather than add a bespoke CSRF token, this leans on the same mechanism the core UI mutations rely on: the API server's auth cookie is `SameSite=Lax`, so it isn't sent on a cross-site `POST` — moving the mutation off `GET` is what closes the CSRF hole.
- **Trusted identifiers only.** The link now carries only Airflow identifiers — the run's launch `task_id`, and for a single-task repair the target `task_id`. `databricks_conn_id` and the Databricks run id are derived server-side from the launch task's `WorkflowRunMetadata` XCom, never from the request, so a crafted link can't point the repair at an arbitrary connection or run. Task keys are reconstructed from the serialized Dag.
- **Open redirect.** The redirect target is now a same-site relative path built from validated identifiers, which resolves both CodeQL findings.
- **Version gating.** Registration is gated to Airflow 3.1+. I confirmed `resolve_user_from_token` is absent on 3.0.6, and that 3.0.6's `get_user` reads only the bearer header (not the `_token` cookie), so cookie-based link navigation couldn't work there regardless — 3.1+ is the right floor.
- **Error hygiene.** Databricks failures now return a generic `502` with no upstream exception text, and task clearing happens only after a successful repair call.
- **Output encoding.** Confirmation-page values are HTML-escaped so identifiers can't become reflected markup.
- **`related: #52280`** added to the PR body.
Tests cover the GET-doesn't-mutate behaviour, confirmation-page escaping, server-side identifier resolution on the POST, and the generic 502. Appreciate the offer to prepare a patch branch — I think the current revision covers the items raised, but happy to iterate if anything still looks off.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3726071448
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,260 @@
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ from urllib.parse import urlencode
+
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _serialized_task_key(dag_id: str, task: Any) -> str:
+ """
+ Reproduce an operator's ``databricks_task_key`` from a serialized task.
+
+ Serialized tasks don't expose the operator property, so mirror its default: an explicit key
+ if one was set, else ``md5(dag_id__task_id)``.
+ """
+ import hashlib
+
+ return (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, run_id: str, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from sqlalchemy import select
+
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.get_task_instances(session=session) if ti.task_id in target_task_ids]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ request: Request,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Same-site relative action; SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = f"{request.url.path}?{request.url.query}"
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ # Redirect to a same-site relative path with the identifiers percent-encoded, so the
+ # target can never be steered to another host or scheme (CodeQL open-redirect).
+ return_url = f"/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
+
+ metadata = _read_launch_metadata(dag_id, run_id, launch_task_id, session)
+
+ if repair_all:
+ repaired_task_ids: list[str] = [] # resolved from live Databricks state below
+ else:
+ if task_id is None or not dag.has_task(task_id):
+ raise HTTPException(status_code=404, detail="Task not found in Dag.")
+ repaired_task_ids = [task_id]
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a
+ # generic error to the UI without leaking the upstream exception text.
+ try:
+ if repair_all:
+ hook = DatabricksHook(databricks_conn_id=metadata.conn_id)
+ task_keys = hook.get_run_failed_task_keys(metadata.run_id)
+ key_to_task_id = {_serialized_task_key(dag_id, t): t.task_id for t in dag.tasks}
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ else:
+ task_keys = [_serialized_task_key(dag_id, dag.get_task(repaired_task_ids[0]))]
+
+ 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)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/620)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4228601216
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +509,45 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks import (
+ DatabricksNotebookOperator,
+ DatabricksTaskOperator,
+ )
+
+ return [DatabricksNotebookOperator, DatabricksTaskOperator]
+
def get_link( # type: ignore[override] # Signature intentionally kept this way for Airflow 2.x compatibility
self,
operator,
dttm=None,
*,
ti_key: TaskInstanceKey | None = None,
) -> str:
+ if AIRFLOW_V_3_0_PLUS:
+ # The Airflow-3 repair backend is only registered on 3.1.1+ (see
+ # DatabricksWorkflowPlugin), so render no link below that.
Review Comment:
Why only on 3.1.1? I can't see the reason that is needed?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "Vamsi-klu (via GitHub)" <gi...@apache.org>.
Vamsi-klu commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5014524555
@PrakshiGoyal10 , Some CI signals are failing. PTAL and fix those. Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3918221750
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +382,33 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
Review Comment:
a) Is this really true
b) why do we need to not use XComOperatorLink? That is the important thing to include in the comment
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +382,33 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
+ # XComOperatorLink. Lazy import avoids a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks_workflow import (
+ _CreateDatabricksWorkflowOperator,
+ )
+
+ return [_CreateDatabricksWorkflowOperator]
+
def get_link( # type: ignore[override] # Signature intentionally kept this way for Airflow 2.x compatibility
self,
operator,
dttm=None,
*,
ti_key: TaskInstanceKey | None = None,
) -> str:
+ if AIRFLOW_V_3_0_PLUS:
+ if not AIRFLOW_V_3_1_PLUS or ti_key is None:
Review Comment:
This conditionall is hard to read. But why do we need to check both, isn't it enough to just do:
```suggestion
if ti_key is not None:
```
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,39 @@ def get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED", "MAXIMUM_CONCURRENT_RUNS_REACHED"}
+
+ # ``get_run_tasks`` returns one entry per attempt and is not ordered by attempt, so a
+ # retried or already-repaired task appears several times under the same ``task_key``. Keep
+ # only the latest attempt per key (sort by ``start_time``, same idiom as
+ # ``DatabricksTaskBaseOperator._get_current_databricks_task``) before judging its state, so
+ # the result never contains duplicate keys — Databricks rejects duplicates in
+ # ``rerun_tasks`` — and a task whose latest attempt succeeded is not reported as failed.
+ # Never-started sub-tasks omit ``start_time`` or send null; treat those as 0.
+ sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda task: task.get("start_time") or 0)
+ latest_by_key = {task["task_key"]: task for task in sorted_tasks}
+
+ failed_task_keys = []
+ for task_key, task in latest_by_key.items():
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
Review Comment:
Is there no way to do this filtering on the API side?
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py:
##########
@@ -49,11 +49,15 @@ class WorkflowRunMetadata:
:param run_id: The ID of the Databricks workflow run.
:param job_id: The ID of the Databricks workflow job.
:param conn_id: The connection ID used to connect to Databricks.
+ :param task_key_map: Airflow ``task_id`` → Databricks ``task_key`` for the workflow's tasks.
+ Optional and defaulted for backward compatibility with runs launched before it was added
+ (those XComs carry only conn_id/job_id/run_id).
"""
conn_id: str
job_id: int
run_id: int
+ task_key_map: dict[str, str] = field(default_factory=dict)
Review Comment:
Why did we (have/choose to) add this here? It wasn't needed in Airflow 2 -- is this really part of the same feature?
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
Review Comment:
I'm surprised we need to unquote here -- I would have thought FastAPI would do that for us
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
Review Comment:
Do we do anything with this object? Is it really worth an extra query to handle this case vs just hitting the "dag_run is None" path below.
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
Review Comment:
What in this block is actually depending on Airflow 3.1?
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
Review Comment:
Is there really not a function already to clear a list of task_ids? IT feels like we shouldn't have to write this
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
Review Comment:
This feels very very wrong. You shouldn't be parsing the HTTP headers yourself, nor making any assumptions what so ever about how the Auth Manager maps request to the user or their perms.
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
Review Comment:
Isn't there an XComModel.get_one()?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4228580941
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +385,40 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
Review Comment:
I think this is fine, but I also don't think the logic holds.
Given the plugin includes a custom app, couldn't the link could be hard-coded to "the URL of that custom repair endpoint with task_id etc pre-filled"?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4228624866
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
Review Comment:
The only place this appears to be used is to check that the task exists in the dag.
That would be much better done just looking for the TaskInstance row instead. It tells you the same, and doesn't need to deserialize a dag to do it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4248475684
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
Review Comment:
Follow-up: `find_relevant_relatives` doesn't exist on Airflow 3.1, the oldest version this backend supports, so on 3.1.x the repair succeeded on Databricks but the clear then 500'd. I found it end-to-end testing on 3.1.8. Switched to `dag.partial_subset(..., include_downstream=True, exclude_original=True)`, which is what the 3.1 core clear route uses, plus a DB-backed test for the clear. Verified the full repair → clear → resume flow on 3.1.8 against a real Databricks workspace.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "moomindani (via GitHub)" <gi...@apache.org>.
moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3877530721
##########
providers/databricks/tests/unit/databricks/plugins/test_databricks_workflow.py:
##########
@@ -251,21 +257,58 @@ def test_workflow_job_repair_single_failed_link_airflow2():
ti_key.run_id = "run_id"
ti_key.try_number = 1
- with patch(
- "airflow.providers.databricks.plugins.databricks_workflow.get_task_instance"
- ) as mock_get_task_instance:
- with patch(
+ with (
+ patch(
+ "airflow.providers.databricks.plugins.databricks_workflow.get_task_instance"
+ ) as mock_get_task_instance,
+ patch(
"airflow.providers.databricks.plugins.databricks_workflow.get_xcom_result"
- ) as mock_get_xcom_result:
- with patch(
- "airflow.providers.databricks.plugins.databricks_workflow.DagBag.get_dag"
- ) as mock_get_dag:
- mock_get_task_instance.return_value = Mock(key=ti_key)
- mock_get_xcom_result.return_value = Mock(conn_id="conn_id", run_id=1)
- mock_get_dag.return_value.get_task = Mock(return_value=Mock(task_id="task_id"))
+ ) as mock_get_xcom_result,
+ patch("airflow.providers.databricks.plugins.databricks_workflow.DagBag.get_dag") as mock_get_dag,
+ ):
+ mock_get_task_instance.return_value = Mock(key=ti_key)
+ mock_get_xcom_result.return_value = Mock(conn_id="conn_id", run_id=1)
+ mock_get_dag.return_value.get_task = Mock(return_value=Mock(task_id="task_id"))
- result = link.get_link(operator, ti_key=ti_key)
- assert result.startswith("http://localhost/repair_databricks_job")
+ result = link.get_link(operator, ti_key=ti_key)
+ assert result.startswith("http://localhost/repair_databricks_job")
+
+
+@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow < 3.0")
+@pytest.mark.skipif(
+ RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES, reason="Web plugin test doesn't work when not against sources"
+)
Review Comment:
Worth knowing that this test never runs, in CI or locally: the two guards cover every environment between them.
`RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES` is `True` whenever `USE_AIRFLOW_VERSION` is set *or* `<repo root>/airflow/__init__.py` is missing (`devel-common/src/tests_common/__init__.py:27`) — and in the monorepo layout that file no longer exists, so the constant is `True` even on a sources checkout. In the Airflow 2.11 compatibility job the first clause makes it `True` as well, and on Airflow 3 sources `AIRFLOW_V_3_0_PLUS` skips it first. Measured here: the whole `airflow2` selection reports `11 skipped`, this test among them.
That is structural rather than a defect in this PR — the legacy branch's helpers (`_get_dag`, `_get_dagrun`, the FAB views) are defined inside `if not AIRFLOW_V_3_0_PLUS:`, so nothing in the 2.x path can be exercised from an Airflow 3 checkout, and the neighbouring `*_airflow2` tests sit in the same position. No change requested. I am noting it so nobody later reads this test as a live regression guard for the fix.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5353216915
Pulled your patch onto the branch — thanks @Vamsi-klu, this is a clean fix.
Confirming what it does, verified locally against a deserialized Dag rather than mocks:
- Both repair-link classes now declare `operators`, so `_deserialize_operator` keeps the real link objects instead of replacing them with `XComOperatorLink` — @moomindani's suggested route, so `_build_repair_url` runs at request time and no URL is frozen into XCom.
- `_get_launch_task_id_v3` no longer calls `get_child_by_label`; it resolves the launch task via the `.launch` suffix, `upstream_task_ids`, then a group walk, so it works on `SerializedTaskGroup`.
- The redirect and the confirmation form action are same-site relative and now include the `[api] base_url` path prefix, so a non-root deployment doesn't 404 after repair.
- `start_time` missing/null is treated as oldest, and legacy repair-all now logs task keys it can't map back instead of dropping them silently.
The `SerializedDagModel.get_dag` round-trip test is the important part — it goes through the API server's own path and asserts a non-empty URL for a member task, so the "buttons silently stop rendering" break can't come back unseen. All 178 provider unit tests pass locally, ruff and mypy clean.
@moomindani — I think this addresses the two blockers plus the test you asked to treat as non-optional. PTAL when you have a chance.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on PR #69998:
URL: https://github.com/apache/airflow/pull/69998#issuecomment-5344218176
Thanks both — this is a great catch, and the `SerializedDagModel.get_dag` round-trip is exactly the kind of test I should have had; my link tests ran against a `Mock` operator and never exercised the deserialized path, so both blockers were invisible to them.
@Vamsi-klu — yes please, go ahead and take a pass at these on this PR; I'd rather build on your fix than have you open a competing one. A patch I can pull works well.
On approach: I don't have a strong preference between persisting the URLs into XCom (matches the job-run link precedent) and @moomindani's suggestion of declaring `operators` on the two link classes so the real link objects survive deserialization and `_build_repair_url` runs at request time. The latter is appealing because it avoids freezing `base_url` at execute time and actually exercises the branches already in the PR — but I'll defer to whichever you and @moomindani think is cleaner.
Appreciate the thorough reviews.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "moomindani (via GitHub)" <gi...@apache.org>.
moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3835874904
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +499,38 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
+ # XComOperatorLink. Lazy import avoids a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks import (
+ DatabricksNotebookOperator,
+ DatabricksTaskOperator,
+ )
+
+ return [DatabricksNotebookOperator, DatabricksTaskOperator]
Review Comment:
This makes the plugin link match every `DatabricksNotebookOperator` / `DatabricksTaskOperator`, and
plugin links take precedence over the operator's own serialized links ("If OperatorLinks with the same
name exists, Links via Plugin have higher precedence", `serialization/serialized_objects.py:1117-1120`).
That is what fixes the workflow case, but it also overrides `operators/databricks.py:1690-1697`, where
the operator deliberately picks its link set per case:
```python
else:
# Databricks does not support repair for non-workflow tasks, hence do not show the repair link.
self.operator_extra_links = (DatabricksJobRunLink(),)
```
On Airflow 3 that is harmless, and I measured it: a standalone `DatabricksNotebookOperator` gains a
`"Repair a single task"` entry that resolves to `None` — `_get_launch_task_id_v3` finds no launch task
and `get_link` returns `""` — and the UI drops null URLs, so only the `/links` response gains a null key.
On Airflow 2 it is not harmless. The provider still declares `apache-airflow>=2.11.0`
(`pyproject.toml:62`); 2.11's deserializer has the same plugin matching and the same precedence merge
(`serialized_objects.py:1266-1272` and `:1310-1315`); and the pre-3.0 branch of `get_link` degrades
differently from the 3.x branch above. For a top-level task `operator.task_group` is the Dag's root
group, so the `if not task_group` guard does not fire, and `get_launch_task_id` then raises
`AirflowException("No launch task can be found in the task group.")`. Airflow 2.11's extra-links view
catches only `ValueError` (`airflow/www/views.py:3249-3252`), so this surfaces as a 500 behind a button
the UI has already rendered — on the plain notebook operator, which is the common case.
I cannot run 2.11 here, so that last chain is code reading rather than measurement, but every step of it
is in the released 2.11 source. Suggested fix in the next comment.
`WorkflowJobRepairAllFailedLink` does not have this problem — `_CreateDatabricksWorkflowOperator` only
ever exists inside a workflow task group.
---
Drafted-by: Claude Code (Opus 5); reviewed by @moomindani before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +499,38 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def operators(self):
+ # Declared so deserialization keeps this plugin link instead of replacing it with
+ # XComOperatorLink. Lazy import avoids a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks import (
+ DatabricksNotebookOperator,
+ DatabricksTaskOperator,
+ )
+
+ return [DatabricksNotebookOperator, DatabricksTaskOperator]
+
def get_link( # type: ignore[override] # Signature intentionally kept this way for Airflow 2.x compatibility
self,
operator,
dttm=None,
*,
ti_key: TaskInstanceKey | None = None,
) -> str:
+ if AIRFLOW_V_3_0_PLUS:
+ if not AIRFLOW_V_3_1_PLUS or ti_key is None:
+ # The Airflow-3 repair backend requires 3.1+ (see DatabricksWorkflowPlugin).
+ return ""
+ launch_task_id = _get_launch_task_id_v3(operator, ti_key)
+ if not launch_task_id:
+ return ""
Review Comment:
This is the graceful degradation the Airflow 2 branch lacks. Now that the link is attached to
non-workflow tasks as well, the cheapest fix is to mirror this shape below: when the launch task cannot
be resolved, return `""` instead of letting `get_launch_task_id(task_group)` raise — i.e. around
```python
if ".launch" not in ti_key.task_id:
launch_task_id = get_launch_task_id(task_group)
```
catch the "no launch task in this group" case and return `""`, so a standalone Databricks task renders
no URL rather than a 500. That keeps the behaviour of the two branches aligned and leaves the workflow
case untouched.
---
Drafted-by: Claude Code (Opus 5); reviewed by @moomindani before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub)" <gi...@apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4230037127
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -374,13 +385,40 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
name = "Repair All Failed Tasks"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
Review Comment:
Agreed, the comment's reasoning no longer held: the links stopped reading run metadata at render time a few rounds ago, so `operators` was only stopping the swap to `XComOperatorLink`. I've switched to your suggestion. At execute time, the launch task and each workflow task now push the pre-built `/databricks/workflow/repair/...` URL to XCom, the same way `WorkflowJobRunLink` already stores the job-run link. The URL holds Airflow identifiers only. The API server renders it through the stock `XComOperatorLink`, so the `operators` overrides, the render-time launch-task resolution and the Airflow 2 standalone-task guard they forced are all gone. Trade-off: the links only appear for runs executed after the provider upgrade (noted in the docs).
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +509,45 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def operators(self):
+ # On Airflow 3 a plugin extra link that declares no ``operators`` is replaced at
+ # deserialization by an ``XComOperatorLink`` that just returns a URL the task stored in
+ # XCom under ``xcom_key``. This link stores no such URL — it builds the URL at request time
+ # in ``get_link`` from the run's XCom metadata — so it must survive as the real object.
+ # Declaring the operators it applies to keeps it from being swapped out. Lazy import avoids
+ # a circular import with the operator module.
+ from airflow.providers.databricks.operators.databricks import (
+ DatabricksNotebookOperator,
+ DatabricksTaskOperator,
+ )
+
+ return [DatabricksNotebookOperator, DatabricksTaskOperator]
+
def get_link( # type: ignore[override] # Signature intentionally kept this way for Airflow 2.x compatibility
self,
operator,
dttm=None,
*,
ti_key: TaskInstanceKey | None = None,
) -> str:
+ if AIRFLOW_V_3_0_PLUS:
+ # The Airflow-3 repair backend is only registered on 3.1.1+ (see
+ # DatabricksWorkflowPlugin), so render no link below that.
Review Comment:
It's for auth on a link clicked in the browser. Extra links open as plain navigations, so the request carries only the `_token` cookie, not an `Authorization` header. `get_user` (behind `requires_access_dag`) only started reading that cookie in 3.1.1, in #56677. In 3.1.0 it resolves the token from the `oauth2_scheme`/`bearer_scheme` dependencies only, and in 3.0.x from `oauth2_scheme` only, so every click would 401 (apart from the all-admins `SimpleAllAdminMiddleware` dev mode). Rather than parse the cookie in the provider, which is what we moved away from in the previous round, the backend and links are only enabled from 3.1.1. I've also moved the reason to where links are stored and put it in the docs.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +574,311 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server, and the
+# repair links (below) build URLs that point at it.
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+def _get_launch_task_id_v3(operator: BaseOperator, ti_key: TaskInstanceKey) -> str | None:
+ """
+ Resolve the ``task_id`` of the workflow's launch task for an extra-link render.
+
+ Works on both live operators and deserialized ones. ``SerializedTaskGroup`` has no
+ ``get_child_by_label``, so this never calls it. Returns ``None`` when the launch task
+ cannot be found (so the link is not rendered).
+ """
+ if ti_key.task_id.endswith(".launch"):
+ return ti_key.task_id
+
+ for tid in getattr(operator, "upstream_task_ids", ()) or ():
+ if tid.endswith(".launch"):
+ return tid
+
+ task_group = getattr(operator, "task_group", None)
+ while task_group is not None:
+ child_id = getattr(task_group, "child_id", None)
+ children = getattr(task_group, "children", None)
+ if callable(child_id) and children is not None:
+ launch_id = child_id("launch")
+ if launch_id in children:
+ child = children[launch_id]
+ return getattr(child, "task_id", launch_id)
+ task_group = getattr(task_group, "parent_group", None)
+ return None
+
+
+if AIRFLOW_V_3_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
+ from airflow.api_fastapi.core_api.security import resolve_user_from_token
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ async def _resolve_request_user(request: Request):
+ """Authenticate via the bearer header (UI XHR) or the ``_token`` cookie (link navigation)."""
+ token = None
+ auth_header = request.headers.get("Authorization", "")
+ if auth_header.lower().startswith("bearer "):
+ token = auth_header.split(" ", 1)[1]
+ if not token:
+ token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ # resolve_user_from_token raises HTTP 401 for a missing/invalid token.
+ return await resolve_user_from_token(token)
+
+ async def _require_dag_run_edit(dag_id: str, request: Request):
+ from airflow.api_fastapi.app import get_auth_manager
+
+ user = await _resolve_request_user(request)
+ authorized = get_auth_manager().is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.RUN,
+ details=DagDetails(id=dag_id),
+ user=user,
+ )
+ if not authorized:
+ raise HTTPException(status_code=403, detail="Not authorized to repair runs of this Dag.")
+ return user
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
+ import hashlib
+
+ return hashlib.md5(f"{dag_id}__{task_id}".encode()).hexdigest()
+
+ def _read_launch_metadata(dag_id: str, run_id: str, launch_task_id: str, session) -> Any:
+ """
+ Read the launch task's trusted ``WorkflowRunMetadata`` XCom (conn_id, job_id, run_id).
+
+ The Databricks connection and run id come from here — never from the request — so a crafted
+ link cannot redirect the repair at an arbitrary connection or Databricks run.
+ """
+ from airflow.models.xcom import XComModel
+ from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+ result = session.scalars(
+ XComModel.get_many(
+ run_id=run_id,
+ key="return_value",
+ task_ids=launch_task_id,
+ dag_ids=dag_id,
+ limit=1,
+ )
+ ).first()
+ if result is None:
+ raise HTTPException(status_code=404, detail="Databricks workflow run metadata not found.")
+ return WorkflowRunMetadata(**XComModel.deserialize_value(result))
+
+ def _clear_repaired_and_downstream(
+ dag, dag_run: DagRun, task_ids: list[str], session, logger: logging.Logger
+ ) -> None:
+ """
+ Clear the repaired tasks' instances and their downstream instances for this run.
+
+ Runs inside the API server (the DB-facing component), so clearing the repaired tasks plus
+ their downstream lets the upstream-failed dependents resume deterministically when the
+ repaired Databricks sub-runs succeed — without clearing the whole Dag.
+ """
+ from airflow.models.taskinstance import clear_task_instances
+
+ target_task_ids: set[str] = set(task_ids)
+ for task_id in task_ids:
+ target_task_ids.update(dag.get_task(task_id).get_flat_relative_ids(upstream=False))
+
+ tis_to_clear = [
+ ti for ti in dag_run.get_task_instances(session=session) if ti.task_id in target_task_ids
+ ]
+ logger.info("Clearing %s task instances after Databricks repair", len(tis_to_clear))
+ clear_task_instances(tis_to_clear, session)
+
+ def _repair_confirmation_page(dag_id: str, run_id: str, action: str, summary: str) -> HTMLResponse:
+ """Render the read-only confirmation page whose form issues the state-changing POST."""
+ return HTMLResponse(
+ "<!doctype html><html><head><title>Repair Databricks workflow</title></head><body>"
+ "<h2>Repair Databricks workflow tasks</h2>"
+ f"<p>Dag <b>{escape(dag_id)}</b>, run <b>{escape(run_id)}</b>.</p>"
+ f"<p>{escape(summary)}</p>"
+ f'<form method="post" action="{escape(action)}">'
+ '<button type="submit">Repair</button></form>'
+ "</body></html>"
+ )
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow_confirm(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Render a read-only confirmation page; the repair itself happens on the POST below."""
+ run_id = unquote(run_id)
+ summary = (
+ "This will repair all failed tasks of the run and resume their downstream tasks."
+ if repair_all
+ else f"This will repair task '{task_id}' and resume its downstream tasks."
+ )
+ # Rebuild a same-site relative POST target from the validated identifiers rather than
+ # echoing the request URL. SameSite=Lax on the auth cookie means a cross-site POST cannot
+ # carry it, so moving the mutation to POST is what protects it from CSRF.
+ action = _build_repair_url(dag_id, run_id, launch_task_id, repair_all=repair_all, task_id=task_id)
+ return _repair_confirmation_page(dag_id, run_id, action, summary)
+
+ @repair_app.post("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ task_id: str | None = None,
+ repair_all: bool = False,
+ _user=Depends(_require_dag_run_edit),
+ ):
+ """Repair failed Databricks tasks for a workflow run and resume the Airflow run."""
+ run_id = unquote(run_id)
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ dag = SerializedDagModel.get_dag(dag_id, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail="Dag not found.")
Review Comment:
Done. The task check now looks up the run's `TaskInstance` rows, and repair-all maps the Databricks keys back over those same rows instead of `dag.tasks`. The Dag is still deserialized once, but only for what needs its structure: `find_relevant_relatives` + `dag.clear` for the downstream clear. It's now `DBDagBag().get_dag_for_run(dag_run)`, i.e. the version the run was created from (matching the core clear route) rather than the latest serialized Dag. A task that isn't in the run now 404s before the Dag is loaded or Databricks is called.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
Posted by "PrakshiGoyal10 (via GitHub) via commits" <co...@airflow.apache.org>.
PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r4451886045
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, task_id=task_id)
+ ti.xcom_push(key=key, value=link)
+ logger.info("Stored Databricks repair link in XCom: %s", link)
+ except Exception as e:
+ logger.warning("Failed to store Databricks repair link: %s", e)
+
+
+def _get_link_from_xcom(ti_key: TaskInstanceKey | None, key: str, logger: Logger) -> str:
+ if ti_key is None:
+ return ""
+ try:
+ return XCom.get_value(ti_key=ti_key, key=key) or ""
+ except Exception as e:
+ logger.warning("Failed to retrieve Databricks link from XCom: %s", e)
+ return ""
+
+
class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
"""Constructs a link to send a request to repair all failed tasks in the Databricks workflow."""
name = "Repair All Failed Tasks"
+ @property
+ def xcom_key(self) -> str:
+ """XCom key where the link is stored during task execution."""
+ return "databricks_repair_all_failed_link"
Review Comment:
Not meant to be settable. Changed to `xcom_key: ClassVar[str] = "databricks_repair_all_failed_link"` as suggested. Nothing instantiates the class to read it any more, because the function that did is gone (see the `:408` thread).
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -471,13 +531,21 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
name = "Repair a single task"
+ @property
+ def xcom_key(self) -> str:
+ """XCom key where the link is stored during task execution."""
+ return "databricks_repair_single_task_link"
Review Comment:
Done, same `ClassVar` change here.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
Review Comment:
`store_databricks_repair_link` is removed now, so this branch no longer exists (see the `:408` thread). For the record, the two arms did use different keys (`..._all_failed_link` vs `..._single_task_link`), one per link class.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -369,18 +380,67 @@ def store_databricks_job_run_link(
logger.warning("Failed to store Databricks job run link: %s", e)
+def store_databricks_repair_link(
+ context: Context,
+ launch_task_id: str,
+ logger: Logger,
+ *,
+ task_id: str | None = None,
+) -> None:
+ """
+ Store a repair link in XCom during task execution.
+
+ Stores the single-task repair link for ``task_id``, or the repair-all link when ``task_id`` is
+ omitted. The API server renders it from XCom like any other stored link, so it needs no
+ Databricks code to do so.
+ """
+ if not AIRFLOW_V_3_1_1_PLUS:
+ return # The repair backend is only registered on 3.1.1+ (see DatabricksWorkflowPlugin).
+
+ try:
+ ti = context["ti"]
+ if task_id is None:
+ key = WorkflowJobRepairAllFailedLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, repair_all=True)
+ else:
+ key = WorkflowJobRepairSingleTaskLink().xcom_key
+ link = _build_repair_url(ti.dag_id, ti.run_id, launch_task_id, task_id=task_id)
+ ti.xcom_push(key=key, value=link)
Review Comment:
You're right, there was no good reason. I'd missed that the task runner's `finalize()` already calls each link's `get_link()` on the worker and pushes the result under `xcom_key`. That's been true since 3.0.0. My explicit push stored the same value a second time.
I've removed `store_databricks_repair_link` and the two operator call sites. On Airflow 3 (3.1.1+), `get_link` now builds the URL directly on the worker:
- **Repair all:** the operator is the launch task, so it uses `ti_key.task_id`.
- **Single task:** it uses `get_launch_task_id(operator.task_group)`, the same lookup as the Airflow 2 path.
The API server still renders the link from XCom, but there's now a single write. The round-trip test now mirrors `finalize`: it pushes `get_link()` under `xcom_key`, then reads the link back through the serialized Dag's `get_extra_links`.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +583,304 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way
return url_for("RepairDatabricksTasks.repair", **query_params)
+# Airflow-3 repair backend. Flask-AppBuilder was dropped in Airflow 3, so the repair
+# action is re-implemented as a FastAPI sub-application mounted on the API server. The workflow
+# tasks store the repair links pointing at it in XCom (``store_databricks_repair_link``).
+REPAIR_URL_PREFIX = "/databricks/workflow/repair"
+
+
+def _build_repair_url(
+ dag_id: str,
+ run_id: str,
+ launch_task_id: str,
+ *,
+ repair_all: bool = False,
+ task_id: str | None = None,
+) -> str:
+ """
+ Build the URL to the Airflow-3 FastAPI repair confirmation page for a workflow run.
+
+ The URL carries only Airflow identifiers: the run's launch ``task_id`` (from which the
+ endpoint reads the trusted ``WorkflowRunMetadata`` XCom) and, for a single-task repair, the
+ target ``task_id``. The Databricks connection, run id, and task keys are never placed in the
+ link — the endpoint derives them server-side, so the request cannot point the repair at an
+ arbitrary connection or Databricks run.
+
+ Called on the worker at task execution, so the path prefix comes from the worker's
+ ``[api] base_url``.
+ """
+ query: dict[str, Any] = {"launch_task_id": launch_task_id}
+ if repair_all:
+ query["repair_all"] = "true"
+ if task_id:
+ query["task_id"] = task_id
+
+ # Same-site relative path only. Using the full ``[api] base_url`` (scheme + host) would make
+ # the confirmation POST cross-origin when that setting names a different domain than the UI,
+ # and SameSite=Lax would then withhold the auth cookie.
+ return (
+ f"{_api_root_path()}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/"
+ f"{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+def _api_root_path() -> str:
+ """Path prefix from ``[api] base_url``, or empty when the API is mounted at the origin root."""
+ return urlsplit(conf.get("api", "base_url", fallback="") or "").path.rstrip("/")
+
+
+def _ui_run_path(dag_id: str, run_id: str) -> str:
+ """Same-site relative path to the Dag run in the UI, including the API root path if set."""
+ return f"{_api_root_path()}/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+
+if AIRFLOW_V_3_1_1_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException
+ from fastapi.responses import HTMLResponse, RedirectResponse
+ from markupsafe import escape
+
+ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
+ from airflow.api_fastapi.core_api.security import requires_access_dag
+
+ repair_app = FastAPI(
+ title="Databricks Workflow Repair",
+ description="Repair failed tasks of a Databricks workflow run from Airflow.",
+ )
+
+ # Authenticate (bearer header or the UI's ``_token`` cookie) and authorize Dag-run edit access
+ # through the API server's own dependency, rather than parsing the request or assuming how the
+ # auth manager maps requests to users. ``requires_access_dag`` reads the ``dag_id`` path param
+ # and calls the auth manager; ``get_user`` behind it already honours both the bearer header and
+ # the cookie, so a repair link clicked in the browser is authorized. Bound to a module-level
+ # name so tests can override just this dependency.
+ _require_dag_run_edit = requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)
+
+ def _task_id_to_key(dag_id: str, task_id: str, task_key_map: dict[str, str]) -> str:
+ """
+ Resolve a task's Databricks ``task_key`` from the launch task's trusted key map.
+
+ An explicit ``databricks_task_key`` does not survive Dag serialization, so the serialized
+ task can't be trusted to reproduce it. The launch task captured the real keys from the live
+ operators into ``task_key_map``. Runs launched before that map existed fall back to the
+ operator's default derivation, ``md5(dag_id__task_id)`` — correct for any task that did not
+ set an explicit key (the common case).
+ """
+ mapped = task_key_map.get(task_id)
+ if mapped:
+ return mapped
Review Comment:
Done.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @PrakshiGoyal10 before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org