apache-airflow-providers-databricks 6.9.0__tar.gz → 6.10.0__tar.gz

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.9.0
+Version: 6.10.0
 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.9.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.9.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.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.9.0``
+Release: ``6.10.0``
 
 
 `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.9.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.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.9.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.9.0``
+Release: ``6.10.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -55,7 +55,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.9.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0/>`_.
 
 Installation
 ------------
@@ -103,4 +103,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.9.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/__init__.py

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

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.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": 1723970061,
+        "source-date-epoch": 1726860611,
         "versions": [
+            "6.10.0",
             "6.9.0",
             "6.8.0",
             "6.7.0",

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/hooks/databricks.py

@@ -44,6 +44,7 @@ TERMINATE_CLUSTER_ENDPOINT = ("POST", "api/2.0/clusters/delete")
 
 CREATE_ENDPOINT = ("POST", "api/2.1/jobs/create")
 RESET_ENDPOINT = ("POST", "api/2.1/jobs/reset")
+UPDATE_ENDPOINT = ("POST", "api/2.1/jobs/update")
 RUN_NOW_ENDPOINT = ("POST", "api/2.1/jobs/run-now")
 SUBMIT_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/submit")
 GET_RUN_ENDPOINT = ("GET", "api/2.1/jobs/runs/get")
@@ -233,6 +234,15 @@ class DatabricksHook(BaseDatabricksHook):
         """
         self._do_api_call(RESET_ENDPOINT, {"job_id": job_id, "new_settings": json})
 
+    def update_job(self, job_id: str, json: dict) -> None:
+        """
+        Call the ``api/2.1/jobs/update`` endpoint.
+
+        :param job_id: The id of the job to update.
+        :param json: The data used in the new_settings of the request to the ``update`` endpoint.
+        """
+        self._do_api_call(UPDATE_ENDPOINT, {"job_id": job_id, "new_settings": json})
+
     def run_now(self, json: dict) -> int:
         """
         Call the ``api/2.1/jobs/run-now`` endpoint.

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/operators/databricks.py

@@ -100,13 +100,23 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
                         f"and with the error {run_state.state_message}"
                     )
 
-                if isinstance(operator, DatabricksRunNowOperator) and operator.repair_run:
+                should_repair = (
+                    isinstance(operator, DatabricksRunNowOperator)
+                    and operator.repair_run
+                    and (
+                        not operator.databricks_repair_reason_new_settings
+                        or is_repair_reason_match_exist(operator, run_state)
+                    )
+                )
+
+                if should_repair:
                     operator.repair_run = False
                     log.warning(
                         "%s but since repair run is set, repairing the run with all failed tasks",
                         error_message,
                     )
-
+                    job_id = operator.json["job_id"]
+                    update_job_for_repair(operator, hook, job_id, run_state)
                     latest_repair_id = hook.get_latest_repair_id(operator.run_id)
                     repair_json = {"run_id": operator.run_id, "rerun_all_failed_tasks": True}
                     if latest_repair_id is not None:
@@ -123,6 +133,41 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
     log.info("View run status, Spark UI, and logs at %s", run_page_url)
 
 
+def is_repair_reason_match_exist(operator: Any, run_state: RunState) -> bool:
+    """
+    Check if the repair reason matches the run state message.
+
+    :param operator: Databricks operator being handled
+    :param run_state: Run state of the Databricks job
+    :return: True if repair reason matches the run state message, False otherwise
+    """
+    return any(reason in run_state.state_message for reason in operator.databricks_repair_reason_new_settings)
+
+
+def update_job_for_repair(operator: Any, hook: Any, job_id: int, run_state: RunState) -> None:
+    """
+    Update job settings(partial) to repair the run with all failed tasks.
+
+    :param operator: Databricks operator being handled
+    :param hook: Databricks hook
+    :param job_id: Job ID of Databricks
+    :param run_state: Run state of the Databricks job
+    """
+    repair_reason = next(
+        (
+            reason
+            for reason in operator.databricks_repair_reason_new_settings
+            if reason in run_state.state_message
+        ),
+        None,
+    )
+    if repair_reason is not None:
+        new_settings_json = normalise_json_content(
+            operator.databricks_repair_reason_new_settings[repair_reason]
+        )
+        hook.update_job(job_id=job_id, json=new_settings_json)
+
+
 def _handle_deferrable_databricks_operator_execution(operator, hook, log, context) -> None:
     """
     Handle the Airflow + Databricks lifecycle logic for deferrable Databricks operators.
@@ -674,6 +719,7 @@ class DatabricksRunNowOperator(BaseOperator):
         - ``spark_submit_params``
         - ``idempotency_token``
         - ``repair_run``
+        - ``databricks_repair_reason_new_settings``
         - ``cancel_previous_runs``
 
     :param job_id: the job_id of the existing Databricks job.
@@ -764,6 +810,12 @@ class DatabricksRunNowOperator(BaseOperator):
     :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default.
     :param deferrable: Run operator in the deferrable mode.
     :param repair_run: Repair the databricks run in case of failure.
+    :param databricks_repair_reason_new_settings: A dict of reason and new_settings JSON object for which
+            to repair the run. `None` by default. `None` means to repair at all cases with existing job
+            settings otherwise check whether `RunState` state_message contains reason and
+            update job settings as per new_settings using databricks partial job update endpoint
+            (https://docs.databricks.com/api/workspace/jobs/update). If nothing is matched, then repair
+            will not get triggered.
     :param cancel_previous_runs: Cancel all existing running jobs before submitting new one.
     """
 
@@ -796,6 +848,7 @@ class DatabricksRunNowOperator(BaseOperator):
         wait_for_termination: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         repair_run: bool = False,
+        databricks_repair_reason_new_settings: dict[str, Any] | None = None,
         cancel_previous_runs: bool = False,
         **kwargs,
     ) -> None:
@@ -810,6 +863,7 @@ class DatabricksRunNowOperator(BaseOperator):
         self.wait_for_termination = wait_for_termination
         self.deferrable = deferrable
         self.repair_run = repair_run
+        self.databricks_repair_reason_new_settings = databricks_repair_reason_new_settings or {}
         self.cancel_previous_runs = cancel_previous_runs
 
         if job_id is not None:
@@ -870,9 +924,16 @@ class DatabricksRunNowOperator(BaseOperator):
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
         if event:
             _handle_deferrable_databricks_operator_completion(event, self.log)
-            if event["repair_run"]:
+            run_state = RunState.from_json(event["run_state"])
+            should_repair = event["repair_run"] and (
+                not self.databricks_repair_reason_new_settings
+                or is_repair_reason_match_exist(self, run_state)
+            )
+            if should_repair:
                 self.repair_run = False
                 self.run_id = event["run_id"]
+                job_id = self._hook.get_job_id(self.run_id)
+                update_job_for_repair(self, self._hook, job_id, run_state)
                 latest_repair_id = self._hook.get_latest_repair_id(self.run_id)
                 repair_json = {"run_id": self.run_id, "rerun_all_failed_tasks": True}
                 if latest_repair_id is not None:

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/plugins/databricks_workflow.py

@@ -24,7 +24,6 @@ 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
@@ -34,13 +33,11 @@ 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
 
@@ -55,15 +52,6 @@ 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)

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/sensors/databricks_partition.py

@@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence
 
 from databricks.sql.utils import ParamEscaper
 
-from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.sql import fetch_all_handler
 from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
 from airflow.sensors.base import BaseSensorOperator
@@ -182,10 +182,7 @@ class DatabricksPartitionSensor(BaseSensorOperator):
         partition_columns = self._sql_sensor(f"DESCRIBE DETAIL {table_name}")[0][7]
         self.log.debug("Partition columns: %s", partition_columns)
         if len(partition_columns) < 1:
-            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = f"Table {table_name} does not have partitions"
-            if self.soft_fail:
-                raise AirflowSkipException(message)
             raise AirflowException(message)
 
         formatted_opts = ""
@@ -207,17 +204,11 @@ class DatabricksPartitionSensor(BaseSensorOperator):
                             f"""{partition_col}{self.partition_operator}{self.escaper.escape_item(partition_value)}"""
                         )
                 else:
-                    # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
                     message = f"Column {partition_col} not part of table partitions: {partition_columns}"
-                    if self.soft_fail:
-                        raise AirflowSkipException(message)
                     raise AirflowException(message)
         else:
             # Raises exception if the table does not have any partitions.
-            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = "No partitions specified to check with the sensor."
-            if self.soft_fail:
-                raise AirflowSkipException(message)
             raise AirflowException(message)
         formatted_opts = f"{prefix} {joiner_val.join(output_list)} {suffix}"
         self.log.debug("Formatted options: %s", formatted_opts)
@@ -231,8 +222,5 @@ class DatabricksPartitionSensor(BaseSensorOperator):
         if partition_result:
             return True
         else:
-            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = f"Specified partition(s): {self.partitions} were not found."
-            if self.soft_fail:
-                raise AirflowSkipException(message)
             raise AirflowException(message)

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/airflow/providers/databricks/sensors/databricks_sql.py

@@ -23,7 +23,7 @@ from __future__ import annotations
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence
 
-from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.sql import fetch_all_handler
 from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
 from airflow.sensors.base import BaseSensorOperator
@@ -117,13 +117,10 @@ class DatabricksSqlSensor(BaseSensorOperator):
     def _get_results(self) -> bool:
         """Use the Databricks SQL hook and run the specified SQL query."""
         if not (self._http_path or self._sql_warehouse_name):
-            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = (
                 "Databricks SQL warehouse/cluster configuration missing. Please specify either"
                 " http_path or sql_warehouse_name."
             )
-            if self.soft_fail:
-                raise AirflowSkipException(message)
             raise AirflowException(message)
         hook = self.hook
         sql_result = hook.run(

{apache_airflow_providers_databricks-6.9.0 → apache_airflow_providers_databricks-6.10.0}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "6.9.0"
+version = "6.10.0"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -68,8 +68,8 @@ dependencies = [
 ]
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.9.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.9.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.10.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"