apache-airflow-providers-databricks 6.0.0rc1__tar.gz → 6.1.0__tar.gz

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 6.0.0rc1
+Version: 6.1.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>
@@ -21,15 +21,15 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: aiohttp>=3.6.3, <4
-Requires-Dist: apache-airflow-providers-common-sql>=1.10.0.dev0
-Requires-Dist: apache-airflow>=2.6.0.dev0
+Requires-Dist: apache-airflow-providers-common-sql>=1.10.0
+Requires-Dist: apache-airflow>=2.6.0
 Requires-Dist: databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0
 Requires-Dist: requests>=2.27,<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.0.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.0.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.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
@@ -81,7 +81,7 @@ Provides-Extra: sdk
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.0.0.rc1``
+Release: ``6.1.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -94,7 +94,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.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.0/>`_.
 
 Installation
 ------------
@@ -138,4 +138,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.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``6.0.0.rc1``
+Release: ``6.1.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.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.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.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.0/changelog.html>`_.

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/airflow/providers/databricks/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "6.0.0"
+__version__ = "6.1.0"
 
 try:
     from airflow import __version__ as airflow_version

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/airflow/providers/databricks/get_provider_info.py

@@ -27,9 +27,10 @@ def get_provider_info():
         "package-name": "apache-airflow-providers-databricks",
         "name": "Databricks",
         "description": "`Databricks <https://databricks.com/>`__\n",
-        "suspended": False,
-        "source-date-epoch": 1703288125,
+        "state": "ready",
+        "source-date-epoch": 1705912006,
         "versions": [
+            "6.1.0",
             "6.0.0",
             "5.0.1",
             "5.0.0",
@@ -70,6 +71,14 @@ def get_provider_info():
             "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0",
             "aiohttp>=3.6.3, <4",
         ],
+        "additional-extras": [
+            {
+                "name": "sdk",
+                "description": "Install Databricks SDK",
+                "dependencies": ["databricks-sdk==0.10.0"],
+            }
+        ],
+        "devel-dependencies": ["deltalake>=0.12.0"],
         "integrations": [
             {
                 "integration-name": "Databricks",
@@ -153,11 +162,4 @@ def get_provider_info():
             }
         ],
         "extra-links": ["airflow.providers.databricks.operators.databricks.DatabricksJobRunLink"],
-        "additional-extras": [
-            {
-                "name": "sdk",
-                "description": "Install Databricks SDK",
-                "dependencies": ["databricks-sdk==0.10.0"],
-            }
-        ],
     }

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/airflow/providers/databricks/hooks/databricks.py

@@ -519,13 +519,24 @@ class DatabricksHook(BaseDatabricksHook):
         json = {"run_id": run_id}
         self._do_api_call(DELETE_RUN_ENDPOINT, json)
 
-    def repair_run(self, json: dict) -> None:
+    def repair_run(self, json: dict) -> int:
         """
         Re-run one or more tasks.
 
         :param json: repair a job run.
         """
-        self._do_api_call(REPAIR_RUN_ENDPOINT, json)
+        response = self._do_api_call(REPAIR_RUN_ENDPOINT, json)
+        return response["repair_id"]
+
+    def get_latest_repair_id(self, run_id: int) -> int | None:
+        """Get latest repair id if any exist for run_id else None."""
+        json = {"run_id": run_id, "include_history": True}
+        response = self._do_api_call(GET_RUN_ENDPOINT, json)
+        repair_history = response["repair_history"]
+        if len(repair_history) == 1:
+            return None
+        else:
+            return repair_history[-1]["id"]
 
     def get_cluster_state(self, cluster_id: str) -> ClusterState:
         """

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/airflow/providers/databricks/hooks/databricks_sql.py

@@ -174,7 +174,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
             )
         return self._sql_conn
 
-    @overload
+    @overload  # type: ignore[override]
     def run(
         self,
         sql: str | Iterable[str],
@@ -249,7 +249,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
                 self.set_autocommit(conn, autocommit)
 
                 with closing(conn.cursor()) as cur:
-                    self._run_command(cur, sql_statement, parameters)
+                    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
                     if handler is not None:
                         raw_result = handler(cur)
                         if self.return_tuple:
@@ -280,13 +280,15 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
         # instantiated namedtuple, and will never do: https://github.com/python/mypy/issues/848
         if isinstance(result, list):
             rows: list[Row] = result
-            rows_fields = rows[0].__fields__
-            rows_object = namedtuple("Row", rows_fields)  # type: ignore[misc]
+            if not rows:
+                return []
+            rows_fields = tuple(rows[0].__fields__)
+            rows_object = namedtuple("Row", rows_fields, rename=True)  # type: ignore
             return cast(List[tuple], [rows_object(*row) for row in rows])
         else:
             row: Row = result
-            row_fields = row.__fields__
-            row_object = namedtuple("Row", row_fields)  # type: ignore[misc]
+            row_fields = tuple(row.__fields__)
+            row_object = namedtuple("Row", row_fields, rename=True)  # type: ignore
             return cast(tuple, row_object(*row))
 
     def bulk_dump(self, table, tmp_file):

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/airflow/providers/databricks/operators/databricks.py

@@ -88,6 +88,19 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
                             f"{operator.task_id} failed with terminal state: {run_state} "
                             f"and with the error {run_state.state_message}"
                         )
+                    if isinstance(operator, DatabricksRunNowOperator) and operator.repair_run:
+                        operator.repair_run = False
+                        log.warning(
+                            "%s but since repair run is set, repairing the run with all failed tasks",
+                            error_message,
+                        )
+
+                        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:
+                            repair_json["latest_repair_id"] = latest_repair_id
+                        operator.json["latest_repair_id"] = hook.repair_run(operator, repair_json)
+                        _handle_databricks_operator_execution(operator, hook, log, context)
                     raise AirflowException(error_message)
 
             else:
@@ -119,18 +132,24 @@ def _handle_deferrable_databricks_operator_execution(operator, hook, log, contex
     log.info("View run status, Spark UI, and logs at %s", run_page_url)
 
     if operator.wait_for_termination:
-        operator.defer(
-            trigger=DatabricksExecutionTrigger(
-                run_id=operator.run_id,
-                databricks_conn_id=operator.databricks_conn_id,
-                polling_period_seconds=operator.polling_period_seconds,
-                retry_limit=operator.databricks_retry_limit,
-                retry_delay=operator.databricks_retry_delay,
-                retry_args=operator.databricks_retry_args,
-                run_page_url=run_page_url,
-            ),
-            method_name=DEFER_METHOD_NAME,
-        )
+        run_info = hook.get_run(operator.run_id)
+        run_state = RunState(**run_info["state"])
+        if not run_state.is_terminal:
+            operator.defer(
+                trigger=DatabricksExecutionTrigger(
+                    run_id=operator.run_id,
+                    databricks_conn_id=operator.databricks_conn_id,
+                    polling_period_seconds=operator.polling_period_seconds,
+                    retry_limit=operator.databricks_retry_limit,
+                    retry_delay=operator.databricks_retry_delay,
+                    retry_args=operator.databricks_retry_args,
+                    run_page_url=run_page_url,
+                ),
+                method_name=DEFER_METHOD_NAME,
+            )
+        else:
+            if run_state.is_successful:
+                log.info("%s completed successfully.", operator.task_id)
 
 
 def _handle_deferrable_databricks_operator_completion(event: dict, log: Logger) -> None:
@@ -623,6 +642,7 @@ class DatabricksRunNowOperator(BaseOperator):
         - ``jar_params``
         - ``spark_submit_params``
         - ``idempotency_token``
+        - ``repair_run``
 
     :param job_id: the job_id of the existing Databricks job.
         This field will be templated.
@@ -711,6 +731,7 @@ class DatabricksRunNowOperator(BaseOperator):
     :param do_xcom_push: Whether we should push run_id and run_page_url to xcom.
     :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, doesn't work in deferrable mode
     """
 
     # Used in airflow.models.BaseOperator
@@ -741,6 +762,7 @@ class DatabricksRunNowOperator(BaseOperator):
         do_xcom_push: bool = True,
         wait_for_termination: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        repair_run: bool = False,
         **kwargs,
     ) -> None:
         """Create a new ``DatabricksRunNowOperator``."""
@@ -753,6 +775,7 @@ class DatabricksRunNowOperator(BaseOperator):
         self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
         self.deferrable = deferrable
+        self.repair_run = repair_run
 
         if job_id is not None:
             self.json["job_id"] = job_id

{apache_airflow_providers_databricks-6.0.0rc1 → apache_airflow_providers_databricks-6.1.0}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "6.0.0.rc1"
+version = "6.1.0"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -56,15 +56,15 @@ classifiers = [
 requires-python = "~=3.8"
 dependencies = [
     "aiohttp>=3.6.3, <4",
-    "apache-airflow-providers-common-sql>=1.10.0.dev0",
-    "apache-airflow>=2.6.0.dev0",
+    "apache-airflow-providers-common-sql>=1.10.0",
+    "apache-airflow>=2.6.0",
     "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0",
     "requests>=2.27,<3",
 ]
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.0.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.0.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-databricks/6.1.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"