feat: add support for partial success in ListBuckets for grpc (#3418)

Dhriti07 · web-flow · commit 6c4592080119 · 2025-12-09T11:58:09.000+05:30
diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java b/java-storage/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 =
         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);
+
+    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 {
+      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,78 @@ 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 decoded = bucketNameCodec.decode(name);
+                    return BucketInfo.newBuilder(decoded)
+                        .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 +1941,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(
diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/Storage.java b/java-storage/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.
diff --git a/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITListBucketTest.java b/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITListBucketTest.java
@@ -40,18 +40,18 @@
 @RunWith(StorageITRunner.class)
 @CrossRun(
     backends = {Backend.TEST_BENCH},
-    transports = {Transport.HTTP})
+    transports = {Transport.HTTP, Transport.GRPC})
 public class ITListBucketTest {
   @Inject public Storage storage;
 
   @Inject public BucketInfo defaultBucket;
 
+  @Inject public Generator generator;
+
   @Inject
   @BucketFixture(BucketType.HNS)
   public BucketInfo hnsBucket;
 
-  @Inject public Generator generator;
-
   private static final String UNREACHABLE_BUCKET_SUFFIX = ".unreachable";
 
   @Test
