Adding support for retrieving over 1,000 entries from S3

cloudcloud · Aug 23, 2024 · 2004070 · 2004070
1 parent ee47587
commit 2004070
Showing 1 changed file with 25 additions and 12 deletions.
diff --git a/pkg/sources/s3.go b/pkg/sources/s3.go
@@ -119,20 +119,33 @@ func (s *S3) GetRefs() (r []types.Reference) {
 		delim = "/"
 	}
 
-	in := &s3.ListObjectsV2Input{
-		Bucket:    aws.String(s.Bucket),
-		Delimiter: aws.String(delim),
-		MaxKeys:   aws.Int64(1000),
-		Prefix:    aws.String(prefix),
-	}
+	truncated := true
+	token := ""
+	for truncated {
+		in := &s3.ListObjectsV2Input{
+			Bucket:    aws.String(s.Bucket),
+			Delimiter: aws.String(delim),
+			MaxKeys:   aws.Int64(1000),
+			Prefix:    aws.String(prefix),
+		}
+		if token != "" {
+			in.ContinuationToken = aws.String(token)
+		}
 
-	res, err := s.s3.ListObjectsV2(in)
-	if err != nil {
-		s.c.GetLogger().With("error_message", err).Error("Could not retrieve refs from S3.")
-	}
+		res, err := s.s3.ListObjectsV2(in)
+		if err != nil {
+			s.c.GetLogger().With("error_message", err).Error("Could not retrieve refs from S3.")
+			break
+		}
 
-	for _, x := range res.CommonPrefixes {
-		r = append(r, types.Reference{Entry: strings.TrimSuffix(strings.TrimPrefix(*x.Prefix, prefix), delim)})
+		for _, x := range res.CommonPrefixes {
+			r = append(r, types.Reference{Entry: strings.TrimSuffix(strings.TrimPrefix(*x.Prefix, prefix), delim)})
+		}
+
+		truncated = *res.IsTruncated
+		if truncated {
+			token = *res.NextContinuationToken
+		}
 	}
 
 	return