apache-airflow-providers-databricks 6.4.0__tar.gz → 6.5.0__tar.gz

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.4.0
+Version: 6.5.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>
@@ -29,8 +29,8 @@ Requires-Dist: requests>=2.27.0,<3
 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.4.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.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
@@ -82,7 +82,7 @@ Provides-Extra: sdk
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.4.0``
+Release: ``6.5.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -95,7 +95,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.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0/>`_.
 
 Installation
 ------------
@@ -139,4 +139,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.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.4.0``
+Release: ``6.5.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.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0/>`_.
 
 Installation
 ------------
@@ -99,4 +99,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.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/airflow/providers/databricks/__init__.py

@@ -25,14 +25,11 @@ from __future__ import annotations
 
 import packaging.version
 
-__all__ = ["__version__"]
+from airflow import __version__ as airflow_version
 
-__version__ = "6.4.0"
+__all__ = ["__version__"]
 
-try:
-    from airflow import __version__ as airflow_version
-except ImportError:
-    from airflow.version import version as airflow_version
+__version__ = "6.5.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.7.0"

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.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": 1714476154,
+        "source-date-epoch": 1716287262,
         "versions": [
+            "6.5.0",
             "6.4.0",
             "6.3.0",
             "6.2.0",

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/airflow/providers/databricks/hooks/databricks_base.py

@@ -80,6 +80,7 @@ class BaseDatabricksHook(BaseHook):
     :param retry_delay: The number of seconds to wait between retries (it
         might be a floating point number).
     :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+    :param caller: The name of the operator that is calling the hook.
     """
 
     conn_name_attr: str = "databricks_conn_id"

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/airflow/providers/databricks/operators/databricks.py

@@ -67,23 +67,22 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
                     log.info("%s completed successfully.", operator.task_id)
                     log.info("View run status, Spark UI, and logs at %s", run_page_url)
                     return
-
                 if run_state.result_state == "FAILED":
-                    task_run_id = None
+                    failed_tasks = []
                     for task in run_info.get("tasks", []):
                         if task.get("state", {}).get("result_state", "") == "FAILED":
                             task_run_id = task["run_id"]
-                    if task_run_id is not None:
-                        run_output = hook.get_run_output(task_run_id)
-                        if "error" in run_output:
-                            notebook_error = run_output["error"]
-                        else:
-                            notebook_error = run_state.state_message
-                    else:
-                        notebook_error = run_state.state_message
+                            task_key = task["task_key"]
+                            run_output = hook.get_run_output(task_run_id)
+                            if "error" in run_output:
+                                error = run_output["error"]
+                            else:
+                                error = run_state.state_message
+                            failed_tasks.append({"task_key": task_key, "run_id": task_run_id, "error": error})
+
                     error_message = (
                         f"{operator.task_id} failed with terminal state: {run_state} "
-                        f"and with the error {notebook_error}"
+                        f"and with the errors {failed_tasks}"
                     )
                 else:
                     error_message = (
@@ -168,7 +167,7 @@ def _handle_deferrable_databricks_operator_completion(event: dict, log: Logger)
 
     error_message = f"Job run failed with terminal state: {run_state} and with the errors {errors}"
 
-    if event["repair_run"]:
+    if event.get("repair_run"):
         log.warning(
             "%s but since repair run is set, repairing the run with all failed tasks",
             error_message,
@@ -924,9 +923,11 @@ class DatabricksNotebookOperator(BaseOperator):
     :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default.
     :param databricks_conn_id: The name of the Airflow connection to use.
+    :param deferrable: Run operator in the deferrable mode.
     """
 
     template_fields = ("notebook_params",)
+    CALLER = "DatabricksNotebookOperator"
 
     def __init__(
         self,
@@ -943,6 +944,7 @@ class DatabricksNotebookOperator(BaseOperator):
         databricks_retry_args: dict[Any, Any] | None = None,
         wait_for_termination: bool = True,
         databricks_conn_id: str = "databricks_default",
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs: Any,
     ):
         self.notebook_path = notebook_path
@@ -959,11 +961,12 @@ class DatabricksNotebookOperator(BaseOperator):
         self.wait_for_termination = wait_for_termination
         self.databricks_conn_id = databricks_conn_id
         self.databricks_run_id: int | None = None
+        self.deferrable = deferrable
         super().__init__(**kwargs)
 
     @cached_property
     def _hook(self) -> DatabricksHook:
-        return self._get_hook(caller="DatabricksNotebookOperator")
+        return self._get_hook(caller=self.CALLER)
 
     def _get_hook(self, caller: str) -> DatabricksHook:
         return DatabricksHook(
@@ -971,7 +974,7 @@ class DatabricksNotebookOperator(BaseOperator):
             retry_limit=self.databricks_retry_limit,
             retry_delay=self.databricks_retry_delay,
             retry_args=self.databricks_retry_args,
-            caller=caller,
+            caller=self.CALLER,
         )
 
     def _get_task_timeout_seconds(self) -> int:
@@ -1042,6 +1045,19 @@ class DatabricksNotebookOperator(BaseOperator):
         run = self._hook.get_run(self.databricks_run_id)
         run_state = RunState(**run["state"])
         self.log.info("Current state of the job: %s", run_state.life_cycle_state)
+        if self.deferrable and not run_state.is_terminal:
+            self.defer(
+                trigger=DatabricksExecutionTrigger(
+                    run_id=self.databricks_run_id,
+                    databricks_conn_id=self.databricks_conn_id,
+                    polling_period_seconds=self.polling_period_seconds,
+                    retry_limit=self.databricks_retry_limit,
+                    retry_delay=self.databricks_retry_delay,
+                    retry_args=self.databricks_retry_args,
+                    caller=self.CALLER,
+                ),
+                method_name=DEFER_METHOD_NAME,
+            )
         while not run_state.is_terminal:
             time.sleep(self.polling_period_seconds)
             run = self._hook.get_run(self.databricks_run_id)
@@ -1057,9 +1073,7 @@ class DatabricksNotebookOperator(BaseOperator):
             )
         if not run_state.is_successful:
             raise AirflowException(
-                "Task failed. Final state %s. Reason: %s",
-                run_state.result_state,
-                run_state.state_message,
+                f"Task failed. Final state {run_state.result_state}. Reason: {run_state.state_message}"
             )
         self.log.info("Task succeeded. Final state %s.", run_state.result_state)
 
@@ -1067,3 +1081,16 @@ class DatabricksNotebookOperator(BaseOperator):
         self.launch_notebook_job()
         if self.wait_for_termination:
             self.monitor_databricks_job()
+
+    def execute_complete(self, context: dict | None, event: dict) -> None:
+        run_state = RunState.from_json(event["run_state"])
+        if run_state.life_cycle_state != "TERMINATED":
+            raise AirflowException(
+                f"Databricks job failed with state {run_state.life_cycle_state}. "
+                f"Message: {run_state.state_message}"
+            )
+        if not run_state.is_successful:
+            raise AirflowException(
+                f"Task failed. Final state {run_state.result_state}. Reason: {run_state.state_message}"
+            )
+        self.log.info("Task succeeded. Final state %s.", run_state.result_state)

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/airflow/providers/databricks/triggers/databricks.py

@@ -48,6 +48,7 @@ class DatabricksExecutionTrigger(BaseTrigger):
         retry_args: dict[Any, Any] | None = None,
         run_page_url: str | None = None,
         repair_run: bool = False,
+        caller: str = "DatabricksExecutionTrigger",
     ) -> None:
         super().__init__()
         self.run_id = run_id
@@ -63,6 +64,7 @@ class DatabricksExecutionTrigger(BaseTrigger):
             retry_limit=self.retry_limit,
             retry_delay=self.retry_delay,
             retry_args=retry_args,
+            caller=caller,
         )
 
     def serialize(self) -> tuple[str, dict[str, Any]]:

{apache_airflow_providers_databricks-6.4.0 → apache_airflow_providers_databricks-6.5.0}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "6.4.0"
+version = "6.5.0"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -64,8 +64,8 @@ dependencies = [
 ]
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.4.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.5.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"