From b0a25e2434c7749179a80b420482fe2ece5a4be6 Mon Sep 17 00:00:00 2001
From: Andrew Gaul <gaul@google.com>
Date: Sat, 12 Feb 2022 13:17:19 +0900
Subject: [PATCH] feat(storage): allow specifying includeTrailingDelimiter

References fsouza/fake-gcs-server#676.
---
 storage/bucket.go           |  1 +
 storage/integration_test.go | 40 ++++++++++++++++++++++++++++++++++++-
 storage/storage.go          |  8 ++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/storage/bucket.go b/storage/bucket.go
index 9d145f039636..54f4b45d4dd6 100644
--- a/storage/bucket.go
+++ b/storage/bucket.go
@@ -1543,6 +1543,7 @@ func (it *ObjectIterator) fetch(pageSize int, pageToken string) (string, error)
 	req.StartOffset(it.query.StartOffset)
 	req.EndOffset(it.query.EndOffset)
 	req.Versions(it.query.Versions)
+	req.IncludeTrailingDelimiter(it.query.IncludeTrailingDelimiter)
 	if len(it.query.fieldSelection) > 0 {
 		req.Fields("nextPageToken", googleapi.Field(it.query.fieldSelection))
 	}
diff --git a/storage/integration_test.go b/storage/integration_test.go
index 364d454ec7f5..91219360925c 100644
--- a/storage/integration_test.go
+++ b/storage/integration_test.go
@@ -1341,6 +1341,7 @@ func TestIntegration_Objects(t *testing.T) {
 		"obj1",
 		"obj2",
 		"obj/with/slashes",
+		"obj/",
 	}
 	contents := make(map[string][]byte)
 
@@ -1395,6 +1396,43 @@ func TestIntegration_Objects(t *testing.T) {
 			t.Errorf("prefixes = %v, want %v", gotPrefixes, sortedPrefixes)
 		}
 	})
+	t.Run("testObjectsIterateSelectedAttrsDelimiterIncludeTrailingDelimiter", func(t *testing.T) {
+		query := &Query{Prefix: "", Delimiter: "/", IncludeTrailingDelimiter: true}
+		if err := query.SetAttrSelection([]string{"Name"}); err != nil {
+			t.Fatalf("selecting query attrs: %v", err)
+		}
+
+		var gotNames []string
+		var gotPrefixes []string
+		it := bkt.Objects(context.Background(), query)
+		for {
+			attrs, err := it.Next()
+			if err == iterator.Done {
+				break
+			}
+			if err != nil {
+				t.Fatalf("iterator.Next: %v", err)
+			}
+			if attrs.Name != "" {
+				gotNames = append(gotNames, attrs.Name)
+			} else if attrs.Prefix != "" {
+				gotPrefixes = append(gotPrefixes, attrs.Prefix)
+			}
+
+			if attrs.Bucket != "" {
+				t.Errorf("Bucket field not selected, want empty, got = %v", attrs.Bucket)
+			}
+		}
+
+		sortedNames := []string{"obj/", "obj1", "obj2"}
+		if !cmp.Equal(sortedNames, gotNames) {
+			t.Errorf("names = %v, want %v", gotNames, sortedNames)
+		}
+		sortedPrefixes := []string{"obj/"}
+		if !cmp.Equal(sortedPrefixes, gotPrefixes) {
+			t.Errorf("prefixes = %v, want %v", gotPrefixes, sortedPrefixes)
+		}
+	})
 
 	// Test Reader.
 	for _, obj := range objects {
@@ -1872,7 +1910,7 @@ func testObjectsIterateAllSelectedAttrs(t *testing.T, bkt *BucketHandle, objects
 	// verifying the returned results).
 	query := &Query{
 		Prefix:      "",
-		StartOffset: "obj/with/slashes",
+		StartOffset: "obj/",
 		EndOffset:   "obj2",
 	}
 	var selectedAttrs []string
diff --git a/storage/storage.go b/storage/storage.go
index 457d77c8dbc9..e9b3613d0f16 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -1588,6 +1588,14 @@ type Query struct {
 	// which returns all properties. Passing ProjectionNoACL will omit Owner and ACL,
 	// which may improve performance when listing many objects.
 	Projection Projection
+
+	// IncludeTrailingDelimiter controls how objects which end in a single
+	// instance of Delimiter (for example, if Query.Delimiter = "/" and the
+	// object name is "foo/bar/") are included in the results. By default, these
+	// objects only show up as prefixes. If IncludeTrailingDelimiter is set to
+	// true, they will also be included as objects and their metadata will be
+	// populated in the returned ObjectAttrs.
+	IncludeTrailingDelimiter bool
 }
 
 // attrToFieldMap maps the field names of ObjectAttrs to the underlying field