moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3974626224
##########
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:
One correction to the premise here, since it moves where the gate can sit
rather than whether it should.
`get_user` picks up the `_token` cookie from **3.1.1**, not 3.1.0. In 3.1.0
it takes only the two header schemes:
```python
async def get_user(
oauth_token: str | None = Depends(oauth2_scheme),
bearer_credentials: HTTPAuthorizationCredentials | None =
Depends(bearer_scheme),
) -> BaseUser:
```
The `token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)` branch appears
in 3.1.1 and is present in every release I checked after it (3.1.2-3.1.5,
3.1.8, 3.2.x, 3.3.x).
`AIRFLOW_V_3_1_PLUS` is `>= (3, 1, 0)`, so on exactly 3.1.0 a repair link
clicked in the browser carries only the cookie, `requires_access_dag` resolves
no token, and the click 401s — which the hand-rolled resolver did handle. So
the options are to gate at 3.1.1 (patch-level gates are already used elsewhere,
e.g. `AIRFLOW_V_3_1_9_PLUS` in the celery provider) or to keep 3.1 and say in
that comment that browser clicks need 3.1.1+.
Relaxing to `AIRFLOW_V_3_0_PLUS` looks closed off for the same reason:
3.0.6's `get_user` reads only `oauth2_scheme`, so a cookie-only navigation
never authenticates there either.
---
Drafted-by: Claude Code (Opus 5); reviewed by @moomindani 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]
