query-orchestration/internal/document/init/create.go

package documentinit

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"

	docsyncrunner "queryorchestration/api/docSyncRunner"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
)

type Create struct {
	Bucket       string
	Key          objectstore.BucketKey
	Hash         string
	Filename     *string    // Optional filename for batch processing
	OriginalPath *string    // Original path provided during upload (immutable after creation)
	FolderID     *uuid.UUID // Optional folder ID if folder hierarchy was pre-created
}

func (s *Service) Create(ctx context.Context, doc *Create) (uuid.UUID, error) {
	params, err := s.getCreateParams(ctx, doc)
	if err != nil {
		return uuid.Nil, err
	}

	id, err := s.submitCreate(ctx, params)
	if err != nil {
		return uuid.Nil, err
	}

	if params.ID == nil {
		err := s.cfg.SendToQueue(ctx, &queue.SendParams{
			QueueURL: s.cfg.GetDocumentSyncURL(),
			Body: docsyncrunner.Body{
				DocumentID: id,
			},
		})
		if err != nil {
			return uuid.Nil, err
		}
	}

	return id, nil
}

type createDocumentParams struct {
	ID           *uuid.UUID
	Hash         string
	Bucket       string
	Key          objectstore.BucketKey
	Filename     *string
	OriginalPath *string
	FolderID     *uuid.UUID
}

func (s *Service) getCreateParams(ctx context.Context, doc *Create) (*createDocumentParams, error) {
	idbyhash, err := s.cfg.GetDBQueries().GetDocumentIDByHash(ctx, &repository.GetDocumentIDByHashParams{
		Clientid: doc.Key.ClientID,
		Hash:     doc.Hash,
	})
	var docID *uuid.UUID
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, err
	} else if err == nil {
		docID = &idbyhash
	}

	return &createDocumentParams{
		ID:           docID,
		Hash:         doc.Hash,
		Key:          doc.Key,
		Bucket:       doc.Bucket,
		Filename:     doc.Filename,
		OriginalPath: doc.OriginalPath,
		FolderID:     doc.FolderID,
	}, nil
}

func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams) (uuid.UUID, error) {
	var id uuid.UUID

	// Measure file size from S3 (trusted source - we measure it ourselves)
	// The S3 client must be initialized - file size tracking is required
	fileSizeBytes, err := s.measureFileSize(ctx, params.Bucket, params.Key.String())
	if err != nil {
		return uuid.Nil, fmt.Errorf("failed to measure file size: %w", err)
	}

	err = s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
		var dbid uuid.UUID
		if params.ID == nil {
			createid, err := s.cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
				Clientid:      params.Key.ClientID,
				Hash:          params.Hash,
				BatchID:       params.Key.BatchID,
				Filename:      params.Filename,
				Folderid:      params.FolderID,
				Originalpath:  params.OriginalPath,
				FileSizeBytes: fileSizeBytes,
			})
			if err != nil {
				return err
			}
			slog.Debug("document created", "id", createid.String(), "client", params.Key.ClientID)

			dbid = createid
		} else {
			slog.Debug("document exists", "id", params.ID.String(), "client", params.Key.ClientID)

			dbid = *params.ID
		}

		err := s.cfg.GetDBQueries().AddDocumentEntry(ctx, &repository.AddDocumentEntryParams{
			Documentid: dbid,
			Bucket:     params.Bucket,
			Key:        params.Key.String(),
		})
		if err != nil {
			return err
		}

		id = dbid

		return nil
	})
	if err != nil {
		return uuid.Nil, err
	}

	return id, nil
}

// measureFileSize retrieves the file size from S3 using HeadObject.
// Returns the file size in bytes, or an error if the S3 client is not configured
// or if the HeadObject call fails.
//
// Parameters:
//   - ctx: Context for the S3 operation
//   - bucket: S3 bucket name
//   - key: S3 object key
//
// Returns:
//   - *int64: Pointer to file size in bytes
//   - error: Error if S3 client is nil or HeadObject fails
//
// Note: This operation requires s3:GetObject permission on the bucket/key.
func (s *Service) measureFileSize(ctx context.Context, bucket, key string) (*int64, error) {
	storeClient := s.cfg.GetStoreClient()
	if storeClient == nil {
		return nil, errors.New("S3 client not initialized - required for file size measurement")
	}

	headOutput, err := storeClient.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("HeadObject failed for s3://%s/%s: %w", bucket, key, err)
	}

	if headOutput.ContentLength == nil {
		return nil, fmt.Errorf("HeadObject returned nil ContentLength for s3://%s/%s", bucket, key)
	}

	return headOutput.ContentLength, nil
}