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
+}
@@ -0,0 +1,293 @@
+package foldercleanup_test
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/document/batch/foldercleanup"
+	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/test"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type TestConfig struct {
+	serviceconfig.BaseConfig
+}
+
+func createClientRoot(t *testing.T, cfg *TestConfig, clientID string) *repository.Folder {
+	t.Helper()
+	ctx := t.Context()
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client",
+	})
+	require.NoError(t, err)
+	root, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/",
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	return root
+}
+
+func createBatch(t *testing.T, cfg *TestConfig, clientID string, total int32) uuid.UUID {
+	t.Helper()
+	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
+		ClientID:         clientID,
+		OriginalFilename: "batch.zip",
+		TotalDocuments:   total,
+	})
+	require.NoError(t, err)
+	return batchID
+}
+
+func TestFinalizeDeletesNewRejectedFolders(t *testing.T) {
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	ctx := t.Context()
+
+	clientID := "cleanup_rejected_" + uuid.New().String()[:8]
+	root := createClientRoot(t, cfg, clientID)
+	batchID := createBatch(t, cfg, clientID, 1)
+
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/auto",
+		Parentid:  &root.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+		BatchID:       batchID,
+		FolderID:      folder.ID,
+		Path:          folder.Path,
+		ExistedBefore: false,
+	})
+	require.NoError(t, err)
+
+	filename := "duplicate.pdf"
+	err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
+		ID:        uuid.New(),
+		Clientid:  clientID,
+		Bucket:    "bucket",
+		Key:       "key",
+		Part:      1,
+		Createdat: pgtype.Timestamp{Valid: true},
+		Filename:  &filename,
+		BatchID:   &batchID,
+		FolderID:  &folder.ID,
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: filename,
+		Column3:  repository.BatchOutcomeStatusDuplicate,
+	})
+	require.NoError(t, err)
+	require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
+		ID:      batchID,
+		Column2: repository.BatchStatusFailed,
+	}))
+
+	result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID)
+	require.NoError(t, err)
+	require.True(t, result.Ready)
+	assert.Equal(t, int64(1), result.DetachedUploads)
+	assert.Equal(t, int64(1), result.DeletedFolders)
+
+	upload, err := cfg.GetDBQueries().GetDocumentUploadByKey(ctx, &repository.GetDocumentUploadByKeyParams{
+		Bucket: "bucket",
+		Key:    "key",
+	})
+	require.NoError(t, err)
+	assert.Nil(t, upload.FolderID)
+
+	_, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	assert.True(t, errors.Is(err, pgx.ErrNoRows))
+}
+
+func TestFinalizePreservesExistingFolders(t *testing.T) {
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	ctx := t.Context()
+
+	clientID := "cleanup_existing_" + uuid.New().String()[:8]
+	root := createClientRoot(t, cfg, clientID)
+	batchID := createBatch(t, cfg, clientID, 1)
+
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/existing",
+		Parentid:  &root.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+		BatchID:       batchID,
+		FolderID:      folder.ID,
+		Path:          folder.Path,
+		ExistedBefore: true,
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: "duplicate.pdf",
+		Column3:  repository.BatchOutcomeStatusDuplicate,
+	})
+	require.NoError(t, err)
+	require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
+		ID:      batchID,
+		Column2: repository.BatchStatusFailed,
+	}))
+
+	result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID)
+	require.NoError(t, err)
+	require.True(t, result.Ready)
+	assert.Zero(t, result.DeletedFolders)
+
+	persisted, err := cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	require.NoError(t, err)
+	assert.Equal(t, "/existing", persisted.Path)
+}
+
+func TestFinalizeReadyBatchesRetriesReadyCleanup(t *testing.T) {
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	ctx := t.Context()
+
+	clientID := "cleanup_retry_" + uuid.New().String()[:8]
+	root := createClientRoot(t, cfg, clientID)
+	batchID := createBatch(t, cfg, clientID, 1)
+
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/retry",
+		Parentid:  &root.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+		BatchID:       batchID,
+		FolderID:      folder.ID,
+		Path:          folder.Path,
+		ExistedBefore: false,
+	}))
+	require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: "invalid.exe",
+		Column3:  repository.BatchOutcomeStatusInvalidType,
+	}))
+	require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
+		ID:      batchID,
+		Column2: repository.BatchStatusFailed,
+	}))
+
+	service := foldercleanup.New(cfg)
+	finalized, err := service.FinalizeReadyBatches(ctx, 10)
+	require.NoError(t, err)
+	assert.Equal(t, 1, finalized)
+
+	batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID)
+	require.NoError(t, err)
+	assert.True(t, batch.FolderCleanupCompletedAt.Valid)
+	_, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	assert.True(t, errors.Is(err, pgx.ErrNoRows))
+
+	finalized, err = service.FinalizeReadyBatches(ctx, 10)
+	require.NoError(t, err)
+	assert.Zero(t, finalized)
+}
+
+func TestFinalizeDeclinesFailedBatchWithoutOutcomes(t *testing.T) {
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	ctx := t.Context()
+
+	clientID := "cleanup_zero_" + uuid.New().String()[:8]
+	root := createClientRoot(t, cfg, clientID)
+	batchID := createBatch(t, cfg, clientID, 0)
+
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/zero",
+		Parentid:  &root.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+		BatchID:       batchID,
+		FolderID:      folder.ID,
+		Path:          folder.Path,
+		ExistedBefore: false,
+	}))
+	require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
+		ID:      batchID,
+		Column2: repository.BatchStatusFailed,
+	}))
+
+	service := foldercleanup.New(cfg)
+	result, err := service.TryFinalizeBatch(ctx, batchID)
+	require.NoError(t, err)
+	require.False(t, result.Ready)
+
+	batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID)
+	require.NoError(t, err)
+	assert.False(t, batch.FolderCleanupCompletedAt.Valid)
+	_, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	require.NoError(t, err)
+
+	finalized, err := service.FinalizeReadyBatches(ctx, 10)
+	require.NoError(t, err)
+	assert.Zero(t, finalized)
+}
+
+func TestFinalizeStaleFailedBatchWithOutcomeEvidence(t *testing.T) {
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	ctx := t.Context()
+
+	clientID := "cleanup_stale_" + uuid.New().String()[:8]
+	root := createClientRoot(t, cfg, clientID)
+	batchID := createBatch(t, cfg, clientID, 2)
+
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/stale",
+		Parentid:  &root.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+	require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+		BatchID:       batchID,
+		FolderID:      folder.ID,
+		Path:          folder.Path,
+		ExistedBefore: false,
+	}))
+	require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: "stuck.pdf",
+		Column3:  repository.BatchOutcomeStatusSubmitted,
+	}))
+	require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
+		ID:      batchID,
+		Column2: repository.BatchStatusFailed,
+	}))
+	require.NoError(t, cfg.GetDBQueries().SetBatchCompletedAt(ctx, &repository.SetBatchCompletedAtParams{
+		ID:          batchID,
+		CompletedAt: pgtype.Timestamp{Time: time.Now().Add(-20 * time.Minute), Valid: true},
+	}))
+
+	result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID)
+	require.NoError(t, err)
+	require.True(t, result.Ready)
+	assert.Equal(t, int64(1), result.DeletedFolders)
+	assert.Equal(t, int32(1), result.RejectedDocuments)
+}
@@ -78,7 +78,7 @@ func (r *Reporter) Resolve(ctx context.Context, batchID *uuid.UUID, filename str
 // If no outcome row exists for this document_id (non-batch document), this
 // affects 0 rows and is a no-op.
 func (r *Reporter) Update(ctx context.Context, documentID uuid.UUID,
-	outcomeStatus repository.BatchOutcomeStatus, errorDetail *string) error {
+	outcomeStatus repository.BatchOutcomeStatus, errorDetail *string) ([]uuid.UUID, error) {
 	return r.queries.UpdateBatchDocumentOutcomeByDocumentID(ctx, &repository.UpdateBatchDocumentOutcomeByDocumentIDParams{
 		DocumentID:  &documentID,
 		Column2:     outcomeStatus,
@@ -125,8 +125,9 @@ func TestUpdate_UpdatesOutcomeByDocumentID(t *testing.T) {
 	require.NoError(t, err)

 	// Now update to sync_complete
-	err = reporter.Update(ctx, docID, repository.BatchOutcomeStatusSyncComplete, nil)
+	batchIDs, err := reporter.Update(ctx, docID, repository.BatchOutcomeStatusSyncComplete, nil)
 	require.NoError(t, err)
+	require.Equal(t, []uuid.UUID{batchID}, batchIDs)

 	// Verify the outcome was advanced
 	rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
@@ -145,6 +146,7 @@ func TestUpdate_NoOpForNonBatchDocument(t *testing.T) {
 	reporter := outcome.New(cfg.GetDBQueries())

 	// Call Update with a random document_id that has no outcome row
-	err := reporter.Update(ctx, uuid.New(), repository.BatchOutcomeStatusSyncComplete, nil)
+	batchIDs, err := reporter.Update(ctx, uuid.New(), repository.BatchOutcomeStatusSyncComplete, nil)
 	require.NoError(t, err)
+	require.Empty(t, batchIDs)
 }