apache-airflow-providers-databricks 4.1.0rc1__tar.gz → 4.2.0__tar.gz

{apache-airflow-providers-databricks-4.1.0rc1/apache_airflow_providers_databricks.egg-info → apache-airflow-providers-databricks-4.2.0}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 4.1.0rc1
+Version: 4.2.0
 Summary: Provider for Apache Airflow. Implements apache-airflow-providers-databricks package
 Home-page: https://airflow.apache.org/
 Download-URL: https://archive.apache.org/dist/airflow/providers
 Author: Apache Software Foundation
 Author-email: dev@airflow.apache.org
 License: Apache License 2.0
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.1.0/
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
@@ -53,7 +53,7 @@ License-File: NOTICE
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``4.1.0rc1``
+Release: ``4.2.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -66,7 +66,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/4.1.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/>`_.
 
 
 Installation
@@ -84,8 +84,8 @@ Requirements
 =======================================  ===================
 PIP package                              Version required
 =======================================  ===================
-``apache-airflow``                       ``>=2.3.0``
-``apache-airflow-providers-common-sql``  ``>=1.3.1``
+``apache-airflow``                       ``>=2.4.0``
+``apache-airflow-providers-common-sql``  ``>=1.5.0``
 ``requests``                             ``>=2.27,<3``
 ``databricks-sql-connector``             ``>=2.0.0, <3.0.0``
 ``aiohttp``                              ``>=3.6.3, <4``
@@ -136,6 +136,38 @@ Dependent package
 Changelog
 ---------
 
+4.2.0
+.....
+
+.. note::
+  This release of provider is only available for Airflow 2.4+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
+
+Features
+~~~~~~~~
+
+* ``Add conditional output processing in SQL operators (#31136)``
+* ``Add cancel all runs functionality to Databricks hook (#31038)``
+* ``Add retry param in databrics async operator (#30744)``
+* ``Add repair job functionality to databricks hook (#30786)``
+* ``Add 'DatabricksPartitionSensor' (#30980)``
+
+Misc
+~~~~
+
+* ``Bump minimum Airflow version in providers (#30917)``
+* ``Deprecate databricks async operator (#30761)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Move TaskInstanceKey to a separate file (#31033)``
+   * ``Use 'AirflowProviderDeprecationWarning' in providers (#30975)``
+   * ``Add full automation for min Airflow version for providers (#30994)``
+   * ``Add cli cmd to list the provider trigger info (#30822)``
+   * ``Use '__version__' in providers not 'version' (#31393)``
+   * ``Fixing circular import error in providers caused by airflow version check (#31379)``
+   * ``Prepare docs for May 2023 wave of Providers (#31252)``
+
 4.1.0
 .....
 
@@ -203,8 +235,9 @@ Bug Fixes
 3.4.0
 .....
 
-This release of provider is only available for Airflow 2.3+ as explained in the
-`Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/README.md#support-for-providers>`_.
+.. note::
+  This release of provider is only available for Airflow 2.3+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Misc
 ~~~~
@@ -300,8 +333,9 @@ Bug Fixes
 Breaking changes
 ~~~~~~~~~~~~~~~~
 
-* This release of provider is only available for Airflow 2.2+ as explained in the Apache Airflow
-  providers support policy https://github.com/apache/airflow/blob/main/README.md#support-for-providers
+.. note::
+  This release of provider is only available for Airflow 2.2+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Features
 ~~~~~~~~

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/README.rst

@@ -19,7 +19,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``4.1.0rc1``
+Release: ``4.2.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -32,7 +32,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/4.1.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/>`_.
 
 
 Installation
@@ -50,8 +50,8 @@ Requirements
 =======================================  ===================
 PIP package                              Version required
 =======================================  ===================
-``apache-airflow``                       ``>=2.3.0``
-``apache-airflow-providers-common-sql``  ``>=1.3.1``
+``apache-airflow``                       ``>=2.4.0``
+``apache-airflow-providers-common-sql``  ``>=1.5.0``
 ``requests``                             ``>=2.27,<3``
 ``databricks-sql-connector``             ``>=2.0.0, <3.0.0``
 ``aiohttp``                              ``>=3.6.3, <4``
@@ -102,6 +102,38 @@ Dependent package
 Changelog
 ---------
 
+4.2.0
+.....
+
+.. note::
+  This release of provider is only available for Airflow 2.4+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
+
+Features
+~~~~~~~~
+
+* ``Add conditional output processing in SQL operators (#31136)``
+* ``Add cancel all runs functionality to Databricks hook (#31038)``
+* ``Add retry param in databrics async operator (#30744)``
+* ``Add repair job functionality to databricks hook (#30786)``
+* ``Add 'DatabricksPartitionSensor' (#30980)``
+
+Misc
+~~~~
+
+* ``Bump minimum Airflow version in providers (#30917)``
+* ``Deprecate databricks async operator (#30761)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Move TaskInstanceKey to a separate file (#31033)``
+   * ``Use 'AirflowProviderDeprecationWarning' in providers (#30975)``
+   * ``Add full automation for min Airflow version for providers (#30994)``
+   * ``Add cli cmd to list the provider trigger info (#30822)``
+   * ``Use '__version__' in providers not 'version' (#31393)``
+   * ``Fixing circular import error in providers caused by airflow version check (#31379)``
+   * ``Prepare docs for May 2023 wave of Providers (#31252)``
+
 4.1.0
 .....
 
@@ -169,8 +201,9 @@ Bug Fixes
 3.4.0
 .....
 
-This release of provider is only available for Airflow 2.3+ as explained in the
-`Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/README.md#support-for-providers>`_.
+.. note::
+  This release of provider is only available for Airflow 2.3+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Misc
 ~~~~
@@ -266,8 +299,9 @@ Bug Fixes
 Breaking changes
 ~~~~~~~~~~~~~~~~
 
-* This release of provider is only available for Airflow 2.2+ as explained in the Apache Airflow
-  providers support policy https://github.com/apache/airflow/blob/main/README.md#support-for-providers
+.. note::
+  This release of provider is only available for Airflow 2.2+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Features
 ~~~~~~~~

apache-airflow-providers-databricks-4.2.0/airflow/providers/databricks/__init__.py

@@ -0,0 +1,41 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
+# OVERWRITTEN WHEN PREPARING DOCUMENTATION FOR THE PACKAGES.
+#
+# IF YOU WANT TO MODIFY IT, YOU SHOULD MODIFY THE TEMPLATE
+# `PROVIDER__INIT__PY_TEMPLATE.py.jinja2` IN the `dev/provider_packages` DIRECTORY
+#
+from __future__ import annotations
+
+import packaging.version
+
+__all__ = ["__version__"]
+
+__version__ = "4.2.0"
+
+try:
+    from airflow import __version__ as airflow_version
+except ImportError:
+    from airflow.version import version as airflow_version
+
+if packaging.version.parse(airflow_version) < packaging.version.parse("2.4.0"):
+    raise RuntimeError(
+        f"The package `apache-airflow-providers-databricks:{__version__}` requires Apache Airflow 2.4.0+"  # NOQA: E501
+    )

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/airflow/providers/databricks/get_provider_info.py

@@ -29,6 +29,7 @@ def get_provider_info():
         "description": "`Databricks <https://databricks.com/>`__\n",
         "suspended": False,
         "versions": [
+            "4.2.0",
             "4.1.0",
             "4.0.1",
             "4.0.0",
@@ -51,8 +52,8 @@ def get_provider_info():
             "1.0.0",
         ],
         "dependencies": [
-            "apache-airflow>=2.3.0",
-            "apache-airflow-providers-common-sql>=1.3.1",
+            "apache-airflow>=2.4.0",
+            "apache-airflow-providers-common-sql>=1.5.0",
             "requests>=2.27,<3",
             "databricks-sql-connector>=2.0.0, <3.0.0",
             "aiohttp>=3.6.3, <4",
@@ -120,13 +121,18 @@ def get_provider_info():
         "triggers": [
             {
                 "integration-name": "Databricks",
-                "python-modules": ["airflow.providers.databricks.triggers.databricks"],
+                "class-names": [
+                    "airflow.providers.databricks.triggers.databricks.DatabricksExecutionTrigger"
+                ],
             }
         ],
         "sensors": [
             {
                 "integration-name": "Databricks",
-                "python-modules": ["airflow.providers.databricks.sensors.databricks_sql"],
+                "python-modules": [
+                    "airflow.providers.databricks.sensors.databricks_sql",
+                    "airflow.providers.databricks.sensors.databricks_partition",
+                ],
             }
         ],
         "connection-types": [

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/airflow/providers/databricks/hooks/databricks.py

@@ -44,7 +44,9 @@ SUBMIT_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/submit")
 GET_RUN_ENDPOINT = ("GET", "api/2.1/jobs/runs/get")
 CANCEL_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/cancel")
 DELETE_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/delete")
+REPAIR_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/repair")
 OUTPUT_RUNS_JOB_ENDPOINT = ("GET", "api/2.1/jobs/runs/get-output")
+CANCEL_ALL_RUNS_ENDPOINT = ("POST", "api/2.1/jobs/runs/cancel-all")
 
 INSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/install")
 UNINSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/uninstall")
@@ -352,6 +354,15 @@ class DatabricksHook(BaseDatabricksHook):
         json = {"run_id": run_id}
         self._do_api_call(CANCEL_RUN_ENDPOINT, json)
 
+    def cancel_all_runs(self, job_id: int) -> None:
+        """
+        Cancels all active runs of a job. The runs are canceled asynchronously.
+
+        :param job_id: The canonical identifier of the job to cancel all runs of
+        """
+        json = {"job_id": job_id}
+        self._do_api_call(CANCEL_ALL_RUNS_ENDPOINT, json)
+
     def delete_run(self, run_id: int) -> None:
         """
         Deletes a non-active run.
@@ -361,6 +372,14 @@ class DatabricksHook(BaseDatabricksHook):
         json = {"run_id": run_id}
         self._do_api_call(DELETE_RUN_ENDPOINT, json)
 
+    def repair_run(self, json: dict) -> None:
+        """
+        Re-run one or more tasks.
+
+        :param json: repair a job run.
+        """
+        self._do_api_call(REPAIR_RUN_ENDPOINT, json)
+
     def restart_cluster(self, json: dict) -> None:
         """
         Restarts the cluster.

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/airflow/providers/databricks/operators/databricks.py

@@ -19,22 +19,24 @@
 from __future__ import annotations
 
 import time
+import warnings
 from logging import Logger
 from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.compat.functools import cached_property
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator, BaseOperatorLink, XCom
 from airflow.providers.databricks.hooks.databricks import DatabricksHook, RunState
 from airflow.providers.databricks.triggers.databricks import DatabricksExecutionTrigger
 from airflow.providers.databricks.utils.databricks import normalise_json_content, validate_trigger_event
 
 if TYPE_CHECKING:
-    from airflow.models.taskinstance import TaskInstanceKey
+    from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
 
 DEFER_METHOD_NAME = "execute_complete"
 XCOM_RUN_ID_KEY = "run_id"
+XCOM_JOB_ID_KEY = "job_id"
 XCOM_RUN_PAGE_URL_KEY = "run_page_url"
 
 
@@ -103,6 +105,9 @@ def _handle_deferrable_databricks_operator_execution(operator, hook, log, contex
     :param operator: Databricks async operator being handled
     :param context: Airflow context
     """
+    job_id = hook.get_job_id(operator.run_id)
+    if operator.do_xcom_push and context is not None:
+        context["ti"].xcom_push(key=XCOM_JOB_ID_KEY, value=job_id)
     if operator.do_xcom_push and context is not None:
         context["ti"].xcom_push(key=XCOM_RUN_ID_KEY, value=operator.run_id)
     log.info("Run submitted with run_id: %s", operator.run_id)
@@ -118,6 +123,10 @@ def _handle_deferrable_databricks_operator_execution(operator, hook, log, contex
                 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,
         )
@@ -267,6 +276,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
     :param do_xcom_push: Whether we should push run_id and run_page_url to xcom.
     :param git_source: Optional specification of a remote git repository from which
         supported task types are retrieved.
+    :param deferrable: Run operator in the deferrable mode.
 
         .. seealso::
             https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsRunsSubmit
@@ -306,6 +316,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
         access_control_list: list[dict[str, str]] | None = None,
         wait_for_termination: bool = True,
         git_source: dict[str, str] | None = None,
+        deferrable: bool = False,
         **kwargs,
     ) -> None:
         """Creates a new ``DatabricksSubmitRunOperator``."""
@@ -317,6 +328,7 @@ class DatabricksSubmitRunOperator(BaseOperator):
         self.databricks_retry_delay = databricks_retry_delay
         self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
+        self.deferrable = deferrable
         if tasks is not None:
             self.json["tasks"] = tasks
         if spark_jar_task is not None:
@@ -373,7 +385,10 @@ class DatabricksSubmitRunOperator(BaseOperator):
     def execute(self, context: Context):
         json_normalised = normalise_json_content(self.json)
         self.run_id = self._hook.submit_run(json_normalised)
-        _handle_databricks_operator_execution(self, self._hook, self.log, context)
+        if self.deferrable:
+            _handle_deferrable_databricks_operator_execution(self, self._hook, self.log, context)
+        else:
+            _handle_databricks_operator_execution(self, self._hook, self.log, context)
 
     def on_kill(self):
         if self.run_id:
@@ -384,10 +399,23 @@ class DatabricksSubmitRunOperator(BaseOperator):
         else:
             self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
 
+    def execute_complete(self, context: dict | None, event: dict):
+        _handle_deferrable_databricks_operator_completion(event, self.log)
+
 
 class DatabricksSubmitRunDeferrableOperator(DatabricksSubmitRunOperator):
     """Deferrable version of ``DatabricksSubmitRunOperator``"""
 
+    def __init__(self, *args, **kwargs):
+        warnings.warn(
+            "`DatabricksSubmitRunDeferrableOperator` has been deprecated. "
+            "Please use `airflow.providers.databricks.operators.DatabricksSubmitRunOperator` with "
+            "`deferrable=True` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(deferrable=True, *args, **kwargs)
+
     def execute(self, context):
         hook = self._get_hook(caller="DatabricksSubmitRunDeferrableOperator")
         json_normalised = normalise_json_content(self.json)
@@ -549,6 +577,7 @@ class DatabricksRunNowOperator(BaseOperator):
     :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
     :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.
     """
 
     # Used in airflow.models.BaseOperator
@@ -578,6 +607,7 @@ class DatabricksRunNowOperator(BaseOperator):
         databricks_retry_args: dict[Any, Any] | None = None,
         do_xcom_push: bool = True,
         wait_for_termination: bool = True,
+        deferrable: bool = False,
         **kwargs,
     ) -> None:
         """Creates a new ``DatabricksRunNowOperator``."""
@@ -589,6 +619,7 @@ class DatabricksRunNowOperator(BaseOperator):
         self.databricks_retry_delay = databricks_retry_delay
         self.databricks_retry_args = databricks_retry_args
         self.wait_for_termination = wait_for_termination
+        self.deferrable = deferrable
 
         if job_id is not None:
             self.json["job_id"] = job_id
@@ -636,7 +667,10 @@ class DatabricksRunNowOperator(BaseOperator):
             self.json["job_id"] = job_id
             del self.json["job_name"]
         self.run_id = hook.run_now(self.json)
-        _handle_databricks_operator_execution(self, hook, self.log, context)
+        if self.deferrable:
+            _handle_deferrable_databricks_operator_execution(self, hook, self.log, context)
+        else:
+            _handle_databricks_operator_execution(self, hook, self.log, context)
 
     def on_kill(self):
         if self.run_id:
@@ -651,6 +685,16 @@ class DatabricksRunNowOperator(BaseOperator):
 class DatabricksRunNowDeferrableOperator(DatabricksRunNowOperator):
     """Deferrable version of ``DatabricksRunNowOperator``"""
 
+    def __init__(self, *args, **kwargs):
+        warnings.warn(
+            "`DatabricksRunNowDeferrableOperator` has been deprecated. "
+            "Please use `airflow.providers.databricks.operators.DatabricksRunNowOperator` with "
+            "`deferrable=True` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(deferrable=True, *args, **kwargs)
+
     def execute(self, context):
         hook = self._get_hook(caller="DatabricksRunNowDeferrableOperator")
         self.run_id = hook.run_now(self.json)

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/airflow/providers/databricks/operators/databricks_sql.py

@@ -120,6 +120,9 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
         }
         return DatabricksSqlHook(self.databricks_conn_id, **hook_params)
 
+    def _should_run_output_processing(self) -> bool:
+        return self.do_xcom_push or bool(self._output_path)
+
     def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequence] | None]) -> list[Any]:
         if not self._output_path:
             return list(zip(descriptions, results))

apache-airflow-providers-databricks-4.2.0/airflow/providers/databricks/sensors/databricks_partition.py

@@ -0,0 +1,228 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+"""This module contains Databricks sensors."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Callable, Sequence
+
+from databricks.sql.utils import ParamEscaper
+
+from airflow.compat.functools import cached_property
+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
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class DatabricksPartitionSensor(BaseSensorOperator):
+    """
+    Sensor to detect the presence of table partitions in Databricks.
+
+    :param databricks_conn_id: Reference to :ref:`Databricks
+        connection id<howto/connection:databricks>` (templated), defaults to
+        DatabricksSqlHook.default_conn_name.
+    :param sql_warehouse_name: Optional name of Databricks SQL warehouse. If not specified, ``http_path``
+        must be provided as described below, defaults to None
+    :param http_path: Optional string specifying HTTP path of Databricks SQL warehouse or All Purpose cluster.
+        If not specified, it should be either specified in the Databricks connection's
+        extra parameters, or ``sql_warehouse_name`` must be specified.
+    :param session_configuration: An optional dictionary of Spark session parameters. If not specified,
+        it could be specified in the Databricks connection's extra parameters, defaults to None
+    :param http_headers: An optional list of (k, v) pairs
+        that will be set as HTTP headers on every request. (templated).
+    :param catalog: An optional initial catalog to use.
+        Requires Databricks Runtime version 9.0+ (templated), defaults to ""
+    :param schema: An optional initial schema to use.
+        Requires Databricks Runtime version 9.0+ (templated), defaults to "default"
+    :param table_name: Name of the table to check partitions.
+    :param partitions: Name of the partitions to check.
+        Example: {"date": "2023-01-03", "name": ["abc", "def"]}
+    :param partition_operator: Optional comparison operator for partitions, such as >=.
+    :param handler: Handler for DbApiHook.run() to return results, defaults to fetch_all_handler
+    :param client_parameters: Additional parameters internal to Databricks SQL connector parameters.
+    """
+
+    template_fields: Sequence[str] = (
+        "databricks_conn_id",
+        "catalog",
+        "schema",
+        "table_name",
+        "partitions",
+        "http_headers",
+    )
+
+    template_ext: Sequence[str] = (".sql",)
+    template_fields_renderers = {"sql": "sql"}
+
+    def __init__(
+        self,
+        *,
+        databricks_conn_id: str = DatabricksSqlHook.default_conn_name,
+        http_path: str | None = None,
+        sql_warehouse_name: str | None = None,
+        session_configuration=None,
+        http_headers: list[tuple[str, str]] | None = None,
+        catalog: str = "",
+        schema: str = "default",
+        table_name: str,
+        partitions: dict,
+        partition_operator: str = "=",
+        handler: Callable[[Any], Any] = fetch_all_handler,
+        client_parameters: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        self.databricks_conn_id = databricks_conn_id
+        self._http_path = http_path
+        self._sql_warehouse_name = sql_warehouse_name
+        self.session_config = session_configuration
+        self.http_headers = http_headers
+        self.catalog = catalog
+        self.schema = schema
+        self.caller = "DatabricksPartitionSensor"
+        self.partitions = partitions
+        self.partition_operator = partition_operator
+        self.table_name = table_name
+        self.client_parameters = client_parameters or {}
+        self.hook_params = kwargs.pop("hook_params", {})
+        self.handler = handler
+        self.escaper = ParamEscaper()
+        super().__init__(**kwargs)
+
+    def _sql_sensor(self, sql):
+        """Executes the supplied SQL statement using the hook object."""
+        hook = self._get_hook
+        sql_result = hook.run(
+            sql,
+            handler=self.handler if self.do_xcom_push else None,
+        )
+        self.log.debug("SQL result: %s", sql_result)
+        return sql_result
+
+    @cached_property
+    def _get_hook(self) -> DatabricksSqlHook:
+        """Creates and returns a DatabricksSqlHook object."""
+        return DatabricksSqlHook(
+            self.databricks_conn_id,
+            self._http_path,
+            self._sql_warehouse_name,
+            self.session_config,
+            self.http_headers,
+            self.catalog,
+            self.schema,
+            self.caller,
+            **self.client_parameters,
+            **self.hook_params,
+        )
+
+    def _check_table_partitions(self) -> list:
+        """
+        The method performs the following:
+        * Generates the fully qualified table name.
+        * Calls the generate partition query.
+        * Based on the result returned by the partition generation method,
+        the _sql_sensor method is called.
+        """
+        if self.table_name.split(".")[0] == "delta":
+            _fully_qualified_table_name = self.table_name
+        else:
+            _fully_qualified_table_name = str(self.catalog + "." + self.schema + "." + self.table_name)
+        self.log.debug("Table name generated from arguments: %s", _fully_qualified_table_name)
+        _joiner_val = " AND "
+        _prefix = f"SELECT 1 FROM {_fully_qualified_table_name} WHERE"
+        _suffix = " LIMIT 1"
+
+        partition_sql = self._generate_partition_query(
+            prefix=_prefix,
+            suffix=_suffix,
+            joiner_val=_joiner_val,
+            opts=self.partitions,
+            table_name=_fully_qualified_table_name,
+            escape_key=False,
+        )
+        return self._sql_sensor(partition_sql)
+
+    def _generate_partition_query(
+        self,
+        prefix: str,
+        suffix: str,
+        joiner_val: str,
+        table_name: str,
+        opts: dict[str, str] | None = None,
+        escape_key: bool = False,
+    ) -> str:
+        """
+        Queries the table for available partitions.
+        Generates the SQL query based on the partition data types.
+            * For a list, it prepares the SQL in the format:
+                column_name in (value1, value2,...)
+            * For a numeric type, it prepares the format:
+                column_name =(or other provided operator such as >=) value
+            * For a date type, it prepares the format:
+                column_name =(or other provided operator such as >=) value
+        Once the filter predicates have been generated like above, the query
+        is prepared to be executed using the prefix and suffix supplied, which are:
+        "SELECT 1 FROM {_fully_qualified_table_name} WHERE" and "LIMIT 1".
+        """
+        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:
+            raise AirflowException(f"Table {table_name} does not have partitions")
+        formatted_opts = ""
+        if opts is not None and len(opts) > 0:
+            output_list = []
+            for partition_col, partition_value in opts.items():
+                if escape_key:
+                    partition_col = self.escaper.escape_item(partition_col)
+                if partition_col in partition_columns:
+                    if isinstance(partition_value, list):
+                        output_list.append(f"""{partition_col} in {tuple(partition_value)}""")
+                        self.log.debug("List formatting for partitions: %s", output_list)
+                    if isinstance(partition_value, (int, float, complex)):
+                        output_list.append(
+                            f"""{partition_col}{self.partition_operator}{self.escaper.escape_item(partition_value)}"""
+                        )
+                    if isinstance(partition_value, (str, datetime)):
+                        output_list.append(
+                            f"""{partition_col}{self.partition_operator}{self.escaper.escape_item(partition_value)}"""
+                        )
+                else:
+                    raise AirflowException(
+                        f"Column {partition_col} not part of table partitions: {partition_columns}"
+                    )
+        else:
+            # Raises exception if the table does not have any partitions.
+            raise AirflowException("No partitions specified to check with the sensor.")
+        formatted_opts = f"{prefix} {joiner_val.join(output_list)} {suffix}"
+        self.log.debug("Formatted options: %s", formatted_opts)
+
+        return formatted_opts.strip()
+
+    def poke(self, context: Context) -> bool:
+        """Checks the table partitions and returns the results."""
+        partition_result = self._check_table_partitions()
+        self.log.debug("Partition sensor result: %s", partition_result)
+        if len(partition_result) >= 1:
+            return True
+        else:
+            raise AirflowException(f"Specified partition(s): {self.partitions} were not found.")

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/airflow/providers/databricks/triggers/databricks.py

@@ -32,14 +32,36 @@ class DatabricksExecutionTrigger(BaseTrigger):
     :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
     :param polling_period_seconds: Controls the rate of the poll for the result of this run.
         By default, the trigger will poll every 30 seconds.
+    :param retry_limit: The number of times to retry the connection in case of service outages.
+    :param retry_delay: The number of seconds to wait between retries.
+    :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+    :param run_page_url: The run page url.
     """
 
-    def __init__(self, run_id: int, databricks_conn_id: str, polling_period_seconds: int = 30) -> None:
+    def __init__(
+        self,
+        run_id: int,
+        databricks_conn_id: str,
+        polling_period_seconds: int = 30,
+        retry_limit: int = 3,
+        retry_delay: int = 10,
+        retry_args: dict[Any, Any] | None = None,
+        run_page_url: str | None = None,
+    ) -> None:
         super().__init__()
         self.run_id = run_id
         self.databricks_conn_id = databricks_conn_id
         self.polling_period_seconds = polling_period_seconds
-        self.hook = DatabricksHook(databricks_conn_id)
+        self.retry_limit = retry_limit
+        self.retry_delay = retry_delay
+        self.retry_args = retry_args
+        self.run_page_url = run_page_url
+        self.hook = DatabricksHook(
+            databricks_conn_id,
+            retry_limit=self.retry_limit,
+            retry_delay=self.retry_delay,
+            retry_args=retry_args,
+        )
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
@@ -48,22 +70,31 @@ class DatabricksExecutionTrigger(BaseTrigger):
                 "run_id": self.run_id,
                 "databricks_conn_id": self.databricks_conn_id,
                 "polling_period_seconds": self.polling_period_seconds,
+                "retry_limit": self.retry_limit,
+                "retry_delay": self.retry_delay,
+                "retry_args": self.retry_args,
+                "run_page_url": self.run_page_url,
             },
         )
 
     async def run(self):
         async with self.hook:
-            run_page_url = await self.hook.a_get_run_page_url(self.run_id)
             while True:
                 run_state = await self.hook.a_get_run_state(self.run_id)
                 if run_state.is_terminal:
                     yield TriggerEvent(
                         {
                             "run_id": self.run_id,
+                            "run_page_url": self.run_page_url,
                             "run_state": run_state.to_json(),
-                            "run_page_url": run_page_url,
                         }
                     )
-                    break
+                    return
                 else:
+                    self.log.info(
+                        "run-id %s in run state %s. sleeping for %s seconds",
+                        self.run_id,
+                        run_state,
+                        self.polling_period_seconds,
+                    )
                     await asyncio.sleep(self.polling_period_seconds)

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0/apache_airflow_providers_databricks.egg-info}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-databricks
-Version: 4.1.0rc1
+Version: 4.2.0
 Summary: Provider for Apache Airflow. Implements apache-airflow-providers-databricks package
 Home-page: https://airflow.apache.org/
 Download-URL: https://archive.apache.org/dist/airflow/providers
 Author: Apache Software Foundation
 Author-email: dev@airflow.apache.org
 License: Apache License 2.0
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.1.0/
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
@@ -53,7 +53,7 @@ License-File: NOTICE
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``4.1.0rc1``
+Release: ``4.2.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -66,7 +66,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/4.1.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/>`_.
 
 
 Installation
@@ -84,8 +84,8 @@ Requirements
 =======================================  ===================
 PIP package                              Version required
 =======================================  ===================
-``apache-airflow``                       ``>=2.3.0``
-``apache-airflow-providers-common-sql``  ``>=1.3.1``
+``apache-airflow``                       ``>=2.4.0``
+``apache-airflow-providers-common-sql``  ``>=1.5.0``
 ``requests``                             ``>=2.27,<3``
 ``databricks-sql-connector``             ``>=2.0.0, <3.0.0``
 ``aiohttp``                              ``>=3.6.3, <4``
@@ -136,6 +136,38 @@ Dependent package
 Changelog
 ---------
 
+4.2.0
+.....
+
+.. note::
+  This release of provider is only available for Airflow 2.4+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
+
+Features
+~~~~~~~~
+
+* ``Add conditional output processing in SQL operators (#31136)``
+* ``Add cancel all runs functionality to Databricks hook (#31038)``
+* ``Add retry param in databrics async operator (#30744)``
+* ``Add repair job functionality to databricks hook (#30786)``
+* ``Add 'DatabricksPartitionSensor' (#30980)``
+
+Misc
+~~~~
+
+* ``Bump minimum Airflow version in providers (#30917)``
+* ``Deprecate databricks async operator (#30761)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Move TaskInstanceKey to a separate file (#31033)``
+   * ``Use 'AirflowProviderDeprecationWarning' in providers (#30975)``
+   * ``Add full automation for min Airflow version for providers (#30994)``
+   * ``Add cli cmd to list the provider trigger info (#30822)``
+   * ``Use '__version__' in providers not 'version' (#31393)``
+   * ``Fixing circular import error in providers caused by airflow version check (#31379)``
+   * ``Prepare docs for May 2023 wave of Providers (#31252)``
+
 4.1.0
 .....
 
@@ -203,8 +235,9 @@ Bug Fixes
 3.4.0
 .....
 
-This release of provider is only available for Airflow 2.3+ as explained in the
-`Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/README.md#support-for-providers>`_.
+.. note::
+  This release of provider is only available for Airflow 2.3+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Misc
 ~~~~
@@ -300,8 +333,9 @@ Bug Fixes
 Breaking changes
 ~~~~~~~~~~~~~~~~
 
-* This release of provider is only available for Airflow 2.2+ as explained in the Apache Airflow
-  providers support policy https://github.com/apache/airflow/blob/main/README.md#support-for-providers
+.. note::
+  This release of provider is only available for Airflow 2.2+ as explained in the
+  `Apache Airflow providers support policy <https://github.com/apache/airflow/blob/main/PROVIDERS.rst#minimum-supported-version-of-airflow-for-community-managed-providers>`_.
 
 Features
 ~~~~~~~~

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/apache_airflow_providers_databricks.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ airflow/providers/databricks/operators/databricks.py
 airflow/providers/databricks/operators/databricks_repos.py
 airflow/providers/databricks/operators/databricks_sql.py
 airflow/providers/databricks/sensors/__init__.py
+airflow/providers/databricks/sensors/databricks_partition.py
 airflow/providers/databricks/sensors/databricks_sql.py
 airflow/providers/databricks/triggers/__init__.py
 airflow/providers/databricks/triggers/databricks.py

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/apache_airflow_providers_databricks.egg-info/requires.txt

@@ -1,6 +1,6 @@
 aiohttp<4,>=3.6.3
-apache-airflow-providers-common-sql>=1.3.1.dev0
-apache-airflow>=2.3.0.dev0
+apache-airflow-providers-common-sql>=1.5.0
+apache-airflow>=2.4.0
 databricks-sql-connector<3.0.0,>=2.0.0
 requests<3,>=2.27
 

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/pyproject.toml

@@ -46,11 +46,14 @@ extend-select = [
     "UP", # Pyupgrade
     "RUF100", # Unused noqa (auto-fixable)
 
+    # implicit single-line string concatenation
+    "ISC001",
     # We ignore more pydocstyle than we enable, so be more selective at what we enable
     "D101",
     "D106",
     "D2",
     "D3",
+    # "D400", WIP: see #31135
     # "D401", # Not enabled by ruff, but we don't want it
     "D402",
     "D403",

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/setup.cfg

@@ -27,7 +27,7 @@ classifiers =
 	Programming Language :: Python :: 3.10
 	Topic :: System :: Monitoring
 project_urls = 
-	Documentation=https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.1.0/
+	Documentation=https://airflow.apache.org/docs/apache-airflow-providers-databricks/4.2.0/
 	Bug Tracker=https://github.com/apache/airflow/issues
 	Source Code=https://github.com/apache/airflow
 	Slack Chat=https://s.apache.org/airflow-slack
@@ -47,8 +47,8 @@ setup_requires =
 	wheel
 install_requires = 
 	aiohttp>=3.6.3, <4
-	apache-airflow-providers-common-sql>=1.3.1.dev0
-	apache-airflow>=2.3.0.dev0
+	apache-airflow-providers-common-sql>=1.5.0
+	apache-airflow>=2.4.0
 	databricks-sql-connector>=2.0.0, <3.0.0
 	requests>=2.27,<3
 
@@ -60,6 +60,6 @@ apache_airflow_provider =
 packages = airflow.providers.databricks
 
 [egg_info]
-tag_build = rc1
+tag_build = 
 tag_date = 0
 

{apache-airflow-providers-databricks-4.1.0rc1 → apache-airflow-providers-databricks-4.2.0}/setup.py

@@ -26,7 +26,7 @@
 
 from setuptools import find_namespace_packages, setup
 
-version = "4.1.0"
+version = "4.2.0"
 
 
 def do_setup():

apache-airflow-providers-databricks-4.1.0rc1/airflow/providers/databricks/triggers/__init__.py

@@ -1,17 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.