storage/s3.go

package storage

import (
	"context"
	"io"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

// Uploader is the interface for uploading objects to S3 using the upload manager
type Uploader interface {
	Upload(ctx context.Context, input *s3.PutObjectInput, opts ...func(*manager.Uploader)) (*manager.UploadOutput, error)
}

// DirectUploader is the interface for uploading objects directly to S3
type DirectUploader interface {
	PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
}

// Presigner is the interface for generating presigned URLs
type Presigner interface {
	PresignGetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error)
}

// S3 provides storage operations for AWS S3
type S3 struct {
	bucket          string
	svc             Uploader
	directSvc       DirectUploader
	presigner       Presigner
	useDirectUpload bool
}

// Store uploads content to S3 and returns a presigned URL valid for 15 minutes
func (s *S3) Store(path string, content io.Reader, contentType string) (string, error) {
	if s.useDirectUpload {
		return s.storeWithDirectUpload(path, content, contentType)
	}
	return s.storeWithManager(path, content, contentType)
}

func (s *S3) storeWithManager(path string, content io.Reader, contentType string) (string, error) {
	out, err := s.svc.Upload(context.Background(), &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(path),
		Body:   content,
	})
	if err != nil {
		return "", err
	}
	presignedUrl, err := s.presigner.PresignGetObject(context.Background(),
		&s3.GetObjectInput{
			Bucket:              aws.String(s.bucket),
			Key:                 aws.String(*out.Key),
			ResponseContentType: aws.String(contentType),
		},
		s3.WithPresignExpires(time.Minute*15),
	)
	if err != nil {
		return "", err
	}

	return presignedUrl.URL, nil
}

func (s *S3) storeWithDirectUpload(path string, content io.Reader, contentType string) (string, error) {
	// Upload file to S3
	_, err := s.directSvc.PutObject(context.Background(), &s3.PutObjectInput{
		Bucket:      aws.String(s.bucket),
		Key:         aws.String(path),
		Body:        content,
		ContentType: aws.String(contentType),
	})
	if err != nil {
		return "", err
	}

	// Generate presigned URL valid for 15 minutes
	req, err := s.presigner.PresignGetObject(context.Background(), &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(path),
	}, s3.WithPresignExpires(15*time.Minute))
	if err != nil {
		return "", err
	}

	return req.URL, nil
}

// New creates a new S3 storage instance using the upload manager
// This loads AWS config from the default locations and is suitable for most use cases
func New(bucket string) (*S3, error) {
	cfg, err := config.LoadDefaultConfig(context.Background())
	if err != nil {
		return nil, err
	}
	client := s3.NewFromConfig(cfg)
	uploader := manager.NewUploader(client, func(u *manager.Uploader) {
		u.PartSize = 5 * 1024 * 1024
	})
	presignClient := s3.NewPresignClient(client)
	return &S3{
		bucket:          bucket,
		svc:             uploader,
		presigner:       presignClient,
		useDirectUpload: false,
	}, nil
}

// NewS3 creates a new S3 storage instance using direct PutObject
// This is useful when you want more control over the AWS configuration
func NewS3(cfg aws.Config, bucket string) *S3 {
	client := s3.NewFromConfig(cfg)
	return &S3{
		bucket:          bucket,
		directSvc:       client,
		presigner:       s3.NewPresignClient(client),
		useDirectUpload: true,
	}
}