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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
