query-orchestration/internal/document/batch/foldercleanup/service.go

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
}