apache-airflow-providers-databricks 7.8.3rc1__tar.gz → 7.9.0rc1__tar.gz

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-databricks
-Version: 7.8.3rc1
+Version: 7.9.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>
@@ -23,7 +23,7 @@ Classifier: Topic :: System :: Monitoring
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: apache-airflow>=2.11.0rc1
-Requires-Dist: apache-airflow-providers-common-compat>=1.12.0rc1
+Requires-Dist: apache-airflow-providers-common-compat>=1.13.0rc1
 Requires-Dist: apache-airflow-providers-common-sql>=1.27.0rc1
 Requires-Dist: requests>=2.32.0,<3
 Requires-Dist: databricks-sql-connector>=4.0.0
@@ -33,21 +33,26 @@ 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; python_version < '3.13'
 Requires-Dist: pyarrow>=18.0.0; python_version >= '3.13'
+Requires-Dist: fastavro>=1.9.0 ; extra == "avro"
+Requires-Dist: fastavro>=1.10.0 ; extra == "avro" and (python_version>="3.12")
 Requires-Dist: azure-identity>=1.3.1 ; extra == "azure-identity"
 Requires-Dist: apache-airflow-providers-fab>=2.2.0rc1 ; extra == "fab" and ( python_version < '3.13')
+Requires-Dist: apache-airflow-providers-google>=10.24.0rc1 ; extra == "google"
 Requires-Dist: apache-airflow-providers-openlineage>=2.3.0rc1 ; extra == "openlineage"
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Requires-Dist: databricks-sqlalchemy>=1.0.2 ; extra == "sqlalchemy"
 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.8.3/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.8.3
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.9.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.9.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: avro
 Provides-Extra: azure-identity
 Provides-Extra: fab
+Provides-Extra: google
 Provides-Extra: openlineage
 Provides-Extra: sdk
 Provides-Extra: sqlalchemy
@@ -78,7 +83,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.8.3``
+Release: ``7.9.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -91,7 +96,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.8.3/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.9.0/>`_.
 
 Installation
 ------------
@@ -109,7 +114,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.10.1``
+``apache-airflow-providers-common-compat``  ``>=1.13.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
 ``requests``                                ``>=2.32.0,<3``
 ``databricks-sql-connector``                ``>=4.0.0``
@@ -139,6 +144,7 @@ Dependent package
 ==================================================================================================================  =================
 `apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_  ``common.compat``
 `apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_        ``common.sql``
+`apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_                ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
@@ -154,8 +160,10 @@ Extra               Dependencies
 ``standard``        ``apache-airflow-providers-standard``
 ``openlineage``     ``apache-airflow-providers-openlineage>=2.3.0``
 ``sqlalchemy``      ``databricks-sqlalchemy>=1.0.2``
+``google``          ``apache-airflow-providers-google>=10.24.0``
+``avro``            ``fastavro>=1.9.0``, ``fastavro>=1.10.0;python_version>="3.12"``
 ==================  ================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.8.3/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.9.0/changelog.html>`_.
 

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.8.3``
+Release: ``7.9.0``
 
 
 `Databricks <https://databricks.com/>`__
@@ -36,7 +36,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.8.3/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.9.0/>`_.
 
 Installation
 ------------
@@ -54,7 +54,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.10.1``
+``apache-airflow-providers-common-compat``  ``>=1.13.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
 ``requests``                                ``>=2.32.0,<3``
 ``databricks-sql-connector``                ``>=4.0.0``
@@ -84,6 +84,7 @@ Dependent package
 ==================================================================================================================  =================
 `apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_  ``common.compat``
 `apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_        ``common.sql``
+`apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_                ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
@@ -99,7 +100,9 @@ Extra               Dependencies
 ``standard``        ``apache-airflow-providers-standard``
 ``openlineage``     ``apache-airflow-providers-openlineage>=2.3.0``
 ``sqlalchemy``      ``databricks-sqlalchemy>=1.0.2``
+``google``          ``apache-airflow-providers-google>=10.24.0``
+``avro``            ``fastavro>=1.9.0``, ``fastavro>=1.10.0;python_version>="3.12"``
 ==================  ================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.8.3/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.9.0/changelog.html>`_.

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/docs/changelog.rst

@@ -26,6 +26,30 @@
 Changelog
 ---------
 
+7.9.0
+.....
+
+Features
+~~~~~~~~
+
+* ``Add direct GCS export to DatabricksSqlOperator with Parquet/Avro support #55128 (#60543)``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``Fix missing fastavro after PR #60732 (#60797)``
+* ``Pin fastavro to minimum 1.10.0 for Python 3.13 compatibility (#60732)``
+* ``Updating Databricks API endpoints to appropriate versions (#60647)``
+
+Misc
+~~~~
+
+* ``Define 'TaskInstanceKey' in task-sdk to support client server separation (#60776)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Revert exclusion of deltalake 1.3.1 as aarch64 binaries are available now (#60611)``
+
 7.8.3
 .....
 

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/docs/index.rst

@@ -78,7 +78,7 @@ apache-airflow-providers-databricks package
 `Databricks <https://databricks.com/>`__
 
 
-Release: 7.8.3
+Release: 7.9.0
 
 Provider package
 ----------------
@@ -102,7 +102,7 @@ The minimum Apache Airflow version supported by this provider distribution is ``
 PIP package                                 Version required
 ==========================================  ======================================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.12.0``
+``apache-airflow-providers-common-compat``  ``>=1.13.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
 ``requests``                                ``>=2.32.0,<3``
 ``databricks-sql-connector``                ``>=4.0.0``
@@ -132,6 +132,7 @@ Dependent package
 ==================================================================================================================  =================
 `apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_  ``common.compat``
 `apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_        ``common.sql``
+`apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_                ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
@@ -141,5 +142,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.8.3 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3.tar.gz.sha512>`__)
-* `The apache-airflow-providers-databricks 7.8.3 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.8.3-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-databricks 7.9.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0.tar.gz.sha512>`__)
+* `The apache-airflow-providers-databricks 7.9.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_databricks-7.9.0-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/provider.yaml

@@ -22,12 +22,13 @@ description: |
     `Databricks <https://databricks.com/>`__
 
 state: ready
-source-date-epoch: 1768334769
+source-date-epoch: 1769460991
 # 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.9.0
   - 7.8.3
   - 7.8.2
   - 7.8.1

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-databricks"
-version = "7.8.3rc1"
+version = "7.9.0rc1"
 description = "Provider package apache-airflow-providers-databricks for Apache Airflow"
 readme = "README.rst"
 license = "Apache-2.0"
@@ -59,7 +59,7 @@ requires-python = ">=3.10"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.11.0rc1",
-    "apache-airflow-providers-common-compat>=1.12.0rc1",
+    "apache-airflow-providers-common-compat>=1.13.0rc1",
     "apache-airflow-providers-common-sql>=1.27.0rc1",
     "requests>=2.32.0,<3",
     "databricks-sql-connector>=4.0.0",
@@ -93,6 +93,13 @@ dependencies = [
 "sqlalchemy" = [
     "databricks-sqlalchemy>=1.0.2",
 ]
+"google" = [
+    "apache-airflow-providers-google>=10.24.0rc1"
+]
+"avro" = [
+    "fastavro>=1.9.0",
+    'fastavro>=1.10.0;python_version>="3.12"'  # Need to pin to this version for Python 3.13 compatibility
+]
 
 [dependency-groups]
 dev = [
@@ -101,10 +108,11 @@ dev = [
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-compat",
     "apache-airflow-providers-common-sql",
+    "apache-airflow-providers-google",
     "apache-airflow-providers-openlineage",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    # Limit deltalake to avoid issue with missing linux ARM wheels: https://github.com/delta-io/delta-rs/issues/4041
-    "deltalake>=1.1.3,!=1.3.0,!=1.3.1",
+    # Need to exclude 1.3.0 due to missing aarch64 binaries, fixed with 1.3.1++
+    "deltalake>=1.1.3,!=1.3.0",
     "apache-airflow-providers-fab>=2.2.0; python_version < '3.13'",
     "apache-airflow-providers-microsoft-azure",
     "apache-airflow-providers-common-sql[pandas,polars]",
@@ -138,8 +146,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.8.3"
-"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.8.3/changelog.html"
+"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.9.0"
+"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.9.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.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/src/airflow/providers/databricks/__init__.py

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

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/src/airflow/providers/databricks/hooks/databricks.py

@@ -37,10 +37,10 @@ from requests import exceptions as requests_exceptions
 from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook
 
-GET_CLUSTER_ENDPOINT = ("GET", "2.2/clusters/get")
-RESTART_CLUSTER_ENDPOINT = ("POST", "2.2/clusters/restart")
-START_CLUSTER_ENDPOINT = ("POST", "2.2/clusters/start")
-TERMINATE_CLUSTER_ENDPOINT = ("POST", "2.2/clusters/delete")
+GET_CLUSTER_ENDPOINT = ("GET", "2.1/clusters/get")
+RESTART_CLUSTER_ENDPOINT = ("POST", "2.1/clusters/restart")
+START_CLUSTER_ENDPOINT = ("POST", "2.1/clusters/start")
+TERMINATE_CLUSTER_ENDPOINT = ("POST", "2.1/clusters/delete")
 
 CREATE_ENDPOINT = ("POST", "2.2/jobs/create")
 RESET_ENDPOINT = ("POST", "2.2/jobs/reset")
@@ -54,20 +54,20 @@ REPAIR_RUN_ENDPOINT = ("POST", "2.2/jobs/runs/repair")
 OUTPUT_RUNS_JOB_ENDPOINT = ("GET", "2.2/jobs/runs/get-output")
 CANCEL_ALL_RUNS_ENDPOINT = ("POST", "2.2/jobs/runs/cancel-all")
 
-INSTALL_LIBS_ENDPOINT = ("POST", "2.2/libraries/install")
-UNINSTALL_LIBS_ENDPOINT = ("POST", "2.2/libraries/uninstall")
-UPDATE_REPO_ENDPOINT = ("PATCH", "2.2/repos/")
-DELETE_REPO_ENDPOINT = ("DELETE", "2.2/repos/")
-CREATE_REPO_ENDPOINT = ("POST", "2.2/repos")
+INSTALL_LIBS_ENDPOINT = ("POST", "2.0/libraries/install")
+UNINSTALL_LIBS_ENDPOINT = ("POST", "2.0/libraries/uninstall")
+UPDATE_REPO_ENDPOINT = ("PATCH", "2.0/repos/")
+DELETE_REPO_ENDPOINT = ("DELETE", "2.0/repos/")
+CREATE_REPO_ENDPOINT = ("POST", "2.0/repos")
 
 LIST_JOBS_ENDPOINT = ("GET", "2.2/jobs/list")
-LIST_PIPELINES_ENDPOINT = ("GET", "2.2/pipelines")
-LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "2.2/sql/endpoints")
+LIST_PIPELINES_ENDPOINT = ("GET", "2.0/pipelines")
+LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "2.0/sql/warehouses")
 
-WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "2.2/workspace/get-status")
+WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "2.0/workspace/get-status")
 
-SPARK_VERSIONS_ENDPOINT = ("GET", "2.2/clusters/spark-versions")
-SQL_STATEMENTS_ENDPOINT = "2.2/sql/statements"
+SPARK_VERSIONS_ENDPOINT = ("GET", "2.1/clusters/spark-versions")
+SQL_STATEMENTS_ENDPOINT = "2.0/sql/statements"
 
 
 class RunLifeCycleState(Enum):
@@ -717,7 +717,7 @@ class DatabricksHook(BaseDatabricksHook):
         """
         Install libraries on the cluster.
 
-        Utility function to call the ``2.2/libraries/install`` endpoint.
+        Utility function to call the ``2.0/libraries/install`` endpoint.
 
         :param json: json dictionary containing cluster_id and an array of library
         """
@@ -727,7 +727,7 @@ class DatabricksHook(BaseDatabricksHook):
         """
         Uninstall libraries on the cluster.
 
-        Utility function to call the ``2.2/libraries/uninstall`` endpoint.
+        Utility function to call the ``2.0/libraries/uninstall`` endpoint.
 
         :param json: json dictionary containing cluster_id and an array of library
         """
@@ -790,7 +790,7 @@ class DatabricksHook(BaseDatabricksHook):
         :param json: payload
         :return: json containing permission specification
         """
-        return self._do_api_call(("PATCH", f"2.2/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.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/src/airflow/providers/databricks/operators/databricks.py

@@ -53,7 +53,7 @@ from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMix
 from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
-    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.common.compat.sdk import TaskInstanceKey
     from airflow.providers.databricks.operators.databricks_workflow import (
         DatabricksWorkflowTaskGroup,
     )

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/src/airflow/providers/databricks/operators/databricks_sql.py

@@ -21,13 +21,20 @@ from __future__ import annotations
 
 import csv
 import json
+import os
 from collections.abc import Sequence
 from functools import cached_property
+from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Any, ClassVar
+from urllib.parse import urlparse
 
 from databricks.sql.utils import ParamEscaper
 
-from airflow.providers.common.compat.sdk import AirflowException, BaseOperator
+from airflow.providers.common.compat.sdk import (
+    AirflowException,
+    AirflowOptionalProviderFeatureException,
+    BaseOperator,
+)
 from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
 
@@ -62,13 +69,27 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
     :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ (templated)
     :param schema: An optional initial schema to use. Requires DBR version 9.0+ (templated)
     :param output_path: optional string specifying the file to which write selected data. (templated)
-    :param output_format: format of output data if ``output_path` is specified.
-        Possible values are ``csv``, ``json``, ``jsonl``. Default is ``csv``.
+        Supports local file paths and GCS URIs (e.g., ``gs://bucket/path/file.parquet``).
+        When using GCS URIs, requires the ``apache-airflow-providers-google`` package.
+    :param output_format: format of output data if ``output_path`` is specified.
+        Possible values are ``csv``, ``json``, ``jsonl``, ``parquet``, ``avro``. Default is ``csv``.
     :param csv_params: parameters that will be passed to the ``csv.DictWriter`` class used to write CSV data.
+    :param gcp_conn_id: The connection ID to use for connecting to Google Cloud when using GCS output path.
+        Default is ``google_cloud_default``.
+    :param gcs_impersonation_chain: Optional service account to impersonate using short-term
+        credentials for GCS upload, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request. (templated)
     """
 
     template_fields: Sequence[str] = tuple(
-        {"_output_path", "schema", "catalog", "http_headers", "databricks_conn_id"}
+        {
+            "_output_path",
+            "schema",
+            "catalog",
+            "http_headers",
+            "databricks_conn_id",
+            "_gcs_impersonation_chain",
+        }
         | set(SQLExecuteQueryOperator.template_fields)
     )
 
@@ -90,6 +111,8 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
         output_format: str = "csv",
         csv_params: dict[str, Any] | None = None,
         client_parameters: dict[str, Any] | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        gcs_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
         super().__init__(conn_id=databricks_conn_id, **kwargs)
@@ -105,6 +128,8 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
         self.http_headers = http_headers
         self.catalog = catalog
         self.schema = schema
+        self._gcp_conn_id = gcp_conn_id
+        self._gcs_impersonation_chain = gcs_impersonation_chain
 
     @cached_property
     def _hook(self) -> DatabricksSqlHook:
@@ -127,41 +152,151 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator):
     def _should_run_output_processing(self) -> bool:
         return self.do_xcom_push or bool(self._output_path)
 
+    @property
+    def _is_gcs_output(self) -> bool:
+        """Check if the output path is a GCS URI."""
+        return self._output_path.startswith("gs://") if self._output_path else False
+
+    def _parse_gcs_path(self, path: str) -> tuple[str, str]:
+        """Parse a GCS URI into bucket and object name."""
+        parsed = urlparse(path)
+        bucket = parsed.netloc
+        object_name = parsed.path.lstrip("/")
+        return bucket, object_name
+
+    def _upload_to_gcs(self, local_path: str, gcs_path: str) -> None:
+        """Upload a local file to GCS."""
+        try:
+            from airflow.providers.google.cloud.hooks.gcs import GCSHook
+        except ImportError:
+            raise AirflowOptionalProviderFeatureException(
+                "The 'apache-airflow-providers-google' package is required for GCS output. "
+                "Install it with: pip install apache-airflow-providers-google"
+            )
+
+        bucket, object_name = self._parse_gcs_path(gcs_path)
+        hook = GCSHook(
+            gcp_conn_id=self._gcp_conn_id,
+            impersonation_chain=self._gcs_impersonation_chain,
+        )
+        hook.upload(
+            bucket_name=bucket,
+            object_name=object_name,
+            filename=local_path,
+        )
+        self.log.info("Uploaded output to %s", gcs_path)
+
+    def _write_parquet(self, file_path: str, field_names: list[str], rows: list[Any]) -> None:
+        """Write data to a Parquet file."""
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        data: dict[str, list] = {name: [] for name in field_names}
+        for row in rows:
+            row_dict = row._asdict()
+            for name in field_names:
+                data[name].append(row_dict[name])
+
+        table = pa.Table.from_pydict(data)
+        pq.write_table(table, file_path)
+
+    def _write_avro(self, file_path: str, field_names: list[str], rows: list[Any]) -> None:
+        """Write data to an Avro file using fastavro."""
+        try:
+            from fastavro import writer
+        except ImportError:
+            raise AirflowOptionalProviderFeatureException(
+                "The 'fastavro' package is required for Avro output. Install it with: pip install fastavro"
+            )
+
+        data: dict[str, list] = {name: [] for name in field_names}
+        for row in rows:
+            row_dict = row._asdict()
+            for name in field_names:
+                data[name].append(row_dict[name])
+
+        schema_fields = []
+        for name in field_names:
+            sample_val = next(
+                (data[name][i] for i in range(len(data[name])) if data[name][i] is not None), None
+            )
+            if sample_val is None:
+                avro_type = ["null", "string"]
+            elif isinstance(sample_val, bool):
+                avro_type = ["null", "boolean"]
+            elif isinstance(sample_val, int):
+                avro_type = ["null", "long"]
+            elif isinstance(sample_val, float):
+                avro_type = ["null", "double"]
+            else:
+                avro_type = ["null", "string"]
+            schema_fields.append({"name": name, "type": avro_type})
+
+        avro_schema = {
+            "type": "record",
+            "name": "QueryResult",
+            "fields": schema_fields,
+        }
+
+        records = [row._asdict() for row in rows]
+        with open(file_path, "wb") as f:
+            writer(f, avro_schema, records)
+
     def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequence] | None]) -> list[Any]:
         if not self._output_path:
             return list(zip(descriptions, results))
         if not self._output_format:
             raise AirflowException("Output format should be specified!")
-        # Output to a file only the result of last query
+
         last_description = descriptions[-1]
         last_results = results[-1]
         if last_description is None:
-            raise AirflowException("There is missing description present for the output file. .")
+            raise AirflowException("There is missing description present for the output file.")
         field_names = [field[0] for field in last_description]
-        if self._output_format.lower() == "csv":
-            with open(self._output_path, "w", newline="") as file:
-                if self._csv_params:
-                    csv_params = self._csv_params
-                else:
-                    csv_params = {}
-                write_header = csv_params.get("header", True)
-                if "header" in csv_params:
-                    del csv_params["header"]
-                writer = csv.DictWriter(file, fieldnames=field_names, **csv_params)
-                if write_header:
-                    writer.writeheader()
-                for row in last_results:
-                    writer.writerow(row._asdict())
-        elif self._output_format.lower() == "json":
-            with open(self._output_path, "w") as file:
-                file.write(json.dumps([row._asdict() for row in last_results]))
-        elif self._output_format.lower() == "jsonl":
-            with open(self._output_path, "w") as file:
-                for row in last_results:
-                    file.write(json.dumps(row._asdict()))
-                    file.write("\n")
+
+        if self._is_gcs_output:
+            suffix = f".{self._output_format.lower()}"
+            tmp_file = NamedTemporaryFile(mode="w", suffix=suffix, delete=False, newline="")
+            local_path = tmp_file.name
+            tmp_file.close()
         else:
-            raise AirflowException(f"Unsupported output format: '{self._output_format}'")
+            local_path = self._output_path
+
+        try:
+            output_format = self._output_format.lower()
+            if output_format == "csv":
+                with open(local_path, "w", newline="") as file:
+                    if self._csv_params:
+                        csv_params = self._csv_params.copy()
+                    else:
+                        csv_params = {}
+                    write_header = csv_params.pop("header", True)
+                    writer = csv.DictWriter(file, fieldnames=field_names, **csv_params)
+                    if write_header:
+                        writer.writeheader()
+                    for row in last_results:
+                        writer.writerow(row._asdict())
+            elif output_format == "json":
+                with open(local_path, "w") as file:
+                    file.write(json.dumps([row._asdict() for row in last_results]))
+            elif output_format == "jsonl":
+                with open(local_path, "w") as file:
+                    for row in last_results:
+                        file.write(json.dumps(row._asdict()))
+                        file.write("\n")
+            elif output_format == "parquet":
+                self._write_parquet(local_path, field_names, last_results)
+            elif output_format == "avro":
+                self._write_avro(local_path, field_names, last_results)
+            else:
+                raise ValueError(f"Unsupported output format: '{self._output_format}'")
+
+            if self._is_gcs_output:
+                self._upload_to_gcs(local_path, self._output_path)
+        finally:
+            if self._is_gcs_output and os.path.exists(local_path):
+                os.unlink(local_path)
+
         return list(zip(descriptions, results))
 
 

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/src/airflow/providers/databricks/plugins/databricks_workflow.py

@@ -298,7 +298,7 @@ class WorkflowJobRunLink(BaseOperatorLink, LoggingMixin):
         """XCom key where the link is stored during task execution."""
         return "databricks_job_run_link"
 
-    def get_link(
+    def get_link(  # type: ignore[override]  # Signature intentionally kept this way for Airflow 2.x compatibility
         self,
         operator: BaseOperator,
         dttm=None,
@@ -374,7 +374,7 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
 
     name = "Repair All Failed Tasks"
 
-    def get_link(
+    def get_link(  # type: ignore[override]  # Signature intentionally kept this way for Airflow 2.x compatibility
         self,
         operator,
         dttm=None,
@@ -471,7 +471,7 @@ class WorkflowJobRepairSingleTaskLink(BaseOperatorLink, LoggingMixin):
 
     name = "Repair a single task"
 
-    def get_link(
+    def get_link(  # type: ignore[override]  # Signature intentionally kept this way for Airflow 2.x compatibility
         self,
         operator,
         dttm=None,

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/tests/unit/databricks/hooks/test_databricks.py

@@ -216,42 +216,42 @@ def get_cluster_endpoint(host):
     """
     Utility function to generate the get run endpoint given the host.
     """
-    return f"https://{host}/api/2.2/clusters/get"
+    return f"https://{host}/api/2.1/clusters/get"
 
 
 def start_cluster_endpoint(host):
     """
     Utility function to generate the get run endpoint given the host.
     """
-    return f"https://{host}/api/2.2/clusters/start"
+    return f"https://{host}/api/2.1/clusters/start"
 
 
 def restart_cluster_endpoint(host):
     """
     Utility function to generate the get run endpoint given the host.
     """
-    return f"https://{host}/api/2.2/clusters/restart"
+    return f"https://{host}/api/2.1/clusters/restart"
 
 
 def terminate_cluster_endpoint(host):
     """
     Utility function to generate the get run endpoint given the host.
     """
-    return f"https://{host}/api/2.2/clusters/delete"
+    return f"https://{host}/api/2.1/clusters/delete"
 
 
 def install_endpoint(host):
     """
     Utility function to generate the install endpoint given the host.
     """
-    return f"https://{host}/api/2.2/libraries/install"
+    return f"https://{host}/api/2.0/libraries/install"
 
 
 def uninstall_endpoint(host):
     """
     Utility function to generate the uninstall endpoint given the host.
     """
-    return f"https://{host}/api/2.2/libraries/uninstall"
+    return f"https://{host}/api/2.0/libraries/uninstall"
 
 
 def list_jobs_endpoint(host):
@@ -265,19 +265,19 @@ def list_pipelines_endpoint(host):
     """
     Utility function to generate the list jobs endpoint given the host
     """
-    return f"https://{host}/api/2.2/pipelines"
+    return f"https://{host}/api/2.0/pipelines"
 
 
 def list_spark_versions_endpoint(host):
     """Utility function to generate the list spark versions endpoint given the host"""
-    return f"https://{host}/api/2.2/clusters/spark-versions"
+    return f"https://{host}/api/2.1/clusters/spark-versions"
 
 
 def permissions_endpoint(host, job_id):
     """
     Utility function to generate the permissions endpoint given the host
     """
-    return f"https://{host}/api/2.2/permissions/jobs/{job_id}"
+    return f"https://{host}/api/2.0/permissions/jobs/{job_id}"
 
 
 def create_valid_response_mock(content):
@@ -289,7 +289,7 @@ def create_valid_response_mock(content):
 
 def sql_statements_endpoint(host):
     """Utility function to generate the sql statements endpoint given the host."""
-    return f"https://{host}/api/2.2/sql/statements"
+    return f"https://{host}/api/2.0/sql/statements"
 
 
 def create_successful_response_mock(content):
@@ -1302,7 +1302,7 @@ class TestDatabricksHook:
         self.hook.update_job_permission(1, ACCESS_CONTROL_DICT)
 
         mock_requests.patch.assert_called_once_with(
-            f"https://{HOST}/api/2.2/permissions/jobs/1",
+            f"https://{HOST}/api/2.0/permissions/jobs/1",
             json=utils.normalise_json_content(ACCESS_CONTROL_DICT),
             params=None,
             auth=HTTPBasicAuth(LOGIN, PASSWORD),

{apache_airflow_providers_databricks-7.8.3rc1 → apache_airflow_providers_databricks-7.9.0rc1}/tests/unit/databricks/operators/test_databricks_sql.py

@@ -314,3 +314,142 @@ def test_hook_is_cached():
     hook = op.get_db_hook()
     hook2 = op.get_db_hook()
     assert hook is hook2
+
+
+def test_exec_write_parquet_file(tmp_path):
+    """Test writing output to Parquet format."""
+    with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class:
+        path = tmp_path / "testfile.parquet"
+        op = DatabricksSqlOperator(
+            task_id=TASK_ID,
+            sql="select * from dummy",
+            output_path=os.fspath(path),
+            output_format="parquet",
+            return_last=True,
+            do_xcom_push=True,
+            split_statements=False,
+        )
+        db_mock = db_mock_class.return_value
+        db_mock.run.return_value = [SerializableRow(1, "value1"), SerializableRow(2, "value2")]
+        db_mock.descriptions = [[("id",), ("value",)]]
+
+        op.execute(None)
+
+        import pyarrow.parquet as pq
+
+        table = pq.read_table(path)
+        assert table.num_rows == 2
+        assert table.column_names == ["id", "value"]
+        assert table.column("id").to_pylist() == [1, 2]
+        assert table.column("value").to_pylist() == ["value1", "value2"]
+
+
+def test_exec_write_avro_file_with_fastavro(tmp_path):
+    """Test writing output to Avro format using fastavro."""
+    pytest.importorskip("fastavro")
+    with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class:
+        path = tmp_path / "testfile.avro"
+        op = DatabricksSqlOperator(
+            task_id=TASK_ID,
+            sql="select * from dummy",
+            output_path=os.fspath(path),
+            output_format="avro",
+            return_last=True,
+            do_xcom_push=True,
+            split_statements=False,
+        )
+        db_mock = db_mock_class.return_value
+        db_mock.run.return_value = [SerializableRow(1, "value1"), SerializableRow(2, "value2")]
+        db_mock.descriptions = [[("id",), ("value",)]]
+
+        op.execute(None)
+
+        from fastavro import reader
+
+        with open(path, "rb") as f:
+            records = list(reader(f))
+        assert len(records) == 2
+        assert records[0] == {"id": 1, "value": "value1"}
+        assert records[1] == {"id": 2, "value": "value2"}
+
+
+def test_exec_write_gcs_output(tmp_path):
+    """Test writing output to GCS."""
+    with (
+        patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class,
+        patch("airflow.providers.google.cloud.hooks.gcs.GCSHook") as gcs_mock_class,
+    ):
+        op = DatabricksSqlOperator(
+            task_id=TASK_ID,
+            sql="select * from dummy",
+            output_path="gs://my-bucket/path/to/output.csv",
+            output_format="csv",
+            return_last=True,
+            do_xcom_push=True,
+            split_statements=False,
+            gcp_conn_id="my_gcp_conn",
+        )
+        db_mock = db_mock_class.return_value
+        db_mock.run.return_value = [SerializableRow(1, "value1"), SerializableRow(2, "value2")]
+        db_mock.descriptions = [[("id",), ("value",)]]
+
+        op.execute(None)
+
+        gcs_mock_class.assert_called_once_with(
+            gcp_conn_id="my_gcp_conn",
+            impersonation_chain=None,
+        )
+        gcs_mock_class.return_value.upload.assert_called_once()
+        call_kwargs = gcs_mock_class.return_value.upload.call_args[1]
+        assert call_kwargs["bucket_name"] == "my-bucket"
+        assert call_kwargs["object_name"] == "path/to/output.csv"
+
+
+def test_exec_write_gcs_parquet_output(tmp_path):
+    """Test writing Parquet output to GCS."""
+    with (
+        patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class,
+        patch("airflow.providers.google.cloud.hooks.gcs.GCSHook") as gcs_mock_class,
+    ):
+        op = DatabricksSqlOperator(
+            task_id=TASK_ID,
+            sql="select * from dummy",
+            output_path="gs://my-bucket/data/results.parquet",
+            output_format="parquet",
+            return_last=True,
+            do_xcom_push=True,
+            split_statements=False,
+        )
+        db_mock = db_mock_class.return_value
+        db_mock.run.return_value = [SerializableRow(1, "value1"), SerializableRow(2, "value2")]
+        db_mock.descriptions = [[("id",), ("value",)]]
+
+        op.execute(None)
+
+        gcs_mock_class.return_value.upload.assert_called_once()
+        call_kwargs = gcs_mock_class.return_value.upload.call_args[1]
+        assert call_kwargs["bucket_name"] == "my-bucket"
+        assert call_kwargs["object_name"] == "data/results.parquet"
+
+
+def test_is_gcs_output():
+    """Test _is_gcs_output property."""
+    op_gcs = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1", output_path="gs://bucket/path")
+    assert op_gcs._is_gcs_output is True
+
+    op_local = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1", output_path="/local/path")
+    assert op_local._is_gcs_output is False
+
+    op_s3 = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1", output_path="s3://bucket/path")
+    assert op_s3._is_gcs_output is False
+
+    op_none = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1")
+    assert op_none._is_gcs_output is False
+
+
+def test_parse_gcs_path():
+    """Test _parse_gcs_path method."""
+    op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1")
+    bucket, object_name = op._parse_gcs_path("gs://my-bucket/path/to/file.parquet")
+    assert bucket == "my-bucket"
+    assert object_name == "path/to/file.parquet"