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