query-orchestration/internal/serviceconfig/objectstore/config.go

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/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
	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) SetStoreClient(ctx context.Context) error {
	cfg, err := awsc.GetAWSConfig(ctx)
	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(2 * time.Second)
	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
}