apache-airflow-providers-google 10.7.0rc1__py3-none-any.whl → 10.8.0__py3-none-any.whl

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

@@ -19,18 +19,20 @@
 from __future__ import annotations
 
 import time
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
-from google.api_core.retry import Retry
 from google.cloud.compute_v1.services.instance_group_managers import InstanceGroupManagersClient
 from google.cloud.compute_v1.services.instance_templates import InstanceTemplatesClient
 from google.cloud.compute_v1.services.instances import InstancesClient
-from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 from googleapiclient.discovery import build
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+    from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
+
 # Time to sleep between active checks of the operation results
 TIME_TO_SLEEP_IN_SECONDS = 1
 

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

@@ -314,7 +314,7 @@ class ComputeEngineSSHHook(SSHHook):
                 item["value"] = keys
                 break
         else:
-            new_dict = dict(key="ssh-keys", value=keys)
+            new_dict = {"key": "ssh-keys", "value": keys}
             metadata["items"] = [new_dict]
 
         self._compute_hook.set_instance_metadata(

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

@@ -16,10 +16,9 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud import datacatalog
 from google.cloud.datacatalog import (
     CreateTagRequest,
@@ -31,12 +30,15 @@ from google.cloud.datacatalog import (
     TagTemplate,
     TagTemplateField,
 )
-from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow 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
+    from google.protobuf.field_mask_pb2 import FieldMask
+
 
 class CloudDataCatalogHook(GoogleBaseHook):
     """

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

@@ -226,10 +226,7 @@ class _DataflowJobsController(LoggingMixin):
         if not self._jobs:
             return False
 
-        for job in self._jobs:
-            if job["currentState"] not in DataflowJobStatus.TERMINAL_STATES:
-                return True
-        return False
+        return any(job["currentState"] not in DataflowJobStatus.TERMINAL_STATES for job in self._jobs)
 
     def _get_current_jobs(self) -> list[dict]:
         """
@@ -870,7 +867,7 @@ class DataflowHook(GoogleBaseHook):
         """Builds Dataflow job name."""
         base_job_name = str(job_name).replace("_", "-")
 
-        if not re.match(r"^[a-z]([-a-z0-9]*[a-z0-9])?$", base_job_name):
+        if not re.fullmatch(r"[a-z]([-a-z0-9]*[a-z0-9])?", base_job_name):
             raise ValueError(
                 f"Invalid job_name ({base_job_name}); the name must consist of only the characters "
                 f"[-a-z0-9], starting with a letter and ending with a letter or number "
@@ -1236,12 +1233,12 @@ class AsyncDataflowHook(GoogleBaseAsyncHook):
         client = await self.initialize_client(JobsV1Beta3AsyncClient)
 
         request = GetJobRequest(
-            dict(
-                project_id=project_id,
-                job_id=job_id,
-                view=job_view,
-                location=location,
-            )
+            {
+                "project_id": project_id,
+                "job_id": job_id,
+                "view": job_view,
+                "location": location,
+            }
         )
 
         job = await client.get_job(

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

@@ -17,10 +17,9 @@
 from __future__ import annotations
 
 import time
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.dataform_v1beta1 import DataformClient
 from google.cloud.dataform_v1beta1.types import (
     CompilationResult,
@@ -34,6 +33,9 @@ from google.cloud.dataform_v1beta1.types import (
 from airflow import AirflowException
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
 
 class DataformHook(GoogleBaseHook):
     """Hook for Google Cloud DataForm APIs."""

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

@@ -17,6 +17,7 @@
 """This module contains Google DataFusion hook."""
 from __future__ import annotations
 
+import asyncio
 import json
 import os
 from time import monotonic, sleep
@@ -152,7 +153,7 @@ class DataFusionHook(GoogleBaseHook):
         return os.path.join(instance_url, "v3", "namespaces", quote(namespace), "apps")
 
     def _cdap_request(
-        self, url: str, method: str, body: list | dict | None = None
+        self, url: str, method: str, body: list | dict | None = None, params: dict | None = None
     ) -> google.auth.transport.Response:
         headers: dict[str, str] = {"Content-Type": "application/json"}
         request = google.auth.transport.requests.Request()
@@ -162,7 +163,7 @@ class DataFusionHook(GoogleBaseHook):
 
         payload = json.dumps(body) if body else None
 
-        response = request(method=method, url=url, headers=headers, body=payload)
+        response = request(method=method, url=url, headers=headers, body=payload, params=params)
         return response
 
     @staticmethod
@@ -281,6 +282,23 @@ class DataFusionHook(GoogleBaseHook):
         )
         return instance
 
+    def get_instance_artifacts(
+        self, instance_url: str, namespace: str = "default", scope: str = "SYSTEM"
+    ) -> Any:
+        url = os.path.join(
+            instance_url,
+            "v3",
+            "namespaces",
+            quote(namespace),
+            "artifacts",
+        )
+        response = self._cdap_request(url=url, method="GET", params={"scope": scope})
+        self._check_response_status_and_data(
+            response, f"Retrieving an instance artifacts failed with code {response.status}"
+        )
+        content = json.loads(response.data)
+        return content
+
     @GoogleBaseHook.fallback_to_default_project_id
     def patch_instance(
         self,
@@ -370,12 +388,12 @@ class DataFusionHook(GoogleBaseHook):
                 self._check_response_status_and_data(
                     response, f"Deleting a pipeline failed with code {response.status}: {response.data}"
                 )
-                if response.status == 200:
-                    break
             except ConflictException as exc:
                 self.log.info(exc)
                 sleep(time_to_wait)
-                continue
+            else:
+                if response.status == 200:
+                    break
 
     def list_pipelines(
         self,
@@ -528,7 +546,7 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
                     break
                 except Exception as exc:
                     if "404" in str(exc):
-                        sleep(time_to_wait)
+                        await asyncio.sleep(time_to_wait)
                     else:
                         raise
         if pipeline:

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

@@ -18,13 +18,11 @@
 from __future__ import annotations
 
 import time
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.operation import Operation
-from google.api_core.retry import Retry
-from google.cloud.dataplex_v1 import DataplexServiceClient, DataScanServiceClient
+from google.cloud.dataplex_v1 import DataplexServiceClient, DataScanServiceAsyncClient, DataScanServiceClient
 from google.cloud.dataplex_v1.types import (
     Asset,
     DataScan,
@@ -34,11 +32,15 @@ from google.cloud.dataplex_v1.types import (
     Zone,
 )
 from google.protobuf.field_mask_pb2 import FieldMask
-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 GoogleBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+
+if TYPE_CHECKING:
+    from google.api_core.operation import Operation
+    from google.api_core.retry import Retry
+    from googleapiclient.discovery import Resource
 
 PATH_DATA_SCAN = "projects/{project_id}/locations/{region}/dataScans/{data_scan_id}"
 
@@ -857,3 +859,70 @@ class DataplexHook(GoogleBaseHook):
             metadata=metadata,
         )
         return result
+
+
+class DataplexAsyncHook(GoogleBaseAsyncHook):
+    """
+    Asynchronous Hook for Google Cloud Dataplex APIs.
+
+    All the methods in the hook where project_id is used must be called with
+    keyword arguments rather than positional.
+    """
+
+    sync_hook_class = DataplexHook
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+
+        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain)
+
+    async def get_dataplex_data_scan_client(self) -> DataScanServiceAsyncClient:
+        """Returns DataScanServiceAsyncClient."""
+        client_options = ClientOptions(api_endpoint="dataplex.googleapis.com:443")
+
+        return DataScanServiceAsyncClient(
+            credentials=(await self.get_sync_hook()).get_credentials(),
+            client_info=CLIENT_INFO,
+            client_options=client_options,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def get_data_scan_job(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str | None = None,
+        job_id: str | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Any:
+        """
+        Gets a DataScan Job resource.
+
+        :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+        :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+        :param data_scan_id: Required. DataScan identifier.
+        :param job_id: Required. The resource name of the DataScanJob:
+            projects/{project_id}/locations/{region}/dataScans/{data_scan_id}/jobs/{data_scan_job_id}
+        :param retry: A retry object used  to retry requests. If `None` is specified, requests
+            will not be retried.
+        :param timeout: The amount of time, in seconds, to wait for the request to complete.
+            Note that 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_dataplex_data_scan_client()
+
+        name = f"projects/{project_id}/locations/{region}/dataScans/{data_scan_id}/jobs/{job_id}"
+        result = await client.get_data_scan_job(
+            request={"name": name, "view": "FULL"},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        return result

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

@@ -20,15 +20,11 @@ from __future__ import annotations
 
 import time
 import uuid
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 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
-from google.api_core.operation import Operation
-from google.api_core.operation_async import AsyncOperation
-from google.api_core.operations_v1.operations_client import OperationsClient
-from google.api_core.retry import Retry
 from google.cloud.dataproc_v1 import (
     Batch,
     BatchControllerAsyncClient,
@@ -44,14 +40,20 @@ from google.cloud.dataproc_v1 import (
     WorkflowTemplateServiceAsyncClient,
     WorkflowTemplateServiceClient,
 )
-from google.protobuf.duration_pb2 import Duration
-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 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
+    from google.api_core.operations_v1.operations_client import OperationsClient
+    from google.api_core.retry import Retry
+    from google.protobuf.duration_pb2 import Duration
+    from google.protobuf.field_mask_pb2 import FieldMask
+
 
 class DataProcJobBuilder:
     """A helper class for building Dataproc job."""

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

@@ -18,21 +18,23 @@
 """This module contains a Google Cloud Dataproc Metastore hook."""
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.operation import Operation
-from google.api_core.retry import Retry
 from google.cloud.metastore_v1 import DataprocMetastoreClient
-from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
-from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
-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 GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.operation import Operation
+    from google.api_core.retry import Retry
+    from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
+    from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
+    from google.protobuf.field_mask_pb2 import FieldMask
+
 
 class DataprocMetastoreHook(GoogleBaseHook):
     """Hook for Google Cloud Dataproc Metastore APIs."""