feat: add support for partial list buckets (#1606)

Add support for partial list buckets

Author: Pulkit0110

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,

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):