diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py index 74b6061c0..f03807ca7 100644 --- a/google/cloud/storage/client.py +++ b/google/cloud/storage/client.py @@ -64,6 +64,15 @@ _marker = object() +def _buckets_page_start(iterator, page, response): + """Grab unreachable buckets after a :class:`~google.cloud.iterator.Page` started.""" + unreachable = response.get("unreachable", []) + if not isinstance(unreachable, list): + raise TypeError( + f"expected unreachable to be list, but obtained {type(unreachable)}" + ) + page.unreachable = unreachable + class Client(ClientWithProject): """Client to bundle configuration needed for API requests. @@ -1458,6 +1467,7 @@ def list_buckets( retry=DEFAULT_RETRY, *, soft_deleted=None, + return_partial_success=None, ): """Get all buckets in the project associated to the client. @@ -1516,6 +1526,13 @@ def list_buckets( generation number. This parameter can only be used successfully if the bucket has a soft delete policy. See: https://cloud.google.com/storage/docs/soft-delete + :type return_partial_success: bool + :param return_partial_success: + (Optional) If True, the response will also contain a list of + unreachable buckets if the buckets are unavailable. The + unreachable buckets will be available on the ``unreachable`` + attribute of the returned iterator. + :rtype: :class:`~google.api_core.page_iterator.Iterator` :raises ValueError: if both ``project`` is ``None`` and the client's project is also ``None``. @@ -1551,7 +1568,10 @@ def list_buckets( if soft_deleted is not None: extra_params["softDeleted"] = soft_deleted - return self._list_resource( + if return_partial_success is not None: + extra_params["returnPartialSuccess"] = return_partial_success + + iterator = self._list_resource( "/b", _item_to_bucket, page_token=page_token, @@ -1560,7 +1580,9 @@ def list_buckets( page_size=page_size, timeout=timeout, retry=retry, + page_start=_buckets_page_start, ) + return iterator def restore_bucket( self, diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index d5723740e..331f40d66 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2224,7 +2224,7 @@ def test_list_blobs_w_explicit_w_user_project(self): def test_list_buckets_wo_project(self): from google.cloud.exceptions import BadRequest - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start credentials = _make_credentials() client = self._make_one(project=None, credentials=credentials) @@ -2253,10 +2253,11 @@ def test_list_buckets_wo_project(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_wo_project_w_emulator(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start # mock STORAGE_EMULATOR_ENV_VAR is set host = "http://localhost:8080" @@ -2288,10 +2289,11 @@ def test_list_buckets_wo_project_w_emulator(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_w_environ_project_w_emulator(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start # mock STORAGE_EMULATOR_ENV_VAR is set host = "http://localhost:8080" @@ -2327,10 +2329,11 @@ def test_list_buckets_w_environ_project_w_emulator(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_w_custom_endpoint(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start custom_endpoint = "storage-example.p.googleapis.com" client = self._make_one(client_options={"api_endpoint": custom_endpoint}) @@ -2358,10 +2361,11 @@ def test_list_buckets_w_custom_endpoint(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_w_defaults(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start project = "PROJECT" credentials = _make_credentials() @@ -2390,10 +2394,11 @@ def test_list_buckets_w_defaults(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_w_soft_deleted(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start project = "PROJECT" credentials = _make_credentials() @@ -2423,10 +2428,11 @@ def test_list_buckets_w_soft_deleted(self): page_size=expected_page_size, timeout=self._get_default_timeout(), retry=DEFAULT_RETRY, + page_start=_buckets_page_start, ) def test_list_buckets_w_explicit(self): - from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _item_to_bucket, _buckets_page_start project = "foo-bar" other_project = "OTHER_PROJECT" @@ -2476,6 +2482,7 @@ def test_list_buckets_w_explicit(self): page_size=expected_page_size, timeout=timeout, retry=retry, + page_start=_buckets_page_start, ) def test_restore_bucket(self): @@ -3086,6 +3093,57 @@ def test_get_signed_policy_v4_with_access_token_sa_email(self): self.assertEqual(fields["x-goog-signature"], EXPECTED_SIGN) self.assertEqual(fields["policy"], EXPECTED_POLICY) + def test_list_buckets_w_partial_success(self): + from google.cloud.storage.client import _item_to_bucket + from google.cloud.storage.client import _buckets_page_start + + PROJECT = "project" + bucket_name = "bucket-name" + unreachable_bucket = "projects/_/buckets/unreachable-bucket" + + client = self._make_one(project=PROJECT) + + mock_bucket = mock.Mock() + mock_bucket.name = bucket_name + + mock_page = mock.Mock() + mock_page.unreachable = [unreachable_bucket] + mock_page.__iter__ = mock.Mock(return_value=iter([mock_bucket])) + + mock_iterator = mock.Mock() + mock_iterator.pages = iter([mock_page]) + + client._list_resource = mock.Mock(return_value=mock_iterator) + + iterator = client.list_buckets(return_partial_success=True) + + page = next(iterator.pages) + + self.assertEqual(page.unreachable, [unreachable_bucket]) + + buckets = list(page) + self.assertEqual(len(buckets), 1) + self.assertEqual(buckets[0].name, bucket_name) + + expected_path = "/b" + expected_item_to_value = _item_to_bucket + expected_extra_params = { + "project": PROJECT, + "projection": "noAcl", + "returnPartialSuccess": True, + } + + client._list_resource.assert_called_once_with( + expected_path, + expected_item_to_value, + page_token=None, + max_results=None, + extra_params=expected_extra_params, + page_size=None, + timeout=self._get_default_timeout(), + retry=DEFAULT_RETRY, + page_start=_buckets_page_start, + ) class Test__item_to_bucket(unittest.TestCase): def _call_fut(self, iterator, item):