apache-airflow-providers-databricks 7.6.0rc1__tar.gz → 7.7.0rc1__tar.gz

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-databricks
-Version: 7.6.0rc1
+Version: 7.7.0rc1
 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>
 Maintainer-email: Apache Software Foundation <dev@airflow.apache.org>
-Requires-Python: ~=3.10
+Requires-Python: >=3.10
 Description-Content-Type: text/x-rst
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Console
@@ -18,33 +18,34 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: apache-airflow>=2.10.0rc1
 Requires-Dist: apache-airflow-providers-common-compat>=1.6.0rc1
 Requires-Dist: apache-airflow-providers-common-sql>=1.27.0rc1
+Requires-Dist: apache-airflow-providers-openlineage>=2.3.0rc1
 Requires-Dist: requests>=2.32.0,<3
-Requires-Dist: databricks-sql-connector>=3.0.0
+Requires-Dist: databricks-sql-connector>=4.0.0
 Requires-Dist: databricks-sqlalchemy>=1.0.2
 Requires-Dist: aiohttp>=3.9.2, <4
 Requires-Dist: mergedeep>=1.3.4
 Requires-Dist: pandas>=2.1.2; python_version <"3.13"
 Requires-Dist: pandas>=2.2.3; python_version >="3.13"
-Requires-Dist: pyarrow>=16.1.0
+Requires-Dist: pyarrow>=16.1.0; python_version < '3.13'
+Requires-Dist: pyarrow>=18.0.0; python_version >= '3.13'
 Requires-Dist: azure-identity>=1.3.1 ; extra == "azure-identity"
-Requires-Dist: apache-airflow-providers-fab ; extra == "fab"
-Requires-Dist: apache-airflow-providers-openlineage>=2.3.0rc1 ; extra == "openlineage"
+Requires-Dist: apache-airflow-providers-fab>=2.2.0rc1 ; extra == "fab" and ( python_version < '3.13')
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.6.0/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.6.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 Provides-Extra: azure-identity
 Provides-Extra: fab
-Provides-Extra: openlineage
 Provides-Extra: sdk
 Provides-Extra: standard
 
@@ -73,8 +74,9 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.6.0``
+Release: ``7.7.0``
 
+Release Date: ``|PypiReleaseDate|``
 
 `Databricks <https://databricks.com/>`__
 
@@ -86,7 +88,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/7.6.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.0/>`_.
 
 Installation
 ------------
@@ -95,26 +97,28 @@ You can install this package on top of an existing Airflow 2 installation (see `
 for the minimum Airflow version supported) via
 ``pip install apache-airflow-providers-databricks``
 
-The package supports the following python versions: 3.10,3.11,3.12
+The package supports the following python versions: 3.10,3.11,3.12,3.13
 
 Requirements
 ------------
 
-==========================================  =====================================
+==========================================  ======================================
 PIP package                                 Version required
-==========================================  =====================================
+==========================================  ======================================
 ``apache-airflow``                          ``>=2.10.0``
 ``apache-airflow-providers-common-compat``  ``>=1.6.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
+``apache-airflow-providers-openlineage``    ``>=2.3.0``
 ``requests``                                ``>=2.32.0,<3``
-``databricks-sql-connector``                ``>=3.0.0``
+``databricks-sql-connector``                ``>=4.0.0``
 ``databricks-sqlalchemy``                   ``>=1.0.2``
 ``aiohttp``                                 ``>=3.9.2,<4``
 ``mergedeep``                               ``>=1.3.4``
 ``pandas``                                  ``>=2.1.2; python_version < "3.13"``
 ``pandas``                                  ``>=2.2.3; python_version >= "3.13"``
-``pyarrow``                                 ``>=16.1.0``
-==========================================  =====================================
+``pyarrow``                                 ``>=16.1.0; python_version < "3.13"``
+``pyarrow``                                 ``>=18.0.0; python_version >= "3.13"``
+==========================================  ======================================
 
 Cross provider package dependencies
 -----------------------------------
@@ -139,5 +143,5 @@ Dependent package
 ==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.6.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.0/changelog.html>`_.
 

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/README.rst

@@ -23,8 +23,9 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.6.0``
+Release: ``7.7.0``
 
+Release Date: ``|PypiReleaseDate|``
 
 `Databricks <https://databricks.com/>`__
 
@@ -36,7 +37,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/7.6.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.0/>`_.
 
 Installation
 ------------
@@ -45,26 +46,28 @@ You can install this package on top of an existing Airflow 2 installation (see `
 for the minimum Airflow version supported) via
 ``pip install apache-airflow-providers-databricks``
 
-The package supports the following python versions: 3.10,3.11,3.12
+The package supports the following python versions: 3.10,3.11,3.12,3.13
 
 Requirements
 ------------
 
-==========================================  =====================================
+==========================================  ======================================
 PIP package                                 Version required
-==========================================  =====================================
+==========================================  ======================================
 ``apache-airflow``                          ``>=2.10.0``
 ``apache-airflow-providers-common-compat``  ``>=1.6.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
+``apache-airflow-providers-openlineage``    ``>=2.3.0``
 ``requests``                                ``>=2.32.0,<3``
-``databricks-sql-connector``                ``>=3.0.0``
+``databricks-sql-connector``                ``>=4.0.0``
 ``databricks-sqlalchemy``                   ``>=1.0.2``
 ``aiohttp``                                 ``>=3.9.2,<4``
 ``mergedeep``                               ``>=1.3.4``
 ``pandas``                                  ``>=2.1.2; python_version < "3.13"``
 ``pandas``                                  ``>=2.2.3; python_version >= "3.13"``
-``pyarrow``                                 ``>=16.1.0``
-==========================================  =====================================
+``pyarrow``                                 ``>=16.1.0; python_version < "3.13"``
+``pyarrow``                                 ``>=18.0.0; python_version >= "3.13"``
+==========================================  ======================================
 
 Cross provider package dependencies
 -----------------------------------
@@ -89,4 +92,4 @@ Dependent package
 ==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.6.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.0/changelog.html>`_.

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/docs/changelog.rst

@@ -26,6 +26,40 @@
 Changelog
 ---------
 
+7.7.0
+.....
+
+Features
+~~~~~~~~
+
+* ``feat: Refactor Databricks error handling with utility functions (#52704)``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``fix URI Construction in Databricks Hook (#53217)``
+* ``fix: Improve logging and timeouts in OL helpers (#53139)``
+* ``fix: Task Group Deprecation error from plugin (#53813)``
+* ``Remove 'api' prefix from update_job_permission in databricks hook (#53039)``
+
+Misc
+~~~~
+
+* ``Refactor: Consolidate API endpoints for improved security and maintainability (#53214)``
+* ``Exclude deltalake 1.1.1 (#53729)``
+* ``Add Python 3.13 support for Airflow. (#46891)``
+* ``Cleanup mypy ignores in databricks provider where possible (#53265)``
+* ``Remove type ignore across codebase after mypy upgrade (#53243)``
+* ``Remove upper-binding for "python-requires" (#52980)``
+* ``Temporarily switch to use >=,< pattern instead of '~=' (#52967)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Replace 'mock.patch("utcnow")' with time_machine. (#53642)``
+   * ``Added non-Azure DatabricksBaseHook tests (#53286)``
+   * ``Fix pandas FutureWarning (#53236)``
+   * ``Handle ruff PT028 changes (#53235)``
+
 7.6.0
 .....
 

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/docs/index.rst

@@ -78,7 +78,9 @@ apache-airflow-providers-databricks package
 `Databricks <https://databricks.com/>`__
 
 
-Release: 7.6.0
+Release: 7.7.0
+
+Release Date: ``|PypiReleaseDate|``
 
 Provider package
 ----------------
@@ -98,21 +100,23 @@ Requirements
 
 The minimum Apache Airflow version supported by this provider distribution is ``2.10.0``.
 
-==========================================  =====================================
+==========================================  ======================================
 PIP package                                 Version required
-==========================================  =====================================
+==========================================  ======================================
 ``apache-airflow``                          ``>=2.10.0``
 ``apache-airflow-providers-common-compat``  ``>=1.6.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
+``apache-airflow-providers-openlineage``    ``>=2.3.0``
 ``requests``                                ``>=2.32.0,<3``
-``databricks-sql-connector``                ``>=3.0.0``
+``databricks-sql-connector``                ``>=4.0.0``
 ``databricks-sqlalchemy``                   ``>=1.0.2``
 ``aiohttp``                                 ``>=3.9.2,<4``
 ``mergedeep``                               ``>=1.3.4``
 ``pandas``                                  ``>=2.1.2; python_version < "3.13"``
 ``pandas``                                  ``>=2.2.3; python_version >= "3.13"``
-``pyarrow``                                 ``>=16.1.0``
-==========================================  =====================================
+``pyarrow``                                 ``>=16.1.0; python_version < "3.13"``
+``pyarrow``                                 ``>=18.0.0; python_version >= "3.13"``
+==========================================  ======================================
 
 Cross provider package dependencies
 -----------------------------------
@@ -142,5 +146,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-databricks 7.6.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0.tar.gz.sha512>`__)
-* `The apache-airflow-providers-databricks 7.6.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.6.0-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-databricks 7.7.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0.tar.gz.sha512>`__)
+* `The apache-airflow-providers-databricks 7.7.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.7.0-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/provider.yaml

@@ -22,12 +22,13 @@ description: |
     `Databricks <https://databricks.com/>`__
 
 state: ready
-source-date-epoch: 1751473030
+source-date-epoch: 1753690258
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 7.7.0
   - 7.6.0
   - 7.5.0
   - 7.4.0

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "7.6.0rc1"
+version = "7.7.0rc1"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -47,9 +47,10 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: System :: Monitoring",
 ]
-requires-python = "~=3.10"
+requires-python = ">=3.10"
 
 # The dependencies should be modified in place in the generated file.
 # Any change in the dependencies is preserved when the file is regenerated
@@ -59,14 +60,16 @@ dependencies = [
     "apache-airflow>=2.10.0rc1",
     "apache-airflow-providers-common-compat>=1.6.0rc1",
     "apache-airflow-providers-common-sql>=1.27.0rc1",
+    "apache-airflow-providers-openlineage>=2.3.0rc1",
     "requests>=2.32.0,<3",
-    "databricks-sql-connector>=3.0.0",
+    "databricks-sql-connector>=4.0.0",
     "databricks-sqlalchemy>=1.0.2",
     "aiohttp>=3.9.2, <4",
     "mergedeep>=1.3.4",
     'pandas>=2.1.2; python_version <"3.13"',
     'pandas>=2.2.3; python_version >="3.13"',
-    "pyarrow>=16.1.0",
+    "pyarrow>=16.1.0; python_version < '3.13'",
+    "pyarrow>=18.0.0; python_version >= '3.13'",
 ]
 
 # The optional dependencies should be modified in place in the generated file
@@ -80,14 +83,11 @@ dependencies = [
     "azure-identity>=1.3.1",
 ]
 "fab" = [
-    "apache-airflow-providers-fab"
+    "apache-airflow-providers-fab>=2.2.0rc1; python_version < '3.13'"
 ]
 "standard" = [
     "apache-airflow-providers-standard"
 ]
-"openlineage" = [
-    "apache-airflow-providers-openlineage>=2.3.0rc1"
-]
 
 [dependency-groups]
 dev = [
@@ -99,7 +99,9 @@ dev = [
     "apache-airflow-providers-fab",
     "apache-airflow-providers-openlineage",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    "deltalake>=0.12.0",
+    # https://github.com/delta-io/delta-rs/issues/3621
+    "deltalake>=0.12.0,!=1.1.1",
+    "apache-airflow-providers-fab>=2.2.0; python_version < '3.13'",
     "apache-airflow-providers-microsoft-azure",
     "apache-airflow-providers-common-sql[pandas,polars]",
 ]
@@ -130,8 +132,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.6.0"
-"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.6.0/changelog.html"
+"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.0"
+"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.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"

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/__init__.py

@@ -14,4 +14,4 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

{apache_airflow_providers_databricks-7.6.0rc1/tests/system → apache_airflow_providers_databricks-7.7.0rc1/src/airflow/providers}/__init__.py

@@ -14,4 +14,4 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "7.6.0"
+__version__ = "7.7.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/hooks/databricks.py

@@ -62,6 +62,7 @@ CREATE_REPO_ENDPOINT = ("POST", "2.0/repos")
 
 LIST_JOBS_ENDPOINT = ("GET", "2.1/jobs/list")
 LIST_PIPELINES_ENDPOINT = ("GET", "2.0/pipelines")
+LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "2.0/sql/endpoints")
 
 WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "2.0/workspace/get-status")
 
@@ -770,7 +771,7 @@ class DatabricksHook(BaseDatabricksHook):
         :param json: payload
         :return: json containing permission specification
         """
-        return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json)
+        return self._do_api_call(("PATCH", f"2.0/permissions/jobs/{job_id}"), json)
 
     def post_sql_statement(self, json: dict[str, Any]) -> str:
         """

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/hooks/databricks_sql.py

@@ -30,13 +30,15 @@ from typing import (
     overload,
 )
 
-from databricks import sql  # type: ignore[attr-defined]
+from databricks import sql
 from databricks.sql.types import Row
+from sqlalchemy.engine import URL
 
 from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.handlers import return_single_query_results
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.databricks.exceptions import DatabricksSqlExecutionError, DatabricksSqlExecutionTimeout
+from airflow.providers.databricks.hooks.databricks import LIST_SQL_ENDPOINTS_ENDPOINT
 from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook
 
 if TYPE_CHECKING:
@@ -47,9 +49,6 @@ if TYPE_CHECKING:
     from airflow.providers.openlineage.sqlparser import DatabaseInfo
 
 
-LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints")
-
-
 T = TypeVar("T")
 
 
@@ -173,7 +172,38 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
             raise AirflowException("SQL connection is not initialized")
         return cast("AirflowConnection", self._sql_conn)
 
-    @overload  # type: ignore[override]
+    @property
+    def sqlalchemy_url(self) -> URL:
+        """
+        Return a Sqlalchemy.engine.URL object from the connection.
+
+        :return: the extracted sqlalchemy.engine.URL object.
+        """
+        conn = self.get_conn()
+        url_query = {
+            "http_path": self._http_path,
+            "catalog": self.catalog,
+            "schema": self.schema,
+        }
+        url_query = {k: v for k, v in url_query.items() if v is not None}
+        return URL.create(
+            drivername="databricks",
+            username="token",
+            password=conn.password,
+            host=conn.host,
+            port=conn.port,
+            query=url_query,
+        )
+
+    def get_uri(self) -> str:
+        """
+        Extract the URI from the connection.
+
+        :return: the extracted uri.
+        """
+        return self.sqlalchemy_url.render_as_string(hide_password=False)
+
+    @overload
     def run(
         self,
         sql: str | Iterable[str],
@@ -258,7 +288,7 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
 
                     # TODO: adjust this to make testing easier
                     try:
-                        self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
+                        self._run_command(cur, sql_statement, parameters)
                     except Exception as e:
                         if t is None or t.is_alive():
                             raise DatabricksSqlExecutionError(

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/operators/databricks.py

@@ -46,21 +46,32 @@ from airflow.providers.databricks.plugins.databricks_workflow import (
 from airflow.providers.databricks.triggers.databricks import (
     DatabricksExecutionTrigger,
 )
-from airflow.providers.databricks.utils.databricks import normalise_json_content, validate_trigger_event
+from airflow.providers.databricks.utils.databricks import (
+    extract_failed_task_errors,
+    normalise_json_content,
+    validate_trigger_event,
+)
 from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
 from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.databricks.operators.databricks_workflow import (
+        DatabricksWorkflowTaskGroup,
+    )
     from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
-    from airflow.utils.task_group import TaskGroup
+
+    try:
+        from airflow.sdk import TaskGroup
+    except ImportError:
+        from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk import BaseOperatorLink
     from airflow.sdk.execution_time.xcom import XCom
 else:
-    from airflow.models import XCom  # type: ignore[no-redef]
+    from airflow.models import XCom
     from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
 
 DEFER_METHOD_NAME = "execute_complete"
@@ -95,17 +106,7 @@ def _handle_databricks_operator_execution(operator, hook, log, context) -> None:
                     log.info("View run status, Spark UI, and logs at %s", run_page_url)
                     return
                 if run_state.result_state == "FAILED":
-                    failed_tasks = []
-                    for task in run_info.get("tasks", []):
-                        if task.get("state", {}).get("result_state", "") == "FAILED":
-                            task_run_id = task["run_id"]
-                            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})
+                    failed_tasks = extract_failed_task_errors(hook, run_info, run_state)
 
                     error_message = (
                         f"{operator.task_id} failed with terminal state: {run_state} "
@@ -1324,15 +1325,15 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
 
         return self.databricks_run_id
 
-    def _handle_terminal_run_state(self, run_state: RunState) -> None:
+    def _handle_terminal_run_state(self, run_state: RunState, errors: list) -> None:
         """Handle the terminal state of the run."""
         if run_state.life_cycle_state != RunLifeCycleState.TERMINATED.value:
             raise AirflowException(
-                f"Databricks job failed with state {run_state.life_cycle_state}. Message: {run_state.state_message}"
+                f"Databricks job failed with state {run_state.life_cycle_state}. Message: {run_state.state_message}. Errors: {errors}"
             )
         if not run_state.is_successful:
             raise AirflowException(
-                f"Task failed. Final state {run_state.result_state}. Reason: {run_state.state_message}"
+                f"Task failed. Final state {run_state.result_state}. Reason: {run_state.state_message}. Errors: {errors}"
             )
         self.log.info("Task succeeded. Final state %s.", run_state.result_state)
 
@@ -1414,12 +1415,17 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
             time.sleep(self.polling_period_seconds)
             run = self._hook.get_run(current_task_run_id)
             run_state = RunState(**run["state"])
+
             self.log.info(
                 "Current state of the databricks task %s is %s",
                 self.databricks_task_key,
                 run_state.life_cycle_state,
             )
-        self._handle_terminal_run_state(run_state)
+
+        # Extract errors from the run response using utility function
+        errors = extract_failed_task_errors(self._hook, run, run_state)
+
+        self._handle_terminal_run_state(run_state, errors)
 
     def execute(self, context: Context) -> None:
         """Execute the operator. Launch the job and monitor it if wait_for_termination is set to True."""
@@ -1428,9 +1434,7 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
             if not self.workflow_run_metadata:
                 launch_task_id = next(task for task in self.upstream_task_ids if task.endswith(".launch"))
                 self.workflow_run_metadata = context["ti"].xcom_pull(task_ids=launch_task_id)
-            workflow_run_metadata = WorkflowRunMetadata(  # type: ignore[arg-type]
-                **self.workflow_run_metadata
-            )
+            workflow_run_metadata = WorkflowRunMetadata(**self.workflow_run_metadata)
             self.databricks_run_id = workflow_run_metadata.run_id
             self.databricks_conn_id = workflow_run_metadata.conn_id
 
@@ -1449,7 +1453,8 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC):
 
     def execute_complete(self, context: dict | None, event: dict) -> None:
         run_state = RunState.from_json(event["run_state"])
-        self._handle_terminal_run_state(run_state)
+        errors = event.get("errors", [])
+        self._handle_terminal_run_state(run_state, errors)
 
 
 class DatabricksNotebookOperator(DatabricksTaskBaseOperator):

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/operators/databricks_workflow.py

@@ -33,7 +33,11 @@ from airflow.providers.databricks.plugins.databricks_workflow import (
     store_databricks_job_run_link,
 )
 from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
-from airflow.utils.task_group import TaskGroup
+
+try:
+    from airflow.sdk import TaskGroup
+except ImportError:
+    from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from types import TracebackType

{apache_airflow_providers_databricks-7.6.0rc1 → apache_airflow_providers_databricks-7.7.0rc1}/src/airflow/providers/databricks/plugins/databricks_workflow.py

@@ -41,7 +41,6 @@ else:
     from airflow.www import auth  # type: ignore
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
-from airflow.utils.task_group import TaskGroup
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
@@ -51,11 +50,12 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
+    from airflow.sdk import BaseOperatorLink, TaskGroup
     from airflow.sdk.execution_time.xcom import XCom
 else:
-    from airflow.models import XCom  # type: ignore[no-redef]
+    from airflow.models import XCom
     from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
+    from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
 
 
 REPAIR_WAIT_ATTEMPTS = os.getenv("DATABRICKS_REPAIR_WAIT_ATTEMPTS", 20)