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

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

@@ -19,10 +19,12 @@
 
 from __future__ import annotations
 
+import shlex
+import subprocess
 import time
 import uuid
-from collections.abc import MutableSequence
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import MutableSequence, Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.exceptions import ServerError
@@ -217,19 +219,6 @@ class DataprocHook(GoogleBaseHook):
     keyword arguments rather than positional.
     """
 
-    def __init__(
-        self,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain)
-
     def get_cluster_client(self, region: str | None = None) -> ClusterControllerClient:
         """Create a ClusterControllerClient."""
         client_options = None
@@ -274,6 +263,49 @@ class DataprocHook(GoogleBaseHook):
         """Create a OperationsClient."""
         return self.get_batch_client(region=region).transport.operations_client
 
+    def dataproc_options_to_args(self, options: dict) -> list[str]:
+        """
+        Return a formatted cluster parameters from a dictionary of arguments.
+
+        :param options: Dictionary with options
+        :return: List of arguments
+        """
+        if not options:
+            return []
+
+        args: list[str] = []
+        for attr, value in options.items():
+            if value is None or (isinstance(value, bool) and value):
+                args.append(f"--{attr}")
+            elif isinstance(value, bool) and not value:
+                continue
+            elif isinstance(value, list):
+                args.extend([f"--{attr}={v}" for v in value])
+            else:
+                args.append(f"--{attr}={value}")
+        return args
+
+    def _build_gcloud_command(self, command: list[str], parameters: dict[str, str]) -> list[str]:
+        return [*command, *(self.dataproc_options_to_args(parameters))]
+
+    def _create_dataproc_cluster_with_gcloud(self, cmd: list[str]) -> str:
+        """Create a Dataproc cluster with a gcloud command and return the job's ID."""
+        self.log.info("Executing command: %s", " ".join(shlex.quote(c) for c in cmd))
+        success_code = 0
+
+        with self.provide_authorized_gcloud():
+            proc = subprocess.run(cmd, capture_output=True)
+
+        if proc.returncode != success_code:
+            stderr_last_20_lines = "\n".join(proc.stderr.decode().strip().splitlines()[-20:])
+            raise AirflowException(
+                f"Process exit with non-zero exit code. Exit code: {proc.returncode}. Error Details : "
+                f"{stderr_last_20_lines}"
+            )
+
+        response = proc.stdout.decode().strip()
+        return response
+
     def wait_for_operation(
         self,
         operation: Operation,
@@ -302,7 +334,7 @@ class DataprocHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
-    ) -> Operation:
+    ) -> Operation | str:
         """
         Create a cluster in a specified project.
 
@@ -339,6 +371,34 @@ class DataprocHook(GoogleBaseHook):
             "project_id": project_id,
             "cluster_name": cluster_name,
         }
+
+        if virtual_cluster_config and "kubernetes_cluster_config" in virtual_cluster_config:
+            kube_config = virtual_cluster_config["kubernetes_cluster_config"]["gke_cluster_config"]
+            try:
+                spark_engine_version = virtual_cluster_config["kubernetes_cluster_config"][
+                    "kubernetes_software_config"
+                ]["component_version"]["SPARK"]
+            except KeyError:
+                spark_engine_version = "latest"
+            gke_cluster_name = kube_config["gke_cluster_target"].rsplit("/", 1)[1]
+            gke_pools = kube_config["node_pool_target"][0]
+            gke_pool_name = gke_pools["node_pool"].rsplit("/", 1)[1]
+            gke_pool_role = gke_pools["roles"][0]
+            gke_pool_machine_type = gke_pools["node_pool_config"]["config"]["machine_type"]
+            gcp_flags = {
+                "region": region,
+                "gke-cluster": gke_cluster_name,
+                "spark-engine-version": spark_engine_version,
+                "pools": f"name={gke_pool_name},roles={gke_pool_role.lower()},machineType={gke_pool_machine_type},min=1,max=10",
+                "setup-workload-identity": None,
+            }
+            cmd = self._build_gcloud_command(
+                command=["gcloud", "dataproc", "clusters", "gke", "create", cluster_name],
+                parameters=gcp_flags,
+            )
+            response = self._create_dataproc_cluster_with_gcloud(cmd=cmd)
+            return response
+
         if virtual_cluster_config is not None:
             cluster["virtual_cluster_config"] = virtual_cluster_config  # type: ignore
         if cluster_config is not None:
@@ -1219,12 +1279,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain)
+        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
         self._cached_client: JobControllerAsyncClient | None = None
 
     def get_cluster_client(self, region: str | None = None) -> ClusterControllerAsyncClient:

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -40,14 +41,6 @@ if TYPE_CHECKING:
 class DataprocMetastoreHook(GoogleBaseHook):
     """Hook for Google Cloud Dataproc Metastore APIs."""
 
-    def __init__(self, **kwargs):
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(**kwargs)
-
     def get_dataproc_metastore_client(self) -> DataprocMetastoreClient:
         """Return DataprocMetastoreClient."""
         client_options = ClientOptions(api_endpoint="metastore.googleapis.com:443")

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import time
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 
 from googleapiclient.discovery import Resource, build
 
@@ -44,14 +45,10 @@ class DatastoreHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self.connection = None
         self.api_version = api_version

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

@@ -28,7 +28,8 @@ from __future__ import annotations
 
 import re
 import time
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dlp import DlpServiceClient
@@ -91,14 +92,10 @@ class CloudDLPHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._client: DlpServiceClient | None = None
 

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import time
-from typing import Sequence
+from collections.abc import Sequence
 
 import requests
 from googleapiclient.discovery import build
@@ -49,14 +49,10 @@ class CloudFunctionsHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self.api_version = api_version
 

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

@@ -26,11 +26,12 @@ import os
 import shutil
 import time
 import warnings
+from collections.abc import Generator, Sequence
 from contextlib import contextmanager
 from functools import partial
 from io import BytesIO
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Sequence, TypeVar, cast, overload
+from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
 from gcloud.aio.storage import Storage
@@ -151,23 +152,6 @@ class GCSHook(GoogleBaseHook):
 
     _conn: storage.Client | None = None
 
-    def __init__(
-        self,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-
-        super().__init__(
-            gcp_conn_id=gcp_conn_id,
-            impersonation_chain=impersonation_chain,
-        )
-
     def get_conn(self) -> storage.Client:
         """Return a Google Cloud Storage service object."""
         if not self._conn:

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

@@ -17,7 +17,7 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import Any
 
 from googleapiclient.discovery import Resource, build
 
@@ -32,22 +32,6 @@ class GoogleDeploymentManagerHook(GoogleBaseHook):
     This allows for scheduled and programmatic inspection and deletion of resources managed by GDM.
     """
 
-    def __init__(
-        self,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(
-            gcp_conn_id=gcp_conn_id,
-            impersonation_chain=impersonation_chain,
-        )
-
     def get_conn(self) -> Resource:
         """Return a Google Deployment Manager service object."""
         http_authorized = self._authorize()

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import base64
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.kms_v1 import KeyManagementServiceClient
@@ -63,14 +64,10 @@ class CloudKMSHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._conn: KeyManagementServiceClient | None = None
 

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

@@ -22,14 +22,15 @@ from __future__ import annotations
 import contextlib
 import json
 import time
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.auth.transport import requests as google_requests
 
 # not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
-from google.cloud import container_v1, exceptions  # type: ignore[attr-defined]
+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
@@ -38,11 +39,10 @@ from kubernetes_asyncio import client as async_client
 from kubernetes_asyncio.config.kube_config import FileOrData
 
 from airflow import version
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook
 from airflow.providers.cncf.kubernetes.kube_client import _enable_tcp_keepalive
 from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
     GoogleBaseAsyncHook,
@@ -118,14 +118,10 @@ class GKEHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._client: ClusterManagerClient | None = None
         self.location = location
@@ -136,26 +132,6 @@ class GKEHook(GoogleBaseHook):
             self._client = ClusterManagerClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
-    # To preserve backward compatibility
-    # TODO: remove one day
-    @deprecated(
-        planned_removal_date="November 01, 2024",
-        use_instead="get_cluster_manager_client",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_conn(self) -> container_v1.ClusterManagerClient:
-        return self.get_cluster_manager_client()
-
-    # To preserve backward compatibility
-    # TODO: remove one day
-    @deprecated(
-        planned_removal_date="November 01, 2024",
-        use_instead="get_cluster_manager_client",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_client(self) -> ClusterManagerClient:
-        return self.get_conn()
-
     def wait_for_operation(self, operation: Operation, project_id: str = PROVIDE_PROJECT_ID) -> Operation:
         """
         Continuously fetch the status from Google Cloud.
@@ -388,10 +364,12 @@ class GKEAsyncHook(GoogleBaseAsyncHook):
         gcp_conn_id: str = "google_cloud_default",
         location: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
     ) -> None:
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._client: ClusterManagerAsyncClient | None = None
         self.location = location
@@ -576,71 +554,3 @@ class GKEKubernetesAsyncHook(GoogleBaseAsyncHook, AsyncKubernetesHook):
             ).as_file(),
         )
         return configuration
-
-
-@deprecated(
-    planned_removal_date="October 01, 2024",
-    use_instead="GKEKubernetesHook",
-    category=AirflowProviderDeprecationWarning,
-)
-class GKEDeploymentHook(GKEKubernetesHook):
-    """Google Kubernetes Engine Deployment APIs."""
-
-
-@deprecated(
-    planned_removal_date="October 01, 2024",
-    use_instead="GKEKubernetesHook",
-    category=AirflowProviderDeprecationWarning,
-)
-class GKECustomResourceHook(GKEKubernetesHook):
-    """Google Kubernetes Engine Custom Resource APIs."""
-
-
-@deprecated(
-    planned_removal_date="October 01, 2024",
-    use_instead="GKEKubernetesHook",
-    category=AirflowProviderDeprecationWarning,
-)
-class GKEPodHook(GKEKubernetesHook):
-    """Google Kubernetes Engine pod APIs."""
-
-    def __init__(
-        self,
-        cluster_url: str,
-        ssl_ca_cert: str,
-        disable_tcp_keepalive: bool | None = None,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ):
-        super().__init__(
-            gcp_conn_id=gcp_conn_id,
-            impersonation_chain=impersonation_chain,
-            cluster_url=cluster_url,
-            ssl_ca_cert=ssl_ca_cert,
-            **kwargs,
-        )
-        self.enable_tcp_keepalive = not bool(disable_tcp_keepalive)
-
-
-@deprecated(
-    planned_removal_date="October 01, 2024",
-    use_instead="GKEKubernetesHook",
-    category=AirflowProviderDeprecationWarning,
-)
-class GKEJobHook(GKEKubernetesHook):
-    """Google Kubernetes Engine Job APIs."""
-
-
-@deprecated(
-    planned_removal_date="October 01, 2024",
-    use_instead="GKEKubernetesAsyncHook",
-    category=AirflowProviderDeprecationWarning,
-)
-class GKEPodAsyncHook(GKEKubernetesAsyncHook):
-    """
-    Google Kubernetes Engine pods APIs asynchronously.
-
-    :param cluster_url: The URL pointed to the cluster.
-    :param ssl_ca_cert: SSL certificate used for authentication to the pod.
-    """

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import time
-from typing import Sequence
+from collections.abc import Sequence
 
 import google.api_core.path_template
 from googleapiclient.discovery import build
@@ -73,14 +73,10 @@ class LifeSciencesHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self.api_version = api_version
 

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

@@ -42,8 +42,9 @@ class LookerHook(BaseHook):
     def __init__(
         self,
         looker_conn_id: str,
+        **kwargs,
     ) -> None:
-        super().__init__()
+        super().__init__(**kwargs)
         self.looker_conn_id = looker_conn_id
         # source is used to track origin of the requests
         self.source = f"airflow:{version}"

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

@@ -92,14 +92,6 @@ class MLEngineHook(GoogleBaseHook):
     keyword arguments rather than positional.
     """
 
-    def __init__(self, **kwargs):
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(**kwargs)
-
     def get_conn(self) -> Resource:
         """
         Retrieve the connection to MLEngine.

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.language_v1 import EncodingType, LanguageServiceClient
@@ -62,14 +63,10 @@ class CloudNaturalLanguageHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._conn: LanguageServiceClient | None = None
 

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

@@ -24,7 +24,8 @@ OS Login hooks.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient
@@ -50,14 +51,10 @@ class OSLoginHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._conn: OsLoginServiceClient | None = None
 

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

@@ -28,8 +28,9 @@ from __future__ import annotations
 
 import warnings
 from base64 import b64decode
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
@@ -83,14 +84,10 @@ class PubSubHook(GoogleBaseHook):
         enable_message_ordering: bool = False,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self.enable_message_ordering = enable_message_ordering
         self._client = None