feat: add support for partial success in ListBuckets for grpc

google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java

@@ -454,26 +454,37 @@ public Page<Bucket> list(BucketListOption... options) {
     Opts<BucketListOpt> opts = Opts.unwrap(options).prepend(defaultOpts).prepend(ALL_BUCKET_FIELDS);
     GrpcCallContext grpcCallContext =
         opts.grpcMetadataMapper().apply(GrpcCallContext.createDefault());
-    ListBucketsRequest request =
+    ListBucketsRequest.Builder builder =
         defaultProjectId
             .get()
             .listBuckets()
             .andThen(opts.listBucketsRequest())
-            .apply(ListBucketsRequest.newBuilder())
-            .build();
-    try {
-      GrpcCallContext merge = Utils.merge(grpcCallContext, Retrying.newCallContext());
-      return retrier.run(
-          retryAlgorithmManager.getFor(request),
-          () -> storageClient.listBucketsPagedCallable().call(request, merge),
-          resp ->
-              new TransformingPageDecorator<>(
-                  resp.getPage(),
-                  syntaxDecoders.bucket.andThen(opts.clearBucketFields()),
-                  retrier,
-                  retryAlgorithmManager.getFor(request)));
-    } catch (Exception e) {
-      throw StorageException.coalesce(e);
+            .apply(ListBucketsRequest.newBuilder());
+
+    final ListBucketsRequest request = builder.build();
+    if (!request.getReturnPartialSuccess()) {
+      try {
+        GrpcCallContext merge = Utils.merge(grpcCallContext, Retrying.newCallContext());
+        return retrier.run(
+            retryAlgorithmManager.getFor(request),
+            () -> storageClient.listBucketsPagedCallable().call(request, merge),
+            resp ->
+                new TransformingPageDecorator<>(
+                    resp.getPage(),
+                    syntaxDecoders.bucket.andThen(opts.clearBucketFields()),
+                    retrier,
+                    retryAlgorithmManager.getFor(request)));
+      } catch (Exception e) {
+        throw StorageException.coalesce(e);
+      }
+    } else {
+      // New logic for partial success
+      try {
+        com.google.storage.v2.ListBucketsResponse response = listBuckets(grpcCallContext, request);
+        return new ListBucketsWithPartialSuccessPage(grpcCallContext, request, response, opts);
+      } catch (Exception e) {
+        throw StorageException.coalesce(e);
+      }
     }
   }
 
@@ -1619,6 +1630,79 @@ public Iterable<Blob> getValues() {
     }
   }
 
+  private final class ListBucketsWithPartialSuccessPage implements Page<Bucket> {
+
+    private final GrpcCallContext ctx;
+    private final ListBucketsRequest req;
+    private final com.google.storage.v2.ListBucketsResponse resp;
+    private final Opts<BucketListOpt> opts;
+
+    private ListBucketsWithPartialSuccessPage(
+        GrpcCallContext ctx,
+        ListBucketsRequest req,
+        com.google.storage.v2.ListBucketsResponse resp,
+        Opts<BucketListOpt> opts) {
+      this.ctx = ctx;
+      this.req = req;
+      this.resp = resp;
+      this.opts = opts;
+    }
+
+    @Override
+    public boolean hasNextPage() {
+      return !resp.getNextPageToken().isEmpty();
+    }
+
+    @Override
+    public String getNextPageToken() {
+      return resp.getNextPageToken();
+    }
+
+    @Override
+    public Page<Bucket> getNextPage() {
+      if (!hasNextPage()) {
+        return null;
+      }
+      ListBucketsRequest nextPageReq =
+          req.toBuilder().setPageToken(resp.getNextPageToken()).build();
+      try {
+        com.google.storage.v2.ListBucketsResponse nextPageResp =
+            listBuckets(ctx, nextPageReq);
+        return new ListBucketsWithPartialSuccessPage(ctx, nextPageReq, nextPageResp, opts);
+      } catch (Exception e) {
+        throw StorageException.coalesce(e);
+      }
+    }
+
+    @Override
+    public Iterable<Bucket> getValues() {
+      Decoder<com.google.storage.v2.Bucket, Bucket> bucketDecoder =
+          syntaxDecoders.bucket.andThen(opts.clearBucketFields());
+      Stream<Bucket> reachable = resp.getBucketsList().stream().map(bucketDecoder::decode);
+      Stream<Bucket> unreachable =
+          resp.getUnreachableList().stream()
+              .map(
+                  name -> {
+                    String encoded = bucketNameCodec.encode(name);
+                    return BucketInfo.newBuilder(encoded)
+                        .setIsUnreachable(true)
+                        .build()
+                        .asBucket(GrpcStorageImpl.this);
+                  });
+      return Streams.concat(reachable, unreachable).collect(ImmutableList.toImmutableList());
+    }
+
+    @Override
+    public Iterable<Bucket> iterateAll() {
+      Page<Bucket> curr = this;
+      return () ->
+          streamIterate(curr, p -> p != null && p.hasNextPage(), Page::getNextPage)
+              .filter(Objects::nonNull)
+              .flatMap(p -> StreamSupport.stream(p.getValues().spliterator(), false))
+              .iterator();
+    }
+  }
+
   static final class TransformingPageDecorator<
           RequestT,
           ResponseT,
@@ -1858,6 +1942,15 @@ private SourceObject sourceObjectEncode(SourceBlob from) {
     return to.build();
   }
 
+  private com.google.storage.v2.ListBucketsResponse listBuckets(
+      GrpcCallContext grpcCallContext, ListBucketsRequest request) {
+    GrpcCallContext merge = Utils.merge(grpcCallContext, Retrying.newCallContext());
+    return retrier.run(
+        retryAlgorithmManager.getFor(request),
+        () -> storageClient.listBucketsCallable().call(request, merge),
+        Decoder.identity());
+  }
+
   private com.google.storage.v2.Bucket getBucketWithDefaultAcls(String bucketName) {
     Fields fields =
         UnifiedOpts.fields(

google-cloud-storage/src/main/java/com/google/cloud/storage/Storage.java

@@ -2470,17 +2470,17 @@ public static BucketListOption pageSize(long pageSize) {
       return new BucketListOption(UnifiedOpts.pageSize(pageSize));
     }
 
+    @TransportCompatibility({Transport.HTTP, Transport.GRPC})
+    public static BucketListOption returnPartialSuccess(boolean returnPartialSuccess) {
+      return new BucketListOption(UnifiedOpts.returnPartialSuccess(returnPartialSuccess));
+    }
+
     /** Returns an option to specify the page token from which to start listing buckets. */
     @TransportCompatibility({Transport.HTTP, Transport.GRPC})
     public static BucketListOption pageToken(@NonNull String pageToken) {
       return new BucketListOption(UnifiedOpts.pageToken(pageToken));
     }
 
-    @TransportCompatibility({Transport.HTTP})
-    public static BucketListOption returnPartialSuccess(boolean returnPartialSuccess) {
-      return new BucketListOption(UnifiedOpts.returnPartialSuccess(returnPartialSuccess));
-    }
-
     /**
      * Returns an option to set a prefix to filter results to buckets whose names begin with this
      * prefix.

google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITListBucketTest.java

@@ -17,119 +17,85 @@
 package com.google.cloud.storage.it;
 
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.assertTrue;
 
 import com.google.api.gax.paging.Page;
 import com.google.cloud.storage.Bucket;
 import com.google.cloud.storage.BucketInfo;
 import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.Storage.BucketListOption;
 import com.google.cloud.storage.TransportCompatibility.Transport;
 import com.google.cloud.storage.it.runner.StorageITRunner;
 import com.google.cloud.storage.it.runner.annotations.Backend;
+import com.google.cloud.storage.it.runner.annotations.BucketFixture;
+import com.google.cloud.storage.it.runner.annotations.BucketType;
 import com.google.cloud.storage.it.runner.annotations.CrossRun;
 import com.google.cloud.storage.it.runner.annotations.Inject;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
-import java.util.stream.StreamSupport;
-import org.junit.After;
-import org.junit.Before;
+import com.google.cloud.storage.it.runner.registry.Generator;
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+import java.util.stream.Collectors;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @RunWith(StorageITRunner.class)
 @CrossRun(
     backends = {Backend.TEST_BENCH},
-    transports = {Transport.HTTP})
+    transports = {Transport.HTTP, Transport.GRPC})
 public class ITListBucketTest {
   @Inject public Storage storage;
 
-  private static final String NORMAL_BUCKET_NAME = "normal_bucket";
-  // For testing purposes, the TESTBENCH considers a bucket to be unreachable if the bucket name
-  // contains "unreachable"
-  private static final String UNREACHABLE_BUCKET_NAME_1 = "unreachable_bucket_1";
-  private static final String UNREACHABLE_BUCKET_NAME_2 = "unreachable_bucket_2";
+  @Inject public BucketInfo defaultBucket;
 
-  // The unreachable buckets are returned as a list of bucket resource names in string form. (e.g.
-  // "projects/_/buckets/bucket1")
-  private static final String EXPECTED_UNREACHABLE_BUCKET_NAME_1 =
-      "projects/_/buckets/" + UNREACHABLE_BUCKET_NAME_1;
-  private static final String EXPECTED_UNREACHABLE_BUCKET_NAME_2 =
-      "projects/_/buckets/" + UNREACHABLE_BUCKET_NAME_2;
+  @Inject
+  @BucketFixture(BucketType.HNS)
+  public BucketInfo hnsBucket;
 
-  @Before
-  public void setup() {
-    Bucket normalBucket = storage.create(BucketInfo.of(NORMAL_BUCKET_NAME));
-    Bucket unreachableBucket = storage.create(BucketInfo.of(UNREACHABLE_BUCKET_NAME_1));
-  }
-
-  @After
-  public void tearDown() {
-    BucketCleaner.doCleanup(NORMAL_BUCKET_NAME, storage);
-    BucketCleaner.doCleanup(UNREACHABLE_BUCKET_NAME_1, storage);
-  }
+  @Inject public Generator generator;
 
   @Test
-  public void testListBucketWithPartialSuccess() {
-    Page<Bucket> page = storage.list(Storage.BucketListOption.returnPartialSuccess(true));
-    Iterable<Bucket> allBuckets = page.getValues();
-
-    Bucket actualNormalBucket =
-        Iterables.getOnlyElement(
-            Iterables.filter(allBuckets, b -> b.getName().equals(NORMAL_BUCKET_NAME)));
-
-    Bucket actualUnreachableBucket =
-        Iterables.getOnlyElement(
-            Iterables.filter(allBuckets, b -> b.getName().contains(UNREACHABLE_BUCKET_NAME_1)));
-
-    assertThat(actualNormalBucket.getName()).isEqualTo(NORMAL_BUCKET_NAME);
-    assertThat(actualUnreachableBucket.getName()).isEqualTo(EXPECTED_UNREACHABLE_BUCKET_NAME_1);
-    assertTrue(
-        "The unreachable bucket must have the isUnreachable flag set to true",
-        actualUnreachableBucket.isUnreachable());
+  public void testListBucketWithPartialSuccess() throws Exception {
+    doTest(Reachability.Unreachable, BucketListOption.returnPartialSuccess(true));
   }
 
   @Test
-  public void testMultipleUnreachableBuckets() {
-    Bucket unreachableBucket2 = storage.create(BucketInfo.of(UNREACHABLE_BUCKET_NAME_2));
-
-    try {
-      Page<Bucket> page = storage.list(Storage.BucketListOption.returnPartialSuccess(true));
-      Iterable<Bucket> allBuckets = page.getValues();
-
-      Bucket actualNormalBucket =
-          Iterables.getOnlyElement(
-              Iterables.filter(allBuckets, b -> b.getName().equals(NORMAL_BUCKET_NAME)));
-
-      Bucket actualUnreachableBucket1 =
-          Iterables.getOnlyElement(
-              Iterables.filter(allBuckets, b -> b.getName().contains(UNREACHABLE_BUCKET_NAME_1)));
-
-      Bucket actualUnreachableBucket2 =
-          Iterables.getOnlyElement(
-              Iterables.filter(allBuckets, b -> b.getName().contains(UNREACHABLE_BUCKET_NAME_2)));
+  public void testListBucketWithoutPartialSuccess() throws Exception {
+    doTest(Reachability.Reachable);
+  }
 
-      assertThat(actualNormalBucket.getName()).isEqualTo(NORMAL_BUCKET_NAME);
-      assertThat(actualUnreachableBucket1.getName()).isEqualTo(EXPECTED_UNREACHABLE_BUCKET_NAME_1);
-      assertTrue(
-          "The unreachable bucket 1 must have the isUnreachable flag set to true",
-          actualUnreachableBucket1.isUnreachable());
-      assertThat(actualUnreachableBucket2.getName()).isEqualTo(EXPECTED_UNREACHABLE_BUCKET_NAME_2);
-      assertTrue(
-          "The unreachable bucket 2 must have the isUnreachable flag set to true",
-          actualUnreachableBucket2.isUnreachable());
-    } finally {
-      BucketCleaner.doCleanup(UNREACHABLE_BUCKET_NAME_2, storage);
+  private void doTest(
+      Reachability expectedReachabilityOfUnreachableBucket, BucketListOption... bucketListOption)
+      throws Exception {
+    // TESTBENCH considers a bucket to be unreachable if the bucket name contains "unreachable"
+    String name = generator.randomBucketName() + ".unreachable";
+    BucketInfo info = BucketInfo.of(name);
+    try (TemporaryBucket tmpBucket =
+        TemporaryBucket.newBuilder().setBucketInfo(info).setStorage(storage).build()) {
+      // bucket name to unreachable status
+      Map<String, Reachability> expected =
+          ImmutableMap.of(
+              defaultBucket.getName(), Reachability.Reachable,
+              hnsBucket.getName(), Reachability.Reachable,
+              tmpBucket.getBucket().getName(), expectedReachabilityOfUnreachableBucket);
+
+      Page<Bucket> page = storage.list(bucketListOption);
+
+      Map<String, Reachability> actual =
+          page.streamAll().collect(Collectors.toMap(BucketInfo::getName, Reachability::forBucket));
+
+      assertThat(actual).containsAtLeastEntriesIn(expected);
     }
   }
 
-  @Test
-  public void testListBucketWithoutPartialSuccess() {
-    Page<Bucket> page = storage.list();
-    ImmutableList<String> bucketNames =
-        StreamSupport.stream(page.iterateAll().spliterator(), false)
-            .map(Bucket::getName)
-            .collect(ImmutableList.toImmutableList());
-    assertThat(bucketNames).contains(NORMAL_BUCKET_NAME);
-    assertThat(bucketNames).doesNotContain(EXPECTED_UNREACHABLE_BUCKET_NAME_1);
+  private enum Reachability {
+    Reachable,
+    Unreachable;
+
+    static Reachability forBucket(BucketInfo b) {
+      if (b.isUnreachable() != null && b.isUnreachable()) {
+        return Unreachable;
+      } else {
+        return Reachable;
+      }
+    }
   }
-}
+}