package objectstore import ( "context" "crypto/md5" "encoding/hex" "fmt" "io" "log/slog" "time" awsc "queryorchestration/internal/serviceconfig/aws" "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/s3" ) type ObjectStoreConfig struct { AWSEndpointUrlS3 string `env:"AWS_ENDPOINT_URL_S3"` AWSS3UsePathStyle bool `env:"AWS_S3_USE_PATH_STYLE" envDefault:"false"` StoreClient S3Client ChunkSize int MaxDirSize int64 } type ConfigProvider interface { awsc.ConfigProvider SetStoreClient(context.Context) error SetStoreClientFromS3Cfg(S3Client) SetStoreClientAndPingByName(context.Context, string) error GetStoreClient() S3Client PingStoreByName(context.Context, string) error SetS3Endpoint(string) GetS3Endpoint() string SetS3UsePathStyle(bool) CalculateETag(context.Context, io.Reader) (string, error) GetDirectoryPart(uint16, int64) uint16 } func (c *ObjectStoreConfig) SetStoreClientAndPingByName(ctx context.Context, name string) error { err := c.SetStoreClient(ctx) if err != nil { return err } return c.PingStoreByName(ctx, name) } func (c *ObjectStoreConfig) SetStoreClientFromS3Cfg(cli S3Client) { c.StoreClient = cli } func (c *ObjectStoreConfig) SetStoreClient(ctx context.Context) error { cfg, err := awsc.GetAWSConfigWithOpts(ctx, func(lo *config.LoadOptions) error { if c.AWSEndpointUrlS3 != "" { lo.BaseEndpoint = c.AWSEndpointUrlS3 } return nil }) if err != nil { return err } c.StoreClient = s3.NewFromConfig(cfg, func(o *s3.Options) { o.UsePathStyle = c.AWSS3UsePathStyle }) return nil } func (c *ObjectStoreConfig) GetStoreClient() S3Client { return c.StoreClient } func (c *ObjectStoreConfig) PingStoreByName(ctx context.Context, name string) error { timeout := time.After(30 * time.Second) tick := time.NewTicker(500 * time.Millisecond) defer tick.Stop() input := &s3.HeadBucketInput{ Bucket: &name, } for { select { case <-timeout: return fmt.Errorf("unable to ping bucket by name: %s", name) case <-tick.C: _, err := c.StoreClient.HeadBucket(ctx, input) if err == nil { slog.Info("Successful bucket ping", "name", name) return nil } slog.Debug("Attempted bucket ping", "name", name, "address", c.GetS3Endpoint()) case <-ctx.Done(): return ctx.Err() } } } func (c *ObjectStoreConfig) GetS3Endpoint() string { return c.AWSEndpointUrlS3 } func (c *ObjectStoreConfig) SetS3Endpoint(e string) { c.AWSEndpointUrlS3 = e } func (c *ObjectStoreConfig) SetS3UsePathStyle(e bool) { c.AWSS3UsePathStyle = e } func (c *ObjectStoreConfig) CalculateETag(ctx context.Context, doc io.Reader) (string, error) { if c.ChunkSize <= 0 { c.ChunkSize = 5 * 1024 * 1024 } buffer := make([]byte, c.ChunkSize) var partHashes [][]byte var totalBytes int for { bytesRead, err := doc.Read(buffer) if bytesRead == 0 { if err != nil && err != io.EOF { return "", fmt.Errorf("error reading chunk: %w", err) } break } hash := md5.Sum(buffer[:bytesRead]) partHashes = append(partHashes, hash[:]) totalBytes += bytesRead if err == io.EOF { break } if err != nil { return "", fmt.Errorf("error reading chunk: %w", err) } } if totalBytes == 0 { emptyHash := md5.Sum([]byte{}) return fmt.Sprintf(`"%s"`, hex.EncodeToString(emptyHash[:])), nil } if len(partHashes) == 1 && totalBytes <= c.ChunkSize { return fmt.Sprintf(`"%s"`, hex.EncodeToString(partHashes[0])), nil } finalHash := md5.New() for _, hash := range partHashes { finalHash.Write(hash) } return fmt.Sprintf(`"%s-%d"`, hex.EncodeToString(finalHash.Sum(nil)), len(partHashes)), nil } func (c *ObjectStoreConfig) GetDirectoryPart(part uint16, count int64) uint16 { if c.MaxDirSize <= 0 { c.MaxDirSize = 1000 } if count >= c.MaxDirSize { return part + 1 } return part }