to #246 feat: use transfer manager in object_storage_client.go

go.mod

@@ -10,7 +10,9 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.30.4
 	github.com/aws/aws-sdk-go-v2/config v1.27.31
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.30
+	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.15
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.60.1
+	github.com/aws/smithy-go v1.20.4
 	github.com/cockroachdb/apd/v3 v3.2.1
 	github.com/cockroachdb/cockroachdb-parser v0.23.2
 	github.com/cockroachdb/errors v1.9.0
@@ -57,7 +59,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sso v1.22.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.30.5 // indirect
-	github.com/aws/smithy-go v1.20.4 // indirect
 	github.com/bazelbuild/rules_go v0.46.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/biogo/store v0.0.0-20201120204734-aad293a2328f // indirect

go.sum

@@ -73,6 +73,8 @@ github.com/aws/aws-sdk-go-v2/credentials v1.17.30 h1:aau/oYFtibVovr2rDt8FHlU17BT
 github.com/aws/aws-sdk-go-v2/credentials v1.17.30/go.mod h1:BPJ/yXV92ZVq6G8uYvbU0gSl8q94UB63nMT5ctNO38g=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12 h1:yjwoSyDZF8Jth+mUk5lSPJCkMC0lMy6FaCD51jm6ayE=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12/go.mod h1:fuR57fAgMk7ot3WcNQfb6rSEn+SUffl7ri+aa8uKysI=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.15 h1:ijB7hr56MngOiELJe0C5aQRaBQ11LveNgWFyG02AUto=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.15/go.mod h1:0QEmQSSWMVfiAk93l1/ayR9DQ9+jwni7gHS2NARZXB0=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16 h1:TNyt/+X43KJ9IJJMjKfa3bNTiZbUP7DeCxfbTROESwY=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16/go.mod h1:2DwJF39FlNAUiX5pAc0UNeiz16lK2t7IaFcm0LFHEgc=
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16 h1:jYfy8UPmd+6kJW5YhY0L1/KftReOGxI/4NtVSTh9O/I=

storage/object_storage_client.go

@@ -5,16 +5,24 @@ import (
 	"errors"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
-	"github.com/aws/smithy-go"
 	"os"
-	"time"
 )
 
+// References:
+// - S3 client:        https://docs.aws.amazon.com/code-library/latest/ug/go_2_s3_code_examples.html
+// - Transfer Manager: https://aws.github.io/aws-sdk-go-v2/docs/sdk-utilities/s3/
+
 type BucketBasics struct {
 	S3Client *s3.Client
 }
 
+const (
+	uploadPartSize   = 5 * 1024 * 1024 // 5 MiB
+	downloadPartSize = 5 * 1024 * 1024 // 5 MiB
+)
+
 func NewBucketBasics(config *aws.Config) *BucketBasics {
 	return &BucketBasics{S3Client: s3.NewFromConfig(*config)}
 }
@@ -26,29 +34,50 @@ func (basics *BucketBasics) UploadFile(ctx context.Context, bucketName string, o
 	} else {
 		defer file.Close()
 
-		_, err = basics.S3Client.PutObject(ctx, &s3.PutObjectInput{
+		uploader := manager.NewUploader(basics.S3Client, func(u *manager.Uploader) {
+			u.PartSize = uploadPartSize // 5 MiB
+		})
+
+		_, err := uploader.Upload(ctx, &s3.PutObjectInput{
 			Bucket: aws.String(bucketName),
 			Key:    aws.String(objectKey),
 			Body:   file,
 		})
 
 		if err != nil {
-			var apiErr smithy.APIError
-			if errors.As(err, &apiErr) && apiErr.ErrorCode() == "EntityTooLarge" {
-				return fmt.Errorf("Error while uploading object to %s. The object is too large.\n"+
-					"To upload objects larger than 5GB, use the S3 console (160GB max)\n"+
-					"or the multipart upload API (5TB max).", bucketName)
+			var mu manager.MultiUploadFailure
+			if errors.As(err, &mu) {
+				return fmt.Errorf("Error while uploading object to %s.\n"+
+					"The UploadId is %s. Error is %v\n", bucketName, mu.UploadID(), err)
 			} else {
-				return fmt.Errorf("Couldn't upload file %v to %v:%v. Here's why: %v\n",
-					fileName, bucketName, objectKey, err)
-			}
-		} else {
-			err = s3.NewObjectExistsWaiter(basics.S3Client).Wait(
-				ctx, &s3.HeadObjectInput{Bucket: aws.String(bucketName), Key: aws.String(objectKey)}, time.Minute)
-			if err != nil {
-				return fmt.Errorf("Failed attempt to wait for object %s to exist.\n", objectKey)
+				return fmt.Errorf("Error while uploading object to %s.\n"+
+					"Error is %v\n", bucketName, err)
 			}
 		}
+		return nil
+	}
+}
+
+func (basics *BucketBasics) DownloadFile(ctx context.Context, bucketName string, objectKey string, fileName string) error {
+	f, err := os.Create(fileName)
+	if err != nil {
+		return fmt.Errorf("failed to create file %q, %v", fileName, err)
+	}
+	defer f.Close()
+
+	downloader := manager.NewDownloader(basics.S3Client, func(u *manager.Downloader) {
+		u.PartSize = downloadPartSize // 5 MiB
+	})
+
+	numBytes, err := downloader.Download(ctx, f, &s3.GetObjectInput{
+		Bucket: aws.String(bucketName),
+		Key:    aws.String(objectKey),
+	})
+
+	if err != nil {
+		return fmt.Errorf("failed to download file, %v", err)
 	}
-	return err
+
+	fmt.Printf("Successfully downloaded %d bytes from s3://%s/%s to %s\n", numBytes, bucketName, objectKey, fileName)
+	return nil
 }