Enable repair of failed Databricks Workflow tasks from Airflow 3 UI

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 <https://docs.databricks.com/api/workspace/jobs/repairrun>`_
+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.

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.

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(),)

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,