Skip to content

Commit eca8ce8

Browse files
YikaiHuYikaiHu
committed
feat: add the support for getting prefix list file from third bucket and support configu sse encrypthion during putting object
2 parents b9ff267 + 3553479 commit eca8ce8

6 files changed

Lines changed: 112 additions & 21 deletions

File tree

‎Dockerfile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ ENV SINGLE_PART_TABLE_NAME ''
2424
ENV SRC_BUCKET ''
2525
ENV SRC_PREFIX ''
2626
ENV SRC_PREFIX_LIST ''
27+
ENV SRC_PREFIX_LIST_BUCKET ''
2728
ENV SRC_REGION ''
2829
ENV SRC_ENDPOINT ''
2930
ENV SRC_CREDENTIALS ''

‎cmd/root.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,14 @@ func initConfig() {
7272
viper.SetDefault("destStorageClass", "STANDARD")
7373
viper.SetDefault("srcPrefix", "")
7474
viper.SetDefault("srcPrefixList", "")
75+
viper.SetDefault("srcPrefixListBucket", "")
7576
viper.SetDefault("srcCredential", "")
7677
viper.SetDefault("srcEndpoint", "")
7778
viper.SetDefault("destPrefix", "")
7879
viper.SetDefault("destCredential", "")
7980
viper.SetDefault("destAcl", "bucket-owner-full-control")
81+
viper.SetDefault("destSSEType", "None")
82+
viper.SetDefault("destSSEKMSKeyId", "")
8083

8184
viper.SetDefault("options.chunkSize", dth.DefaultChunkSize)
8285
viper.SetDefault("options.multipartThreshold", dth.DefaultMultipartThreshold)
@@ -92,6 +95,7 @@ func initConfig() {
9295
viper.BindEnv("srcBucket", "SRC_BUCKET")
9396
viper.BindEnv("srcPrefix", "SRC_PREFIX")
9497
viper.BindEnv("srcPrefixList", "SRC_PREFIX_LIST")
98+
viper.BindEnv("srcPrefixListBucket", "SRC_PREFIX_LIST_BUCKET")
9599
viper.BindEnv("srcRegion", "SRC_REGION")
96100
viper.BindEnv("srcEndpoint", "SRC_ENDPOINT")
97101
viper.BindEnv("srcCredential", "SRC_CREDENTIALS")
@@ -106,6 +110,8 @@ func initConfig() {
106110
viper.BindEnv("destInCurrentAccount", "DEST_IN_CURRENT_ACCOUNT")
107111
viper.BindEnv("destStorageClass", "DEST_STORAGE_CLASS")
108112
viper.BindEnv("destAcl", "DEST_ACL")
113+
viper.BindEnv("destSSEType", "DEST_SSE_TYPE")
114+
viper.BindEnv("destSSEKMSKeyId", "DEST_SSE_KMS_KEY_ID")
109115

110116
viper.BindEnv("jobTableName", "JOB_TABLE_NAME")
111117
viper.BindEnv("jobQueueName", "JOB_QUEUE_NAME")
@@ -156,6 +162,7 @@ func initConfig() {
156162
SrcBucket: viper.GetString("srcBucket"),
157163
SrcPrefix: viper.GetString("srcPrefix"),
158164
SrcPrefixList: viper.GetString("srcPrefixList"),
165+
SrcPrefixListBucket: viper.GetString("srcPrefixListBucket"),
159166
SrcRegion: viper.GetString("srcRegion"),
160167
SrcEndpoint: viper.GetString("srcEndpoint"),
161168
SrcCredential: viper.GetString("srcCredential"),
@@ -168,6 +175,8 @@ func initConfig() {
168175
DestCredential: viper.GetString("destCredential"),
169176
DestStorageClass: viper.GetString("destStorageClass"),
170177
DestAcl: viper.GetString("destAcl"),
178+
DestSSEType: viper.GetString("destSSEType"),
179+
DestSSEKMSKeyId: viper.GetString("destSSEKMSKeyId"),
171180
DestInCurrentAccount: viper.GetBool("destInCurrentAccount"),
172181
JobTableName: viper.GetString("jobTableName"),
173182
JobQueueName: viper.GetString("jobQueueName"),

‎config-example.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
srcType: Amazon_S3
22
srcBucket: src-bucket
33
srcPrefix:
4+
srcPrefixListBucket:
45
srcRegion: us-west-2
56
srcEndpoint:
67
srcCredential: src
@@ -15,6 +16,8 @@ destCredential:
1516
destStorageClass: STANDARD
1617
destInCurrentAccount: true
1718
destAcl: bucket-owner-full-control
19+
destSSEType: None
20+
destSSEKMSKeyId:
1821

1922
jobTableName: test-table
2023
jobQueueName: test-queue

‎dth/client.go‎

Lines changed: 70 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,11 @@ type Client interface {
4747
ListParts(ctx context.Context, key, uploadID *string) (parts map[int]*Part)
4848
GetUploadID(ctx context.Context, key *string) (uploadID *string)
4949
ListSelectedPrefixes(ctx context.Context, key *string) (prefixes []*string)
50+
ListSelectedPrefixesFromThirdBucket(ctx context.Context, bucket *string, key *string) (prefixes []*string)
5051

5152
// WRITE
52-
PutObject(ctx context.Context, key *string, body []byte, storageClass, acl *string, meta *Metadata) (etag *string, err error)
53-
CreateMultipartUpload(ctx context.Context, key, storageClass, acl *string, meta *Metadata) (uploadID *string, err error)
53+
PutObject(ctx context.Context, key *string, body []byte, storageClass, acl *string, sseType *string, sseKMSKeyId *string, meta *Metadata) (etag *string, err error)
54+
CreateMultipartUpload(ctx context.Context, key, storageClass, acl *string, sseType *string, sseKMSKeyId *string, meta *Metadata) (uploadID *string, err error)
5455
CompleteMultipartUpload(ctx context.Context, key, uploadID *string, parts []*Part) (etag *string, err error)
5556
UploadPart(ctx context.Context, key *string, body []byte, uploadID *string, partNumber int) (etag *string, err error)
5657
AbortMultipartUpload(ctx context.Context, key, uploadID *string) (err error)
@@ -166,6 +167,19 @@ func NewS3Client(ctx context.Context, bucket, prefix, prefixList, endpoint, regi
166167

167168
}
168169

170+
// NewS3ClientWithEC2Role creates a S3Client instance which uses EC2 Role to access S3
171+
func NewS3ClientWithEC2Role(ctx context.Context, bucket, prefixList string) *S3Client {
172+
cfg := loadDefaultConfig(ctx)
173+
174+
client := s3.NewFromConfig(cfg)
175+
176+
return &S3Client{
177+
bucket: bucket,
178+
prefixList: prefixList,
179+
client: client,
180+
}
181+
}
182+
169183
// GetObject is a function to get (download) object from Amazon S3
170184
func (c *S3Client) GetObject(ctx context.Context, key *string, size, start, chunkSize int64, version string) ([]byte, error) {
171185
// log.Printf("S3> Downloading %s with %d bytes start from %d\n", key, size, start)
@@ -459,8 +473,45 @@ func (c *S3Client) ListSelectedPrefixes(ctx context.Context, key *string) (prefi
459473
return
460474
}
461475

476+
// ListSelectedPrefixesFromThirdBucket is a function to list prefixes from a list file in a specific bucket.
477+
func (c *S3Client) ListSelectedPrefixesFromThirdBucket(ctx context.Context, bucket *string, key *string) (prefixes []*string) {
478+
downloader := manager.NewDownloader(c.client)
479+
getBuf := manager.NewWriteAtBuffer([]byte{})
480+
481+
input := &s3.GetObjectInput{
482+
Bucket: bucket,
483+
Key: key,
484+
}
485+
486+
downloadStart := time.Now()
487+
log.Printf("Start downloading the Prefix List File from bucket: %s", *bucket)
488+
_, err := downloader.Download(ctx, getBuf, input)
489+
downloadEnd := time.Since(downloadStart)
490+
if err != nil {
491+
log.Printf("Error downloading the Prefix List File: %s", err)
492+
return nil
493+
} else {
494+
log.Printf("Download the Prefix List File Completed in %v\n", downloadEnd)
495+
}
496+
497+
start := time.Now()
498+
prefixesValue := make([]string, 0)
499+
500+
for i, line := range strings.Split(string(getBuf.Bytes()), "\n") {
501+
if len(line) > 0 {
502+
prefixesValue = append(prefixesValue, line)
503+
prefixes = append(prefixes, &prefixesValue[i])
504+
}
505+
}
506+
507+
end := time.Since(start)
508+
log.Printf("Got %d prefixes from the customized list file in %v", len(prefixes), end)
509+
return
510+
}
511+
512+
462513
// PutObject is a function to put (upload) an object to Amazon S3
463-
func (c *S3Client) PutObject(ctx context.Context, key *string, body []byte, storageClass, acl *string, meta *Metadata) (etag *string, err error) {
514+
func (c *S3Client) PutObject(ctx context.Context, key *string, body []byte, storageClass, acl *string, sseType *string, sseKMSKeyId *string, meta *Metadata) (etag *string, err error) {
464515
// log.Printf("S3> Uploading object %s to bucket %s\n", key, c.bucket)
465516

466517
md5Bytes := md5.Sum(body)
@@ -481,6 +532,14 @@ func (c *S3Client) PutObject(ctx context.Context, key *string, body []byte, stor
481532
StorageClass: types.StorageClass(*storageClass),
482533
ACL: types.ObjectCannedACL(*acl),
483534
}
535+
switch *sseType {
536+
case "AES256":
537+
input.ServerSideEncryption = types.ServerSideEncryptionAes256
538+
case "AWS_KMS":
539+
input.ServerSideEncryption = types.ServerSideEncryptionAwsKms
540+
input.SSEKMSKeyId = sseKMSKeyId
541+
}
542+
484543
if meta != nil {
485544
input.ContentType = meta.ContentType
486545
input.ContentEncoding = meta.ContentEncoding
@@ -522,7 +581,7 @@ func (c *S3Client) DeleteObject(ctx context.Context, key *string) (err error) {
522581
// CreateMultipartUpload is a function to initilize a multipart upload process.
523582
// This func returns an upload ID used to indicate the multipart upload.
524583
// All parts will be uploaded with this upload ID, after that, all parts by this ID will be combined to create the full object.
525-
func (c *S3Client) CreateMultipartUpload(ctx context.Context, key, storageClass, acl *string, meta *Metadata) (uploadID *string, err error) {
584+
func (c *S3Client) CreateMultipartUpload(ctx context.Context, key, storageClass, acl *string, sseType *string, sseKMSKeyId *string, meta *Metadata) (uploadID *string, err error) {
526585
// log.Printf("S3> Create Multipart Upload for %s\n", *key)
527586
if *acl == "" {
528587
*acl = string(types.ObjectCannedACLBucketOwnerFullControl)
@@ -534,6 +593,13 @@ func (c *S3Client) CreateMultipartUpload(ctx context.Context, key, storageClass,
534593
StorageClass: types.StorageClass(*storageClass),
535594
ACL: types.ObjectCannedACL(*acl),
536595
}
596+
switch *sseType {
597+
case "AES256":
598+
input.ServerSideEncryption = types.ServerSideEncryptionAes256
599+
case "AWS_KMS":
600+
input.ServerSideEncryption = types.ServerSideEncryptionAwsKms
601+
input.SSEKMSKeyId = sseKMSKeyId
602+
}
537603
if meta != nil {
538604
input.ContentType = meta.ContentType
539605
input.ContentEncoding = meta.ContentEncoding

‎dth/config.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,10 +70,10 @@ type JobOptions struct {
7070

7171
// JobConfig is General Job Info
7272
type JobConfig struct {
73-
SrcType, SrcBucket, SrcPrefix, SrcPrefixList, SrcRegion, SrcEndpoint, SrcCredential string
74-
DestBucket, DestPrefix, DestRegion, DestCredential, DestStorageClass, DestAcl string
75-
JobTableName, JobQueueName, SinglePartTableName, SfnArn string
76-
SrcInCurrentAccount, DestInCurrentAccount, SkipCompare, PayerRequest bool
73+
SrcType, SrcBucket, SrcPrefix, SrcPrefixList, SrcPrefixListBucket, SrcRegion, SrcEndpoint, SrcCredential string
74+
DestBucket, DestPrefix, DestRegion, DestCredential, DestStorageClass, DestAcl, DestSSEType, DestSSEKMSKeyId string
75+
JobTableName, JobQueueName, SinglePartTableName, SfnArn string
76+
SrcInCurrentAccount, DestInCurrentAccount, SkipCompare, PayerRequest bool
7777
*JobOptions
7878
}
7979

‎dth/job.go‎

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -54,10 +54,10 @@ type Job interface {
5454
// Finder is an implemenation of Job interface
5555
// Finder compares the differences of source and destination and sends the delta to SQS
5656
type Finder struct {
57-
srcClient, desClient Client
58-
sqs *SqsService
59-
cfg *JobConfig
60-
sfn *SfnService
57+
srcClient, desClient, ec2RoleS3Client Client
58+
sqs *SqsService
59+
cfg *JobConfig
60+
sfn *SfnService
6161
}
6262

6363
// Worker is an implemenation of Job interface
@@ -148,6 +148,8 @@ func NewFinder(ctx context.Context, cfg *JobConfig) (f *Finder) {
148148
srcClient := NewS3Client(ctx, cfg.SrcBucket, cfg.SrcPrefix, cfg.SrcPrefixList, cfg.SrcEndpoint, cfg.SrcRegion, cfg.SrcType, srcCred)
149149
desClient := NewS3Client(ctx, cfg.DestBucket, cfg.DestPrefix, "", "", cfg.DestRegion, "Amazon_S3", desCred)
150150

151+
ec2RoleS3Client := NewS3ClientWithEC2Role(ctx, cfg.SrcPrefixListBucket, cfg.SrcPrefixList)
152+
151153
if srcClient != nil {
152154
srcClient.isSrcClient = true
153155
}
@@ -160,11 +162,12 @@ func NewFinder(ctx context.Context, cfg *JobConfig) (f *Finder) {
160162
DST_CRED = desCred
161163

162164
f = &Finder{
163-
srcClient: srcClient,
164-
desClient: desClient,
165-
sfn: sfn,
166-
sqs: sqs,
167-
cfg: cfg,
165+
srcClient: srcClient,
166+
desClient: desClient,
167+
ec2RoleS3Client: ec2RoleS3Client,
168+
sfn: sfn,
169+
sqs: sqs,
170+
cfg: cfg,
168171
}
169172
return
170173
}
@@ -227,7 +230,9 @@ func (f *Finder) Run(ctx context.Context) {
227230
log.Printf("Enable Payer Request Mode")
228231
}
229232

230-
if len(f.cfg.SrcPrefixList) > 0 {
233+
if f.cfg.SrcPrefixListBucket != "" && len(f.cfg.SrcPrefixList) > 0 {
234+
prefixes = f.ec2RoleS3Client.ListSelectedPrefixesFromThirdBucket(ctx, &f.cfg.SrcPrefixListBucket, &f.cfg.SrcPrefixList)
235+
} else if len(f.cfg.SrcPrefixList) > 0 {
231236
prefixes = f.srcClient.ListSelectedPrefixes(ctx, &f.cfg.SrcPrefixList)
232237
} else {
233238
prefixes = f.srcClient.ListCommonPrefixes(ctx, f.cfg.FinderDepth, f.cfg.MaxKeys)
@@ -843,7 +848,9 @@ func (w *Worker) generateMultiPartTransferJobs(ctx context.Context, obj *Object,
843848
meta = w.srcClient.HeadObject(ctx, &obj.Key)
844849
}
845850

846-
uploadID, err = w.desClient.CreateMultipartUpload(ctx, destKey, &w.cfg.DestStorageClass, &w.cfg.DestAcl, meta)
851+
uploadID, err = w.desClient.CreateMultipartUpload(
852+
ctx, destKey, &w.cfg.DestStorageClass, &w.cfg.DestAcl, &w.cfg.DestSSEType, &w.cfg.DestSSEKMSKeyId, meta,
853+
)
847854
if err != nil {
848855
log.Printf("Failed to create upload ID - %s for %s\n", err.Error(), *destKey)
849856
return 0, err
@@ -949,7 +956,9 @@ func (w *Worker) migrateBigFile(ctx context.Context, obj *Object, destKey *strin
949956
meta = w.srcClient.HeadObject(ctx, &obj.Key)
950957
}
951958

952-
uploadID, err = w.desClient.CreateMultipartUpload(ctx, destKey, &w.cfg.DestStorageClass, &w.cfg.DestAcl, meta)
959+
uploadID, err = w.desClient.CreateMultipartUpload(
960+
ctx, destKey, &w.cfg.DestStorageClass, &w.cfg.DestAcl, &w.cfg.DestSSEType, &w.cfg.DestSSEKMSKeyId, meta,
961+
)
953962
if err != nil {
954963
log.Printf("Failed to create upload ID - %s for %s\n", err.Error(), *destKey)
955964
return &TransferResult{
@@ -1100,7 +1109,10 @@ func (w *Worker) transfer(ctx context.Context, obj *Object, destKey *string, sta
11001109

11011110
} else {
11021111
log.Printf("----->Uploading %d Bytes to %s/%s\n", chunkSize, w.cfg.DestBucket, *destKey)
1103-
etag, err = w.desClient.PutObject(ctx, destKey, body, &w.cfg.DestStorageClass, &w.cfg.DestAcl, meta)
1112+
etag, err = w.desClient.PutObject(
1113+
ctx, destKey, body, &w.cfg.DestStorageClass, &w.cfg.DestAcl,
1114+
&w.cfg.DestSSEType, &w.cfg.DestSSEKMSKeyId, meta,
1115+
)
11041116
}
11051117

11061118
body = nil // release memory

0 commit comments

Comments
 (0)