diff --git a/providers/databricks/docs/operators/workflow.rst b/providers/databricks/docs/operators/workflow.rst index 42cb0b20f9666..d62941087ad9a 100644 --- a/providers/databricks/docs/operators/workflow.rst +++ b/providers/databricks/docs/operators/workflow.rst @@ -68,3 +68,22 @@ To minimize update conflicts, we recommend that you keep parameters in the ``not ``DatabricksWorkflowTaskGroup`` and not in the ``DatabricksNotebookOperator`` whenever possible. This is because, tasks in the ``DatabricksWorkflowTaskGroup`` are passed in on the job trigger time and do not modify the job definition. + +Repairing a failed workflow run +------------------------------- + +When a Databricks Workflow run fails, each task instance of the task group exposes repair links in the +Airflow UI: + +* **Repair a single task** (on a notebook/task operator) re-runs that one Databricks task. +* **Repair All Failed Tasks** (on the ``launch`` task) re-runs every failed task of the run. + +Clicking a repair link calls the Databricks `repair_run `_ +API on the existing run (continuing the repair chain via ``latest_repair_id``) and reruns the dependent +tasks, then clears the corresponding Airflow task instances and their downstream tasks so the run resumes +without having to clear the whole Dag. + +On Airflow 3 the repair action is served by a FastAPI endpoint registered by the +``DatabricksWorkflowPlugin``; the set of failed tasks is resolved from the live Databricks run state. On +Airflow 2 it is served by the legacy Flask-AppBuilder view. In both cases the repair links are +authorized with Dag-run edit access. diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py index d780816f9faac..709ef6d02c8fd 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py @@ -563,6 +563,28 @@ 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"} + failed_task_keys = [] + for task in self.get_run_tasks(run_id): + state = task.get("state", {}) + if ( + state.get("result_state") in failed_result_states + or state.get("life_cycle_state") == "INTERNAL_ERROR" + ): + failed_task_keys.append(task["task_key"]) + return failed_task_keys + def get_run(self, run_id: int) -> dict[str, Any]: """ Retrieve run information. diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index 2d17a21587f2a..ed339fef53079 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -1688,16 +1688,10 @@ def __init__( super().__init__(**kwargs) if self._databricks_workflow_task_group is not None: - # Conditionally set operator_extra_links based on Airflow version. In Airflow 3, only show the job run link. - # In Airflow 2, show the job run link and the repair link. - # TODO: Once we expand the plugin functionality in Airflow 3.1, this can be re-evaluated on how to handle the repair link. - if AIRFLOW_V_3_0_PLUS: - self.operator_extra_links = (WorkflowJobRunLink(),) - else: - self.operator_extra_links = ( - WorkflowJobRunLink(), - WorkflowJobRepairSingleTaskLink(), - ) + self.operator_extra_links = ( + WorkflowJobRunLink(), + WorkflowJobRepairSingleTaskLink(), + ) else: # Databricks does not support repair for non-workflow tasks, hence do not show the repair link. self.operator_extra_links = (DatabricksJobRunLink(),) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py index 317bd444911c9..0350327976d31 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_workflow.py @@ -120,16 +120,10 @@ class _CreateDatabricksWorkflowOperator(BaseOperator): "spark_submit_params", ) caller = "_CreateDatabricksWorkflowOperator" - # Conditionally set operator_extra_links based on Airflow version - if AIRFLOW_V_3_0_PLUS: - # In Airflow 3, disable "Repair All Failed Tasks" since we can't pre-determine failed tasks - operator_extra_links = (WorkflowJobRunLink(),) - else: - # In Airflow 2.x, keep both links - operator_extra_links = ( # type: ignore[assignment] - WorkflowJobRunLink(), - WorkflowJobRepairAllFailedLink(), - ) + operator_extra_links = ( + WorkflowJobRunLink(), + WorkflowJobRepairAllFailedLink(), + ) def __init__( self, diff --git a/providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py b/providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py index 756bbf6bdeef4..dd1679de8dd38 100644 --- a/providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py +++ b/providers/databricks/src/airflow/providers/databricks/plugins/databricks_workflow.py @@ -17,8 +17,9 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING, Any -from urllib.parse import unquote +from urllib.parse import quote, unquote from airflow.exceptions import TaskInstanceNotFound from airflow.models.dagrun import DagRun @@ -30,9 +31,10 @@ BaseOperatorLink, TaskGroup, XCom, + conf, ) from airflow.providers.databricks.hooks.databricks import DatabricksHook -from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.state import TaskInstanceState @@ -44,6 +46,8 @@ from airflow.providers.databricks.operators.databricks import DatabricksTaskBaseOperator from airflow.sdk.types import Logger +log = logging.getLogger(__name__) + def get_databricks_task_ids( group_id: str, task_map: dict[str, DatabricksTaskBaseOperator], log: Logger @@ -67,6 +71,50 @@ def get_databricks_task_ids( return task_ids +def _repair_task( + databricks_conn_id: str, + databricks_run_id: int, + tasks_to_repair: list[str], + logger: Logger | logging.Logger, +) -> int: + """ + Repair a Databricks task using the Databricks API. + + This function allows the Airflow repair buttons to create a repair job for Databricks. + It uses the Databricks API to get the latest repair ID before sending the repair query. + + :param databricks_conn_id: The Databricks connection ID. + :param databricks_run_id: The Databricks run ID. + :param tasks_to_repair: A list of Databricks task IDs to repair. + :param logger: The logger to use for logging. + :return: the repair id returned by the Databricks API. + """ + hook = DatabricksHook(databricks_conn_id=databricks_conn_id) + + repair_history_id = hook.get_latest_repair_id(databricks_run_id) + logger.debug("Latest repair ID is %s", repair_history_id) + logger.debug( + "Sending repair query for tasks %s on run %s", + tasks_to_repair, + databricks_run_id, + ) + + run_data = hook.get_run(databricks_run_id) + repair_json = { + "run_id": databricks_run_id, + "latest_repair_id": repair_history_id, + "rerun_tasks": tasks_to_repair, + # Also rerun dependents so upstream-failed downstream tasks resume rather than + # staying skipped after the repaired task succeeds. + "rerun_dependent_tasks": True, + } + + if "overriding_parameters" in run_data: + repair_json["overriding_parameters"] = run_data["overriding_parameters"] + + return hook.repair_run(repair_json) + + # TODO: Need to re-think on how to support the currently unavailable repair functionality in Airflow 3. Probably a # good time to re-evaluate this would be once the plugin functionality is expanded in Airflow 3.1. if not AIRFLOW_V_3_0_PLUS: @@ -195,46 +243,6 @@ def get_task_instance(operator: BaseOperator, dttm, *, session: Session = NEW_SE raise TaskInstanceNotFound("Task instance not found") return ti - def _repair_task( - databricks_conn_id: str, - databricks_run_id: int, - tasks_to_repair: list[str], - logger: Logger, - ) -> int: - """ - Repair a Databricks task using the Databricks API. - - This function allows the Airflow retry function to create a repair job for Databricks. - It uses the Databricks API to get the latest repair ID before sending the repair query. - - :param databricks_conn_id: The Databricks connection ID. - :param databricks_run_id: The Databricks run ID. - :param tasks_to_repair: A list of Databricks task IDs to repair. - :param logger: The logger to use for logging. - :return: None - """ - hook = DatabricksHook(databricks_conn_id=databricks_conn_id) - - repair_history_id = hook.get_latest_repair_id(databricks_run_id) - logger.debug("Latest repair ID is %s", repair_history_id) - logger.debug( - "Sending repair query for tasks %s on run %s", - tasks_to_repair, - databricks_run_id, - ) - - run_data = hook.get_run(databricks_run_id) - repair_json = { - "run_id": databricks_run_id, - "latest_repair_id": repair_history_id, - "rerun_tasks": tasks_to_repair, - } - - if "overriding_parameters" in run_data: - repair_json["overriding_parameters"] = run_data["overriding_parameters"] - - return hook.repair_run(repair_json) - def get_launch_task_id(task_group: TaskGroup) -> str: """ @@ -381,6 +389,16 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way *, ti_key: TaskInstanceKey | None = None, ) -> str: + if AIRFLOW_V_3_0_PLUS: + if not AIRFLOW_V_3_1_PLUS or ti_key is None: + # The Airflow-3 repair backend requires 3.1+ (see DatabricksWorkflowPlugin). + return "" + launch_task_id = _get_launch_task_id_v3(operator, ti_key) + if not launch_task_id: + return "" + # The set of failed tasks is resolved from the live Databricks run by the endpoint. + return _build_repair_url(ti_key.dag_id, ti_key.run_id, launch_task_id, repair_all=True) + if not ti_key: ti = get_task_instance(operator, dttm) ti_key = ti.key @@ -478,6 +496,20 @@ def get_link( # type: ignore[override] # Signature intentionally kept this way *, ti_key: TaskInstanceKey | None = None, ) -> str: + if AIRFLOW_V_3_0_PLUS: + if not AIRFLOW_V_3_1_PLUS or ti_key is None: + # The Airflow-3 repair backend requires 3.1+ (see DatabricksWorkflowPlugin). + return "" + launch_task_id = _get_launch_task_id_v3(operator, ti_key) + if not launch_task_id: + return "" + return _build_repair_url( + ti_key.dag_id, + ti_key.run_id, + launch_task_id, + task_id=operator.task_id, + ) + if not ti_key: ti = get_task_instance(operator, dttm) ti_key = ti.key @@ -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( + "Repair Databricks workflow" + "

Repair Databricks workflow tasks

" + f"

Dag {escape(dag_id)}, run {escape(run_id)}.

" + f"

{escape(summary)}

" + f'
' + '
' + "" + ) + + @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) + + class DatabricksWorkflowPlugin(AirflowPlugin): """ Databricks Workflows plugin for Airflow. @@ -526,19 +812,25 @@ class DatabricksWorkflowPlugin(AirflowPlugin): name = "databricks_workflow" - # Conditionally set operator_extra_links based on Airflow version - if AIRFLOW_V_3_0_PLUS: - # In Airflow 3, disable the links for repair functionality until it is figured out it can be supported - operator_extra_links = [ - WorkflowJobRunLink(), - ] - else: - # In Airflow 2.x, keep all links including repair all failed tasks - operator_extra_links = [ - WorkflowJobRepairAllFailedLink(), - WorkflowJobRepairSingleTaskLink(), - WorkflowJobRunLink(), + operator_extra_links = [ + WorkflowJobRepairAllFailedLink(), + WorkflowJobRepairSingleTaskLink(), + WorkflowJobRunLink(), + ] + + if AIRFLOW_V_3_1_PLUS: + # Airflow 3.1+: repair is served by a FastAPI sub-application on the API server. The app + # relies on cookie-or-bearer auth resolution (`resolve_user_from_token`) that is only + # available from 3.1, so on 3.0.x the repair backend and its links are not registered. + fastapi_apps = [ + { + "app": repair_app, + "name": "Databricks Workflow Repair", + "url_prefix": REPAIR_URL_PREFIX, + } ] + elif not AIRFLOW_V_3_0_PLUS: + # Airflow 2.x: repair is served by a Flask-AppBuilder view. repair_databricks_view = RepairDatabricksTasks() repair_databricks_package = { "view": repair_databricks_view, diff --git a/providers/databricks/src/airflow/providers/databricks/version_compat.py b/providers/databricks/src/airflow/providers/databricks/version_compat.py index 0956edd21112f..209e8b63f35dc 100644 --- a/providers/databricks/src/airflow/providers/databricks/version_compat.py +++ b/providers/databricks/src/airflow/providers/databricks/version_compat.py @@ -33,7 +33,9 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) +AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0) __all__ = [ "AIRFLOW_V_3_0_PLUS", + "AIRFLOW_V_3_1_PLUS", ] diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py index 9086e29198635..6ea430affae8f 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py @@ -749,6 +749,20 @@ def test_get_run_tasks_success_multiple_pages(self, mock_requests): assert len(tasks) == 2 assert tasks == GET_RUN_RESPONSE["tasks"] * 2 + def test_get_run_failed_task_keys(self): + run_tasks = [ + {"task_key": "ok", "state": {"result_state": "SUCCESS", "life_cycle_state": "TERMINATED"}}, + {"task_key": "boom", "state": {"result_state": "FAILED", "life_cycle_state": "TERMINATED"}}, + {"task_key": "slow", "state": {"result_state": "TIMEDOUT", "life_cycle_state": "TERMINATED"}}, + {"task_key": "killed", "state": {"result_state": "CANCELED", "life_cycle_state": "SKIPPED"}}, + {"task_key": "crashed", "state": {"life_cycle_state": "INTERNAL_ERROR"}}, + {"task_key": "running", "state": {"life_cycle_state": "RUNNING"}}, + ] + with mock.patch.object(self.hook, "get_run_tasks", return_value=run_tasks): + failed = self.hook.get_run_failed_task_keys(RUN_ID) + + assert failed == ["boom", "slow", "killed", "crashed"] + @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") def test_cancel_run(self, mock_requests): mock_requests.post.return_value.json.return_value = GET_RUN_RESPONSE diff --git a/providers/databricks/tests/unit/databricks/plugins/test_databricks_workflow.py b/providers/databricks/tests/unit/databricks/plugins/test_databricks_workflow.py index 0a6e6c617160a..0838208c478c8 100644 --- a/providers/databricks/tests/unit/databricks/plugins/test_databricks_workflow.py +++ b/providers/databricks/tests/unit/databricks/plugins/test_databricks_workflow.py @@ -30,22 +30,25 @@ from airflow.models.taskinstance import TaskInstanceKey from airflow.providers.common.compat.sdk import AirflowException, AirflowPlugin from airflow.providers.databricks.plugins.databricks_workflow import ( + REPAIR_URL_PREFIX, DatabricksWorkflowPlugin, + WorkflowJobRepairAllFailedLink, WorkflowJobRepairSingleTaskLink, WorkflowJobRunLink, + _build_repair_url, _get_launch_task_key, + _repair_task, get_databricks_task_ids, get_launch_task_id, store_databricks_job_run_link, ) from tests_common import RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS if not AIRFLOW_V_3_0_PLUS: from airflow.providers.databricks.plugins.databricks_workflow import ( RepairDatabricksTasks, - _repair_task, ) DAG_ID = "test_dag" @@ -84,7 +87,6 @@ def test_get_dagrun_airflow2(): assert isinstance(result, DagRun) -@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow < 3.0") @patch("airflow.providers.databricks.plugins.databricks_workflow.DatabricksHook") def test_repair_task(mock_databricks_hook): mock_hook_instance = mock_databricks_hook.return_value @@ -97,9 +99,10 @@ def test_repair_task(mock_databricks_hook): assert result == 200 mock_hook_instance.get_latest_repair_id.assert_called_once_with(DATABRICKS_RUN_ID) mock_hook_instance.repair_run.assert_called_once() + # downstream dependents must also be rerun so upstream-failed tasks resume + assert mock_hook_instance.repair_run.call_args[0][0]["rerun_dependent_tasks"] is True -@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow < 3.0") @patch("airflow.providers.databricks.plugins.databricks_workflow.DatabricksHook") def test_repair_task_with_params(mock_databricks_hook): mock_hook_instance = mock_databricks_hook.return_value @@ -119,6 +122,7 @@ def test_repair_task_with_params(mock_databricks_hook): "run_id": DATABRICKS_RUN_ID, "rerun_tasks": tasks_to_repair, "latest_repair_id": 100, + "rerun_dependent_tasks": True, "overriding_parameters": { "key1": "value1", "key2": "value2", @@ -292,28 +296,193 @@ def test_appbuilder_views_airflow2(plugin): assert repair_view.default_view == "repair" -@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3.0+") +@pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="Airflow-3 repair backend requires 3.1+") class TestDatabricksWorkflowPluginAirflow3: - """Test Databricks Workflow Plugin functionality specific to Airflow 3.x.""" + """Test Databricks Workflow Plugin functionality specific to Airflow 3.1+.""" - def test_plugin_operator_extra_links_limited_functionality(self): - """Test that operator_extra_links are limited in Airflow 3.x (only job run link).""" + def test_plugin_operator_extra_links_include_repair(self): + """All three extra links (incl. repair) are registered in Airflow 3.x.""" plugin = DatabricksWorkflowPlugin() - # In Airflow 3, only WorkflowJobRunLink should be present - assert len(plugin.operator_extra_links) == 1 - assert isinstance(plugin.operator_extra_links[0], WorkflowJobRunLink) - - # Verify repair links are not present - link_types = [type(link).__name__ for link in plugin.operator_extra_links] - assert not any("Repair" in link_type for link_type in link_types) + link_types = {type(link).__name__ for link in plugin.operator_extra_links} + assert link_types == { + "WorkflowJobRunLink", + "WorkflowJobRepairAllFailedLink", + "WorkflowJobRepairSingleTaskLink", + } - def test_plugin_no_appbuilder_views(self): - """Test that appbuilder_views are not configured in Airflow 3.x.""" + def test_plugin_registers_fastapi_repair_app(self): + """Repair is served by a FastAPI app (not Flask-AppBuilder) in Airflow 3.1+.""" plugin = DatabricksWorkflowPlugin() - # In Airflow 3, appbuilder_views should not be set (repair functionality disabled) assert not getattr(plugin, "appbuilder_views", []) + assert len(plugin.fastapi_apps) == 1 + app = plugin.fastapi_apps[0] + assert app["url_prefix"] == REPAIR_URL_PREFIX + assert app["app"] is not None + + def test_build_repair_url_single_task(self): + """The URL carries only Airflow identifiers — never the Databricks conn/run/task keys.""" + url = _build_repair_url("my_dag", "run 1", "grp.launch", task_id="grp.nb") + assert f"{REPAIR_URL_PREFIX}/my_dag/run%201" in url + assert "launch_task_id=grp.launch" in url + assert "task_id=grp.nb" in url + assert "repair_all" not in url + assert "databricks_conn_id" not in url + assert "databricks_run_id" not in url + + def test_build_repair_url_repair_all(self): + url = _build_repair_url("my_dag", "run1", "grp.launch", repair_all=True) + assert "repair_all=true" in url + assert "launch_task_id=grp.launch" in url + assert "&task_id=" not in url + + def test_repair_single_task_link_uses_endpoint_url(self): + """The single-task repair link points at the FastAPI endpoint, naming Airflow ids only.""" + link = WorkflowJobRepairSingleTaskLink() + operator = Mock(task_id="grp.nb") + ti_key = TaskInstanceKey(dag_id="my_dag", task_id="grp.nb", run_id="run1", try_number=1) + with patch( + "airflow.providers.databricks.plugins.databricks_workflow._get_launch_task_id_v3", + return_value="grp.launch", + ): + url = link.get_link(operator, ti_key=ti_key) + assert REPAIR_URL_PREFIX in url + assert "launch_task_id=grp.launch" in url + assert "task_id=grp.nb" in url + + def test_repair_all_link_uses_endpoint_url(self): + link = WorkflowJobRepairAllFailedLink() + operator = Mock() + ti_key = TaskInstanceKey(dag_id="my_dag", task_id="grp.nb", run_id="run1", try_number=1) + with patch( + "airflow.providers.databricks.plugins.databricks_workflow._get_launch_task_id_v3", + return_value="grp.launch", + ): + url = link.get_link(operator, ti_key=ti_key) + assert REPAIR_URL_PREFIX in url + assert "repair_all=true" in url + assert "launch_task_id=grp.launch" in url + + def test_repair_link_returns_empty_without_launch_task(self): + """When the launch task can't be resolved, the link renders empty (no crash).""" + link = WorkflowJobRepairSingleTaskLink() + ti_key = TaskInstanceKey(dag_id="my_dag", task_id="grp.nb", run_id="run1", try_number=1) + with patch( + "airflow.providers.databricks.plugins.databricks_workflow._get_launch_task_id_v3", + return_value=None, + ): + assert link.get_link(Mock(task_id="grp.nb"), ti_key=ti_key) == "" + + def test_get_confirmation_page_does_not_mutate(self): + """GET renders a read-only confirmation form and never repairs or clears.""" + from fastapi.testclient import TestClient + + from airflow.providers.databricks.plugins import databricks_workflow as m + + m.repair_app.dependency_overrides[m._require_dag_run_edit] = lambda: Mock() + try: + with ( + patch.object(m, "_repair_task") as mock_repair, + patch.object(m, "_clear_repaired_and_downstream") as mock_clear, + ): + client = TestClient(m.repair_app) + resp = client.get( + "/my_dag/run1", + params={"launch_task_id": "grp.launch", "repair_all": "true"}, + follow_redirects=False, + ) + assert resp.status_code == 200 + assert "x"}, + follow_redirects=False, + ) + assert resp.status_code == 200 + assert "" not in resp.text + assert "<script>" in resp.text + finally: + m.repair_app.dependency_overrides.clear() + + def test_post_repairs_clears_and_redirects(self): + """POST resolves failed tasks from live Databricks state, repairs, clears, and redirects.""" + from fastapi.testclient import TestClient + + from airflow.providers.databricks.plugins import databricks_workflow as m + + m.repair_app.dependency_overrides[m._require_dag_run_edit] = lambda: Mock() + dag = Mock() + task = Mock(task_id="grp.nb", databricks_task_key="k1") + dag.tasks = [task] + dag.dag_id = "my_dag" + try: + with ( + patch("airflow.models.serialized_dag.SerializedDagModel.get_dag", return_value=dag), + patch("airflow.utils.session.create_session"), + patch.object(m, "_read_launch_metadata", return_value=Mock(conn_id="c", run_id=999)), + patch.object(m, "DatabricksHook") as mock_hook, + patch.object(m, "_repair_task") as mock_repair, + patch.object(m, "_clear_repaired_and_downstream") as mock_clear, + ): + mock_hook.return_value.get_run_failed_task_keys.return_value = ["k1"] + client = TestClient(m.repair_app) + resp = client.post( + "/my_dag/run1", + params={"launch_task_id": "grp.launch", "repair_all": "true"}, + follow_redirects=False, + ) + assert resp.status_code == 303 + assert resp.headers["location"] == "/dags/my_dag/runs/run1" + assert mock_repair.call_args.kwargs["tasks_to_repair"] == ["k1"] + assert mock_repair.call_args.kwargs["databricks_run_id"] == 999 + mock_clear.assert_called_once() + finally: + m.repair_app.dependency_overrides.clear() + + def test_post_returns_502_on_databricks_error_without_leaking_detail(self): + """A Databricks/hook failure surfaces as a clean 502 that does not echo the exception.""" + from fastapi.testclient import TestClient + + from airflow.providers.databricks.plugins import databricks_workflow as m + + m.repair_app.dependency_overrides[m._require_dag_run_edit] = lambda: Mock() + dag = Mock(dag_id="my_dag", tasks=[]) + try: + with ( + patch("airflow.models.serialized_dag.SerializedDagModel.get_dag", return_value=dag), + patch("airflow.utils.session.create_session"), + patch.object(m, "_read_launch_metadata", return_value=Mock(conn_id="c", run_id=1)), + patch.object(m, "DatabricksHook") as mock_hook, + patch.object(m, "_clear_repaired_and_downstream") as mock_clear, + ): + mock_hook.return_value.get_run_failed_task_keys.side_effect = Exception("secret token abc") + client = TestClient(m.repair_app, raise_server_exceptions=False) + resp = client.post( + "/my_dag/run1", + params={"launch_task_id": "grp.launch", "repair_all": "true"}, + follow_redirects=False, + ) + assert resp.status_code == 502 + assert "secret token abc" not in resp.text + mock_clear.assert_not_called() + finally: + m.repair_app.dependency_overrides.clear() def test_store_databricks_job_run_link_function_works(self): """Test that store_databricks_job_run_link works correctly in Airflow 3.x."""