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