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)
 }
@@ -2,6 +2,7 @@ package documentupload

 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"mime"
@@ -14,6 +15,7 @@ import (

 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
 )

@@ -80,7 +82,7 @@ func parsePath(path string) (folderPath string, filename string) {
 // Returns the UUID of the leaf folder (/contracts/2025/Q1 in this example).
 // If folderPath is empty, returns the root folder ID (documents without paths go to root).
 // The root folder "/" must exist (created when the client was created).
-func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queries, clientID string, folderPath string, createdBy string) (uuid.UUID, error) {
+func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queries, clientID string, folderPath string, createdBy string, batchID *uuid.UUID) (uuid.UUID, error) {
 	// Get the root folder for this client (must exist - created when client was created)
 	rootFolder, err := q.GetFolderByPath(ctx, &repository.GetFolderByPathParams{
 		Clientid: clientID,
@@ -119,16 +121,20 @@ func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queri
 			currentPath = currentPath + "/" + part
 		}

-		// Upsert the folder (create if not exists, return existing if exists)
-		folder, err := q.UpsertFolder(ctx, &repository.UpsertFolderParams{
-			Path:      currentPath,
-			Parentid:  parentID,
-			Clientid:  clientID,
-			Createdby: createdBy,
-		})
+		folder, existedBefore, err := s.createOrGetFolder(ctx, q, clientID, currentPath, parentID, createdBy, batchID)
 		if err != nil {
 			return uuid.Nil, err
 		}
+		if batchID != nil {
+			if err := q.RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{
+				BatchID:       *batchID,
+				FolderID:      folder.ID,
+				Path:          currentPath,
+				ExistedBefore: existedBefore,
+			}); err != nil {
+				return uuid.Nil, err
+			}
+		}

 		// This folder becomes the parent for the next level
 		lastFolderID = folder.ID
@@ -138,6 +144,48 @@ func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queri
 	return lastFolderID, nil
 }

+func (s *Service) createOrGetFolder(
+	ctx context.Context,
+	q *repository.Queries,
+	clientID string,
+	currentPath string,
+	parentID *uuid.UUID,
+	createdBy string,
+	batchID *uuid.UUID,
+) (*repository.Folder, bool, error) {
+	if batchID == nil {
+		folder, err := q.UpsertFolder(ctx, &repository.UpsertFolderParams{
+			Path:      currentPath,
+			Parentid:  parentID,
+			Clientid:  clientID,
+			Createdby: createdBy,
+		})
+		return folder, true, err
+	}
+
+	folder, err := q.CreateFolderIfAbsent(ctx, &repository.CreateFolderIfAbsentParams{
+		Path:      currentPath,
+		ParentID:  parentID,
+		ClientID:  clientID,
+		CreatedBy: createdBy,
+	})
+	if err == nil {
+		return folder, false, nil
+	}
+	if !errors.Is(err, pgx.ErrNoRows) {
+		return nil, false, err
+	}
+
+	folder, err = q.GetFolderByPath(ctx, &repository.GetFolderByPathParams{
+		Clientid: clientID,
+		Path:     currentPath,
+	})
+	if err != nil {
+		return nil, false, err
+	}
+	return folder, true, nil
+}
+
 func (s *Service) Upload(ctx context.Context, file File) error {
 	_, err := s.UploadWithResult(ctx, file)
 	return err
@@ -197,7 +245,7 @@ func (s *Service) UploadWithResult(ctx context.Context, file File) (UploadResult

 		// Create folder hierarchy if folder path is provided
 		// Use system email as createdBy since this is an automated upload process
-		folderID, err := s.createFolderHierarchy(ctx, q, file.ClientID, folderPath, "system@doczy.local")
+		folderID, err := s.createFolderHierarchy(ctx, q, file.ClientID, folderPath, "system@doczy.local", file.BatchID)
 		if err != nil {
 			return err
 		}
@@ -14,6 +14,7 @@ import (
 	"queryorchestration/internal/test"

 	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -144,6 +145,99 @@ func TestUploadWithPath(t *testing.T) {
 	assert.Len(t, os.Contents, 1)
 }

+func TestUploadWithBatchRecordsFolderCandidates(t *testing.T) {
+	t.Parallel()
+	cfg := &Config{}
+	test.CreateDB(t, cfg)
+	acfg := test.CreateAWSContainer(t, cfg)
+
+	test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
+	test.CreateBucket(t, cfg)
+
+	clientID := "client_batch_candidates"
+	createTestClient(t, cfg, clientID, "client")
+
+	rootFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{
+		Clientid: clientID,
+		Path:     "/",
+	})
+	require.NoError(t, err)
+	parentFolder, err := cfg.GetDBQueries().CreateFolder(t.Context(), &repository.CreateFolderParams{
+		Path:      "/existing",
+		Parentid:  &rootFolder.ID,
+		Clientid:  clientID,
+		Createdby: "test",
+	})
+	require.NoError(t, err)
+
+	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
+		ClientID:         clientID,
+		OriginalFilename: "batch.zip",
+		TotalDocuments:   1,
+	})
+	require.NoError(t, err)
+
+	svc := documentupload.New(cfg)
+	_, err = svc.UploadWithResult(t.Context(), documentupload.File{
+		ClientID: clientID,
+		Content:  strings.NewReader("batch content"),
+		BatchID:  &batchID,
+		Filename: "agreement.pdf",
+		Path:     "existing/new-child/agreement.pdf",
+	})
+	require.NoError(t, err)
+
+	candidates, err := cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID)
+	require.NoError(t, err)
+	require.Len(t, candidates, 2)
+	assert.Equal(t, "/existing", candidates[0].Path)
+	assert.Equal(t, parentFolder.ID, candidates[0].FolderID)
+	assert.True(t, candidates[0].ExistedBefore)
+	assert.Equal(t, "/existing/new-child", candidates[1].Path)
+	assert.False(t, candidates[1].ExistedBefore)
+
+	_, err = svc.UploadWithResult(t.Context(), documentupload.File{
+		ClientID: clientID,
+		Content:  strings.NewReader("batch content retry"),
+		BatchID:  &batchID,
+		Filename: "agreement-2.pdf",
+		Path:     "existing/new-child/agreement-2.pdf",
+	})
+	require.NoError(t, err)
+
+	candidates, err = cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID)
+	require.NoError(t, err)
+	require.Len(t, candidates, 2)
+	assert.False(t, candidates[1].ExistedBefore, "new folder candidates stay marked as newly created across retries")
+}
+
+func TestUploadWithPathDoesNotRecordCandidatesForNonBatch(t *testing.T) {
+	t.Parallel()
+	cfg := &Config{}
+	test.CreateDB(t, cfg)
+	acfg := test.CreateAWSContainer(t, cfg)
+
+	test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
+	test.CreateBucket(t, cfg)
+
+	clientID := "client_non_batch_candidates"
+	createTestClient(t, cfg, clientID, "client")
+
+	batchID := uuid.New()
+	svc := documentupload.New(cfg)
+	_, err := svc.UploadWithResult(t.Context(), documentupload.File{
+		ClientID: clientID,
+		Content:  strings.NewReader("non batch content"),
+		Filename: "agreement.pdf",
+		Path:     "contracts/agreement.pdf",
+	})
+	require.NoError(t, err)
+
+	candidates, err := cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID)
+	require.NoError(t, err)
+	assert.Empty(t, candidates)
+}
+
 func TestUploadWithPathNoFolder(t *testing.T) {
 	t.Parallel()
 	cfg := &Config{}