apache-airflow-providers-google 11.0.0rc1__tar.gz → 12.0.0rc1__tar.gz

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: apache-airflow-providers-google
-Version: 11.0.0rc1
+Version: 12.0.0rc1
 Summary: Provider package apache-airflow-providers-google for Apache Airflow
 Keywords: airflow-provider,google,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -21,9 +21,9 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyOpenSSL>=23.0.0
-Requires-Dist: apache-airflow-providers-common-compat>=1.2.1rc0
+Requires-Dist: apache-airflow-providers-common-compat>=1.3.0rc0
 Requires-Dist: apache-airflow-providers-common-sql>=1.20.0rc0
-Requires-Dist: apache-airflow>=2.8.0rc0
+Requires-Dist: apache-airflow>=2.9.0rc0
 Requires-Dist: asgiref>=3.5.2
 Requires-Dist: dill>=0.2.3
 Requires-Dist: gcloud-aio-auth>=5.2.0
@@ -36,7 +36,8 @@ Requires-Dist: google-api-core>=2.11.0,!=2.16.0,!=2.18.0
 Requires-Dist: google-api-python-client>=2.0.2
 Requires-Dist: google-auth-httplib2>=0.0.1
 Requires-Dist: google-auth>=2.29.0
-Requires-Dist: google-cloud-aiplatform>=1.70.0
+Requires-Dist: google-cloud-aiplatform>=1.73.0
+Requires-Dist: google-cloud-alloydb
 Requires-Dist: google-cloud-automl>=2.12.0
 Requires-Dist: google-cloud-batch>=0.13.0
 Requires-Dist: google-cloud-bigquery-datatransfer>=3.13.0
@@ -91,9 +92,7 @@ Requires-Dist: apache-airflow-providers-amazon>=2.6.0rc0 ; extra == "amazon"
 Requires-Dist: apache-airflow-providers-apache-beam ; extra == "apache-beam"
 Requires-Dist: apache-beam[gcp] ; extra == "apache-beam"
 Requires-Dist: apache-airflow-providers-apache-cassandra ; extra == "apache-cassandra"
-Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0rc0 ; extra == "cncf-kubernetes"
-Requires-Dist: apache-airflow-providers-common-compat ; extra == "common-compat"
-Requires-Dist: apache-airflow-providers-common-sql ; extra == "common-sql"
+Requires-Dist: apache-airflow-providers-cncf-kubernetes>=10.1.0rc0 ; extra == "cncf-kubernetes"
 Requires-Dist: apache-airflow-providers-facebook>=2.2.0rc0 ; extra == "facebook"
 Requires-Dist: plyvel>=1.5.1 ; extra == "leveldb"
 Requires-Dist: apache-airflow-providers-microsoft-azure ; extra == "microsoft-azure"
@@ -108,18 +107,16 @@ Requires-Dist: apache-airflow-providers-sftp ; extra == "sftp"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-trino ; extra == "trino"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
-Project-URL: Twitter, https://twitter.com/ApacheAirflow
+Project-URL: Twitter, https://x.com/ApacheAirflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 Provides-Extra: amazon
 Provides-Extra: apache-beam
 Provides-Extra: apache-cassandra
 Provides-Extra: cncf-kubernetes
-Provides-Extra: common-compat
-Provides-Extra: common-sql
 Provides-Extra: facebook
 Provides-Extra: leveldb
 Provides-Extra: microsoft-azure
@@ -178,7 +175,7 @@ Provides-Extra: trino
 
 Package ``apache-airflow-providers-google``
 
-Release: ``11.0.0.rc1``
+Release: ``12.0.0.rc1``
 
 
 Google services including:
@@ -198,7 +195,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0/>`_.
 
 Installation
 ------------
@@ -215,8 +212,8 @@ Requirements
 ==========================================  =========================================
 PIP package                                 Version required
 ==========================================  =========================================
-``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.2.1``
+``apache-airflow``                          ``>=2.9.0``
+``apache-airflow-providers-common-compat``  ``>=1.3.0``
 ``apache-airflow-providers-common-sql``     ``>=1.20.0``
 ``asgiref``                                 ``>=3.5.2``
 ``dill``                                    ``>=0.2.3``
@@ -230,7 +227,8 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
-``google-cloud-aiplatform``                 ``>=1.70.0``
+``google-cloud-aiplatform``                 ``>=1.73.0``
+``google-cloud-alloydb``
 ``google-cloud-automl``                     ``>=2.12.0``
 ``google-cloud-bigquery``                   ``!=3.21.*,!=3.22.0,!=3.23.*,>=3.4.0``
 ``google-cloud-bigquery-datatransfer``      ``>=3.13.0``
@@ -321,4 +319,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0/changelog.html>`_.

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-google``
 
-Release: ``11.0.0.rc1``
+Release: ``12.0.0.rc1``
 
 
 Google services including:
@@ -62,7 +62,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0/>`_.
 
 Installation
 ------------
@@ -79,8 +79,8 @@ Requirements
 ==========================================  =========================================
 PIP package                                 Version required
 ==========================================  =========================================
-``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.2.1``
+``apache-airflow``                          ``>=2.9.0``
+``apache-airflow-providers-common-compat``  ``>=1.3.0``
 ``apache-airflow-providers-common-sql``     ``>=1.20.0``
 ``asgiref``                                 ``>=3.5.2``
 ``dill``                                    ``>=0.2.3``
@@ -94,7 +94,8 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
-``google-cloud-aiplatform``                 ``>=1.70.0``
+``google-cloud-aiplatform``                 ``>=1.73.0``
+``google-cloud-alloydb``
 ``google-cloud-automl``                     ``>=2.12.0``
 ``google-cloud-bigquery``                   ``!=3.21.*,!=3.22.0,!=3.23.*,>=3.4.0``
 ``google-cloud-bigquery-datatransfer``      ``>=3.13.0``
@@ -185,4 +186,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/11.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/12.0.0/changelog.html>`_.

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/airflow/providers/google/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "11.0.0"
+__version__ = "12.0.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.8.0"
+    "2.9.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.8.0+"
+        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.9.0+"
     )

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/airflow/providers/google/assets/gcs.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
@@ -26,16 +27,9 @@ if TYPE_CHECKING:
     from airflow.providers.common.compat.assets import Asset
     from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
 else:
-    # TODO: Remove this try-exception block after bumping common provider to 1.3.0
-    # This is due to common provider AssetDetails import error handling
     try:
         from airflow.providers.common.compat.assets import Asset
     except ImportError:
-        from packaging.version import Version
-
-        from airflow import __version__ as AIRFLOW_VERSION
-
-        AIRFLOW_V_3_0_PLUS = Version(Version(AIRFLOW_VERSION).base_version) >= Version("3.0.0")
         if AIRFLOW_V_3_0_PLUS:
             from airflow.sdk.definitions.asset import Asset
         else:

apache_airflow_providers_google-12.0.0rc1/airflow/providers/google/cloud/hooks/alloy_db.py

@@ -0,0 +1,289 @@
+#
+# 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.
+"""Module contains a Google Alloy DB Hook."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from copy import deepcopy
+from typing import TYPE_CHECKING
+
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud import alloydb_v1
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+if TYPE_CHECKING:
+    import proto
+    from google.api_core.operation import Operation
+    from google.api_core.retry import Retry
+    from google.protobuf.field_mask_pb2 import FieldMask
+
+
+class AlloyDbHook(GoogleBaseHook):
+    """Google Alloy DB Hook."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._client: alloydb_v1.AlloyDBAdminClient | None = None
+
+    def get_alloy_db_admin_client(self) -> alloydb_v1.AlloyDBAdminClient:
+        """Retrieve AlloyDB client."""
+        if not self._client:
+            self._client = alloydb_v1.AlloyDBAdminClient(
+                credentials=self.get_credentials(), client_info=CLIENT_INFO
+            )
+        return self._client
+
+    def wait_for_operation(self, timeout: float | None, operation: Operation) -> proto.Message:
+        """Wait for long-lasting operation to complete."""
+        self.log.info("Waiting for operation to complete...")
+        _timeout: int | None = int(timeout) if timeout else None
+        try:
+            return operation.result(timeout=_timeout)
+        except Exception:
+            error = operation.exception(timeout=_timeout)
+            raise AirflowException(error)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_cluster(
+        self,
+        cluster_id: str,
+        cluster: alloydb_v1.Cluster | dict,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        request_id: str | None = None,
+        validate_only: bool = False,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Create an Alloy DB cluster.
+
+        .. seealso::
+            For more details see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.CreateClusterRequest
+
+        :param cluster_id: Required. ID of the cluster to create.
+        :param cluster: Required. Cluster to create. For more details please see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.Cluster
+        :param location: Required. The ID of the Google Cloud region where the cluster is located.
+        :param project_id: Optional. The ID of the Google Cloud project where the cluster is located.
+        :param request_id: Optional. The ID of an existing request object.
+        :param validate_only: Optional. If set, performs request validation, but does not actually execute
+            the create request.
+        :param retry: Optional. Designation of what errors, if any, should be retried.
+        :param timeout: Optional. The timeout for this request.
+        :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        """
+        client = self.get_alloy_db_admin_client()
+        return client.create_cluster(
+            request={
+                "parent": client.common_location_path(project_id, location),
+                "cluster_id": cluster_id,
+                "cluster": cluster,
+                "request_id": request_id,
+                "validate_only": validate_only,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_secondary_cluster(
+        self,
+        cluster_id: str,
+        cluster: alloydb_v1.Cluster | dict,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        request_id: str | None = None,
+        validate_only: bool = False,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Create a secondary Alloy DB cluster.
+
+        .. seealso::
+            For more details see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.CreateClusterRequest
+
+        :param cluster_id: Required. ID of the cluster to create.
+        :param cluster: Required. Cluster to create. For more details please see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.Cluster
+        :param location: Required. The ID of the Google Cloud region where the cluster is located.
+        :param project_id: Optional. The ID of the Google Cloud project where the cluster is located.
+        :param request_id: Optional. The ID of an existing request object.
+        :param validate_only: Optional. If set, performs request validation, but does not actually execute
+            the create request.
+        :param retry: Optional. Designation of what errors, if any, should be retried.
+        :param timeout: Optional. The timeout for this request.
+        :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        """
+        client = self.get_alloy_db_admin_client()
+        return client.create_secondary_cluster(
+            request={
+                "parent": client.common_location_path(project_id, location),
+                "cluster_id": cluster_id,
+                "cluster": cluster,
+                "request_id": request_id,
+                "validate_only": validate_only,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_cluster(
+        self,
+        cluster_id: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> alloydb_v1.Cluster:
+        """
+        Retrieve an Alloy DB cluster.
+
+        .. seealso::
+            For more details see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.GetClusterRequest
+
+        :param cluster_id: Required. ID of the cluster to create.
+        :param location: Required. The ID of the Google Cloud region where the cluster is located.
+        :param project_id: Optional. The ID of the Google Cloud project where the cluster is located.
+        :param retry: Optional. Designation of what errors, if any, should be retried.
+        :param timeout: Optional. The timeout for this request.
+        :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        """
+        client = self.get_alloy_db_admin_client()
+        return client.get_cluster(
+            request={"name": client.cluster_path(project_id, location, cluster_id)},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_cluster(
+        self,
+        cluster_id: str,
+        cluster: alloydb_v1.Cluster | dict,
+        location: str,
+        update_mask: FieldMask | dict | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+        allow_missing: bool = False,
+        request_id: str | None = None,
+        validate_only: bool = False,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Update an Alloy DB cluster.
+
+        .. seealso::
+            For more details see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.UpdateClusterRequest
+
+        :param cluster_id: Required. ID of the cluster to update.
+        :param cluster: Required. Cluster to create. For more details please see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.Cluster
+        :param location: Required. The ID of the Google Cloud region where the cluster is located.
+        :param update_mask: Optional. Field mask is used to specify the fields to be overwritten in the
+            Cluster resource by the update.
+        :param request_id: Optional. The ID of an existing request object.
+        :param validate_only: Optional. If set, performs request validation, but does not actually execute
+            the create request.
+        :param project_id: Optional. The ID of the Google Cloud project where the cluster is located.
+        :param allow_missing: Optional. If set to true, update succeeds even if cluster is not found.
+            In that case, a new cluster is created and update_mask is ignored.
+        :param retry: Optional. Designation of what errors, if any, should be retried.
+        :param timeout: Optional. The timeout for this request.
+        :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        """
+        client = self.get_alloy_db_admin_client()
+        _cluster = deepcopy(cluster) if isinstance(cluster, dict) else alloydb_v1.Cluster.to_dict(cluster)
+        _cluster["name"] = client.cluster_path(project_id, location, cluster_id)
+        return client.update_cluster(
+            request={
+                "update_mask": update_mask,
+                "cluster": _cluster,
+                "request_id": request_id,
+                "validate_only": validate_only,
+                "allow_missing": allow_missing,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_cluster(
+        self,
+        cluster_id: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        request_id: str | None = None,
+        etag: str | None = None,
+        validate_only: bool = False,
+        force: bool = False,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Delete an Alloy DB cluster.
+
+        .. seealso::
+            For more details see API documentation:
+            https://cloud.google.com/python/docs/reference/alloydb/latest/google.cloud.alloydb_v1.types.DeleteClusterRequest
+
+        :param cluster_id: Required. ID of the cluster to delete.
+        :param location: Required. The ID of the Google Cloud region where the cluster is located.
+        :param project_id: Optional. The ID of the Google Cloud project where the cluster is located.
+        :param request_id: Optional. The ID of an existing request object.
+        :param etag: Optional. The current etag of the Cluster. If an etag is provided and does not match the
+            current etag of the Cluster, deletion will be blocked and an ABORTED error will be returned.
+        :param validate_only: Optional. If set, performs request validation, but does not actually execute
+            the create request.
+        :param force: Optional. Whether to cascade delete child instances for given cluster.
+        :param retry: Optional. Designation of what errors, if any, should be retried.
+        :param timeout: Optional. The timeout for this request.
+        :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        """
+        client = self.get_alloy_db_admin_client()
+        return client.delete_cluster(
+            request={
+                "name": client.cluster_path(project_id, location, cluster_id),
+                "request_id": request_id,
+                "etag": etag,
+                "validate_only": validate_only,
+                "force": force,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/airflow/providers/google/cloud/hooks/cloud_batch.py

@@ -145,12 +145,20 @@ class CloudBatchHook(GoogleBaseHook):
             try:
                 job = client.get_job(name=f"{job_name}")
                 status: JobStatus.State = job.status.state
-                if (
-                    status == JobStatus.State.SUCCEEDED
-                    or status == JobStatus.State.FAILED
-                    or status == JobStatus.State.DELETION_IN_PROGRESS
-                ):
+                if status == JobStatus.State.SUCCEEDED:
                     return job
+                elif status == JobStatus.State.FAILED:
+                    message = (
+                        "Unexpected error in the operation: "
+                        "Batch job with name {job_name} has failed its execution."
+                    )
+                    raise AirflowException(message)
+                elif status == JobStatus.State.DELETION_IN_PROGRESS:
+                    message = (
+                        "Unexpected error in the operation: "
+                        "Batch job with name {job_name} is being deleted."
+                    )
+                    raise AirflowException(message)
                 else:
                     time.sleep(polling_period_seconds)
             except Exception as e:

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/airflow/providers/google/cloud/hooks/dataproc.py

@@ -116,13 +116,17 @@ class DataProcJobBuilder:
         if args is not None:
             self.job["job"][self.job_type]["args"] = args
 
-    def add_query(self, query: str) -> None:
+    def add_query(self, query: str | list[str]) -> None:
         """
-        Set query for Dataproc job.
+        Add query for Dataproc job.
 
         :param query: query for the job.
         """
-        self.job["job"][self.job_type]["query_list"] = {"queries": [query]}
+        queries = self.job["job"][self.job_type].setdefault("query_list", {"queries": []})["queries"]
+        if isinstance(query, str):
+            queries.append(query)
+        elif isinstance(query, list):
+            queries.extend(query)
 
     def add_query_uri(self, query_uri: str) -> None:
         """

{apache_airflow_providers_google-11.0.0rc1 → apache_airflow_providers_google-12.0.0rc1}/airflow/providers/google/cloud/hooks/dataproc_metastore.py

@@ -60,11 +60,16 @@ class DataprocMetastoreHook(GoogleBaseHook):
 
     def wait_for_operation(self, timeout: float | None, operation: Operation):
         """Wait for long-lasting operation to complete."""
+        self.log.info("Waiting for operation (timeout: %s seconds)", timeout)
+
         try:
-            return operation.result(timeout=timeout)
-        except Exception:
+            result = operation.result(timeout=timeout)
+            self.log.info("Operation completed successfully")
+            return result
+        except Exception as e:
+            self.log.error("Operation failed: %s", str(e))
             error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
+            raise AirflowException(f"Operation failed: {error}")
 
     @GoogleBaseHook.fallback_to_default_project_id
     def create_backup(
@@ -669,23 +674,37 @@ class DataprocMetastoreHook(GoogleBaseHook):
         # because dictionaries are ordered since Python 3.7+
         _partitions = list(dict.fromkeys(partition_names)) if partition_names else []
 
-        query = f"""
-                SELECT *
-                FROM PARTITIONS
-                INNER JOIN TBLS
-                ON PARTITIONS.TBL_ID = TBLS.TBL_ID
-                WHERE
-                    TBLS.TBL_NAME = '{table}'"""
         if _partitions:
-            query += f"""
-                    AND PARTITIONS.PART_NAME IN ({', '.join(f"'{p}'" for p in _partitions)})"""
-        query += ";"
-
-        client = self.get_dataproc_metastore_client_v1beta()
-        result = client.query_metadata(
-            request={
-                "service": f"projects/{project_id}/locations/{region}/services/{service_id}",
-                "query": query,
-            }
-        )
-        return result
+            partition_list = ", ".join(f"'{p}'" for p in _partitions)
+            query = f"""
+    SELECT PARTITIONS.*, TBLS.TBL_TYPE, TBLS.TBL_NAME
+    FROM PARTITIONS
+    INNER JOIN TBLS ON PARTITIONS.TBL_ID = TBLS.TBL_ID
+    WHERE TBLS.TBL_NAME = '{table}'
+        AND PARTITIONS.PART_NAME IN ({partition_list});"""
+        else:
+            query = f"""
+    SELECT PARTITIONS.*, TBLS.TBL_TYPE, TBLS.TBL_NAME
+    FROM PARTITIONS
+    INNER JOIN TBLS ON PARTITIONS.TBL_ID = TBLS.TBL_ID
+    WHERE TBLS.TBL_NAME = '{table}';"""
+
+        request = {
+            "service": f"projects/{project_id}/locations/{region}/services/{service_id}",
+            "query": query,
+        }
+
+        self.log.info("Prepared request:")
+        self.log.info(request)
+
+        # Execute query
+        try:
+            self.log.info("Getting Dataproc Metastore client (v1beta)...")
+            client = self.get_dataproc_metastore_client_v1beta()
+            self.log.info("Executing query_metadata...")
+            result = client.query_metadata(request=request)
+            self.log.info("Query executed successfully")
+            return result
+        except Exception as e:
+            self.log.error("Error executing query_metadata: %s", str(e))
+            raise