apache-airflow-providers-google 11.0.0rc1__py3-none-any.whl → 12.0.0__py3-none-any.whl

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+"
     )

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:

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,
+        )

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:

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

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

airflow/providers/google/cloud/hooks/kubernetes_engine.py

@@ -23,7 +23,7 @@ import contextlib
 import json
 import time
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -33,8 +33,7 @@ from google.auth.transport import requests as google_requests
 from google.cloud import exceptions  # type: ignore[attr-defined]
 from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
 from google.cloud.container_v1.types import Cluster, Operation
-from kubernetes import client, utils
-from kubernetes.client.models import V1Deployment
+from kubernetes import client
 from kubernetes_asyncio import client as async_client
 from kubernetes_asyncio.config.kube_config import FileOrData
 
@@ -434,38 +433,9 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
             enable_tcp_keepalive=self.enable_tcp_keepalive,
         ).get_conn()
 
-    def check_kueue_deployment_running(self, name, namespace):
-        timeout = 300
-        polling_period_seconds = 2
-
-        while timeout is None or timeout > 0:
-            try:
-                deployment = self.get_deployment_status(name=name, namespace=namespace)
-                deployment_status = V1Deployment.to_dict(deployment)["status"]
-                replicas = deployment_status["replicas"]
-                ready_replicas = deployment_status["ready_replicas"]
-                unavailable_replicas = deployment_status["unavailable_replicas"]
-                if (
-                    replicas is not None
-                    and ready_replicas is not None
-                    and unavailable_replicas is None
-                    and replicas == ready_replicas
-                ):
-                    return
-                else:
-                    self.log.info("Waiting until Deployment will be ready...")
-                    time.sleep(polling_period_seconds)
-            except Exception as e:
-                self.log.exception("Exception occurred while checking for Deployment status.")
-                raise e
-
-            if timeout is not None:
-                timeout -= polling_period_seconds
-
-        raise AirflowException("Deployment timed out")
-
     def apply_from_yaml_file(
         self,
+        api_client: Any = None,
         yaml_file: str | None = None,
         yaml_objects: list[dict] | None = None,
         verbose: bool = False,
@@ -474,18 +444,17 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
         """
         Perform an action from a yaml file.
 
+        :param api_client: A Kubernetes client application.
         :param yaml_file: Contains the path to yaml file.
         :param yaml_objects: List of YAML objects; used instead of reading the yaml_file.
         :param verbose: If True, print confirmation from create action. Default is False.
         :param namespace: Contains the namespace to create all resources inside. The namespace must
             preexist otherwise the resource creation will fail.
         """
-        k8s_client = self.get_conn()
-
-        utils.create_from_yaml(
-            k8s_client=k8s_client,
-            yaml_objects=yaml_objects,
+        super().apply_from_yaml_file(
+            api_client=api_client or self.get_conn(),
             yaml_file=yaml_file,
+            yaml_objects=yaml_objects,
             verbose=verbose,
             namespace=namespace,
         )