apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -23,13 +23,6 @@ from collections.abc import MutableSequence, Sequence
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
-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,
-    GoogleBaseAsyncHook,
-    GoogleBaseHook,
-)
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dataplex_v1 import (
@@ -53,9 +46,16 @@ from google.cloud.dataplex_v1.types import (
 )
 from google.protobuf.field_mask_pb2 import FieldMask
 
-if TYPE_CHECKING:
-    from googleapiclient.discovery import Resource
+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,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+)
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
 
+if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
     from google.api_core.retry_async import AsyncRetry
@@ -66,6 +66,7 @@ if TYPE_CHECKING:
         ListEntryTypesPager,
         SearchEntriesPager,
     )
+    from googleapiclient.discovery import Resource
 
 PATH_DATA_SCAN = "projects/{project_id}/locations/{region}/dataScans/{data_scan_id}"
 
@@ -78,7 +79,7 @@ class AirflowDataQualityScanResultTimeoutException(AirflowException):
     """Raised when no result found after specified amount of seconds."""
 
 
-class DataplexHook(GoogleBaseHook):
+class DataplexHook(GoogleBaseHook, OperationHelper):
     """
     Hook for Google Dataplex.
 
@@ -136,7 +137,7 @@ class DataplexHook(GoogleBaseHook):
             credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
-    def wait_for_operation(self, timeout: float | None, operation: Operation):
+    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
         """Wait for long-lasting operation to complete."""
         try:
             return operation.result(timeout=timeout)
@@ -514,7 +515,7 @@ class DataplexHook(GoogleBaseHook):
             fields are non-required and omitted in the request body, their values are emptied.
         :param allow_missing: Optional. If set to true and entry doesn't exist, the service will create it.
         :param delete_missing_aspects: Optional. If set to true and the aspect_keys specify aspect
-            ranges, the service deletes any existing aspects from that range that weren't provided
+            ranges, the service deletes any existing aspects from that range that were not provided
             in the request.
         :param aspect_keys: Optional. The map keys of the Aspects which the service should modify.
             It supports the following syntax:

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

@@ -28,10 +28,10 @@ import requests
 from requests import HTTPError
 from tenacity import retry, stop_after_attempt, wait_exponential
 
-from airflow.hooks.base import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 
 
-def _get_field(extras: dict, field_name: str):
+def _get_field(extras: dict, field_name: str) -> str | None:
     """Get field from extra, first checking short name, then for backcompat we check for prefixed name."""
     backcompat_prefix = "extra__dataprep__"
     if field_name.startswith("extra__"):

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

@@ -26,10 +26,6 @@ import uuid
 from collections.abc import MutableSequence, Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-from airflow.version import version as airflow_version
 from google.api_core.client_options import ClientOptions
 from google.api_core.exceptions import ServerError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -49,6 +45,11 @@ from google.cloud.dataproc_v1 import (
     WorkflowTemplateServiceClient,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+from airflow.version import version as airflow_version
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.operation_async import AsyncOperation
@@ -297,7 +298,7 @@ class DataprocHook(GoogleBaseHook):
         success_code = 0
 
         with self.provide_authorized_gcloud():
-            proc = subprocess.run(cmd, capture_output=True)
+            proc = subprocess.run(cmd, check=False, capture_output=True)
 
         if proc.returncode != success_code:
             stderr_last_20_lines = "\n".join(proc.stderr.decode().strip().splitlines()[-20:])
@@ -911,12 +912,15 @@ class DataprocHook(GoogleBaseHook):
         state = None
         start = time.monotonic()
         while state not in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED):
+            self.log.debug("Waiting for job %s to complete", job_id)
             if timeout and start + timeout < time.monotonic():
                 raise AirflowException(f"Timeout: dataproc job {job_id} is not ready after {timeout}s")
+            self.log.debug("Sleeping for %s seconds", wait_time)
             time.sleep(wait_time)
             try:
                 job = self.get_job(project_id=project_id, region=region, job_id=job_id)
                 state = job.status.state
+                self.log.debug("Job %s is in state %s", job_id, state)
             except ServerError as err:
                 self.log.info("Retrying. Dataproc API returned server error when waiting for job: %s", err)
 
@@ -1268,7 +1272,7 @@ class DataprocHook(GoogleBaseHook):
         return all([word in error_msg for word in key_words])
 
 
-class DataprocAsyncHook(GoogleBaseHook):
+class DataprocAsyncHook(GoogleBaseAsyncHook):
     """
     Asynchronous interaction with Google Cloud Dataproc APIs.
 
@@ -1276,6 +1280,8 @@ class DataprocAsyncHook(GoogleBaseHook):
     keyword arguments rather than positional.
     """
 
+    sync_hook_class = DataprocHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",
@@ -1285,53 +1291,90 @@ class DataprocAsyncHook(GoogleBaseHook):
         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:
+    async def get_cluster_client(self, region: str | None = None) -> ClusterControllerAsyncClient:
         """Create a ClusterControllerAsyncClient."""
         client_options = None
         if region and region != "global":
             client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
 
+        sync_hook = await self.get_sync_hook()
         return ClusterControllerAsyncClient(
-            credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+            credentials=sync_hook.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
-    def get_template_client(self, region: str | None = None) -> WorkflowTemplateServiceAsyncClient:
+    async def get_template_client(self, region: str | None = None) -> WorkflowTemplateServiceAsyncClient:
         """Create a WorkflowTemplateServiceAsyncClient."""
         client_options = None
         if region and region != "global":
             client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
 
+        sync_hook = await self.get_sync_hook()
         return WorkflowTemplateServiceAsyncClient(
-            credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+            credentials=sync_hook.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
-    def get_job_client(self, region: str | None = None) -> JobControllerAsyncClient:
+    async def get_job_client(self, region: str | None = None) -> JobControllerAsyncClient:
         """Create a JobControllerAsyncClient."""
         if self._cached_client is None:
             client_options = None
             if region and region != "global":
                 client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
 
+            sync_hook = await self.get_sync_hook()
             self._cached_client = JobControllerAsyncClient(
-                credentials=self.get_credentials(),
+                credentials=sync_hook.get_credentials(),
                 client_info=CLIENT_INFO,
                 client_options=client_options,
             )
         return self._cached_client
 
-    def get_batch_client(self, region: str | None = None) -> BatchControllerAsyncClient:
+    async def get_batch_client(self, region: str | None = None) -> BatchControllerAsyncClient:
         """Create a BatchControllerAsyncClient."""
         client_options = None
         if region and region != "global":
             client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
 
+        sync_hook = await self.get_sync_hook()
         return BatchControllerAsyncClient(
-            credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+            credentials=sync_hook.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
-    def get_operations_client(self, region: str) -> OperationsClient:
+    async def get_operations_client(self, region: str) -> OperationsClient:
         """Create a OperationsClient."""
-        return self.get_template_client(region=region).transport.operations_client
+        template_client = await self.get_template_client(region=region)
+        return template_client.transport.operations_client
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def get_cluster(
+        self,
+        region: str,
+        cluster_name: str,
+        project_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Cluster:
+        """
+        Get a cluster.
+
+        :param region: Cloud Dataproc region in which to handle the request.
+        :param cluster_name: Name of the cluster to get.
+        :param project_id: Google Cloud project ID that the cluster belongs to.
+        :param retry: A retry object used to retry requests. If *None*, requests
+            will not be retried.
+        :param timeout: The amount of time, in seconds, to wait for the request
+            to complete. If *retry* is specified, the timeout applies to each
+            individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+        """
+        client = await self.get_cluster_client(region=region)
+        result = await client.get_cluster(
+            request={"project_id": project_id, "region": region, "cluster_name": cluster_name},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
 
     @GoogleBaseHook.fallback_to_default_project_id
     async def create_cluster(
@@ -1389,7 +1432,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             cluster["config"] = cluster_config  # type: ignore
             cluster["labels"] = labels  # type: ignore
 
-        client = self.get_cluster_client(region=region)
+        client = await self.get_cluster_client(region=region)
         result = await client.create_cluster(
             request={
                 "project_id": project_id,
@@ -1434,7 +1477,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_cluster_client(region=region)
+        client = await self.get_cluster_client(region=region)
         result = await client.delete_cluster(
             request={
                 "project_id": project_id,
@@ -1482,7 +1525,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_cluster_client(region=region)
+        client = await self.get_cluster_client(region=region)
         result = await client.diagnose_cluster(
             request={
                 "project_id": project_id,
@@ -1499,38 +1542,6 @@ class DataprocAsyncHook(GoogleBaseHook):
         )
         return result
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    async def get_cluster(
-        self,
-        region: str,
-        cluster_name: str,
-        project_id: str,
-        retry: AsyncRetry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> Cluster:
-        """
-        Get the resource representation for a cluster in a project.
-
-        :param project_id: Google Cloud project ID that the cluster belongs to.
-        :param region: Cloud Dataproc region to handle the request.
-        :param cluster_name: The cluster name.
-        :param retry: A retry object used to retry requests. If *None*, requests
-            will not be retried.
-        :param timeout: The amount of time, in seconds, to wait for the request
-            to complete. If *retry* is specified, the timeout applies to each
-            individual attempt.
-        :param metadata: Additional metadata that is provided to the method.
-        """
-        client = self.get_cluster_client(region=region)
-        result = await client.get_cluster(
-            request={"project_id": project_id, "region": region, "cluster_name": cluster_name},
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return result
-
     @GoogleBaseHook.fallback_to_default_project_id
     async def list_clusters(
         self,
@@ -1560,7 +1571,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_cluster_client(region=region)
+        client = await self.get_cluster_client(region=region)
         result = await client.list_clusters(
             request={"project_id": project_id, "region": region, "filter": filter_, "page_size": page_size},
             retry=retry,
@@ -1637,7 +1648,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         """
         if region is None:
             raise TypeError("missing 1 required keyword argument: 'region'")
-        client = self.get_cluster_client(region=region)
+        client = await self.get_cluster_client(region=region)
         operation = await client.update_cluster(
             request={
                 "project_id": project_id,
@@ -1679,10 +1690,8 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        if region is None:
-            raise TypeError("missing 1 required keyword argument: 'region'")
         metadata = metadata or ()
-        client = self.get_template_client(region)
+        client = await self.get_template_client(region)
         parent = f"projects/{project_id}/regions/{region}"
         return await client.create_workflow_template(
             request={"parent": parent, "template": template}, retry=retry, timeout=timeout, metadata=metadata
@@ -1724,10 +1733,8 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        if region is None:
-            raise TypeError("missing 1 required keyword argument: 'region'")
         metadata = metadata or ()
-        client = self.get_template_client(region)
+        client = await self.get_template_client(region)
         name = f"projects/{project_id}/regions/{region}/workflowTemplates/{template_name}"
         operation = await client.instantiate_workflow_template(
             request={"name": name, "version": version, "request_id": request_id, "parameters": parameters},
@@ -1766,10 +1773,8 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        if region is None:
-            raise TypeError("missing 1 required keyword argument: 'region'")
         metadata = metadata or ()
-        client = self.get_template_client(region)
+        client = await self.get_template_client(region)
         parent = f"projects/{project_id}/regions/{region}"
         operation = await client.instantiate_inline_workflow_template(
             request={"parent": parent, "template": template, "request_id": request_id},
@@ -1780,7 +1785,8 @@ class DataprocAsyncHook(GoogleBaseHook):
         return operation
 
     async def get_operation(self, region, operation_name):
-        return await self.get_operations_client(region).get_operation(name=operation_name)
+        operations_client = await self.get_operations_client(region)
+        return await operations_client.get_operation(name=operation_name)
 
     @GoogleBaseHook.fallback_to_default_project_id
     async def get_job(
@@ -1805,9 +1811,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        if region is None:
-            raise TypeError("missing 1 required keyword argument: 'region'")
-        client = self.get_job_client(region=region)
+        client = await self.get_job_client(region=region)
         job = await client.get_job(
             request={"project_id": project_id, "region": region, "job_id": job_id},
             retry=retry,
@@ -1844,9 +1848,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        if region is None:
-            raise TypeError("missing 1 required keyword argument: 'region'")
-        client = self.get_job_client(region=region)
+        client = await self.get_job_client(region=region)
         return await client.submit_job(
             request={"project_id": project_id, "region": region, "job": job, "request_id": request_id},
             retry=retry,
@@ -1877,7 +1879,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_job_client(region=region)
+        client = await self.get_job_client(region=region)
 
         job = await client.cancel_job(
             request={"project_id": project_id, "region": region, "job_id": job_id},
@@ -1919,7 +1921,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_batch_client(region)
+        client = await self.get_batch_client(region)
         parent = f"projects/{project_id}/regions/{region}"
 
         result = await client.create_batch(
@@ -1958,7 +1960,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_batch_client(region)
+        client = await self.get_batch_client(region)
         name = f"projects/{project_id}/locations/{region}/batches/{batch_id}"
 
         await client.delete_batch(
@@ -1993,7 +1995,7 @@ class DataprocAsyncHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
-        client = self.get_batch_client(region)
+        client = await self.get_batch_client(region)
         name = f"projects/{project_id}/locations/{region}/batches/{batch_id}"
 
         result = await client.get_batch(
@@ -2038,7 +2040,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         :param filter: Result filters as specified in ListBatchesRequest
         :param order_by: How to order results as specified in ListBatchesRequest
         """
-        client = self.get_batch_client(region)
+        client = await self.get_batch_client(region)
         parent = f"projects/{project_id}/regions/{region}"
 
         result = await client.list_batches(

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

@@ -22,13 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.metastore_v1 import DataprocMetastoreClient
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

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

@@ -31,9 +31,6 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-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
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dlp import DlpServiceClient
 from google.cloud.dlp_v2.types import (
@@ -58,6 +55,10 @@ from google.cloud.dlp_v2.types import (
 )
 from google.protobuf.field_mask_pb2 import FieldMask
 
+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:
     from google.api_core.retry import Retry
 
@@ -267,7 +268,7 @@ class CloudDLPHook(GoogleBaseHook):
 
             if job.state == DlpJob.JobState.DONE:
                 return job
-            elif job.state in [
+            if job.state in [
                 DlpJob.JobState.PENDING,
                 DlpJob.JobState.RUNNING,
                 DlpJob.JobState.JOB_STATE_UNSPECIFIED,