[s3] Implement CommonPrefixes

dzbarsky · Aug 3, 2024 · eb878c2 · eb878c2
1 parent b4dc1e0
commit eb878c2
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 12 deletions.
diff --git a/services/s3/itest/s3_test.go b/services/s3/itest/s3_test.go
@@ -641,7 +641,6 @@ func TestListObjectsV2(t *testing.T) {
 	if resp.ContinuationToken != nil {
 		t.Fatal("expected nil continuation token", resp)
 	}
-
 }
 
 func TestListBuckets(t *testing.T) {
@@ -663,6 +662,110 @@ func TestListBuckets(t *testing.T) {
 	}
 }
 
+func TestListObjectsV2_CommonPrefixes(t *testing.T) {
+	ctx := context.Background()
+	client, srv := makeClientServerPair()
+	defer srv.Shutdown(ctx)
+
+	_, err := client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: &bucket,
+		Key:    aws.String("top"),
+		Body:   strings.NewReader(""),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 0; i < 5; i++ {
+		_, err := client.PutObject(ctx, &s3.PutObjectInput{
+			Bucket: &bucket,
+			Key:    aws.String("nested2/" + strconv.Itoa(i)),
+			Body:   strings.NewReader(""),
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	for i := 0; i < 5; i++ {
+		_, err := client.PutObject(ctx, &s3.PutObjectInput{
+			Bucket: &bucket,
+			Key:    aws.String("nested1/" + strconv.Itoa(i)),
+			Body:   strings.NewReader(""),
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	resp, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
+		Bucket:    &bucket,
+		Delimiter: aws.String("/"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if *resp.Delimiter != "/" {
+		t.Fatal("Should return delimiter")
+	}
+
+	if len(resp.Contents) != 1 {
+		t.Fatal("not 1 contents", resp.Contents)
+	}
+	if *resp.Contents[0].Key != "top" {
+		t.Fatal("incorrect contents", resp.Contents)
+	}
+
+	if len(resp.CommonPrefixes) != 2 {
+		t.Fatal("incorrect commonPrefixes", resp.CommonPrefixes)
+	}
+	if *resp.CommonPrefixes[0].Prefix != "nested1/" {
+		t.Fatal("incorrect commonPrefixes", *resp.CommonPrefixes[0].Prefix)
+	}
+	if *resp.CommonPrefixes[1].Prefix != "nested2/" {
+		t.Fatal("incorrect commonPrefixes", *resp.CommonPrefixes[1].Prefix)
+	}
+
+	resp, err = client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
+		Bucket:    &bucket,
+		Delimiter: aws.String("/"),
+		Prefix:    aws.String("nes"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(resp.Contents) != 0 {
+		t.Fatal("should have no contents", resp.Contents)
+	}
+	if len(resp.CommonPrefixes) != 2 {
+		t.Fatal("incorrect commonPrefixes", resp.CommonPrefixes)
+	}
+	if *resp.CommonPrefixes[0].Prefix != "nested1/" {
+		t.Fatal("incorrect commonPrefixes", *resp.CommonPrefixes[0].Prefix)
+	}
+	if *resp.CommonPrefixes[1].Prefix != "nested2/" {
+		t.Fatal("incorrect commonPrefixes", *resp.CommonPrefixes[1].Prefix)
+	}
+
+	resp, err = client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
+		Bucket:    &bucket,
+		Delimiter: aws.String("/"),
+		Prefix:    aws.String("nested1"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(resp.Contents) != 0 {
+		t.Fatal("should have no contents", resp.Contents)
+	}
+	if len(resp.CommonPrefixes) != 1 {
+		t.Fatal("incorrect commonPrefixes", resp.CommonPrefixes)
+	}
+	if *resp.CommonPrefixes[0].Prefix != "nested1/" {
+		t.Fatal("incorrect commonPrefixes", resp.CommonPrefixes[0])
+	}
+}
+
 func TestHead(t *testing.T) {
 	ctx := context.Background()
 	client, srv := makeClientServerPair()

diff --git a/services/s3/s3.go b/services/s3/s3.go
@@ -20,6 +20,7 @@ import (
 	"time"
 
 	"github.com/gofrs/uuid/v5"
+	"golang.org/x/exp/maps"
 
 	"aws-in-a-box/atomicfile"
 	"aws-in-a-box/awserrors"
@@ -909,14 +910,15 @@ func (s *S3) ListObjectsV2(input ListObjectsV2Input) (*ListObjectsV2Output, *aws
 	// Gather a list of all keys in bucket, sort them.
 	var keysSorted []string
 	for key := range b.objects {
+		if input.Prefix != nil && !strings.HasPrefix(key, *input.Prefix) {
+			continue
+		}
 		keysSorted = append(keysSorted, key)
 	}
 	sort.Strings(keysSorted)
 
-	var maxKeys int
-	if input.MaxKeys == nil {
-		maxKeys = 1000
-	} else {
+	maxKeys := 1000
+	if input.MaxKeys != nil {
 		maxKeys = *input.MaxKeys
 	}
 
@@ -925,6 +927,10 @@ func (s *S3) ListObjectsV2(input ListObjectsV2Input) (*ListObjectsV2Output, *aws
 	continuationToken := ""
 	var keysToInclude []string
 	for _, key := range keysSorted {
+		if input.Delimiter != nil && strings.Contains(key, *input.Delimiter) {
+			continue
+		}
+
 		if len(keysToInclude) >= maxKeys {
 			isTruncated = true
 			continuationToken = key
@@ -935,12 +941,6 @@ func (s *S3) ListObjectsV2(input ListObjectsV2Input) (*ListObjectsV2Output, *aws
 			continue
 		}
 
-		if input.Prefix != nil {
-			if !strings.HasPrefix(key, *input.Prefix) {
-				continue
-			}
-		}
-
 		if input.ContinuationToken != nil && key < *input.ContinuationToken {
 			continue
 		}
@@ -969,8 +969,31 @@ func (s *S3) ListObjectsV2(input ListObjectsV2Input) (*ListObjectsV2Output, *aws
 		NextContinuationToken: continuationToken,
 		Prefix:                input.Prefix,
 		StartAfter:            input.StartAfter,
+		// TODO(zbarsky): Do we just echo this? Docs seem unclear.
+		Delimiter: input.Delimiter,
+	}
+
+	if input.Delimiter != nil {
+		commonPrefixes := make(map[string]struct{})
+		for _, key := range keysSorted {
+			i := strings.Index(key, *input.Delimiter)
+			if i != -1 {
+				commonPrefixes[key[:i+len(*input.Delimiter)]] = struct{}{}
+			}
+		}
+
+		prefixesSorted := maps.Keys(commonPrefixes)
+		sort.Strings(prefixesSorted)
+		for _, p := range prefixesSorted {
+			response.CommonPrefixes = append(response.CommonPrefixes, Prefix{p})
+		}
 	}
 
+	// TODO(zbarsky): Method may be incomplete, here's what AWS docs say:
+	// When you query ListObjectsV2 with a delimiter during in-progress multipart uploads,
+	// the CommonPrefixes response parameter contains the prefixes that are associated with
+	// the in-progress multipart uploads.
+
 	return response, nil
 
 }
diff --git a/services/s3/types.go b/services/s3/types.go
@@ -306,8 +306,8 @@ type ListObjectsV2Input struct {
 	MaxKeys           *int    `s3:"query:max-keys"`
 	Prefix            *string `s3:"query:prefix"`
 	StartAfter        *string `s3:"query:start-after"`
+	Delimiter         *string `s3:"query:delimiter"`
 	// Not supported:
-	// Delimiter
 	// Encoding-Type
 	// Fetch-Owner
 }
@@ -332,6 +332,12 @@ type ListObjectsV2Output struct {
 	NextContinuationToken string
 	Prefix                *string
 	StartAfter            *string
+	Delimiter             *string
+	CommonPrefixes        []Prefix
+}
+
+type Prefix struct {
+	Prefix string
 }
 
 type InvalidRangeError struct {