Skip to content

Commit

Permalink
chore: separate out listobjects functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
@huanjani
huanjani committed Jun 16, 2023
1 parent 55d7a49 commit bfc28c4
Show file tree
Hide file tree
Showing 8 changed files with 76 additions and 74 deletions.
94 changes: 46 additions & 48 deletions internal/pkg/aws/s3/s3.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,50 +195,21 @@ func FormatARN(partition, location string) string {
return fmt.Sprintf("arn:%s:s3:::%s", partition, location)
}

func (s *S3) bucketExists(bucket string) (bool, error) {
input := &s3.HeadBucketInput{
Bucket: aws.String(bucket),
}
_, err := s.s3Client.HeadBucket(input)
if err != nil {
if aerr, ok := err.(awserr.Error); ok && aerr.Code() == errCodeNotFound {
return false, nil
}
return false, err
}

return true, nil
}

// GetBucketTree retrieves the objects in an S3 bucket and creates an ASCII tree representing their folder structure.
func (s *S3) GetBucketTree(bucket string) (string, error) {
exists, err := s.bucketExists(bucket)
// BucketTree creates an ASCII tree representing the folder structure of a bucket's objects.
func (s *S3) BucketTree(bucket string) (string, error) {
outputs, err := s.listObjects(bucket, "/")
if err != nil {
return "", err
}
if !exists {
if outputs == nil {
return "", nil
}
var contents []*s3.Object
var prefixes []*s3.CommonPrefix
listResp := &s3.ListObjectsV2Output{}
for {
listParams := &s3.ListObjectsV2Input{
Bucket: aws.String(bucket),
Delimiter: aws.String(slashDelimiter),
ContinuationToken: listResp.NextContinuationToken,
}
listResp, err = s.s3Client.ListObjectsV2(listParams)
if err != nil {
return "", fmt.Errorf("list objects for bucket %s: %w", bucket, err)
}
contents = append(contents, listResp.Contents...)
prefixes = append(prefixes, listResp.CommonPrefixes...)
if listResp.NextContinuationToken == nil {
break
}
for _, output := range outputs {
contents = append(contents, output.Contents...)
prefixes = append(prefixes, output.CommonPrefixes...)
}

tree := treeprint.New()
// Add top-level files.
for _, object := range contents {
Expand All @@ -251,45 +222,72 @@ func (s *S3) GetBucketTree(bucket string) (string, error) {
return tree.String(), nil
}

// GetBucketSizeAndCount returns the total size and number of objects in an S3 bucket.
func (s *S3) GetBucketSizeAndCount(bucket string) (string, int, error) {
exists, err := s.bucketExists(bucket)
// BucketSizeAndCount returns the total size and number of objects in an S3 bucket.
func (s *S3) BucketSizeAndCount(bucket string) (string, int, error) {
outputs, err := s.listObjects(bucket, "")
if err != nil {
return "", 0, err
}
if !exists {
if outputs == nil {
return "", 0, nil
}
var objects []*s3.Object
var size int64
var number int
for _, output := range outputs {
for _, object := range output.Contents {
size = size + aws.Int64Value(object.Size)
number = number + 1
}
}
return humanize.Bytes(uint64(size)), number, nil
}

func (s *S3) listObjects(bucket, delimiter string) ([]s3.ListObjectsV2Output, error) {
exists, err := s.bucketExists(bucket)
if err != nil {
return nil, err
}
if !exists {
return nil, nil
}
var outputs []s3.ListObjectsV2Output
listResp := &s3.ListObjectsV2Output{}
for {
listParams := &s3.ListObjectsV2Input{
Bucket: aws.String(bucket),
Delimiter: aws.String(delimiter),
ContinuationToken: listResp.NextContinuationToken,
}
listResp, err = s.s3Client.ListObjectsV2(listParams)
if err != nil {
return "", 0, fmt.Errorf("list objects for bucket %s: %w", bucket, err)
return nil, fmt.Errorf("list objects for bucket %s: %w", bucket, err)
}
objects = append(objects, listResp.Contents...)
outputs = append(outputs, *listResp)
if listResp.NextContinuationToken == nil {
break
}
}
for _, object := range objects {
size = size + aws.Int64Value(object.Size)
number = number + 1
return outputs, nil
}

func (s *S3) bucketExists(bucket string) (bool, error) {
input := &s3.HeadBucketInput{
Bucket: aws.String(bucket),
}
_, err := s.s3Client.HeadBucket(input)
if err != nil {
if aerr, ok := err.(awserr.Error); ok && aerr.Code() == errCodeNotFound {
return false, nil
}
return false, err
}
return humanize.Bytes(uint64(size)), number, nil
return true, nil
}

func (s *S3) addNodes(tree treeprint.Tree, prefixes []*s3.CommonPrefix, bucket string) error {
if len(prefixes) == 0 {
return nil
}

listResp := &s3.ListObjectsV2Output{}
var err error
for _, prefix := range prefixes {
Expand Down
14 changes: 9 additions & 5 deletions internal/pkg/aws/s3/s3_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -489,7 +489,7 @@ func TestS3_FormatARN(t *testing.T) {
}
}

func TestS3_GetBucketTree(t *testing.T) {
func TestS3_BucketTree(t *testing.T) {
type s3Mocks struct {
s3API *mocks.Mocks3API
s3ManagerAPI *mocks.Mocks3ManagerAPI
Expand Down Expand Up @@ -689,7 +689,7 @@ func TestS3_GetBucketTree(t *testing.T) {
}
tc.setupMocks(s3mocks)

gotTree, gotErr := service.GetBucketTree(aws.StringValue(mockBucket))
gotTree, gotErr := service.BucketTree(aws.StringValue(mockBucket))
if tc.wantErr != nil {
require.EqualError(t, gotErr, tc.wantErr.Error())
return
Expand All @@ -701,7 +701,7 @@ func TestS3_GetBucketTree(t *testing.T) {
}
}

func TestS3_GetBucketSizeAndCount(t *testing.T) {
func TestS3_BucketSizeAndCount(t *testing.T) {
type s3Mocks struct {
s3API *mocks.Mocks3API
s3ManagerAPI *mocks.Mocks3ManagerAPI
Expand Down Expand Up @@ -742,6 +742,7 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
m.s3API.EXPECT().HeadBucket(&s3.HeadBucketInput{Bucket: mockBucket}).Return(&s3.HeadBucketOutput{}, nil)
m.s3API.EXPECT().ListObjectsV2(&s3.ListObjectsV2Input{
Bucket: mockBucket,
Delimiter: aws.String(""),
ContinuationToken: nil,
Prefix: nil,
}).Return(&resp, nil)
Expand All @@ -754,6 +755,7 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
m.s3API.EXPECT().HeadBucket(&s3.HeadBucketInput{Bucket: mockBucket}).Return(&s3.HeadBucketOutput{}, nil)
m.s3API.EXPECT().ListObjectsV2(&s3.ListObjectsV2Input{
Bucket: mockBucket,
Delimiter: aws.String(""),
ContinuationToken: nil,
Prefix: nil,
}).Return(
Expand All @@ -771,6 +773,7 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
}, nil)
m.s3API.EXPECT().ListObjectsV2(&s3.ListObjectsV2Input{
Bucket: mockBucket,
Delimiter: aws.String(""),
ContinuationToken: &mockContinuationToken,
Prefix: nil,
}).Return(
Expand All @@ -795,6 +798,7 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
m.s3API.EXPECT().HeadBucket(&s3.HeadBucketInput{Bucket: mockBucket}).Return(&s3.HeadBucketOutput{}, nil)
m.s3API.EXPECT().ListObjectsV2(&s3.ListObjectsV2Input{
Bucket: mockBucket,
Delimiter: aws.String(""),
ContinuationToken: nil,
Prefix: nil,
}).Return(&s3.ListObjectsV2Output{
Expand All @@ -821,8 +825,8 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
m.s3API.EXPECT().HeadBucket(&s3.HeadBucketInput{Bucket: mockBucket}).Return(&s3.HeadBucketOutput{}, nil)
m.s3API.EXPECT().ListObjectsV2(&s3.ListObjectsV2Input{
Bucket: mockBucket,
Delimiter: aws.String(""),
ContinuationToken: nil,
Prefix: nil,
}).Return(nil, errors.New("some error"))
},
wantErr: errors.New("list objects for bucket bucketName: some error"),
Expand All @@ -847,7 +851,7 @@ func TestS3_GetBucketSizeAndCount(t *testing.T) {
}
tc.setupMocks(s3mocks)

gotSize, gotCount, gotErr := service.GetBucketSizeAndCount(aws.StringValue(mockBucket))
gotSize, gotCount, gotErr := service.BucketSizeAndCount(aws.StringValue(mockBucket))
if tc.wantErr != nil {
require.EqualError(t, gotErr, tc.wantErr.Error())
return
Expand Down
24 changes: 12 additions & 12 deletions internal/pkg/describe/mocks/mock_service.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions internal/pkg/describe/service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,11 @@ type cwAlarmDescriber interface {
}

type bucketDescriber interface {
GetBucketTree(bucket string) (string, error)
BucketTree(bucket string) (string, error)
}

type bucketDataGetter interface {
GetBucketSizeAndCount(bucket string) (string, int, error)
BucketSizeAndCount(bucket string) (string, int, error)
}

type bucketNameGetter interface {
Expand Down
2 changes: 1 addition & 1 deletion internal/pkg/describe/static_site.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ func (d *StaticSiteDescriber) Describe() (HumanJSONStringer, error) {
if err != nil {
return nil, fmt.Errorf("get bucket name for %q env: %w", env, err)
}
tree, err := bucketDescriber.GetBucketTree(bucketName)
tree, err := bucketDescriber.BucketTree(bucketName)
if err != nil {
return nil, fmt.Errorf("get tree representation of bucket contents: %w", err)
}
Expand Down
6 changes: 3 additions & 3 deletions internal/pkg/describe/static_site_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ func TestStaticSiteDescriber_Describe(t *testing.T) {
"CloudFrontDistributionDomainName": "dut843shvcmvn.cloudfront.net",
}, nil),
m.s3Client.EXPECT().BucketName(mockApp, mockEnv, mockSvc).Return(mockBucket, nil),
m.awsS3Client.EXPECT().GetBucketTree(mockBucket).Return("", nil),
m.awsS3Client.EXPECT().BucketTree(mockBucket).Return("", nil),
)
},
wantedHuman: `About
Expand All @@ -161,7 +161,7 @@ Routes
"CloudFrontDistributionDomainName": "dut843shvcmvn.cloudfront.net",
}, nil),
m.s3Client.EXPECT().BucketName(mockApp, mockEnv, mockSvc).Return(mockBucket, nil),
m.awsS3Client.EXPECT().GetBucketTree(mockBucket).Return("", nil),
m.awsS3Client.EXPECT().BucketTree(mockBucket).Return("", nil),
m.wkldDescriber.EXPECT().StackResources().Return(nil, mockErr),
)
},
Expand All @@ -176,7 +176,7 @@ Routes
"CloudFrontDistributionDomainName": "dut843shvcmvn.cloudfront.net",
}, nil),
m.s3Client.EXPECT().BucketName(mockApp, mockEnv, mockSvc).Return(mockBucket, nil),
m.awsS3Client.EXPECT().GetBucketTree(mockBucket).Return(`.
m.awsS3Client.EXPECT().BucketTree(mockBucket).Return(`.
├── README.md
├── error.html
├── index.html
Expand Down
2 changes: 1 addition & 1 deletion internal/pkg/describe/status_describe.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,7 @@ func (d *staticSiteStatusDescriber) Describe() (HumanJSONStringer, error) {
if err != nil {
return nil, fmt.Errorf("get bucket name for %q Static Site service in %q environment: %w", d.svc, d.env, err)
}
size, count, err := bucketDataGetter.GetBucketSizeAndCount(bucketName)
size, count, err := bucketDataGetter.BucketSizeAndCount(bucketName)
if err != nil {
return nil, fmt.Errorf("get size and count data for %q S3 bucket: %w", bucketName, err)
}
Expand Down
4 changes: 2 additions & 2 deletions internal/pkg/describe/status_describe_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -785,7 +785,7 @@ func TestStaticSiteStatusDescriber_Describe(t *testing.T) {
"success": {
setupMocks: func(m serviceStatusDescriberMocks) {
m.s3Client.EXPECT().BucketName(appName, envName, svcName).Return(mockBucket, nil)
m.bucketDataGetter.EXPECT().GetBucketSizeAndCount(mockBucket).Return("mockSize", 123, nil)
m.bucketDataGetter.EXPECT().BucketSizeAndCount(mockBucket).Return("mockSize", 123, nil)
},
wantedContent: &staticSiteServiceStatus{
BucketName: mockBucket,
Expand All @@ -802,7 +802,7 @@ func TestStaticSiteStatusDescriber_Describe(t *testing.T) {
"error getting bucket size and count": {
setupMocks: func(m serviceStatusDescriberMocks) {
m.s3Client.EXPECT().BucketName(appName, envName, svcName).Return(mockBucket, nil)
m.bucketDataGetter.EXPECT().GetBucketSizeAndCount(mockBucket).Return("", 0, mockError)
m.bucketDataGetter.EXPECT().BucketSizeAndCount(mockBucket).Return("", 0, mockError)
},
wantedError: fmt.Errorf(`get size and count data for "jimmyBuckets" S3 bucket: %w`, mockError),
},
Expand Down

0 comments on commit bfc28c4

Please sign in to comment.