apache-airflow-providers-databricks 6.7.0rc1__py3-none-any.whl → 6.8.0rc1__py3-none-any.whl

airflow/providers/databricks/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "6.7.0"
+__version__ = "6.8.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.7.0"

airflow/providers/databricks/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Databricks",
         "description": "`Databricks <https://databricks.com/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1720422668,
+        "source-date-epoch": 1722663644,
         "versions": [
+            "6.8.0",
             "6.7.0",
             "6.6.0",
             "6.5.0",
@@ -189,5 +190,11 @@ def get_provider_info():
                 "connection-type": "databricks",
             }
         ],
+        "plugins": [
+            {
+                "name": "databricks_workflow",
+                "plugin-class": "airflow.providers.databricks.plugins.databricks_workflow.DatabricksWorkflowPlugin",
+            }
+        ],
         "extra-links": ["airflow.providers.databricks.operators.databricks.DatabricksJobRunLink"],
     }

airflow/providers/databricks/operators/databricks.py

@@ -35,8 +35,12 @@ from airflow.providers.databricks.operators.databricks_workflow import (
     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
+from airflow.providers.databricks.utils.databricks import normalise_json_content, validate_trigger_event
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
@@ -182,17 +186,6 @@ def _handle_deferrable_databricks_operator_completion(event: dict, log: Logger)
     raise AirflowException(error_message)
 
 
-def _handle_overridden_json_params(operator):
-    for key, value in operator.overridden_json_params.items():
-        if value is not None:
-            operator.json[key] = value
-
-
-def normalise_json_content(operator):
-    if operator.json:
-        operator.json = _normalise_json_content(operator.json)
-
-
 class DatabricksJobRunLink(BaseOperatorLink):
     """Constructs a link to monitor a Databricks Job Run."""
 
@@ -296,21 +289,34 @@ class DatabricksCreateJobsOperator(BaseOperator):
         self.databricks_retry_limit = databricks_retry_limit
         self.databricks_retry_delay = databricks_retry_delay
         self.databricks_retry_args = databricks_retry_args
-        self.overridden_json_params = {
-            "name": name,
-            "description": description,
-            "tags": tags,
-            "tasks": tasks,
-            "job_clusters": job_clusters,
-            "email_notifications": email_notifications,
-            "webhook_notifications": webhook_notifications,
-            "notification_settings": notification_settings,
-            "timeout_seconds": timeout_seconds,
-            "schedule": schedule,
-            "max_concurrent_runs": max_concurrent_runs,
-            "git_source": git_source,
-            "access_control_list": access_control_list,
-        }
+        if name is not None:
+            self.json["name"] = name
+        if description is not None:
+            self.json["description"] = description
+        if tags is not None:
+            self.json["tags"] = tags
+        if tasks is not None:
+            self.json["tasks"] = tasks
+        if job_clusters is not None:
+            self.json["job_clusters"] = job_clusters
+        if email_notifications is not None:
+            self.json["email_notifications"] = email_notifications
+        if webhook_notifications is not None:
+            self.json["webhook_notifications"] = webhook_notifications
+        if notification_settings is not None:
+            self.json["notification_settings"] = notification_settings
+        if timeout_seconds is not None:
+            self.json["timeout_seconds"] = timeout_seconds
+        if schedule is not None:
+            self.json["schedule"] = schedule
+        if max_concurrent_runs is not None:
+            self.json["max_concurrent_runs"] = max_concurrent_runs
+        if git_source is not None:
+            self.json["git_source"] = git_source
+        if access_control_list is not None:
+            self.json["access_control_list"] = access_control_list
+        if self.json:
+            self.json = normalise_json_content(self.json)
 
     @cached_property
     def _hook(self):
@@ -322,24 +328,16 @@ class DatabricksCreateJobsOperator(BaseOperator):
             caller="DatabricksCreateJobsOperator",
         )
 
-    def _setup_and_validate_json(self):
-        _handle_overridden_json_params(self)
-
+    def execute(self, context: Context) -> int:
         if "name" not in self.json:
             raise AirflowException("Missing required parameter: name")
-
-        normalise_json_content(self)
-
-    def execute(self, context: Context) -> int:
-        self._setup_and_validate_json()
-
         job_id = self._hook.find_job_id_by_name(self.json["name"])
         if job_id is None:
             return self._hook.create_job(self.json)
         self._hook.reset_job(str(job_id), self.json)
         if (access_control_list := self.json.get("access_control_list")) is not None:
             acl_json = {"access_control_list": access_control_list}
-            self._hook.update_job_permission(job_id, _normalise_json_content(acl_json))
+            self._hook.update_job_permission(job_id, normalise_json_content(acl_json))
 
         return job_id
 
@@ -512,23 +510,43 @@ class DatabricksSubmitRunOperator(BaseOperator):
         self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
         self.deferrable = deferrable
-        self.overridden_json_params = {
-            "tasks": tasks,
-            "spark_jar_task": spark_jar_task,
-            "notebook_task": notebook_task,
-            "spark_python_task": spark_python_task,
-            "spark_submit_task": spark_submit_task,
-            "pipeline_task": pipeline_task,
-            "dbt_task": dbt_task,
-            "new_cluster": new_cluster,
-            "existing_cluster_id": existing_cluster_id,
-            "libraries": libraries,
-            "run_name": run_name,
-            "timeout_seconds": timeout_seconds,
-            "idempotency_token": idempotency_token,
-            "access_control_list": access_control_list,
-            "git_source": git_source,
-        }
+        if tasks is not None:
+            self.json["tasks"] = tasks
+        if spark_jar_task is not None:
+            self.json["spark_jar_task"] = spark_jar_task
+        if notebook_task is not None:
+            self.json["notebook_task"] = notebook_task
+        if spark_python_task is not None:
+            self.json["spark_python_task"] = spark_python_task
+        if spark_submit_task is not None:
+            self.json["spark_submit_task"] = spark_submit_task
+        if pipeline_task is not None:
+            self.json["pipeline_task"] = pipeline_task
+        if dbt_task is not None:
+            self.json["dbt_task"] = dbt_task
+        if new_cluster is not None:
+            self.json["new_cluster"] = new_cluster
+        if existing_cluster_id is not None:
+            self.json["existing_cluster_id"] = existing_cluster_id
+        if libraries is not None:
+            self.json["libraries"] = libraries
+        if run_name is not None:
+            self.json["run_name"] = run_name
+        if timeout_seconds is not None:
+            self.json["timeout_seconds"] = timeout_seconds
+        if "run_name" not in self.json:
+            self.json["run_name"] = run_name or kwargs["task_id"]
+        if idempotency_token is not None:
+            self.json["idempotency_token"] = idempotency_token
+        if access_control_list is not None:
+            self.json["access_control_list"] = access_control_list
+        if git_source is not None:
+            self.json["git_source"] = git_source
+
+        if "dbt_task" in self.json and "git_source" not in self.json:
+            raise AirflowException("git_source is required for dbt_task")
+        if pipeline_task is not None and "pipeline_id" in pipeline_task and "pipeline_name" in pipeline_task:
+            raise AirflowException("'pipeline_name' is not allowed in conjunction with 'pipeline_id'")
 
         # This variable will be used in case our task gets killed.
         self.run_id: int | None = None
@@ -547,25 +565,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
             caller=caller,
         )
 
-    def _setup_and_validate_json(self):
-        _handle_overridden_json_params(self)
-
-        if "run_name" not in self.json or self.json["run_name"] is None:
-            self.json["run_name"] = self.task_id
-
-        if "dbt_task" in self.json and "git_source" not in self.json:
-            raise AirflowException("git_source is required for dbt_task")
-        if (
-            "pipeline_task" in self.json
-            and "pipeline_id" in self.json["pipeline_task"]
-            and "pipeline_name" in self.json["pipeline_task"]
-        ):
-            raise AirflowException("'pipeline_name' is not allowed in conjunction with 'pipeline_id'")
-
-        normalise_json_content(self)
-
     def execute(self, context: Context):
-        self._setup_and_validate_json()
         if (
             "pipeline_task" in self.json
             and self.json["pipeline_task"].get("pipeline_id") is None
@@ -575,7 +575,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
             pipeline_name = self.json["pipeline_task"]["pipeline_name"]
             self.json["pipeline_task"]["pipeline_id"] = self._hook.find_pipeline_id_by_name(pipeline_name)
             del self.json["pipeline_task"]["pipeline_name"]
-        json_normalised = _normalise_json_content(self.json)
+        json_normalised = normalise_json_content(self.json)
         self.run_id = self._hook.submit_run(json_normalised)
         if self.deferrable:
             _handle_deferrable_databricks_operator_execution(self, self._hook, self.log, context)
@@ -611,7 +611,7 @@ class DatabricksSubmitRunDeferrableOperator(DatabricksSubmitRunOperator):
 
     def execute(self, context):
         hook = self._get_hook(caller="DatabricksSubmitRunDeferrableOperator")
-        json_normalised = _normalise_json_content(self.json)
+        json_normalised = normalise_json_content(self.json)
         self.run_id = hook.submit_run(json_normalised)
         _handle_deferrable_databricks_operator_execution(self, hook, self.log, context)
 
@@ -811,16 +811,27 @@ class DatabricksRunNowOperator(BaseOperator):
         self.deferrable = deferrable
         self.repair_run = repair_run
         self.cancel_previous_runs = cancel_previous_runs
-        self.overridden_json_params = {
-            "job_id": job_id,
-            "job_name": job_name,
-            "notebook_params": notebook_params,
-            "python_params": python_params,
-            "python_named_params": python_named_params,
-            "jar_params": jar_params,
-            "spark_submit_params": spark_submit_params,
-            "idempotency_token": idempotency_token,
-        }
+
+        if job_id is not None:
+            self.json["job_id"] = job_id
+        if job_name is not None:
+            self.json["job_name"] = job_name
+        if "job_id" in self.json and "job_name" in self.json:
+            raise AirflowException("Argument 'job_name' is not allowed with argument 'job_id'")
+        if notebook_params is not None:
+            self.json["notebook_params"] = notebook_params
+        if python_params is not None:
+            self.json["python_params"] = python_params
+        if python_named_params is not None:
+            self.json["python_named_params"] = python_named_params
+        if jar_params is not None:
+            self.json["jar_params"] = jar_params
+        if spark_submit_params is not None:
+            self.json["spark_submit_params"] = spark_submit_params
+        if idempotency_token is not None:
+            self.json["idempotency_token"] = idempotency_token
+        if self.json:
+            self.json = normalise_json_content(self.json)
         # This variable will be used in case our task gets killed.
         self.run_id: int | None = None
         self.do_xcom_push = do_xcom_push
@@ -838,16 +849,7 @@ class DatabricksRunNowOperator(BaseOperator):
             caller=caller,
         )
 
-    def _setup_and_validate_json(self):
-        _handle_overridden_json_params(self)
-
-        if "job_id" in self.json and "job_name" in self.json:
-            raise AirflowException("Argument 'job_name' is not allowed with argument 'job_id'")
-
-        normalise_json_content(self)
-
     def execute(self, context: Context):
-        self._setup_and_validate_json()
         hook = self._hook
         if "job_name" in self.json:
             job_id = hook.find_job_id_by_name(self.json["job_name"])
@@ -958,6 +960,15 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
 
         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 +1027,17 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
             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 +1056,15 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
         """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 +1149,7 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
             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 mergedeep import merge
 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.

airflow/providers/databricks/plugins/databricks_workflow.py

@@ -0,0 +1,477 @@
+# 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.
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import unquote
+
+from flask import current_app, flash, redirect, request, url_for
+from flask_appbuilder.api import expose
+from packaging.version import Version
+
+from airflow.exceptions import AirflowException, TaskInstanceNotFound
+from airflow.models import BaseOperator, BaseOperatorLink
+from airflow.models.dag import DAG, clear_task_instances
+from airflow.models.dagrun import DagRun
+from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
+from airflow.models.xcom import XCom
+from airflow.plugins_manager import AirflowPlugin
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.security import permissions
+from airflow.utils.airflow_flask_app import AirflowApp
+from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.state import TaskInstanceState
+from airflow.utils.task_group import TaskGroup
+from airflow.version import version
+from airflow.www import auth
+from airflow.www.views import AirflowBaseView
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm.session import Session
+
+
+REPAIR_WAIT_ATTEMPTS = os.getenv("DATABRICKS_REPAIR_WAIT_ATTEMPTS", 20)
+REPAIR_WAIT_DELAY = os.getenv("DATABRICKS_REPAIR_WAIT_DELAY", 0.5)
+
+airflow_app = cast(AirflowApp, current_app)
+
+
+def get_auth_decorator():
+    # TODO: remove this if block when min_airflow_version is set to higher than 2.8.0
+    if Version(version) < Version("2.8"):
+        return auth.has_access(
+            [
+                (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG),
+                (permissions.ACTION_CAN_CREATE, permissions.RESOURCE_DAG_RUN),
+            ]
+        )
+
+    from airflow.auth.managers.models.resource_details import DagAccessEntity
+
+    return auth.has_access_dag("POST", DagAccessEntity.RUN)
+
+
+def _get_databricks_task_id(task: BaseOperator) -> str:
+    """
+    Get the databricks task ID using dag_id and task_id. removes illegal characters.
+
+    :param task: The task to get the databricks task ID for.
+    :return: The databricks task ID.
+    """
+    return f"{task.dag_id}__{task.task_id.replace('.', '__')}"
+
+
+def get_databricks_task_ids(
+    group_id: str, task_map: dict[str, BaseOperator], log: logging.Logger
+) -> list[str]:
+    """
+    Return a list of all Databricks task IDs for a dictionary of Airflow tasks.
+
+    :param group_id: The task group ID.
+    :param task_map: A dictionary mapping task IDs to BaseOperator instances.
+    :param log: The logger to use for logging.
+    :return: A list of Databricks task IDs for the given task group.
+    """
+    task_ids = []
+    log.debug("Getting databricks task ids for group %s", group_id)
+    for task_id, task in task_map.items():
+        if task_id == f"{group_id}.launch":
+            continue
+        databricks_task_id = _get_databricks_task_id(task)
+        log.debug("databricks task id for task %s is %s", task_id, databricks_task_id)
+        task_ids.append(databricks_task_id)
+    return task_ids
+
+
+@provide_session
+def _get_dagrun(dag: DAG, run_id: str, session: Session | None = None) -> DagRun:
+    """
+    Retrieve the DagRun object associated with the specified DAG and run_id.
+
+    :param dag: The DAG object associated with the DagRun to retrieve.
+    :param run_id: The run_id associated with the DagRun to retrieve.
+    :param session: The SQLAlchemy session to use for the query. If None, uses the default session.
+    :return: The DagRun object associated with the specified DAG and run_id.
+    """
+    if not session:
+        raise AirflowException("Session not provided.")
+
+    return session.query(DagRun).filter(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id).first()
+
+
+@provide_session
+def _clear_task_instances(
+    dag_id: str, run_id: str, task_ids: list[str], log: logging.Logger, session: Session | None = None
+) -> None:
+    dag = airflow_app.dag_bag.get_dag(dag_id)
+    log.debug("task_ids %s to clear", str(task_ids))
+    dr: DagRun = _get_dagrun(dag, run_id, session=session)
+    tis_to_clear = [ti for ti in dr.get_task_instances() if _get_databricks_task_id(ti) in task_ids]
+    clear_task_instances(tis_to_clear, session)
+
+
+def _repair_task(
+    databricks_conn_id: str,
+    databricks_run_id: int,
+    tasks_to_repair: list[str],
+    logger: logging.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,
+    )
+
+    repair_json = {
+        "run_id": databricks_run_id,
+        "latest_repair_id": repair_history_id,
+        "rerun_tasks": tasks_to_repair,
+    }
+
+    return hook.repair_run(repair_json)
+
+
+def get_launch_task_id(task_group: TaskGroup) -> str:
+    """
+    Retrieve the launch task ID from the current task group or a parent task group, recursively.
+
+    :param task_group: Task Group to be inspected
+    :return: launch Task ID
+    """
+    try:
+        launch_task_id = task_group.get_child_by_label("launch").task_id  # type: ignore[attr-defined]
+    except KeyError as e:
+        if not task_group.parent_group:
+            raise AirflowException("No launch task can be found in the task group.") from e
+        launch_task_id = get_launch_task_id(task_group.parent_group)
+
+    return launch_task_id
+
+
+def _get_launch_task_key(current_task_key: TaskInstanceKey, task_id: str) -> TaskInstanceKey:
+    """
+    Return the task key for the launch task.
+
+    This allows us to gather databricks Metadata even if the current task has failed (since tasks only
+    create xcom values if they succeed).
+
+    :param current_task_key: The task key for the current task.
+    :param task_id: The task ID for the current task.
+    :return: The task key for the launch task.
+    """
+    if task_id:
+        return TaskInstanceKey(
+            dag_id=current_task_key.dag_id,
+            task_id=task_id,
+            run_id=current_task_key.run_id,
+            try_number=current_task_key.try_number,
+        )
+
+    return current_task_key
+
+
+@provide_session
+def get_task_instance(operator: BaseOperator, dttm, session: Session = NEW_SESSION) -> TaskInstance:
+    dag_id = operator.dag.dag_id
+    dag_run = DagRun.find(dag_id, execution_date=dttm)[0]
+    ti = (
+        session.query(TaskInstance)
+        .filter(
+            TaskInstance.dag_id == dag_id,
+            TaskInstance.run_id == dag_run.run_id,
+            TaskInstance.task_id == operator.task_id,
+        )
+        .one_or_none()
+    )
+    if not ti:
+        raise TaskInstanceNotFound("Task instance not found")
+    return ti
+
+
+def get_xcom_result(
+    ti_key: TaskInstanceKey,
+    key: str,
+) -> Any:
+    result = XCom.get_value(
+        ti_key=ti_key,
+        key=key,
+    )
+    from airflow.providers.databricks.operators.databricks_workflow import WorkflowRunMetadata
+
+    return WorkflowRunMetadata(**result)
+
+
+class WorkflowJobRunLink(BaseOperatorLink, LoggingMixin):
+    """Constructs a link to monitor a Databricks Job Run."""
+
+    name = "See Databricks Job Run"
+
+    def get_link(
+        self,
+        operator: BaseOperator,
+        dttm=None,
+        *,
+        ti_key: TaskInstanceKey | None = None,
+    ) -> str:
+        if not ti_key:
+            ti = get_task_instance(operator, dttm)
+            ti_key = ti.key
+        task_group = operator.task_group
+
+        if not task_group:
+            raise AirflowException("Task group is required for generating Databricks Workflow Job Run Link.")
+
+        dag = airflow_app.dag_bag.get_dag(ti_key.dag_id)
+        dag.get_task(ti_key.task_id)
+        self.log.info("Getting link for task %s", ti_key.task_id)
+        if ".launch" not in ti_key.task_id:
+            self.log.debug("Finding the launch task for job run metadata %s", 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)
+        metadata = get_xcom_result(ti_key, "return_value")
+
+        hook = DatabricksHook(metadata.conn_id)
+        return f"https://{hook.host}/#job/{metadata.job_id}/run/{metadata.run_id}"
+
+
+class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
+    """Constructs a link to send a request to repair all failed tasks in the Databricks workflow."""
+
+    name = "Repair All Failed Tasks"
+
+    def get_link(
+        self,
+        operator,
+        dttm=None,
+        *,
+        ti_key: TaskInstanceKey | None = None,
+    ) -> str:
+        if not ti_key:
+            ti = get_task_instance(operator, dttm)
+            ti_key = ti.key
+        task_group = operator.task_group
+        self.log.debug(
+            "Creating link to repair all tasks for databricks job run %s",
+            task_group.group_id,
+        )
+
+        metadata = get_xcom_result(ti_key, "return_value")
+
+        tasks_str = self.get_tasks_to_run(ti_key, operator, self.log)
+        self.log.debug("tasks to rerun: %s", tasks_str)
+
+        query_params = {
+            "dag_id": ti_key.dag_id,
+            "databricks_conn_id": metadata.conn_id,
+            "databricks_run_id": metadata.run_id,
+            "run_id": ti_key.run_id,
+            "tasks_to_repair": tasks_str,
+        }
+
+        return url_for("RepairDatabricksTasks.repair", **query_params)
+
+    @classmethod
+    def get_task_group_children(cls, task_group: TaskGroup) -> dict[str, BaseOperator]:
+        """
+        Given a TaskGroup, return children which are Tasks, inspecting recursively any TaskGroups within.
+
+        :param task_group: An Airflow TaskGroup
+        :return: Dictionary that contains Task IDs as keys and Tasks as values.
+        """
+        children: dict[str, Any] = {}
+        for child_id, child in task_group.children.items():
+            if isinstance(child, TaskGroup):
+                child_children = cls.get_task_group_children(child)
+                children = {**children, **child_children}
+            else:
+                children[child_id] = child
+        return children
+
+    def get_tasks_to_run(self, ti_key: TaskInstanceKey, operator: BaseOperator, log: logging.Logger) -> str:
+        task_group = operator.task_group
+        if not task_group:
+            raise AirflowException("Task group is required for generating repair link.")
+        if not task_group.group_id:
+            raise AirflowException("Task group ID is required for generating repair link.")
+        dag = airflow_app.dag_bag.get_dag(ti_key.dag_id)
+        dr = _get_dagrun(dag, ti_key.run_id)
+        log.debug("Getting failed and skipped tasks for dag run %s", dr.run_id)
+        task_group_sub_tasks = self.get_task_group_children(task_group).items()
+        failed_and_skipped_tasks = self._get_failed_and_skipped_tasks(dr)
+        log.debug("Failed and skipped tasks: %s", failed_and_skipped_tasks)
+
+        tasks_to_run = {ti: t for ti, t in task_group_sub_tasks if ti in failed_and_skipped_tasks}
+
+        return ",".join(get_databricks_task_ids(task_group.group_id, tasks_to_run, log))
+
+    @staticmethod
+    def _get_failed_and_skipped_tasks(dr: DagRun) -> list[str]:
+        """
+        Return a list of task IDs for tasks that have failed or have been skipped in the given DagRun.
+
+        :param dr: The DagRun object for which to retrieve failed and skipped tasks.
+
+        :return: A list of task IDs for tasks that have failed or have been skipped.
+        """
+        return [
+            t.task_id
+            for t in dr.get_task_instances(
+                state=[
+                    TaskInstanceState.FAILED,
+                    TaskInstanceState.SKIPPED,
+                    TaskInstanceState.UP_FOR_RETRY,
+                    TaskInstanceState.UPSTREAM_FAILED,
+                    None,
+                ],
+            )
+        ]
+
+
+class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
+    """Construct a link to send a repair request for a single databricks task."""
+
+    name = "Repair a single task"
+
+    def get_link(
+        self,
+        operator,
+        dttm=None,
+        *,
+        ti_key: TaskInstanceKey | None = None,
+    ) -> str:
+        if not ti_key:
+            ti = get_task_instance(operator, dttm)
+            ti_key = ti.key
+
+        task_group = operator.task_group
+        if not task_group:
+            raise AirflowException("Task group is required for generating repair link.")
+
+        self.log.info(
+            "Creating link to repair a single task for databricks job run %s task %s",
+            task_group.group_id,
+            ti_key.task_id,
+        )
+        dag = airflow_app.dag_bag.get_dag(ti_key.dag_id)
+        task = dag.get_task(ti_key.task_id)
+
+        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)
+        metadata = get_xcom_result(ti_key, "return_value")
+
+        query_params = {
+            "dag_id": ti_key.dag_id,
+            "databricks_conn_id": metadata.conn_id,
+            "databricks_run_id": metadata.run_id,
+            "run_id": ti_key.run_id,
+            "tasks_to_repair": _get_databricks_task_id(task),
+        }
+        return url_for("RepairDatabricksTasks.repair", **query_params)
+
+
+class RepairDatabricksTasks(AirflowBaseView, LoggingMixin):
+    """Repair databricks tasks from Airflow."""
+
+    default_view = "repair"
+
+    @expose("/repair_databricks_job/<string:dag_id>/<string:run_id>", methods=("GET",))
+    @get_auth_decorator()
+    def repair(self, dag_id: str, run_id: str):
+        return_url = self._get_return_url(dag_id, run_id)
+
+        tasks_to_repair = request.values.get("tasks_to_repair")
+        self.log.info("Tasks to repair: %s", tasks_to_repair)
+        if not tasks_to_repair:
+            flash("No tasks to repair. Not sending repair request.")
+            return redirect(return_url)
+
+        databricks_conn_id = request.values.get("databricks_conn_id")
+        databricks_run_id = request.values.get("databricks_run_id")
+
+        if not databricks_conn_id:
+            flash("No Databricks connection ID provided. Cannot repair tasks.")
+            return redirect(return_url)
+
+        if not databricks_run_id:
+            flash("No Databricks run ID provided. Cannot repair tasks.")
+            return redirect(return_url)
+
+        self.log.info("Repairing databricks job %s", databricks_run_id)
+        res = _repair_task(
+            databricks_conn_id=databricks_conn_id,
+            databricks_run_id=int(databricks_run_id),
+            tasks_to_repair=tasks_to_repair.split(","),
+            logger=self.log,
+        )
+        self.log.info("Repairing databricks job query for run %s sent", databricks_run_id)
+
+        self.log.info("Clearing tasks to rerun in airflow")
+
+        run_id = unquote(run_id)
+        _clear_task_instances(dag_id, run_id, tasks_to_repair.split(","), self.log)
+        flash(f"Databricks repair job is starting!: {res}")
+        return redirect(return_url)
+
+    @staticmethod
+    def _get_return_url(dag_id: str, run_id: str) -> str:
+        return url_for("Airflow.grid", dag_id=dag_id, dag_run_id=run_id)
+
+
+repair_databricks_view = RepairDatabricksTasks()
+
+repair_databricks_package = {
+    "view": repair_databricks_view,
+}
+
+
+class DatabricksWorkflowPlugin(AirflowPlugin):
+    """
+    Databricks Workflows plugin for Airflow.
+
+    .. seealso::
+        For more information on how to use this plugin, take a look at the guide:
+        :ref:`howto/plugin:DatabricksWorkflowPlugin`
+    """
+
+    name = "databricks_workflow"
+    operator_extra_links = [
+        WorkflowJobRepairAllFailedLink(),
+        WorkflowJobRepairSingleTaskLink(),
+        WorkflowJobRunLink(),
+    ]
+    appbuilder_views = [repair_databricks_package]

airflow/providers/databricks/utils/databricks.py

@@ -21,7 +21,7 @@ from airflow.exceptions import AirflowException
 from airflow.providers.databricks.hooks.databricks import RunState
 
 
-def _normalise_json_content(content, json_path: str = "json") -> str | bool | list | dict:
+def normalise_json_content(content, json_path: str = "json") -> str | bool | list | dict:
     """
     Normalize content or all values of content if it is a dict to a string.
 
@@ -33,7 +33,7 @@ def _normalise_json_content(content, json_path: str = "json") -> str | bool | li
     The only one exception is when we have boolean values, they can not be converted
     to string type because databricks does not understand 'True' or 'False' values.
     """
-    normalise = _normalise_json_content
+    normalise = normalise_json_content
     if isinstance(content, (str, bool)):
         return content
     elif isinstance(content, (int, float)):

{apache_airflow_providers_databricks-6.7.0rc1.dist-info → apache_airflow_providers_databricks-6.8.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.7.0rc1
+Version: 6.8.0rc1
 Summary: Provider package apache-airflow-providers-databricks for Apache Airflow
 Keywords: airflow-provider,databricks,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -34,8 +34,8 @@ Requires-Dist: azure-identity>=1.3.1 ; extra == "azure-identity"
 Requires-Dist: apache-airflow-providers-common-sql ; extra == "common.sql"
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.7.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.7.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.8.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.8.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -88,7 +88,7 @@ Provides-Extra: sdk
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.7.0.rc1``
+Release: ``6.8.0.rc1``
 
 
 `Databricks <https://databricks.com/>`__
@@ -101,7 +101,7 @@ This is a provider package for ``databricks`` provider. All classes for this pro
 are in ``airflow.providers.databricks`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.7.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.8.0/>`_.
 
 Installation
 ------------
@@ -149,4 +149,4 @@ Dependent package
 ============================================================================================================  ==============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.7.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.8.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.7.0rc1.dist-info → apache_airflow_providers_databricks-6.8.0rc1.dist-info}/RECORD

@@ -1,23 +1,25 @@
 airflow/providers/databricks/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
-airflow/providers/databricks/__init__.py,sha256=AQT8sipR-7rzOdJMQXLB7NINAOwTchfT9aLbkByAmiI,1497
-airflow/providers/databricks/get_provider_info.py,sha256=JyuncYtVVz2MguMPFpXeQKEzdGIYQXFxsIAPaC6FUek,7430
+airflow/providers/databricks/__init__.py,sha256=Gt2Qw2mCeamiuJ-Q4g-FmO_GFJC7dSaQ-8yIbEWFS5o,1497
+airflow/providers/databricks/get_provider_info.py,sha256=8xrnG7qAd4SNmRpCuqew3TC7HOez7CjK9yAhER_FoHo,7675
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/hooks/databricks.py,sha256=yBjtCJLN77Pd6VLGMS9sDjr3S431Z3AOwkACKIYUU6Q,24353
 airflow/providers/databricks/hooks/databricks_base.py,sha256=7zXCcIAdB8OpluU6a78pCkIT3UZXfjIWvvQhcgRo1dQ,28795
 airflow/providers/databricks/hooks/databricks_sql.py,sha256=sRT5gryPrFhCdjcXLzQ73OxtWUSLINwwDeBJMwRZERY,12569
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=8ZXgbD9ZSQmNKtEaPNVpysYup9_BCaz50Atz2aJV4-M,63784
+airflow/providers/databricks/operators/databricks.py,sha256=nxS_ktoAIMp_JvgXbjdYby6CHtg99YmPuG1Irny_2CM,65935
 airflow/providers/databricks/operators/databricks_repos.py,sha256=I1z2ppGfM_oPxR8BM6Nk1i6JuUOWB40dLbRNEZUvccA,13093
 airflow/providers/databricks/operators/databricks_sql.py,sha256=ABqM3aEeZELcVly2qUc4vKR1qmSW32yWYrS6seRwHi0,16800
-airflow/providers/databricks/operators/databricks_workflow.py,sha256=R0PdsjL30pfLeioT8RtrFyDG3omeHBUX8wLWXD-3iRI,12989
+airflow/providers/databricks/operators/databricks_workflow.py,sha256=qAGVx2skx1dNHXJ-JcDBG8TN0LfKiOFDNb1YDhKlZzI,13206
+airflow/providers/databricks/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/databricks/plugins/databricks_workflow.py,sha256=03dspUdWZW3AtvLjyG-A1QGff51pF8t4z-Cw5qBNguI,17216
 airflow/providers/databricks/sensors/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/sensors/databricks_partition.py,sha256=TRZVjGEYzIbN4nZDPJEVavXdcegpyj0BVKXINMRFlCM,10605
 airflow/providers/databricks/sensors/databricks_sql.py,sha256=8qSfbzpWgU2_oZU9pS9SA_MSYhyIHgYZjTYfDkDH84Y,5771
 airflow/providers/databricks/triggers/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/triggers/databricks.py,sha256=xk9aEfdZnG33a4WSFfg6SZF4FfROV8B4HOyBYBvZR_Q,5104
 airflow/providers/databricks/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/databricks/utils/databricks.py,sha256=VlowoI8DSxOQmS8n064wu-aWYfdhsO0jmgM872LOB_U,2892
-apache_airflow_providers_databricks-6.7.0rc1.dist-info/entry_points.txt,sha256=1WxGXTFDb107eV5Zmrt3p12J4LHYk56-ZKlvpOK7vg4,106
-apache_airflow_providers_databricks-6.7.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_databricks-6.7.0rc1.dist-info/METADATA,sha256=fW70HJBgkYlvtf1B5E91gVS9dwEEiDCp4k4ct-F6fTY,7087
-apache_airflow_providers_databricks-6.7.0rc1.dist-info/RECORD,,
+airflow/providers/databricks/utils/databricks.py,sha256=EICTPZTD0R0dy9UGKgv8srkrBTgzCQrcYNL9oBWuhzk,2890
+apache_airflow_providers_databricks-6.8.0rc1.dist-info/entry_points.txt,sha256=hjmZm3ab2cteTR4t9eE28oKixHwNIKtLCThd6sx3XRQ,227
+apache_airflow_providers_databricks-6.8.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_databricks-6.8.0rc1.dist-info/METADATA,sha256=aSv7hcbgFcaYzrsJ9FTyyopie8JZbpdnVrKc-cvVA3Q,7087
+apache_airflow_providers_databricks-6.8.0rc1.dist-info/RECORD,,

apache_airflow_providers_databricks-6.8.0rc1.dist-info/entry_points.txt

@@ -0,0 +1,6 @@
+[airflow.plugins]
+databricks_workflow=airflow.providers.databricks.plugins.databricks_workflow:DatabricksWorkflowPlugin
+
+[apache_airflow_provider]
+provider_info=airflow.providers.databricks.get_provider_info:get_provider_info
+

apache_airflow_providers_databricks-6.7.0rc1.dist-info/entry_points.txt

@@ -1,3 +0,0 @@
-[apache_airflow_provider]
-provider_info=airflow.providers.databricks.get_provider_info:get_provider_info
-