github-advanced-security[bot] commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3608706108
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,193 @@
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,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ *,
+ repair_all: bool = False,
+ tasks_to_repair: list[str] | None = None,
+) -> str:
+ """Build the URL to the Airflow-3 FastAPI repair endpoint for a workflow
run."""
+ from urllib.parse import quote, urlencode
+
+ from airflow.configuration import conf
+
+ query: dict[str, Any] = {
+ "databricks_conn_id": databricks_conn_id,
+ "databricks_run_id": databricks_run_id,
+ }
+ if repair_all:
+ query["repair_all"] = "true"
+ if tasks_to_repair:
+ query["tasks_to_repair"] = ",".join(tasks_to_repair)
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return f"{base_url}{REPAIR_URL_PREFIX}/{dag_id}/{quote(run_id,
safe='')}?{urlencode(query)}"
+
+
+def _get_launch_metadata_v3(operator: BaseOperator, ti_key: TaskInstanceKey)
-> Any:
+ """
+ Read the launch task's ``WorkflowRunMetadata`` (conn_id, job_id, run_id)
from XCom.
+
+ Uses only the public XCom API keyed by ``ti_key`` — no metadata-DB access
— which is the
+ only mechanism available to extra links on Airflow 3.
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" not in ti_key.task_id:
+ launch_task_id = get_launch_task_id(task_group)
+ ti_key = _get_launch_task_key(ti_key, task_id=launch_task_id)
+ result = XCom.get_value(ti_key=ti_key, key="return_value")
+ if not result:
+ return None
+ from airflow.providers.databricks.operators.databricks_workflow import
WorkflowRunMetadata
+
+ return WorkflowRunMetadata(**result)
+
+
+if AIRFLOW_V_3_0_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import RedirectResponse
+
+ 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 _clear_repaired_and_downstream(
+ dag_id: str, run_id: str, task_keys: list[str], 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.
+ """
+ import hashlib
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.models.taskinstance import clear_task_instances
+ 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=f"Dag {dag_id} not
found.")
+
+ # Serialized tasks don't expose the operator's
``databricks_task_key`` property, so
+ # reproduce its default: an explicit key if present, else
md5(dag_id__task_id).
+ key_to_task_id = {}
+ for task in dag.tasks:
+ task_key = (
+ getattr(task, "databricks_task_key", None)
+ or
hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+ key_to_task_id[task_key] = task.task_id
+
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in
key_to_task_id]
+ target_task_ids: set[str] = set(repaired_task_ids)
+ for task_id in repaired_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_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)
+ session.commit()
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ tasks_to_repair: 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)
+ from airflow.configuration import conf
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return_url = f"{base_url}/dags/{dag_id}/runs/{run_id}"
+
+ # Databricks API calls can fail (e.g. expired/invalid connection
token); surface a clear
+ # error to the UI instead of a bare 500.
+ try:
+ hook = DatabricksHook(databricks_conn_id=databricks_conn_id)
+ if repair_all:
+ task_keys = hook.get_run_failed_task_keys(databricks_run_id)
+ elif tasks_to_repair:
+ task_keys = tasks_to_repair.split(",")
+ else:
+ task_keys = []
+
+ if not task_keys:
+ log.info("No failed Databricks tasks to repair for run %s",
databricks_run_id)
+ return RedirectResponse(return_url, status_code=303)
+
+ log.info("Repairing Databricks run %s tasks %s",
databricks_run_id, task_keys)
+ _repair_task(
+ databricks_conn_id=databricks_conn_id,
+ databricks_run_id=databricks_run_id,
+ tasks_to_repair=task_keys,
+ logger=log,
+ )
+ except HTTPException:
+ raise
+ except Exception as e:
+ log.exception("Databricks repair failed for run %s",
databricks_run_id)
+ raise HTTPException(status_code=502, detail=f"Databricks repair
request failed: {e}") from e
+
+ _clear_repaired_and_downstream(dag_id, run_id, task_keys, log)
+ 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).
[Show more
details](https://github.com/apache/airflow/security/code-scanning/618)
##########
providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py:
##########
@@ -515,6 +547,193 @@
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,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ *,
+ repair_all: bool = False,
+ tasks_to_repair: list[str] | None = None,
+) -> str:
+ """Build the URL to the Airflow-3 FastAPI repair endpoint for a workflow
run."""
+ from urllib.parse import quote, urlencode
+
+ from airflow.configuration import conf
+
+ query: dict[str, Any] = {
+ "databricks_conn_id": databricks_conn_id,
+ "databricks_run_id": databricks_run_id,
+ }
+ if repair_all:
+ query["repair_all"] = "true"
+ if tasks_to_repair:
+ query["tasks_to_repair"] = ",".join(tasks_to_repair)
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return f"{base_url}{REPAIR_URL_PREFIX}/{dag_id}/{quote(run_id,
safe='')}?{urlencode(query)}"
+
+
+def _get_launch_metadata_v3(operator: BaseOperator, ti_key: TaskInstanceKey)
-> Any:
+ """
+ Read the launch task's ``WorkflowRunMetadata`` (conn_id, job_id, run_id)
from XCom.
+
+ Uses only the public XCom API keyed by ``ti_key`` — no metadata-DB access
— which is the
+ only mechanism available to extra links on Airflow 3.
+ """
+ task_group = operator.task_group
+ if not task_group:
+ return None
+ if ".launch" not in ti_key.task_id:
+ launch_task_id = get_launch_task_id(task_group)
+ ti_key = _get_launch_task_key(ti_key, task_id=launch_task_id)
+ result = XCom.get_value(ti_key=ti_key, key="return_value")
+ if not result:
+ return None
+ from airflow.providers.databricks.operators.databricks_workflow import
WorkflowRunMetadata
+
+ return WorkflowRunMetadata(**result)
+
+
+if AIRFLOW_V_3_0_PLUS:
+ from fastapi import Depends, FastAPI, HTTPException, Request
+ from fastapi.responses import RedirectResponse
+
+ 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 _clear_repaired_and_downstream(
+ dag_id: str, run_id: str, task_keys: list[str], 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.
+ """
+ import hashlib
+
+ from sqlalchemy import select
+
+ from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.models.taskinstance import clear_task_instances
+ 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=f"Dag {dag_id} not
found.")
+
+ # Serialized tasks don't expose the operator's
``databricks_task_key`` property, so
+ # reproduce its default: an explicit key if present, else
md5(dag_id__task_id).
+ key_to_task_id = {}
+ for task in dag.tasks:
+ task_key = (
+ getattr(task, "databricks_task_key", None)
+ or
hashlib.md5(f"{dag_id}__{task.task_id}".encode()).hexdigest()
+ )
+ key_to_task_id[task_key] = task.task_id
+
+ repaired_task_ids = [key_to_task_id[k] for k in task_keys if k in
key_to_task_id]
+ target_task_ids: set[str] = set(repaired_task_ids)
+ for task_id in repaired_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_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)
+ session.commit()
+
+ @repair_app.get("/{dag_id}/{run_id}")
+ def repair_databricks_workflow(
+ dag_id: str,
+ run_id: str,
+ databricks_conn_id: str,
+ databricks_run_id: int,
+ tasks_to_repair: 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)
+ from airflow.configuration import conf
+
+ base_url = conf.get("api", "base_url", fallback="").rstrip("/")
+ return_url = f"{base_url}/dags/{dag_id}/runs/{run_id}"
+
+ # Databricks API calls can fail (e.g. expired/invalid connection
token); surface a clear
+ # error to the UI instead of a bare 500.
+ try:
+ hook = DatabricksHook(databricks_conn_id=databricks_conn_id)
+ if repair_all:
+ task_keys = hook.get_run_failed_task_keys(databricks_run_id)
+ elif tasks_to_repair:
+ task_keys = tasks_to_repair.split(",")
+ else:
+ task_keys = []
+
+ if not task_keys:
+ log.info("No failed Databricks tasks to repair for run %s",
databricks_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).
[Show more
details](https://github.com/apache/airflow/security/code-scanning/617)
--
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]
