Merged in jmathison/v2-upload-batch (pull request #224)

Implement v2 upload batch cleanup * Implement v2 upload batch cleanup * Merge remote-tracking branch 'origin/main' into jmathison/v2-upload-batch * Address upload batch review feedback * Raise batch worker coverage * Fix batch cleanup review issues Approved-by: Jay Brown
2026-05-08 21:42:46 +00:00
parent 7b820b18c2
commit a080ca59d8
35 changed files with 2587 additions and 108 deletions
@@ -0,0 +1,194 @@
+package foldercleanup
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/serviceconfig"
+
+	"github.com/google/uuid"
+)
+
+const (
+	DefaultReadyBatchLimit   int64 = 100
+	DefaultFolderDeleteLimit int64 = 500
+	defaultStaleAfter              = 15 * time.Minute
+)
+
+type Service struct {
+	cfg serviceconfig.ConfigProvider
+}
+
+type FinalizeResult struct {
+	Ready             bool
+	AcceptedDocuments int32
+	RejectedDocuments int32
+	DetachedDocuments int64
+	DetachedUploads   int64
+	DeletedFolders    int64
+}
+
+func New(cfg serviceconfig.ConfigProvider) *Service {
+	return &Service{cfg: cfg}
+}
+
+func (s *Service) TryFinalizeBatch(ctx context.Context, batchID uuid.UUID) (*FinalizeResult, error) {
+	result := &FinalizeResult{}
+
+	err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
+		batch, err := q.LockBatchUploadForFolderCleanup(ctx, batchID)
+		if err != nil {
+			return err
+		}
+		if batch.FolderCleanupCompletedAt.Valid {
+			result.Ready = true
+			return nil
+		}
+
+		outcomeCount, err := q.CountBatchOutcomes(ctx, batchID)
+		if err != nil {
+			return err
+		}
+		nonTerminalCount, err := q.CountNonTerminalBatchOutcomes(ctx, batchID)
+		if err != nil {
+			return err
+		}
+		acceptedCount, err := q.CountAcceptedBatchOutcomes(ctx, batchID)
+		if err != nil {
+			return err
+		}
+
+		staleTerminal := isStaleTerminalBatch(batch, defaultStaleAfter) && outcomeCount > 0
+		if !isBatchReadyForFolderCleanup(batch, outcomeCount, nonTerminalCount, staleTerminal) {
+			result.Ready = false
+			return nil
+		}
+
+		includeNonTerminal := staleTerminal && nonTerminalCount > 0
+		result.Ready = true
+		result.AcceptedDocuments = acceptedCount
+		result.RejectedDocuments = outcomeCount - acceptedCount
+
+		detachedDocuments, err := q.DetachRejectedBatchDocumentsFromCandidateFolders(ctx, &repository.DetachRejectedBatchDocumentsFromCandidateFoldersParams{
+			BatchID:            &batchID,
+			IncludeNonterminal: includeNonTerminal,
+		})
+		if err != nil {
+			return err
+		}
+		result.DetachedDocuments = detachedDocuments
+
+		detachedUploads, err := q.DetachRejectedBatchDocumentUploadsFromCandidateFolders(ctx, &repository.DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams{
+			BatchID:            &batchID,
+			IncludeNonterminal: includeNonTerminal,
+		})
+		if err != nil {
+			return err
+		}
+		result.DetachedUploads = detachedUploads
+
+		return q.MarkBatchFolderCleanupCompleted(ctx, batchID)
+	})
+	if err != nil {
+		return nil, err
+	}
+	if !result.Ready {
+		return result, nil
+	}
+
+	deleted, err := s.DeleteSafeAutoCreatedFolders(ctx, DefaultFolderDeleteLimit)
+	if err != nil {
+		return nil, err
+	}
+	result.DeletedFolders = deleted
+	return result, nil
+}
+
+func (s *Service) FinalizeReadyBatches(ctx context.Context, limit int64) (int, error) {
+	if limit <= 0 {
+		limit = DefaultReadyBatchLimit
+	}
+	batchIDs, err := s.cfg.GetDBQueries().ListFolderCleanupReadyBatchIDs(ctx, &repository.ListFolderCleanupReadyBatchIDsParams{
+		StaleAfterSeconds: int32(defaultStaleAfter.Seconds()),
+		LimitCount:        limit,
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	finalized := 0
+	for _, batchID := range batchIDs {
+		result, err := s.TryFinalizeBatch(ctx, batchID)
+		if err != nil {
+			return finalized, fmt.Errorf("finalize batch %s: %w", batchID, err)
+		}
+		if result.Ready {
+			finalized++
+		}
+	}
+	return finalized, nil
+}
+
+func (s *Service) DeleteSafeAutoCreatedFolders(ctx context.Context, limit int64) (int64, error) {
+	if limit <= 0 {
+		limit = DefaultFolderDeleteLimit
+	}
+
+	var deleted int64
+	for deleted < limit {
+		candidates, err := s.cfg.GetDBQueries().ListSafeAutoCreatedFolderCandidates(ctx, limit-deleted)
+		if err != nil {
+			return deleted, err
+		}
+		if len(candidates) == 0 {
+			return deleted, nil
+		}
+
+		deletedThisPass := int64(0)
+		for _, candidate := range candidates {
+			rows, err := s.cfg.GetDBQueries().DeleteSafeAutoCreatedFolder(ctx, candidate.ID)
+			if err != nil {
+				return deleted, err
+			}
+			deleted += rows
+			deletedThisPass += rows
+			if deleted >= limit {
+				break
+			}
+		}
+		if deletedThisPass == 0 {
+			return deleted, nil
+		}
+	}
+	return deleted, nil
+}
+
+func isTerminalStatus(status repository.BatchStatus) bool {
+	return status == repository.BatchStatusCompleted ||
+		status == repository.BatchStatusFailed ||
+		status == repository.BatchStatusCancelled
+}
+
+func isStaleTerminalBatch(batch *repository.LockBatchUploadForFolderCleanupRow, staleAfter time.Duration) bool {
+	if batch.Status != repository.BatchStatusFailed && batch.Status != repository.BatchStatusCancelled {
+		return false
+	}
+	if !batch.CompletedAt.Valid {
+		return false
+	}
+	return time.Since(batch.CompletedAt.Time) >= staleAfter
+}
+
+func isBatchReadyForFolderCleanup(
+	batch *repository.LockBatchUploadForFolderCleanupRow,
+	outcomeCount, nonTerminalCount int32,
+	staleTerminal bool,
+) bool {
+	if !isTerminalStatus(batch.Status) {
+		return false
+	}
+	allOutcomesTerminal := batch.TotalDocuments > 0 && outcomeCount == batch.TotalDocuments && nonTerminalCount == 0
+	return allOutcomesTerminal || staleTerminal
+}