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

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

@@ -47,7 +47,7 @@ from google.cloud.dataproc_v1 import (
 
 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
 from airflow.version import version as airflow_version
 
 if TYPE_CHECKING:
@@ -298,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:])
@@ -912,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)
 
@@ -1269,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.
 
@@ -1277,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",
@@ -1286,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(
@@ -1390,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,
@@ -1435,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,
@@ -1483,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,
@@ -1500,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,
@@ -1561,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,
@@ -1638,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,
@@ -1680,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
@@ -1725,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},
@@ -1767,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},
@@ -1781,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(
@@ -1806,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,
@@ -1845,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,
@@ -1878,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},
@@ -1920,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(
@@ -1959,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(
@@ -1994,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(
@@ -2039,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/gcs.py

@@ -26,24 +26,26 @@ import os
 import shutil
 import time
 import warnings
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from contextlib import contextmanager
+from datetime import datetime
 from functools import partial
 from io import BytesIO
+from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, cast, overload
+from typing import IO, TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
+# Make mypy happy by importing as aliases
+import google.cloud.storage as storage
 from gcloud.aio.storage import Storage
 from google.api_core.exceptions import GoogleAPICallError, NotFound
-
-# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
-from google.cloud import storage  # type: ignore[attr-defined]
 from google.cloud.exceptions import GoogleCloudError
 from google.cloud.storage.retry import DEFAULT_RETRY
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
+from airflow.providers.common.compat.sdk import timezone
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
@@ -51,13 +53,9 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from airflow.typing_compat import ParamSpec
-from airflow.utils import timezone
 from airflow.version import version
 
 if TYPE_CHECKING:
-    from datetime import datetime
-
     from aiohttp import ClientSession
     from google.api_core.retry import Retry
     from google.cloud.storage.blob import Blob
@@ -373,8 +371,7 @@ class GCSHook(GoogleBaseHook):
                         num_max_attempts,
                     )
                     raise
-        else:
-            raise NotImplementedError  # should not reach this, but makes mypy happy
+        raise NotImplementedError  # should not reach this, but makes mypy happy
 
     def download_as_byte_array(
         self,
@@ -549,13 +546,13 @@ class GCSHook(GoogleBaseHook):
         if cache_control:
             blob.cache_control = cache_control
 
-        if filename and data:
+        if filename is not None and data is not None:
             raise ValueError(
                 "'filename' and 'data' parameter provided. Please "
                 "specify a single parameter, either 'filename' for "
                 "local file uploads or 'data' for file content uploads."
             )
-        if filename:
+        if filename is not None:
             if not mime_type:
                 mime_type = "application/octet-stream"
             if gzip:
@@ -575,7 +572,7 @@ class GCSHook(GoogleBaseHook):
             if gzip:
                 os.remove(filename)
             self.log.info("File %s uploaded to %s in %s bucket", filename, object_name, bucket_name)
-        elif data:
+        elif data is not None:
             if not mime_type:
                 mime_type = "text/plain"
             if gzip:
@@ -1251,6 +1248,106 @@ class GCSHook(GoogleBaseHook):
 
         self.log.info("Completed successfully.")
 
+    def _sync_to_local_dir_delete_stale_local_files(self, current_gcs_objects: List[Path], local_dir: Path):
+        current_gcs_keys = {key.resolve() for key in current_gcs_objects}
+
+        for item in local_dir.rglob("*"):
+            if item.is_file():
+                if item.resolve() not in current_gcs_keys:
+                    self.log.debug("Deleting stale local file: %s", item)
+                    item.unlink()
+        # Clean up empty directories
+        for root, dirs, _ in os.walk(local_dir, topdown=False):
+            for d in dirs:
+                dir_path = os.path.join(root, d)
+                if not os.listdir(dir_path):
+                    self.log.debug("Deleting stale empty directory: %s", dir_path)
+                    os.rmdir(dir_path)
+
+    def _sync_to_local_dir_if_changed(self, blob: Blob, local_target_path: Path):
+        should_download = False
+        download_msg = ""
+        if not local_target_path.exists():
+            should_download = True
+            download_msg = f"Local file {local_target_path} does not exist."
+        else:
+            local_stats = local_target_path.stat()
+            # Reload blob to get fresh metadata, including size and updated time
+            blob.reload()
+
+            if blob.size != local_stats.st_size:
+                should_download = True
+                download_msg = (
+                    f"GCS object size ({blob.size}) and local file size ({local_stats.st_size}) differ."
+                )
+
+            gcs_last_modified = blob.updated
+            if (
+                not should_download
+                and gcs_last_modified
+                and local_stats.st_mtime < gcs_last_modified.timestamp()
+            ):
+                should_download = True
+                download_msg = f"GCS object last modified ({gcs_last_modified}) is newer than local file last modified ({datetime.fromtimestamp(local_stats.st_mtime, tz=timezone.utc)})."
+
+        if should_download:
+            self.log.debug("%s Downloading %s to %s", download_msg, blob.name, local_target_path.as_posix())
+            self.download(
+                bucket_name=blob.bucket.name, object_name=blob.name, filename=str(local_target_path)
+            )
+        else:
+            self.log.debug(
+                "Local file %s is up-to-date with GCS object %s. Skipping download.",
+                local_target_path.as_posix(),
+                blob.name,
+            )
+
+    def sync_to_local_dir(
+        self,
+        bucket_name: str,
+        local_dir: str | Path,
+        prefix: str | None = None,
+        delete_stale: bool = False,
+    ) -> None:
+        """
+        Download files from a GCS bucket to a local directory.
+
+        It will download all files from the given ``prefix`` and create the corresponding
+        directory structure in the ``local_dir``.
+
+        If ``delete_stale`` is ``True``, it will delete all local files that do not exist in the GCS bucket.
+
+        :param bucket_name: The name of the GCS bucket.
+        :param local_dir: The local directory to which the files will be downloaded.
+        :param prefix: The prefix of the files to be downloaded.
+        :param delete_stale: If ``True``, deletes local files that don't exist in the bucket.
+        """
+        prefix = prefix or ""
+        local_dir_path = Path(local_dir)
+        self.log.debug("Downloading data from gs://%s/%s to %s", bucket_name, prefix, local_dir_path)
+
+        gcs_bucket = self.get_bucket(bucket_name)
+        local_gcs_objects = []
+
+        for blob in gcs_bucket.list_blobs(prefix=prefix):
+            # GCS lists "directories" as objects ending with a slash. We should skip them.
+            if blob.name.endswith("/"):
+                continue
+
+            blob_path = Path(blob.name)
+            local_target_path = local_dir_path.joinpath(blob_path.relative_to(prefix))
+            if not local_target_path.parent.exists():
+                local_target_path.parent.mkdir(parents=True, exist_ok=True)
+                self.log.debug("Created local directory: %s", local_target_path.parent)
+
+            self._sync_to_local_dir_if_changed(blob=blob, local_target_path=local_target_path)
+            local_gcs_objects.append(local_target_path)
+
+        if delete_stale:
+            self._sync_to_local_dir_delete_stale_local_files(
+                current_gcs_objects=local_gcs_objects, local_dir=local_dir_path
+            )
+
     def sync(
         self,
         source_bucket: str,