apache-airflow-providers-databricks 6.0.0rc1__py3-none-any.whl → 6.1.0__py3-none-any.whl

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

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

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

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):

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

@@ -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.dist-info → apache_airflow_providers_databricks-6.1.0.dist-info}/RECORD

@@ -1,12 +1,12 @@
 airflow/providers/databricks/LICENSE,sha256=ywUBpKZc7Jb96rVt5I3IDbg7dIJAbUSHkuoDcF3jbH4,13569
-airflow/providers/databricks/__init__.py,sha256=UcvoLm-UfmmEOWyZjjBQ8rf5a5aUs9Uoge2HmRJKwi0,1585
-airflow/providers/databricks/get_provider_info.py,sha256=tCIb8m8RY9FOo5YKewgLcpntbi4SWCNppvrx3qBycLo,6105
+airflow/providers/databricks/__init__.py,sha256=9F0qMBigUZVuMKQSrzhcdG9mjNAwMVYFZJ9Mu6fvUsQ,1585
+airflow/providers/databricks/get_provider_info.py,sha256=aAWTy02B1anUO7_gY3v3rtjH38L8ZX9MUnA9C-xFvYg,6177
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/hooks/databricks.py,sha256=GeqiBYFsFIwK4vgBTxde-bbdDliWzQAEbvZcShbyuio,22604
+airflow/providers/databricks/hooks/databricks.py,sha256=VRxltNFCRrZ9-U7yszhv1aCX48vdQcsnR-IHe_ACFa8,23076
 airflow/providers/databricks/hooks/databricks_base.py,sha256=nX-40P0XKZyFyahkbbImvWmoCQ4ic35viAYlFIupUHE,30591
-airflow/providers/databricks/hooks/databricks_sql.py,sha256=afzw9l1VNBFvReEheAUSrGgXc7FWePcwuYw9-qMLJUg,12445
+airflow/providers/databricks/hooks/databricks_sql.py,sha256=y55dj2Z1T-5BaMrCh60QT7IW1C_-JiXJBZRU_FJZpvY,12580
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=XmQuoWzesOYpwAtxn3aB-1CihbnT8W67WoxUoAh0Bnc,40012
+airflow/providers/databricks/operators/databricks.py,sha256=g89WZV9VZaLTGEDSdK9gt3GAZfFTrziGa5O1o6o-qkw,41389
 airflow/providers/databricks/operators/databricks_repos.py,sha256=NUxa0jvvmK16CDKb-7Tbs3wF9XoFi1AVJlKxlsE3r4k,13092
 airflow/providers/databricks/operators/databricks_sql.py,sha256=C_XqR5SN_g2t0XXlH248nPAC7BcYB0P1NAsPJT5qcEg,16806
 airflow/providers/databricks/sensors/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
@@ -16,7 +16,7 @@ airflow/providers/databricks/triggers/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvR
 airflow/providers/databricks/triggers/databricks.py,sha256=Qj9mB0bNYRY_toPEU17gxbxmPkkT3P789kCHu_T64BA,3997
 airflow/providers/databricks/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/utils/databricks.py,sha256=iRzRHvdFETGiFxZccOjfC8NGgDofMfP35Tqp3M5CGr0,2880
-apache_airflow_providers_databricks-6.0.0rc1.dist-info/entry_points.txt,sha256=1WxGXTFDb107eV5Zmrt3p12J4LHYk56-ZKlvpOK7vg4,106
-apache_airflow_providers_databricks-6.0.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_databricks-6.0.0rc1.dist-info/METADATA,sha256=Sd37-KPVkHCMIBl1KPqZAE9xUslg0pKMojaZ6SibGog,6444
-apache_airflow_providers_databricks-6.0.0rc1.dist-info/RECORD,,
+apache_airflow_providers_databricks-6.1.0.dist-info/entry_points.txt,sha256=1WxGXTFDb107eV5Zmrt3p12J4LHYk56-ZKlvpOK7vg4,106
+apache_airflow_providers_databricks-6.1.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_databricks-6.1.0.dist-info/METADATA,sha256=nPsc_67tyRVrEcjehAw3sv98WGWcasE7J01nMipjO8U,6427
+apache_airflow_providers_databricks-6.1.0.dist-info/RECORD,,