-
Notifications
You must be signed in to change notification settings - Fork 90
/
client.go
169 lines (147 loc) · 4.84 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
package aws
import (
"encoding/json"
"fmt"
"path"
"strings"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/aws/aws-sdk-go/service/s3/s3iface"
log "github.com/sirupsen/logrus"
)
const (
// BillingDateFormat is the layout for parsing the AWS date format of 'yyyymmdd'.
BillingDateFormat = "20060102"
// ManifestSuffix is the extension of AWS Usage Data manifests.
ManifestSuffix = ".json"
// defaultS3Region is used to make the API call used to determine a bucket's region.
defaultS3Region = "us-east-1"
// maxS3Keys is the maximum amount of keys to be returned by a single S3
// list objects API response
maxS3Keys = 200
)
type ManifestRetriever interface {
RetrieveManifests() ([]*Manifest, error)
}
type manifestRetriever struct {
logger log.FieldLogger
s3API s3iface.S3API
bucket, prefix string
}
func NewManifestRetriever(logger log.FieldLogger, region, bucket, prefix string) ManifestRetriever {
awsSession := session.Must(session.NewSession())
client := s3.New(awsSession, aws.NewConfig().WithRegion(region))
return &manifestRetriever{
logger: logger,
s3API: client,
bucket: bucket,
prefix: prefix,
}
}
// RetrieveManifests downloads the billing manifest for the given bucket and
// prefix. It includes only the top level manifest files, and ignores manifest
// files that are within assemblyId subdirectories, as the top level manifest
// points to the directory containing the most up to date billing report data.
func (r *manifestRetriever) RetrieveManifests() ([]*Manifest, error) {
// ensure that there is a slash at end of location
prefix := r.prefix
if len(prefix) == 0 {
prefix = "/"
} else if prefix[len(prefix)-1] != '/' {
prefix += "/"
}
logger := r.logger.WithFields(log.Fields{
"bucket": r.bucket,
"prefix": prefix,
})
var (
manifests []*Manifest
manifestErr error
page int
)
pageFn := func(out *s3.ListObjectsV2Output, lastPage bool) bool {
page++
keys := r.filterObjects(prefix, out.Contents)
if len(keys) == 0 {
logger.Debugf("page %d had no manifests", page)
return true
}
for _, key := range keys {
logger.WithField("key", key).Debugf("retrieving manifest")
manifest, err := retrieveManifest(r.s3API, r.bucket, key)
if err != nil {
manifestErr = fmt.Errorf("can't get manifest from bucket '%s' with key '%s': %v", r.bucket, key, err)
return false
}
manifests = append(manifests, manifest)
}
return true
}
// list all in <report-prefix>/<report-name>/ of bucket
err := r.s3API.ListObjectsV2Pages(&s3.ListObjectsV2Input{
Bucket: aws.String(r.bucket),
Prefix: aws.String(prefix),
MaxKeys: aws.Int64(maxS3Keys),
}, pageFn)
if err != nil {
return nil, fmt.Errorf("could not list retrieve AWS billing report keys: %v", err)
}
if manifestErr != nil {
return nil, manifestErr
}
return manifests, nil
}
func (r *manifestRetriever) filterObjects(prefix string, objects []*s3.Object) []string {
var keys []string
for _, obj := range objects {
key := *obj.Key
// only look for manifest files
if !strings.HasSuffix(key, ManifestSuffix) {
continue
}
// We're looking for the top-level manifest for a given time range.
// These manifests are copies of manifests within the assemblyId
// directories, and are updated everytime a report for the given time
// period is run. We use these to determine the "most up to date" set
// of data.
// We're looking for manifests in the following format:
// <report-prefix>/<report-name>/YYYYMMDD-YYYYMMDD/<report-name>-Manifest.json
// We ignore the following manifests:
// <report-prefix>/<report-name>/YYYYMMDD-YYYYMMDD/<assemblyId>/<report-name>-Manifest.json
// Strip off <report-prefix>/<report-name>
trimmedPath := strings.TrimPrefix(key, prefix)
// manifestDir will be <YYYYMMDD-YYYYMMDD>/<assemblyId> or <YYYYMMDD-YYYYMMDD>
// The latter is what we're looking for (without the assemblyId subdir)
manifestDir := path.Dir(trimmedPath)
// assemblyDir will be empty if manifestDir is without an assemblyId
// subdirectory: <YYYYMMDD-YYYYMMDD>
assemblyDir, _ := path.Split(manifestDir)
// If there's another directory, it isn't the top-level manifest.
if assemblyDir != "" {
continue
}
// If we've gotten this far, then "key" is a top-level manifest that we
// care about.
keys = append(keys, key)
}
return keys
}
// retrieveManifest retrieves a manifest from the given bucket and key.
func retrieveManifest(client s3iface.S3API, bucket, key string) (*Manifest, error) {
obj, err := client.GetObject(&s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
})
if err != nil {
return nil, err
}
defer obj.Body.Close()
decoder := json.NewDecoder(obj.Body)
var manifest Manifest
err = decoder.Decode(&manifest)
if err != nil {
return nil, err
}
return &manifest, err
}