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

package objectstore

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"regexp"
	"strconv"
	"time"

	"queryorchestration/internal/client"
	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
}

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)
}

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) {
	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
}

type BucketLocation string
type BucketFilename = string
type BucketKey struct {
	ClientID  string
	Location  BucketLocation
	CreatedAt time.Time
	Part      *uint8
	Filename  BucketFilename
}

const (
	Import       BucketLocation = "import"
	TextTextract BucketLocation = "text/textract"
	TextOut      BucketLocation = "text/out"
	Export       BucketLocation = "export"
)

func (c *BucketKey) String() string {
	dateStr := c.CreatedAt.Format("20060102")
	partStr := ""
	if c.Part != nil {
		partStr = fmt.Sprintf("/%d", *c.Part)
	}

	return fmt.Sprintf("%s/%s/%s%s/%s", c.ClientID, c.Location, dateStr, partStr, c.Filename)
}

func ParseBucketKey(key string) (BucketKey, error) {
	locations := fmt.Sprintf("%s|%s|%s|%s", Import, TextTextract, TextOut, Export)
	pattern := fmt.Sprintf(`^(%s)/(%s)/([0-9]{8})/(.+)$`, client.CLIENT_ID_REGEX, locations)
	reg := regexp.MustCompile(pattern)
	matches := reg.FindStringSubmatch(key)
	if len(matches) != 5 {
		return BucketKey{}, errors.New("does not match expectation")
	}

	formattedKey := BucketKey{}

	formattedKey.ClientID = matches[1]
	formattedKey.Location = BucketLocation(matches[2])

	parsedTime, err := time.Parse("20060102", matches[3])
	if err != nil {
		return BucketKey{}, fmt.Errorf("failed to parse date: %w", err)
	}
	formattedKey.CreatedAt = parsedTime

	filename := `[a-zA-Z0-9\_\.\-]+`
	reg = regexp.MustCompile(fmt.Sprintf(`^([0-9]+)/(%s)$`, filename))
	subMatches := reg.FindStringSubmatch(matches[4])
	if len(subMatches) == 3 {
		part, err := strconv.ParseUint(subMatches[1], 10, 8)
		if err != nil {
			return BucketKey{}, err
		}
		safePart := uint8(part)
		formattedKey.Part = &safePart

		formattedKey.Filename = subMatches[2]
	} else {
		reg = regexp.MustCompile(fmt.Sprintf(`^(%s)$`, filename))
		subMatches := reg.FindStringSubmatch(matches[4])
		if len(subMatches) != 2 {
			return BucketKey{}, errors.New("does not match expectation")
		}

		formattedKey.Filename = subMatches[1]
	}

	return formattedKey, nil
}