You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org> on 2026/08/05 13:54:46 UTC
Re: [PR] Enable repair of failed Databricks Workflow tasks from Airflow 3 UI [airflow]
github-advanced-security[bot] commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3721149126
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,260 @@ 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.
+ """
+ from urllib.parse import urlencode
+
+ 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
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+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.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+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 _serialized_task_key(dag_id: str, task: Any) -> str:
+ """
+ Reproduce an operator's ``databricks_task_key`` from a serialized task.
+
+ Serialized tasks don't expose the operator property, so mirror its default: an explicit key
+ if one was set, else ``md5(dag_id__task_id)``.
+ """
+ import hashlib
+
+ return (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.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' 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 sqlalchemy import select
+
+ 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))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.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,
+ request: Request,
+ 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."
+ )
+ # Same-site relative action; 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 = f"{request.url.path}?{request.url.query}"
+ 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)
+
+ # Redirect to a same-site relative path with the identifiers percent-encoded, so the
+ # target can never be steered to another host or scheme (CodeQL open-redirect).
+ return_url = f"/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+ 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.")
+
+ metadata = _read_launch_metadata(dag_id, run_id, launch_task_id, session)
+
+ if repair_all:
+ repaired_task_ids: list[str] = [] # resolved from live Databricks state below
+ else:
+ if task_id is None or not dag.has_task(task_id):
+ raise HTTPException(status_code=404, detail="Task not found in Dag.")
+ repaired_task_ids = [task_id]
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a
+ # generic error to the UI without leaking the upstream exception text.
+ try:
+ if repair_all:
+ hook = DatabricksHook(databricks_conn_id=metadata.conn_id)
+ task_keys = hook.get_run_failed_task_keys(metadata.run_id)
+ key_to_task_id = {_serialized_task_key(dag_id, t): t.task_id for t in dag.tasks}
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ else:
+ task_keys = [_serialized_task_key(dag_id, dag.get_task(repaired_task_ids[0]))]
+
+ if not task_keys:
+ log.info("No failed Databricks tasks to repair for run %s", metadata.run_id)
+ return RedirectResponse(return_url, status_code=303)
+
+ log.info("Repairing Databricks run %s tasks %s", metadata.run_id, task_keys)
+ _repair_task(
+ databricks_conn_id=metadata.conn_id,
+ databricks_run_id=metadata.run_id,
+ tasks_to_repair=task_keys,
+ logger=log,
+ )
+ except HTTPException:
+ raise
+ except Exception:
+ log.exception("Databricks repair failed for run %s", metadata.run_id)
+ raise HTTPException(status_code=502, detail="Databricks repair request failed.")
+
+ # Clear only after a successful repair call, so a failed repair leaves state untouched.
+ _clear_repaired_and_downstream(dag, run_id, repaired_task_ids, session, log)
+ session.commit()
+
+ return RedirectResponse(return_url, status_code=303)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/3)
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,260 @@ 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.
+ """
+ from urllib.parse import urlencode
+
+ 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
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return (
+ f"{base_url}{REPAIR_URL_PREFIX}/{quote(dag_id, safe='')}/{quote(run_id, safe='')}?{urlencode(query)}"
+ )
+
+
+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.
+
+ The link only needs to name the launch task; the repair endpoint reads that task's trusted
+ ``WorkflowRunMetadata`` XCom server-side. Returns ``None`` when the operator is not part of a
+ Databricks workflow task group (so the link is not rendered).
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" in ti_key.task_id:
+ return ti_key.task_id
+ return get_launch_task_id(task_group)
+
+
+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 _serialized_task_key(dag_id: str, task: Any) -> str:
+ """
+ Reproduce an operator's ``databricks_task_key`` from a serialized task.
+
+ Serialized tasks don't expose the operator property, so mirror its default: an explicit key
+ if one was set, else ``md5(dag_id__task_id)``.
+ """
+ import hashlib
+
+ return (
+ getattr(task, "databricks_task_key", None)
+ or hashlib.md5(f"{dag_id}__{task.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' 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 sqlalchemy import select
+
+ 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))
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id)).one()
+ tis_to_clear = [ti for ti in dr.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,
+ request: Request,
+ 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."
+ )
+ # Same-site relative action; 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 = f"{request.url.path}?{request.url.query}"
+ 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)
+
+ # Redirect to a same-site relative path with the identifiers percent-encoded, so the
+ # target can never be steered to another host or scheme (CodeQL open-redirect).
+ return_url = f"/dags/{quote(dag_id, safe='')}/runs/{quote(run_id, safe='')}"
+
+ 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.")
+
+ metadata = _read_launch_metadata(dag_id, run_id, launch_task_id, session)
+
+ if repair_all:
+ repaired_task_ids: list[str] = [] # resolved from live Databricks state below
+ else:
+ if task_id is None or not dag.has_task(task_id):
+ raise HTTPException(status_code=404, detail="Task not found in Dag.")
+ repaired_task_ids = [task_id]
+
+ # Databricks API calls can fail (e.g. expired/invalid connection token); surface a
+ # generic error to the UI without leaking the upstream exception text.
+ try:
+ if repair_all:
+ hook = DatabricksHook(databricks_conn_id=metadata.conn_id)
+ task_keys = hook.get_run_failed_task_keys(metadata.run_id)
+ key_to_task_id = {_serialized_task_key(dag_id, t): t.task_id for t in dag.tasks}
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in key_to_task_id]
+ else:
+ task_keys = [_serialized_task_key(dag_id, dag.get_task(repaired_task_ids[0]))]
+
+ if not task_keys:
+ log.info("No failed Databricks tasks to repair for run %s", metadata.run_id)
+ return RedirectResponse(return_url, status_code=303)
Review Comment:
## CodeQL / URL redirection from remote source
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
Untrusted URL redirection depends on a [user-provided value](1).
Untrusted URL redirection depends on a [user-provided value](2).
[Show more details](https://github.com/apache/airflow/security/code-scanning/31)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org