init
This commit is contained in:
vet
220
internal/s3migrate/runner.go
Normal file
220
internal/s3migrate/runner.go
Normal file
@@ -0,0 +1,220 @@
|
||||
package s3migrate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
const (
|
||||
collectionName = "s3"
|
||||
presignExpiry = time.Hour
|
||||
)
|
||||
|
||||
type Params struct {
|
||||
SourceEngine string `json:"sourceEngine"`
|
||||
SourceEndpoint string `json:"sourceEndpoint"`
|
||||
SourceBucket string `json:"sourceBucket"`
|
||||
SourceAccessKey string `json:"sourceAccessKey"`
|
||||
SourceSecretKey string `json:"sourceSecretKey"`
|
||||
SourceRegion string `json:"sourceRegion"`
|
||||
|
||||
DestEngine string `json:"destEngine"`
|
||||
DestEndpoint string `json:"destEndpoint"`
|
||||
DestBucket string `json:"destBucket"`
|
||||
DestAccessKey string `json:"destAccessKey"`
|
||||
DestSecretKey string `json:"destSecretKey"`
|
||||
DestRegion string `json:"destRegion"`
|
||||
}
|
||||
|
||||
type objectRecord struct {
|
||||
Name string `bson:"name"`
|
||||
Engine string `bson:"engine"`
|
||||
Key string `bson:"key"`
|
||||
Size int64 `bson:"size"`
|
||||
ContentType string `bson:"content_type"`
|
||||
}
|
||||
|
||||
type LogFunc func(msg string, args ...any)
|
||||
|
||||
func Run(ctx context.Context, db *mongo.Database, params Params, logf LogFunc) error {
|
||||
if err := validateParams(params); err != nil {
|
||||
return err
|
||||
}
|
||||
if params.SourceRegion == "" {
|
||||
params.SourceRegion = "us-east-1"
|
||||
}
|
||||
if params.DestRegion == "" {
|
||||
params.DestRegion = "us-east-1"
|
||||
}
|
||||
|
||||
srcPresign := newPresignClient(params.SourceEndpoint, params.SourceAccessKey, params.SourceSecretKey, params.SourceRegion)
|
||||
dstPresign := newPresignClient(params.DestEndpoint, params.DestAccessKey, params.DestSecretKey, params.DestRegion)
|
||||
|
||||
coll := db.Collection(collectionName)
|
||||
count, err := coll.CountDocuments(ctx, bson.M{"engine": params.SourceEngine})
|
||||
if err != nil {
|
||||
return fmt.Errorf("count source objects: %w", err)
|
||||
}
|
||||
logf("source engine=%s count=%d", params.SourceEngine, count)
|
||||
|
||||
if count == 0 {
|
||||
logf("no objects to migrate")
|
||||
return nil
|
||||
}
|
||||
|
||||
cursor, err := coll.Find(ctx, bson.M{"engine": params.SourceEngine}, options.Find().SetBatchSize(100))
|
||||
if err != nil {
|
||||
return fmt.Errorf("find source objects: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
var migrated, skipped, failed int64
|
||||
for cursor.Next(ctx) {
|
||||
var obj objectRecord
|
||||
if err := cursor.Decode(&obj); err != nil {
|
||||
logf("decode error: %v", err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
exists, err := objectExistsInDest(ctx, coll, params.DestEngine, obj.Name)
|
||||
if err != nil {
|
||||
logf("check dest error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
if exists {
|
||||
skipped++
|
||||
continue
|
||||
}
|
||||
|
||||
if err := copyObject(ctx, srcPresign, params.SourceBucket, dstPresign, params.DestBucket, obj); err != nil {
|
||||
logf("copy error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
if _, err := coll.UpdateOne(ctx,
|
||||
bson.M{"engine": params.SourceEngine, "name": obj.Name},
|
||||
bson.M{"$set": bson.M{"engine": params.DestEngine}},
|
||||
); err != nil {
|
||||
logf("update engine error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
migrated++
|
||||
if migrated%100 == 0 {
|
||||
logf("progress: migrated=%d skipped=%d failed=%d", migrated, skipped, failed)
|
||||
}
|
||||
}
|
||||
|
||||
logf("complete: migrated=%d skipped=%d failed=%d total=%d", migrated, skipped, failed, count)
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateParams(p Params) error {
|
||||
if p.SourceEngine == "" {
|
||||
return fmt.Errorf("sourceEngine is required")
|
||||
}
|
||||
if p.DestEngine == "" {
|
||||
return fmt.Errorf("destEngine is required")
|
||||
}
|
||||
if p.SourceEngine == p.DestEngine {
|
||||
return fmt.Errorf("sourceEngine and destEngine must be different")
|
||||
}
|
||||
if p.SourceEndpoint == "" || p.SourceBucket == "" || p.SourceAccessKey == "" || p.SourceSecretKey == "" {
|
||||
return fmt.Errorf("source S3 config (endpoint, bucket, accessKey, secretKey) is required")
|
||||
}
|
||||
if p.DestEndpoint == "" || p.DestBucket == "" || p.DestAccessKey == "" || p.DestSecretKey == "" {
|
||||
return fmt.Errorf("dest S3 config (endpoint, bucket, accessKey, secretKey) is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func newPresignClient(endpoint, accessKey, secretKey, region string) *s3.PresignClient {
|
||||
cfg := aws.Config{
|
||||
Region: region,
|
||||
Credentials: credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""),
|
||||
BaseEndpoint: &endpoint,
|
||||
}
|
||||
client := s3.NewFromConfig(cfg, func(o *s3.Options) {
|
||||
o.UsePathStyle = true
|
||||
})
|
||||
return s3.NewPresignClient(client)
|
||||
}
|
||||
|
||||
func objectExistsInDest(ctx context.Context, coll *mongo.Collection, destEngine, name string) (bool, error) {
|
||||
count, err := coll.CountDocuments(ctx, bson.M{"engine": destEngine, "name": name}, options.Count().SetLimit(1))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func copyObject(ctx context.Context, srcPresign *s3.PresignClient, srcBucket string, dstPresign *s3.PresignClient, dstBucket string, obj objectRecord) error {
|
||||
getReq, err := srcPresign.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: &srcBucket,
|
||||
Key: &obj.Key,
|
||||
}, s3.WithPresignExpires(presignExpiry))
|
||||
if err != nil {
|
||||
return fmt.Errorf("presign get: %w", err)
|
||||
}
|
||||
|
||||
contentType := obj.ContentType
|
||||
putReq, err := dstPresign.PresignPutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: &dstBucket,
|
||||
Key: &obj.Key,
|
||||
ContentType: &contentType,
|
||||
}, s3.WithPresignExpires(presignExpiry))
|
||||
if err != nil {
|
||||
return fmt.Errorf("presign put: %w", err)
|
||||
}
|
||||
|
||||
downloadResp, err := http.Get(getReq.URL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("download: %w", err)
|
||||
}
|
||||
defer downloadResp.Body.Close()
|
||||
|
||||
if downloadResp.StatusCode == http.StatusNotFound {
|
||||
return fmt.Errorf("source object not found: %s", obj.Key)
|
||||
}
|
||||
if downloadResp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("download status: %s", downloadResp.Status)
|
||||
}
|
||||
|
||||
uploadReq, err := http.NewRequestWithContext(ctx, http.MethodPut, putReq.URL, downloadResp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create upload request: %w", err)
|
||||
}
|
||||
if downloadResp.ContentLength > 0 {
|
||||
uploadReq.ContentLength = downloadResp.ContentLength
|
||||
} else {
|
||||
uploadReq.ContentLength = obj.Size
|
||||
}
|
||||
if contentType != "" {
|
||||
uploadReq.Header.Set("Content-Type", contentType)
|
||||
}
|
||||
|
||||
uploadResp, err := http.DefaultClient.Do(uploadReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("upload: %w", err)
|
||||
}
|
||||
defer func() { _, _ = io.Copy(io.Discard, uploadResp.Body); uploadResp.Body.Close() }()
|
||||
|
||||
if uploadResp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("upload status: %s", uploadResp.Status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user