Add DatabricksWorkflowPlugin (apache#40153)

The DatabricksWorkflowPlugin provides with links in the Airflow
UI for tasks that allow us to see the Databricks job run in the
Databricks workspace, additionally it also provides link to 
repair task(s) in the workflow.

Databricks does not allow repairing jobs with single tasks launched
outside the workflow, hence we just provide the link for the job run.
<img width="1342" alt="Screenshot 2024-06-24 at 4 05 53 PM" src="https://github.com/apache/airflow/assets/10206082/c1ded2b7-90fb-4a3c-980d-8043cc5a459f">

Within the workflow, for each of the task, we provide links to the
job run and repair link for the single task
<img width="1368" alt="Screenshot 2024-06-24 at 5 40 27 PM" src="https://github.com/apache/airflow/assets/10206082/d98081b8-8014-4d18-8747-e5b3382db416">

And at the workflow level, for the job launch task, we provide a
link to repair all failed tasks along with the link for job run in 
the Databricks workspace that can be used to monitor the job
in the Databricks account.
<img width="1421" alt="Screenshot 2024-06-24 at 5 40 56 PM" src="https://github.com/apache/airflow/assets/10206082/9d7ad7ae-9bbf-4fef-aa52-16ac5366edf3">


---------

Co-authored-by: Wei Lee <[email protected]>

airflow/providers/databricks/operators/databricks.py

@@ -35,6 +35,10 @@
     DatabricksWorkflowTaskGroup,
     WorkflowRunMetadata,
 )
+from airflow.providers.databricks.plugins.databricks_workflow import (
+    WorkflowJobRepairSingleTaskLink,
+    WorkflowJobRunLink,
+)
 from airflow.providers.databricks.triggers.databricks import DatabricksExecutionTrigger
 from airflow.providers.databricks.utils.databricks import _normalise_json_content, validate_trigger_event
 
@@ -958,6 +962,15 @@ def __init__(
 
         super().__init__(**kwargs)
 
+        if self._databricks_workflow_task_group is not None:
+            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(),)
+
     @cached_property
     def _hook(self) -> DatabricksHook:
         return self._get_hook(caller=self.caller)
@@ -1016,12 +1029,17 @@ def _get_run_json(self) -> dict[str, Any]:
             raise ValueError("Must specify either existing_cluster_id or new_cluster.")
         return run_json
 
-    def _launch_job(self) -> int:
+    def _launch_job(self, context: Context | None = None) -> int:
         """Launch the job on Databricks."""
         run_json = self._get_run_json()
         self.databricks_run_id = self._hook.submit_run(run_json)
         url = self._hook.get_run_page_url(self.databricks_run_id)
         self.log.info("Check the job run in Databricks: %s", url)
+
+        if self.do_xcom_push and context is not None:
+            context["ti"].xcom_push(key=XCOM_RUN_ID_KEY, value=self.databricks_run_id)
+            context["ti"].xcom_push(key=XCOM_RUN_PAGE_URL_KEY, value=url)
+
         return self.databricks_run_id
 
     def _handle_terminal_run_state(self, run_state: RunState) -> None:
@@ -1040,7 +1058,15 @@ def _get_current_databricks_task(self) -> dict[str, Any]:
         """Retrieve the Databricks task corresponding to the current Airflow task."""
         if self.databricks_run_id is None:
             raise ValueError("Databricks job not yet launched. Please run launch_notebook_job first.")
-        return {task["task_key"]: task for task in self._hook.get_run(self.databricks_run_id)["tasks"]}[
+        tasks = self._hook.get_run(self.databricks_run_id)["tasks"]
+
+        # Because the task_key remains the same across multiple runs, and the Databricks API does not return
+        # tasks sorted by their attempts/start time, we sort the tasks by start time. This ensures that we
+        # map the latest attempt (whose status is to be monitored) of the task run to the task_key while
+        # building the {task_key: task} map below.
+        sorted_task_runs = sorted(tasks, key=lambda x: x["start_time"])
+
+        return {task["task_key"]: task for task in sorted_task_runs}[
             self._get_databricks_task_id(self.task_id)
         ]
 
@@ -1125,7 +1151,7 @@ def execute(self, context: Context) -> None:
             self.databricks_run_id = workflow_run_metadata.run_id
             self.databricks_conn_id = workflow_run_metadata.conn_id
         else:
-            self._launch_job()
+            self._launch_job(context=context)
         if self.wait_for_termination:
             self.monitor_databricks_job()
 

airflow/providers/databricks/operators/databricks_workflow.py

@@ -28,6 +28,10 @@
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.databricks.hooks.databricks import DatabricksHook, RunLifeCycleState
+from airflow.providers.databricks.plugins.databricks_workflow import (
+    WorkflowJobRepairAllFailedLink,
+    WorkflowJobRunLink,
+)
 from airflow.utils.task_group import TaskGroup
 
 if TYPE_CHECKING:
@@ -88,6 +92,7 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
         populated after instantiation using the `add_task` method.
     """
 
+    operator_extra_links = (WorkflowJobRunLink(), WorkflowJobRepairAllFailedLink())
     template_fields = ("notebook_params",)
     caller = "_CreateDatabricksWorkflowOperator"
 

airflow/providers/databricks/plugins/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.