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
@@ -33,6 +33,44 @@ func (q *Queries) AddFailedFilename(ctx context.Context, arg *AddFailedFilenameP
 	return err
 }

+const countAcceptedBatchOutcomes = `-- name: CountAcceptedBatchOutcomes :one
+SELECT COUNT(*)::int
+FROM batch_document_outcomes
+WHERE batch_id = $1
+  AND outcome IN ('clean_passed', 'sync_skipped')
+`
+
+// CountAcceptedBatchOutcomes
+//
+//	SELECT COUNT(*)::int
+//	FROM batch_document_outcomes
+//	WHERE batch_id = $1
+//	  AND outcome IN ('clean_passed', 'sync_skipped')
+func (q *Queries) CountAcceptedBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) {
+	row := q.db.QueryRow(ctx, countAcceptedBatchOutcomes, batchID)
+	var column_1 int32
+	err := row.Scan(&column_1)
+	return column_1, err
+}
+
+const countBatchOutcomes = `-- name: CountBatchOutcomes :one
+SELECT COUNT(*)::int
+FROM batch_document_outcomes
+WHERE batch_id = $1
+`
+
+// CountBatchOutcomes
+//
+//	SELECT COUNT(*)::int
+//	FROM batch_document_outcomes
+//	WHERE batch_id = $1
+func (q *Queries) CountBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) {
+	row := q.db.QueryRow(ctx, countBatchOutcomes, batchID)
+	var column_1 int32
+	err := row.Scan(&column_1)
+	return column_1, err
+}
+
 const countDocumentsByBatchId = `-- name: CountDocumentsByBatchId :one
 SELECT COUNT(*) as count
 FROM documents 
@@ -51,6 +89,26 @@ func (q *Queries) CountDocumentsByBatchId(ctx context.Context, batchID *uuid.UUI
 	return count, err
 }

+const countNonTerminalBatchOutcomes = `-- name: CountNonTerminalBatchOutcomes :one
+SELECT COUNT(*)::int
+FROM batch_document_outcomes
+WHERE batch_id = $1
+  AND outcome IN ('submitted', 'init_complete', 'sync_complete')
+`
+
+// CountNonTerminalBatchOutcomes
+//
+//	SELECT COUNT(*)::int
+//	FROM batch_document_outcomes
+//	WHERE batch_id = $1
+//	  AND outcome IN ('submitted', 'init_complete', 'sync_complete')
+func (q *Queries) CountNonTerminalBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) {
+	row := q.db.QueryRow(ctx, countNonTerminalBatchOutcomes, batchID)
+	var column_1 int32
+	err := row.Scan(&column_1)
+	return column_1, err
+}
+
 const createBatchUpload = `-- name: CreateBatchUpload :one
 INSERT INTO batch_uploads (client_id, original_filename, total_documents)
 VALUES ($1, $2, $3) RETURNING id
@@ -128,6 +186,158 @@ func (q *Queries) CreateBatchUploadWithStorage(ctx context.Context, arg *CreateB
 	return &i, err
 }

+const detachRejectedBatchDocumentUploadsFromCandidateFolders = `-- name: DetachRejectedBatchDocumentUploadsFromCandidateFolders :execrows
+UPDATE documentUploads du
+SET folder_id = NULL
+WHERE du.batch_id = $1
+  AND du.folder_id IS NOT NULL
+  AND EXISTS (
+      SELECT 1
+      FROM batch_folder_candidates bfc
+      WHERE bfc.batch_id = $1
+        AND bfc.folder_id = du.folder_id
+  )
+  AND EXISTS (
+      SELECT 1
+      FROM batch_document_outcomes bdo
+      WHERE bdo.batch_id = $1
+        AND bdo.filename = du.filename
+        AND (
+            bdo.outcome IN (
+                'duplicate',
+                'init_duplicate',
+                'invalid_type',
+                'failed_open',
+                'failed_read',
+                'failed_s3_upload',
+                'failed_upload',
+                'clean_failed'
+            )
+            OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete'))
+        )
+  )
+`
+
+type DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams struct {
+	BatchID            *uuid.UUID `db:"batch_id"`
+	IncludeNonterminal bool       `db:"include_nonterminal"`
+}
+
+// @sqlc-vet-disable
+//
+//	UPDATE documentUploads du
+//	SET folder_id = NULL
+//	WHERE du.batch_id = $1
+//	  AND du.folder_id IS NOT NULL
+//	  AND EXISTS (
+//	      SELECT 1
+//	      FROM batch_folder_candidates bfc
+//	      WHERE bfc.batch_id = $1
+//	        AND bfc.folder_id = du.folder_id
+//	  )
+//	  AND EXISTS (
+//	      SELECT 1
+//	      FROM batch_document_outcomes bdo
+//	      WHERE bdo.batch_id = $1
+//	        AND bdo.filename = du.filename
+//	        AND (
+//	            bdo.outcome IN (
+//	                'duplicate',
+//	                'init_duplicate',
+//	                'invalid_type',
+//	                'failed_open',
+//	                'failed_read',
+//	                'failed_s3_upload',
+//	                'failed_upload',
+//	                'clean_failed'
+//	            )
+//	            OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete'))
+//	        )
+//	  )
+func (q *Queries) DetachRejectedBatchDocumentUploadsFromCandidateFolders(ctx context.Context, arg *DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams) (int64, error) {
+	result, err := q.db.Exec(ctx, detachRejectedBatchDocumentUploadsFromCandidateFolders, arg.BatchID, arg.IncludeNonterminal)
+	if err != nil {
+		return 0, err
+	}
+	return result.RowsAffected(), nil
+}
+
+const detachRejectedBatchDocumentsFromCandidateFolders = `-- name: DetachRejectedBatchDocumentsFromCandidateFolders :execrows
+UPDATE documents d
+SET folderId = NULL
+WHERE d.batch_id = $1
+  AND d.folderId IS NOT NULL
+  AND EXISTS (
+      SELECT 1
+      FROM batch_folder_candidates bfc
+      WHERE bfc.batch_id = $1
+        AND bfc.folder_id = d.folderId
+  )
+  AND EXISTS (
+      SELECT 1
+      FROM batch_document_outcomes bdo
+      WHERE bdo.batch_id = $1
+        AND bdo.document_id = d.id
+        AND (
+            bdo.outcome IN (
+                'duplicate',
+                'init_duplicate',
+                'invalid_type',
+                'failed_open',
+                'failed_read',
+                'failed_s3_upload',
+                'failed_upload',
+                'clean_failed'
+            )
+            OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete'))
+        )
+  )
+`
+
+type DetachRejectedBatchDocumentsFromCandidateFoldersParams struct {
+	BatchID            *uuid.UUID `db:"batch_id"`
+	IncludeNonterminal bool       `db:"include_nonterminal"`
+}
+
+// @sqlc-vet-disable
+//
+//	UPDATE documents d
+//	SET folderId = NULL
+//	WHERE d.batch_id = $1
+//	  AND d.folderId IS NOT NULL
+//	  AND EXISTS (
+//	      SELECT 1
+//	      FROM batch_folder_candidates bfc
+//	      WHERE bfc.batch_id = $1
+//	        AND bfc.folder_id = d.folderId
+//	  )
+//	  AND EXISTS (
+//	      SELECT 1
+//	      FROM batch_document_outcomes bdo
+//	      WHERE bdo.batch_id = $1
+//	        AND bdo.document_id = d.id
+//	        AND (
+//	            bdo.outcome IN (
+//	                'duplicate',
+//	                'init_duplicate',
+//	                'invalid_type',
+//	                'failed_open',
+//	                'failed_read',
+//	                'failed_s3_upload',
+//	                'failed_upload',
+//	                'clean_failed'
+//	            )
+//	            OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete'))
+//	        )
+//	  )
+func (q *Queries) DetachRejectedBatchDocumentsFromCandidateFolders(ctx context.Context, arg *DetachRejectedBatchDocumentsFromCandidateFoldersParams) (int64, error) {
+	result, err := q.db.Exec(ctx, detachRejectedBatchDocumentsFromCandidateFolders, arg.BatchID, arg.IncludeNonterminal)
+	if err != nil {
+		return 0, err
+	}
+	return result.RowsAffected(), nil
+}
+
 const getBatchUpload = `-- name: GetBatchUpload :one
 SELECT id, client_id, original_filename, total_documents, processed_documents, 
       failed_documents, invalid_type_documents, status, progress_percent,
@@ -459,6 +669,45 @@ func (q *Queries) ListBatchDocumentOutcomes(ctx context.Context, batchID uuid.UU
 	return items, nil
 }

+const listBatchFolderCandidates = `-- name: ListBatchFolderCandidates :many
+SELECT batch_id, folder_id, path, existed_before, created_at
+FROM batch_folder_candidates
+WHERE batch_id = $1
+ORDER BY path
+`
+
+// ListBatchFolderCandidates
+//
+//	SELECT batch_id, folder_id, path, existed_before, created_at
+//	FROM batch_folder_candidates
+//	WHERE batch_id = $1
+//	ORDER BY path
+func (q *Queries) ListBatchFolderCandidates(ctx context.Context, batchID uuid.UUID) ([]*BatchFolderCandidate, error) {
+	rows, err := q.db.Query(ctx, listBatchFolderCandidates, batchID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := []*BatchFolderCandidate{}
+	for rows.Next() {
+		var i BatchFolderCandidate
+		if err := rows.Scan(
+			&i.BatchID,
+			&i.FolderID,
+			&i.Path,
+			&i.ExistedBefore,
+			&i.CreatedAt,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, &i)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const listBatchUploads = `-- name: ListBatchUploads :many
 SELECT id, client_id, original_filename, total_documents, processed_documents, 
       failed_documents, invalid_type_documents, status, progress_percent,
@@ -530,6 +779,183 @@ func (q *Queries) ListBatchUploads(ctx context.Context, arg *ListBatchUploadsPar
 	return items, nil
 }

+const listFolderCleanupReadyBatchIDs = `-- name: ListFolderCleanupReadyBatchIDs :many
+SELECT bu.id
+FROM batch_uploads bu
+WHERE bu.folder_cleanup_completed_at IS NULL
+  AND bu.status IN ('completed', 'failed', 'cancelled')
+  AND EXISTS (
+      SELECT 1
+      FROM batch_folder_candidates bfc
+      WHERE bfc.batch_id = bu.id
+  )
+  AND (
+      (
+          bu.total_documents > 0
+          AND
+          (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) = bu.total_documents
+          AND NOT EXISTS (
+              SELECT 1
+              FROM batch_document_outcomes bdo
+              WHERE bdo.batch_id = bu.id
+                AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')
+          )
+      )
+      OR (
+          bu.status IN ('failed', 'cancelled')
+          AND (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) > 0
+          AND bu.completed_at IS NOT NULL
+          AND bu.completed_at <= NOW() - ($1::int * INTERVAL '1 second')
+      )
+  )
+ORDER BY bu.created_at ASC
+LIMIT $2
+`
+
+type ListFolderCleanupReadyBatchIDsParams struct {
+	StaleAfterSeconds int32 `db:"stale_after_seconds"`
+	LimitCount        int64 `db:"limit_count"`
+}
+
+// @sqlc-vet-disable
+//
+//	SELECT bu.id
+//	FROM batch_uploads bu
+//	WHERE bu.folder_cleanup_completed_at IS NULL
+//	  AND bu.status IN ('completed', 'failed', 'cancelled')
+//	  AND EXISTS (
+//	      SELECT 1
+//	      FROM batch_folder_candidates bfc
+//	      WHERE bfc.batch_id = bu.id
+//	  )
+//	  AND (
+//	      (
+//	          bu.total_documents > 0
+//	          AND
+//	          (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) = bu.total_documents
+//	          AND NOT EXISTS (
+//	              SELECT 1
+//	              FROM batch_document_outcomes bdo
+//	              WHERE bdo.batch_id = bu.id
+//	                AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')
+//	          )
+//	      )
+//	      OR (
+//	          bu.status IN ('failed', 'cancelled')
+//	          AND (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) > 0
+//	          AND bu.completed_at IS NOT NULL
+//	          AND bu.completed_at <= NOW() - ($1::int * INTERVAL '1 second')
+//	      )
+//	  )
+//	ORDER BY bu.created_at ASC
+//	LIMIT $2
+func (q *Queries) ListFolderCleanupReadyBatchIDs(ctx context.Context, arg *ListFolderCleanupReadyBatchIDsParams) ([]uuid.UUID, error) {
+	rows, err := q.db.Query(ctx, listFolderCleanupReadyBatchIDs, arg.StaleAfterSeconds, arg.LimitCount)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := []uuid.UUID{}
+	for rows.Next() {
+		var id uuid.UUID
+		if err := rows.Scan(&id); err != nil {
+			return nil, err
+		}
+		items = append(items, id)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
+const lockBatchUploadForFolderCleanup = `-- name: LockBatchUploadForFolderCleanup :one
+SELECT id, status, total_documents, created_at, completed_at, folder_cleanup_completed_at
+FROM batch_uploads
+WHERE id = $1
+FOR UPDATE
+`
+
+type LockBatchUploadForFolderCleanupRow struct {
+	ID                       uuid.UUID        `db:"id"`
+	Status                   BatchStatus      `db:"status"`
+	TotalDocuments           int32            `db:"total_documents"`
+	CreatedAt                pgtype.Timestamp `db:"created_at"`
+	CompletedAt              pgtype.Timestamp `db:"completed_at"`
+	FolderCleanupCompletedAt pgtype.Timestamp `db:"folder_cleanup_completed_at"`
+}
+
+// LockBatchUploadForFolderCleanup
+//
+//	SELECT id, status, total_documents, created_at, completed_at, folder_cleanup_completed_at
+//	FROM batch_uploads
+//	WHERE id = $1
+//	FOR UPDATE
+func (q *Queries) LockBatchUploadForFolderCleanup(ctx context.Context, batchID uuid.UUID) (*LockBatchUploadForFolderCleanupRow, error) {
+	row := q.db.QueryRow(ctx, lockBatchUploadForFolderCleanup, batchID)
+	var i LockBatchUploadForFolderCleanupRow
+	err := row.Scan(
+		&i.ID,
+		&i.Status,
+		&i.TotalDocuments,
+		&i.CreatedAt,
+		&i.CompletedAt,
+		&i.FolderCleanupCompletedAt,
+	)
+	return &i, err
+}
+
+const markBatchFolderCleanupCompleted = `-- name: MarkBatchFolderCleanupCompleted :exec
+UPDATE batch_uploads
+SET folder_cleanup_completed_at = NOW()
+WHERE id = $1
+  AND folder_cleanup_completed_at IS NULL
+`
+
+// MarkBatchFolderCleanupCompleted
+//
+//	UPDATE batch_uploads
+//	SET folder_cleanup_completed_at = NOW()
+//	WHERE id = $1
+//	  AND folder_cleanup_completed_at IS NULL
+func (q *Queries) MarkBatchFolderCleanupCompleted(ctx context.Context, batchID uuid.UUID) error {
+	_, err := q.db.Exec(ctx, markBatchFolderCleanupCompleted, batchID)
+	return err
+}
+
+const recordBatchFolderCandidate = `-- name: RecordBatchFolderCandidate :exec
+INSERT INTO batch_folder_candidates (batch_id, folder_id, path, existed_before)
+VALUES ($1, $2, $3, $4)
+ON CONFLICT (batch_id, folder_id) DO UPDATE
+SET path = EXCLUDED.path,
+    existed_before = batch_folder_candidates.existed_before AND EXCLUDED.existed_before
+`
+
+type RecordBatchFolderCandidateParams struct {
+	BatchID       uuid.UUID `db:"batch_id"`
+	FolderID      uuid.UUID `db:"folder_id"`
+	Path          string    `db:"path"`
+	ExistedBefore bool      `db:"existed_before"`
+}
+
+// Records a non-root folder touched by a batch upload. Once a batch observes a
+// folder as newly created, retries must preserve that fact.
+//
+//	INSERT INTO batch_folder_candidates (batch_id, folder_id, path, existed_before)
+//	VALUES ($1, $2, $3, $4)
+//	ON CONFLICT (batch_id, folder_id) DO UPDATE
+//	SET path = EXCLUDED.path,
+//	    existed_before = batch_folder_candidates.existed_before AND EXCLUDED.existed_before
+func (q *Queries) RecordBatchFolderCandidate(ctx context.Context, arg *RecordBatchFolderCandidateParams) error {
+	_, err := q.db.Exec(ctx, recordBatchFolderCandidate,
+		arg.BatchID,
+		arg.FolderID,
+		arg.Path,
+		arg.ExistedBefore,
+	)
+	return err
+}
+
 const resolveBatchDocumentOutcome = `-- name: ResolveBatchDocumentOutcome :exec
 UPDATE batch_document_outcomes
 SET document_id = $3,
@@ -568,6 +994,27 @@ func (q *Queries) ResolveBatchDocumentOutcome(ctx context.Context, arg *ResolveB
 	return err
 }

+const setBatchCompletedAt = `-- name: SetBatchCompletedAt :exec
+UPDATE batch_uploads
+SET completed_at = $1
+WHERE id = $2
+`
+
+type SetBatchCompletedAtParams struct {
+	CompletedAt pgtype.Timestamp `db:"completed_at"`
+	ID          uuid.UUID        `db:"id"`
+}
+
+// SetBatchCompletedAt
+//
+//	UPDATE batch_uploads
+//	SET completed_at = $1
+//	WHERE id = $2
+func (q *Queries) SetBatchCompletedAt(ctx context.Context, arg *SetBatchCompletedAtParams) error {
+	_, err := q.db.Exec(ctx, setBatchCompletedAt, arg.CompletedAt, arg.ID)
+	return err
+}
+
 const setBatchTotalDocuments = `-- name: SetBatchTotalDocuments :exec
 UPDATE batch_uploads SET total_documents = $2 WHERE id = $1
 `
@@ -585,13 +1032,14 @@ func (q *Queries) SetBatchTotalDocuments(ctx context.Context, arg *SetBatchTotal
 	return err
 }

-const updateBatchDocumentOutcomeByDocumentID = `-- name: UpdateBatchDocumentOutcomeByDocumentID :exec
+const updateBatchDocumentOutcomeByDocumentID = `-- name: UpdateBatchDocumentOutcomeByDocumentID :many
 UPDATE batch_document_outcomes
 SET outcome = $2::batch_outcome_status,
    error_detail = $3,
    updated_at = NOW()
 WHERE document_id = $1
  AND outcome IN ('submitted', 'init_complete', 'sync_complete')
+RETURNING batch_id
 `

 type UpdateBatchDocumentOutcomeByDocumentIDParams struct {
@@ -611,9 +1059,25 @@ type UpdateBatchDocumentOutcomeByDocumentIDParams struct {
 //	    updated_at = NOW()
 //	WHERE document_id = $1
 //	  AND outcome IN ('submitted', 'init_complete', 'sync_complete')
-func (q *Queries) UpdateBatchDocumentOutcomeByDocumentID(ctx context.Context, arg *UpdateBatchDocumentOutcomeByDocumentIDParams) error {
-	_, err := q.db.Exec(ctx, updateBatchDocumentOutcomeByDocumentID, arg.DocumentID, arg.Column2, arg.ErrorDetail)
-	return err
+//	RETURNING batch_id
+func (q *Queries) UpdateBatchDocumentOutcomeByDocumentID(ctx context.Context, arg *UpdateBatchDocumentOutcomeByDocumentIDParams) ([]uuid.UUID, error) {
+	rows, err := q.db.Query(ctx, updateBatchDocumentOutcomeByDocumentID, arg.DocumentID, arg.Column2, arg.ErrorDetail)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := []uuid.UUID{}
+	for rows.Next() {
+		var batch_id uuid.UUID
+		if err := rows.Scan(&batch_id); err != nil {
+			return nil, err
+		}
+		items = append(items, batch_id)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
 }

 const updateBatchProgress = `-- name: UpdateBatchProgress :exec