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 }