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