From 721f5d8b6d03880b480366385e11ad908bd91d9d Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Fri, 14 Jun 2024 11:00:47 -0400 Subject: [PATCH] Fix internal test failure caused by PR 31550 --- sdks/python/apache_beam/io/gcp/gcsio.py | 7 ++++--- sdks/python/apache_beam/options/pipeline_options.py | 2 ++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py index b2f8bd4a2da7a..de6d902d87a98 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio.py +++ b/sdks/python/apache_beam/io/gcp/gcsio.py @@ -156,10 +156,10 @@ def get_project_number(self, bucket): return self.bucket_to_project_number.get(bucket, None) - def get_bucket(self, bucket_name): + def get_bucket(self, bucket_name, **kwargs): """Returns an object bucket from its name, or None if it does not exist.""" try: - return self.client.lookup_bucket(bucket_name) + return self.client.lookup_bucket(bucket_name, **kwargs) except NotFound: return None @@ -532,7 +532,8 @@ def _updated_to_seconds(updated): def is_soft_delete_enabled(self, gcs_path): try: bucket_name, _ = parse_gcs_path(gcs_path) - bucket = self.get_bucket(bucket_name) + # set retry timeout to 5 seconds when checking soft delete policy + bucket = self.get_bucket(bucket_name, retry=DEFAULT_RETRY.with_timeout(5)) if (bucket.soft_delete_policy is not None and bucket.soft_delete_policy.retention_duration_seconds > 0): return True diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 42aee47a957e8..55fb2f1703fc6 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -945,6 +945,8 @@ def _handle_temp_and_staging_locations(self, validator): self._warn_if_soft_delete_policy_enabled('temp_location') return [] elif not staging_errors and not temp_errors: + self._warn_if_soft_delete_policy_enabled('temp_location') + self._warn_if_soft_delete_policy_enabled('staging_location') return [] # Both staging and temp locations are bad, try to use default bucket. else: