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