feat(storage): add support for partial success in ListBuckets

storage/bucket.go

@@ -2343,12 +2343,22 @@ func (c *Client) Buckets(ctx context.Context, projectID string) *BucketIterator
 type BucketIterator struct {
 	// Prefix restricts the iterator to buckets whose names begin with it.
 	Prefix string
-
-	ctx       context.Context
-	projectID string
-	buckets   []*BucketAttrs
-	pageInfo  *iterator.PageInfo
-	nextFunc  func() error
+	// If true, the response will contain a list of unreachable buckets if the buckets are unavailable.
+	// Use the Unreachable() method to retrieve the list.
+	ReturnPartialSuccess bool
+
+	ctx         context.Context
+	projectID   string
+	buckets     []*BucketAttrs
+	unreachable []string
+	pageInfo    *iterator.PageInfo
+	nextFunc    func() error
+}
+
+// Unreachable returns a list of bucket names that could not be reached
+// during the iteration if ReturnPartialSuccess was set to true.
+func (it *BucketIterator) Unreachable() []string {
+	return it.unreachable
 }
 
 // Next returns the next result. Its second return value is iterator.Done if

storage/bucket_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	gax "github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/googleapi"
+	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
 	raw "google.golang.org/api/storage/v1"
 	"google.golang.org/protobuf/proto"
@@ -1674,3 +1675,107 @@ func TestDefaultSignBlobRetry(t *testing.T) {
 		t.Fatalf("BucketHandle.SignedURL: %v", err)
 	}
 }
+
+func TestBucketIterator_PartialSuccess(t *testing.T) {
+	ctx := context.Background()
+	const projectID = "test-project"
+
+	// Mocked JSON responses from the GCS API.
+	partialSuccessListBucketsResponse := `{
+		"kind": "storage#buckets",
+		"nextPageToken": "",
+		"items": [
+			{
+				"kind": "storage#bucket",
+				"id": "bucket-1",
+				"name": "bucket-1"
+			},
+			{
+				"kind": "storage#bucket",
+				"id": "bucket-2",
+				"name": "bucket-2"
+			}
+		],
+		"unreachable": [
+			"projects/_/buckets/unreachable-1",
+			"projects/_/buckets/unreachable-2"
+		]
+	}`
+
+	listBucketsResponse := `{
+		"kind": "storage#buckets",
+		"nextPageToken": "",
+		"items": [
+			{
+				"kind": "storage#bucket",
+				"id": "bucket-1",
+				"name": "bucket-1"
+			},
+			{
+				"kind": "storage#bucket",
+				"id": "bucket-2",
+				"name": "bucket-2"
+			}
+		]
+	}`
+
+	testCases := []struct {
+		name                 string
+		returnPartialSuccess bool
+		jsonResponse         string
+		wantUnreachable      []string
+	}{
+		{
+			name:                 "ReturnPartialSuccess is true",
+			returnPartialSuccess: true,
+			jsonResponse:         partialSuccessListBucketsResponse,
+			wantUnreachable:      []string{"projects/_/buckets/unreachable-1", "projects/_/buckets/unreachable-2"},
+		},
+		{
+			name:                 "ReturnPartialSuccess is false",
+			returnPartialSuccess: false,
+			jsonResponse:         listBucketsResponse,
+			wantUnreachable:      nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mt := mockTransport{}
+			mt.addResult(&http.Response{
+				StatusCode: 200,
+				Body:       bodyReader(tc.jsonResponse),
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
+			}, nil)
+
+			client, err := NewClient(ctx, option.WithHTTPClient(&http.Client{Transport: &mt}))
+			if err != nil {
+				t.Fatalf("NewClient: %v", err)
+			}
+
+			it := client.Buckets(ctx, projectID)
+			it.ReturnPartialSuccess = tc.returnPartialSuccess
+
+			var buckets []*BucketAttrs
+			for {
+				b, err := it.Next()
+				if err == iterator.Done {
+					break
+				}
+				if err != nil {
+					t.Fatalf("it.Next() got err: %v", err)
+				}
+				buckets = append(buckets, b)
+			}
+
+			if got, want := len(buckets), 2; got != want {
+				t.Errorf("got %d buckets, want %d", got, want)
+			}
+
+			unreachable := it.Unreachable()
+			if diff := cmp.Diff(unreachable, tc.wantUnreachable); diff != "" {
+				t.Errorf("Unreachable() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

storage/grpc_client.go

@@ -245,8 +245,9 @@ func (c *grpcStorageClient) ListBuckets(ctx context.Context, project string, opt
 			// BucketIterator is returned to them from the veneer.
 			if pageToken == "" {
 				req := &storagepb.ListBucketsRequest{
-					Parent: toProjectResource(it.projectID),
-					Prefix: it.Prefix,
+					Parent:               toProjectResource(it.projectID),
+					Prefix:               it.Prefix,
+					ReturnPartialSuccess: it.ReturnPartialSuccess,
 				}
 				gitr = c.raw.ListBuckets(ctx, req, s.gax...)
 			}
@@ -262,6 +263,9 @@ func (c *grpcStorageClient) ListBuckets(ctx context.Context, project string, opt
 			it.buckets = append(it.buckets, b)
 		}
 
+		if resp, ok := gitr.Response.(*storagepb.ListBucketsResponse); ok {
+			it.unreachable = resp.Unreachable
+		}
 		return next, nil
 	}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(

storage/http_client.go

@@ -224,6 +224,7 @@ func (c *httpStorageClient) ListBuckets(ctx context.Context, project string, opt
 		req.Projection("full")
 		req.Prefix(it.Prefix)
 		req.PageToken(pageToken)
+		req.ReturnPartialSuccess(it.ReturnPartialSuccess)
 		if pageSize > 0 {
 			req.MaxResults(int64(pageSize))
 		}
@@ -242,6 +243,7 @@ func (c *httpStorageClient) ListBuckets(ctx context.Context, project string, opt
 			}
 			it.buckets = append(it.buckets, b)
 		}
+		it.unreachable = resp.Unreachable
 		return resp.NextPageToken, nil
 	}
 

storage/integration_test.go

@@ -6305,6 +6305,48 @@ func TestIntegration_ReaderCancel(t *testing.T) {
 	})
 }
 
+func TestIntegration_ListBuckets(t *testing.T) {
+	ctx := skipExtraReadAPIs(context.Background(), "no reads in test")
+	multiTransportTest(ctx, t, func(t *testing.T, ctx context.Context, _ string, prefix string, client *Client) {
+		h := testHelper{t}
+		projectID := testutil.ProjID()
+
+		newBucketName := prefix + uidSpace.New()
+		bkt := client.Bucket(newBucketName)
+
+		h.mustCreate(bkt, projectID, nil)
+		t.Cleanup(func() { h.mustDeleteBucket(bkt) })
+
+		for _, partialSuccess := range []bool{true, false} {
+			t.Run(fmt.Sprintf("partialSuccess=%v", partialSuccess), func(t *testing.T) {
+				it := client.Buckets(ctx, projectID)
+				it.ReturnPartialSuccess = partialSuccess
+
+				var found bool
+				for {
+					attrs, err := it.Next()
+					if err == iterator.Done {
+						break
+					}
+					if err != nil {
+						t.Fatalf("it.Next: %v", err)
+					}
+					if attrs.Name == newBucketName {
+						found = true
+					}
+				}
+
+				if !found {
+					t.Errorf("created bucket %q not found in list", newBucketName)
+				}
+				if len(it.Unreachable()) > 0 {
+					t.Errorf("got unreachable buckets %v, want none", it.Unreachable())
+				}
+			})
+		}
+	})
+}
+
 // Ensures that a file stored with a:
 // * Content-Encoding of "gzip"
 // * Content-Type of "text/plain"