apache-airflow-providers-databricks 6.4.0rc1__py3-none-any.whl → 6.5.0__py3-none-any.whl

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"

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",

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"

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)

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.0rc1.dist-info → apache_airflow_providers_databricks-6.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.4.0rc1
+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>
@@ -22,15 +22,15 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: aiohttp>=3.9.2, <4
-Requires-Dist: apache-airflow-providers-common-sql>=1.10.0rc0
-Requires-Dist: apache-airflow>=2.7.0rc0
+Requires-Dist: apache-airflow-providers-common-sql>=1.10.0
+Requires-Dist: apache-airflow>=2.7.0
 Requires-Dist: databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0
 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.rc1``
+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.0rc1.dist-info → apache_airflow_providers_databricks-6.5.0.dist-info}/RECORD

@@ -1,22 +1,22 @@
 airflow/providers/databricks/LICENSE,sha256=ywUBpKZc7Jb96rVt5I3IDbg7dIJAbUSHkuoDcF3jbH4,13569
-airflow/providers/databricks/__init__.py,sha256=ZJ0GA08oUAP30uPjov_SNDpKfoOR5bQ017KXjsPXIqk,1585
-airflow/providers/databricks/get_provider_info.py,sha256=Jxx0xdxa7i4-Gxug9kbVNN0sZjPuNFDVP1Ewux3aLfQ,6330
+airflow/providers/databricks/__init__.py,sha256=JsYXdg4453ofY8-fk8U0XwsqkJDG10zVtOPYdrpl9-Y,1497
+airflow/providers/databricks/get_provider_info.py,sha256=CZqPGkTq9yxBbzG5iAGX3Y21HKDIRhHn5KnJeKfizY8,6351
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/hooks/databricks.py,sha256=Pq42FzDGCZ2B0AaZw1zOD74uuXwJWHvwh2YkcQ8G-vk,23715
-airflow/providers/databricks/hooks/databricks_base.py,sha256=uMHbJSL8rdAgb2PKyQnj7F__Fji5lSJgszJrVsUowig,30592
+airflow/providers/databricks/hooks/databricks_base.py,sha256=XtLF6qK5QSMwVkhorniHZbos_GA4XHIxI8urBs8zdOQ,30662
 airflow/providers/databricks/hooks/databricks_sql.py,sha256=8LxnnVEKfk9nIJam3LR8E8uiOqNauMP0WgFa9eKX3Go,12564
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=WW_QJ14qzDmOWOANSZ8mgdgGBylKOiflB7tyszFqd-w,51327
+airflow/providers/databricks/operators/databricks.py,sha256=OCWUA_UzR8LyEMpbfcPejAgHD8rtDD9ILAD4GggWD9E,52819
 airflow/providers/databricks/operators/databricks_repos.py,sha256=I1z2ppGfM_oPxR8BM6Nk1i6JuUOWB40dLbRNEZUvccA,13093
 airflow/providers/databricks/operators/databricks_sql.py,sha256=ABqM3aEeZELcVly2qUc4vKR1qmSW32yWYrS6seRwHi0,16800
 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=LEj3YMr_yIS3z5RP0TzZ0XlRBE1ognPzcbYaiEJT8G8,5025
+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=EICTPZTD0R0dy9UGKgv8srkrBTgzCQrcYNL9oBWuhzk,2890
-apache_airflow_providers_databricks-6.4.0rc1.dist-info/entry_points.txt,sha256=1WxGXTFDb107eV5Zmrt3p12J4LHYk56-ZKlvpOK7vg4,106
-apache_airflow_providers_databricks-6.4.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_databricks-6.4.0rc1.dist-info/METADATA,sha256=2tX3_hVSvXOED2O3s42IBN_7saKuQjPCZ6kVIwVA-PE,6500
-apache_airflow_providers_databricks-6.4.0rc1.dist-info/RECORD,,
+apache_airflow_providers_databricks-6.5.0.dist-info/entry_points.txt,sha256=1WxGXTFDb107eV5Zmrt3p12J4LHYk56-ZKlvpOK7vg4,106
+apache_airflow_providers_databricks-6.5.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_databricks-6.5.0.dist-info/METADATA,sha256=L7_iBfAW2fu9RRcnKaCHSbc1BPCXh8Mb_m4W8MUxvmQ,6487
+apache_airflow_providers_databricks-6.5.0.dist-info/RECORD,,