first commit
This commit is contained in:
Beyhan Oğur
246
framework/objectstore/s3.go
Normal file
246
framework/objectstore/s3.go
Normal file
@@ -0,0 +1,246 @@
|
||||
package objectstore
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
awsconfig "github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||||
"github.com/aws/aws-sdk-go-v2/service/sts"
|
||||
"github.com/maximhq/bifrost/core/schemas"
|
||||
)
|
||||
|
||||
// S3ObjectStore implements ObjectStore using an S3-compatible backend.
|
||||
type S3ObjectStore struct {
|
||||
client *s3.Client
|
||||
bucket string
|
||||
compress bool
|
||||
logger schemas.Logger
|
||||
}
|
||||
|
||||
// NewS3ObjectStore creates a new S3-compatible object store from the given config.
|
||||
func NewS3ObjectStore(ctx context.Context, cfg *Config, logger schemas.Logger) (*S3ObjectStore, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("objectstore: config is nil")
|
||||
}
|
||||
|
||||
bucket := cfg.Bucket.GetValue()
|
||||
if bucket == "" {
|
||||
return nil, fmt.Errorf("objectstore: s3 bucket is required")
|
||||
}
|
||||
|
||||
// Validate static credential fields: reject half-configured credentials.
|
||||
if (cfg.AccessKeyID != nil) != (cfg.SecretAccessKey != nil) {
|
||||
return nil, fmt.Errorf("objectstore: access_key_id and secret_access_key must be set together")
|
||||
}
|
||||
if cfg.AccessKeyID != nil && (cfg.AccessKeyID.GetValue() == "" || cfg.SecretAccessKey.GetValue() == "") {
|
||||
return nil, fmt.Errorf("objectstore: access_key_id and secret_access_key must resolve to non-empty values")
|
||||
}
|
||||
if cfg.SessionToken != nil && cfg.SessionToken.GetValue() != "" &&
|
||||
(cfg.AccessKeyID == nil || cfg.SecretAccessKey == nil || cfg.AccessKeyID.GetValue() == "" || cfg.SecretAccessKey.GetValue() == "") {
|
||||
return nil, fmt.Errorf("objectstore: session_token requires access_key_id and secret_access_key")
|
||||
}
|
||||
|
||||
var opts []func(*awsconfig.LoadOptions) error
|
||||
if cfg.Region != nil && cfg.Region.GetValue() != "" {
|
||||
opts = append(opts, awsconfig.WithRegion(cfg.Region.GetValue()))
|
||||
}
|
||||
|
||||
// Static credentials if provided; otherwise default chain (IAM role, env vars, etc.)
|
||||
hasStaticConfig := cfg.AccessKeyID != nil || cfg.SecretAccessKey != nil || cfg.SessionToken != nil
|
||||
if hasStaticConfig {
|
||||
if cfg.AccessKeyID == nil || cfg.AccessKeyID.GetValue() == "" ||
|
||||
cfg.SecretAccessKey == nil || cfg.SecretAccessKey.GetValue() == "" {
|
||||
return nil, fmt.Errorf("objectstore: access_key_id and secret_access_key must both be set when using static credentials")
|
||||
}
|
||||
sessionToken := ""
|
||||
if cfg.SessionToken != nil {
|
||||
sessionToken = cfg.SessionToken.GetValue()
|
||||
}
|
||||
opts = append(opts, awsconfig.WithCredentialsProvider(
|
||||
credentials.NewStaticCredentialsProvider(
|
||||
cfg.AccessKeyID.GetValue(),
|
||||
cfg.SecretAccessKey.GetValue(),
|
||||
sessionToken,
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
awsCfg, err := awsconfig.LoadDefaultConfig(ctx, opts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("objectstore: failed to load AWS config: %w", err)
|
||||
}
|
||||
|
||||
// If a role ARN is configured, assume that role using STS.
|
||||
// Works on top of either static credentials or the default chain (instance role, env vars, etc.).
|
||||
if cfg.RoleARN != nil && cfg.RoleARN.GetValue() != "" {
|
||||
stsClient := sts.NewFromConfig(awsCfg)
|
||||
awsCfg.Credentials = aws.NewCredentialsCache(
|
||||
stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN.GetValue()),
|
||||
)
|
||||
}
|
||||
|
||||
s3Opts := func(o *s3.Options) {
|
||||
if cfg.Endpoint != nil && cfg.Endpoint.GetValue() != "" {
|
||||
o.BaseEndpoint = aws.String(cfg.Endpoint.GetValue())
|
||||
}
|
||||
if cfg.ForcePathStyle {
|
||||
o.UsePathStyle = true
|
||||
}
|
||||
}
|
||||
|
||||
client := s3.NewFromConfig(awsCfg, s3Opts)
|
||||
|
||||
return &S3ObjectStore{
|
||||
client: client,
|
||||
bucket: bucket,
|
||||
compress: cfg.Compress,
|
||||
logger: logger,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Put uploads data with optional S3 object tags. When compression is enabled,
|
||||
// data is gzip-compressed before upload.
|
||||
func (s *S3ObjectStore) Put(ctx context.Context, key string, data []byte, tags map[string]string) error {
|
||||
body := data
|
||||
if s.compress {
|
||||
compressed, err := gzipCompress(data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("objectstore: gzip compress: %w", err)
|
||||
}
|
||||
body = compressed
|
||||
}
|
||||
|
||||
input := &s3.PutObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: aws.String(key),
|
||||
Body: bytes.NewReader(body),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
if s.compress {
|
||||
input.ContentEncoding = aws.String("gzip")
|
||||
}
|
||||
|
||||
if len(tags) > 0 {
|
||||
input.Tagging = aws.String(encodeTags(tags))
|
||||
}
|
||||
|
||||
_, err := s.client.PutObject(ctx, input)
|
||||
if err != nil {
|
||||
return fmt.Errorf("objectstore: put object %s: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get retrieves and decompresses an object by key.
|
||||
func (s *S3ObjectStore) Get(ctx context.Context, key string) ([]byte, error) {
|
||||
output, err := s.client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("objectstore: get object %s: %w", key, err)
|
||||
}
|
||||
defer output.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(output.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("objectstore: read body %s: %w", key, err)
|
||||
}
|
||||
|
||||
// Only attempt decompression when the object was stored with gzip encoding.
|
||||
if aws.ToString(output.ContentEncoding) == "gzip" {
|
||||
decompressed, err := gzipDecompress(body)
|
||||
if err != nil {
|
||||
s.logger.Warn("objectstore: gzip decompress failed for %s: %v, returning raw bytes", key, err)
|
||||
return body, nil
|
||||
}
|
||||
return decompressed, nil
|
||||
}
|
||||
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// Delete removes a single object by key.
|
||||
func (s *S3ObjectStore) Delete(ctx context.Context, key string) error {
|
||||
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("objectstore: delete object %s: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteBatch removes multiple objects. It uses the S3 DeleteObjects API
|
||||
// which supports up to 1000 keys per call.
|
||||
func (s *S3ObjectStore) DeleteBatch(ctx context.Context, keys []string) error {
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
const maxBatchSize = 1000
|
||||
for i := 0; i < len(keys); i += maxBatchSize {
|
||||
end := i + maxBatchSize
|
||||
if end > len(keys) {
|
||||
end = len(keys)
|
||||
}
|
||||
batch := keys[i:end]
|
||||
|
||||
objects := make([]types.ObjectIdentifier, len(batch))
|
||||
for j, key := range batch {
|
||||
objects[j] = types.ObjectIdentifier{Key: aws.String(key)}
|
||||
}
|
||||
|
||||
output, err := s.client.DeleteObjects(ctx, &s3.DeleteObjectsInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Delete: &types.Delete{
|
||||
Objects: objects,
|
||||
Quiet: aws.Bool(true),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("objectstore: delete objects batch starting at index %d: %w", i, err)
|
||||
}
|
||||
if len(output.Errors) > 0 {
|
||||
return fmt.Errorf("objectstore: %d objects failed to delete in batch starting at index %d", len(output.Errors), i)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Ping checks connectivity by performing a HeadBucket call.
|
||||
// Note: HeadBucket requires the s3:ListBucket IAM permission on the bucket resource.
|
||||
func (s *S3ObjectStore) Ping(ctx context.Context) error {
|
||||
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("objectstore: head bucket %s: %w", s.bucket, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close is a no-op for S3 (no persistent connections to release).
|
||||
func (s *S3ObjectStore) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// encodeTags encodes a tag map into the S3 URL-encoded tagging format.
|
||||
// Format: "key1=value1&key2=value2"
|
||||
func encodeTags(tags map[string]string) string {
|
||||
parts := make([]string, 0, len(tags))
|
||||
for k, v := range tags {
|
||||
parts = append(parts, url.QueryEscape(k)+"="+url.QueryEscape(v))
|
||||
}
|
||||
return strings.Join(parts, "&")
|
||||
}
|
||||
Reference in New Issue
Block a user