From 481ef8816d63e798eeb617478d3ab11ec819393f Mon Sep 17 00:00:00 2001 From: Yuan Chuan Kee <1683885+kylase@users.noreply.github.com> Date: Fri, 14 Feb 2025 10:01:48 +0000 Subject: [PATCH 1/3] Fix `exists` method to support using Requester Pays --- .../providers/google/cloud/hooks/gcs.py | 18 +++++++++++------- .../tests/unit/google/cloud/hooks/test_gcs.py | 6 +++--- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index d6807b3f2997a..430c521fd2597 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py @@ -598,7 +598,9 @@ def _call_with_retry(f: Callable[[], None]) -> None: context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name} ) - def exists(self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETRY) -> bool: + def exists( + self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETRY, user_project: str = None + ) -> bool: """ Check for the existence of a file in Google Cloud Storage. @@ -606,9 +608,11 @@ def exists(self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETR :param object_name: The name of the blob_name to check in the Google cloud storage bucket. :param retry: (Optional) How to retry the RPC + :param user_project: The identifier of the Google Cloud project to bill for the request. + Required for Requester Pays buckets. """ client = self.get_conn() - bucket = client.bucket(bucket_name) + bucket = client.bucket(bucket_name, user_project=user_project) blob = bucket.blob(blob_name=object_name) return blob.exists(retry=retry) @@ -625,7 +629,7 @@ def get_blob_update_time(self, bucket_name: str, object_name: str): def is_updated_after(self, bucket_name: str, object_name: str, ts: datetime) -> bool: """ - Check if an blob_name is updated in Google Cloud Storage. + Check if a blob_name is updated in Google Cloud Storage. :param bucket_name: The Google Cloud Storage bucket where the object is. :param object_name: The name of the object to check in the Google cloud @@ -645,7 +649,7 @@ def is_updated_between( self, bucket_name: str, object_name: str, min_ts: datetime, max_ts: datetime ) -> bool: """ - Check if an blob_name is updated in Google Cloud Storage. + Check if a blob_name is updated in Google Cloud Storage. :param bucket_name: The Google Cloud Storage bucket where the object is. :param object_name: The name of the object to check in the Google cloud @@ -666,7 +670,7 @@ def is_updated_between( def is_updated_before(self, bucket_name: str, object_name: str, ts: datetime) -> bool: """ - Check if an blob_name is updated before given time in Google Cloud Storage. + Check if a blob_name is updated before given time in Google Cloud Storage. :param bucket_name: The Google Cloud Storage bucket where the object is. :param object_name: The name of the object to check in the Google cloud @@ -979,7 +983,7 @@ def list_by_timespan( break return ids - def _get_blob(self, bucket_name: str, object_name: str) -> Blob: + def _get_blob(self, bucket_name: str, object_name: str, user_project: str | None = None) -> Blob: """ Get a blob object in Google Cloud Storage. @@ -989,7 +993,7 @@ def _get_blob(self, bucket_name: str, object_name: str) -> Blob: """ client = self.get_conn() - bucket = client.bucket(bucket_name) + bucket = client.bucket(bucket_name, user_project=user_project) blob = bucket.get_blob(blob_name=object_name) if blob is None: diff --git a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py index 429f10003b0c8..f2cafa605a813 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py @@ -206,11 +206,11 @@ def test_exists(self, mock_service): exists_method.return_value = True # When - response = self.gcs_hook.exists(bucket_name=test_bucket, object_name=test_object) + response = self.gcs_hook.exists(bucket_name=test_bucket, object_name=test_object, user_project=None) # Then assert response - bucket_mock.assert_called_once_with(test_bucket) + bucket_mock.assert_called_once_with(test_bucket, user_project=None) blob_object.assert_called_once_with(blob_name=test_object) exists_method.assert_called_once_with(retry=DEFAULT_RETRY) @@ -226,7 +226,7 @@ def test_exists_nonexisting_object(self, mock_service): exists_method.return_value = False # When - response = self.gcs_hook.exists(bucket_name=test_bucket, object_name=test_object) + response = self.gcs_hook.exists(bucket_name=test_bucket, object_name=test_object, user_project=None) # Then assert not response From 4bacfcace88f5ec11f34a84334fdab1a40594e39 Mon Sep 17 00:00:00 2001 From: Yuan Chuan Kee <1683885+kylase@users.noreply.github.com> Date: Fri, 14 Feb 2025 12:00:53 +0000 Subject: [PATCH 2/3] Fix wrong type for user_project --- .../google/src/airflow/providers/google/cloud/hooks/gcs.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index 430c521fd2597..3179311cee16f 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py @@ -599,7 +599,11 @@ def _call_with_retry(f: Callable[[], None]) -> None: ) def exists( - self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETRY, user_project: str = None + self, + bucket_name: str, + object_name: str, + retry: Retry = DEFAULT_RETRY, + user_project: str | None = None, ) -> bool: """ Check for the existence of a file in Google Cloud Storage. From bcbfab33d908896615caae088bc9cd7d7416dfb1 Mon Sep 17 00:00:00 2001 From: Yuan Chuan Kee <1683885+kylase@users.noreply.github.com> Date: Fri, 14 Feb 2025 12:07:56 +0000 Subject: [PATCH 3/3] Revert change for _get_blob --- .../google/src/airflow/providers/google/cloud/hooks/gcs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index 3179311cee16f..116f5c4501ff7 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py @@ -987,7 +987,7 @@ def list_by_timespan( break return ids - def _get_blob(self, bucket_name: str, object_name: str, user_project: str | None = None) -> Blob: + def _get_blob(self, bucket_name: str, object_name: str) -> Blob: """ Get a blob object in Google Cloud Storage. @@ -997,7 +997,7 @@ def _get_blob(self, bucket_name: str, object_name: str, user_project: str | None """ client = self.get_conn() - bucket = client.bucket(bucket_name, user_project=user_project) + bucket = client.bucket(bucket_name) blob = bucket.get_blob(blob_name=object_name) if blob is None: