Files
query-orchestration/internal/document/init/create.go
T
Jay Brown aafe7d5b5f Merged in feature/batch-status (pull request #215)
Implement and test the batch status feature

* working
2026-03-12 18:53:42 +00:00

205 lines
5.7 KiB
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
}
// CreateResult contains the result of document creation.
//
// Fields:
// - ID: the document's UUID (new or existing)
// - IsDuplicate: true if a document with the same hash already existed
// - BatchID: the batch this document belongs to (nil for non-batch uploads)
// - Filename: the original filename (nil for non-batch uploads)
type CreateResult struct {
ID uuid.UUID
IsDuplicate bool
BatchID *uuid.UUID
Filename *string
}
// Create processes a new document upload by checking for duplicates, creating
// database records, and forwarding new documents to the sync queue.
//
// Parameters:
// - ctx: request context
// - doc: document creation parameters (bucket, key, hash, filename, etc.)
//
// Returns:
// - *CreateResult: result containing document ID, duplicate status, and batch info
// - error: if any infrastructure operation (DB, S3, queue) fails
func (s *Service) Create(ctx context.Context, doc *Create) (*CreateResult, error) {
params, err := s.getCreateParams(ctx, doc)
if err != nil {
return nil, err
}
id, err := s.submitCreate(ctx, params)
if err != nil {
return 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 nil, err
}
}
return &CreateResult{
ID: id,
IsDuplicate: params.ID != nil,
BatchID: doc.Key.BatchID,
Filename: doc.Filename,
}, 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
}