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
@@ -14,18 +14,30 @@ import (

 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/document/batch"
+	"queryorchestration/internal/document/batch/foldercleanup"

 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/google/uuid"
 )

+type batchWorkerService interface {
+	AddFailedFilename(ctx context.Context, batchID uuid.UUID, filename string) error
+	CheckDuplicate(ctx context.Context, clientID string, hash string) (*uuid.UUID, error)
+	MarkCompleted(ctx context.Context, batchID uuid.UUID) error
+	MarkFailed(ctx context.Context, batchID uuid.UUID) error
+	RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string, outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error
+	SetTotalDocuments(ctx context.Context, batchID uuid.UUID, total int32) error
+	UpdateProgress(ctx context.Context, clientID string, batchID uuid.UUID, processed int32, failed int32, invalid int32) error
+}
+
 // Config key constants for batch worker
 const (
 	ConfigKeyBatchService  = "batchService"
 	ConfigKeyS3Client      = "s3Client"
 	ConfigKeyBucket        = "bucket"
 	ConfigKeyUploadHandler = "uploadHandler"
+	ConfigKeyFolderCleanup = "folderCleanupService"
 )

 // acceptedExtensions defines the file extensions allowed in batch uploads.
@@ -115,6 +127,15 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin
 		return fmt.Errorf("upload handler not found in config")
 	}

+	cleanupService, _ := config[ConfigKeyFolderCleanup].(*foldercleanup.Service)
+	if cleanupService != nil {
+		if finalized, err := cleanupService.FinalizeReadyBatches(ctx, foldercleanup.DefaultReadyBatchLimit); err != nil {
+			logger.Error("Failed to finalize ready batch folders", "error", err)
+		} else if finalized > 0 {
+			logger.Info("Finalized ready batch folders", "count", finalized)
+		}
+	}
+
 	// Query unprocessed batches
 	batches, err := batchService.ListUnprocessed(ctx)
 	if err != nil {
@@ -136,7 +157,7 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin
 			continue
 		}

-		if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails); err != nil {
+		if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, cleanupService); err != nil {
 			logger.Error("Failed to process batch",
 				"batch_id", batchDetails.ID,
 				"client_id", batchDetails.ClientID,
@@ -146,6 +167,14 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin
 		}
 	}

+	if cleanupService != nil {
+		if deleted, err := cleanupService.DeleteSafeAutoCreatedFolders(ctx, foldercleanup.DefaultFolderDeleteLimit); err != nil {
+			logger.Error("Failed to delete safe auto-created folders", "error", err)
+		} else if deleted > 0 {
+			logger.Info("Deleted safe auto-created folders", "count", deleted)
+		}
+	}
+
 	return nil
 }

@@ -153,11 +182,12 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin
 func processSingleBatch(
 	ctx context.Context,
 	logger *slog.Logger,
-	batchService *batch.Service,
+	batchService batchWorkerService,
 	s3Client *s3.Client,
 	bucket string,
 	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
 	batchInfo *batch.BatchUploadDetails,
+	cleanupService *foldercleanup.Service,
 ) error {
 	logger.Info("Processing batch",
 		"batch_id", batchInfo.ID,
@@ -171,6 +201,7 @@ func processSingleBatch(
 	if err != nil {
 		// Mark batch as failed if we can't download the zip
 		_ = batchService.MarkFailed(ctx, batchInfo.ID)
+		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
 		return fmt.Errorf("failed to download zip from S3: %w", err)
 	}

@@ -188,17 +219,34 @@ func processSingleBatch(
 	zipReader, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
 	if err != nil {
 		_ = batchService.MarkFailed(ctx, batchInfo.ID)
+		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
 		return fmt.Errorf("failed to open zip archive: %w", err)
 	}

 	// Process each file in the zip
 	for _, file := range zipReader.File {
-		processed, failed, invalid := processZipFile(ctx, logger, batchService, s3Client, bucket, extractPath, uploadHandler, batchInfo, file)
+		processed, failed, invalid, err := processZipFile(ctx, logger, batchService, s3Client, bucket, extractPath, uploadHandler, batchInfo, file)
+		if err != nil {
+			_ = batchService.MarkFailed(ctx, batchInfo.ID)
+			tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
+			return fmt.Errorf("failed to process zip file %q: %w", file.Name, err)
+		}
 		processedCount += processed
 		failedCount += failed
 		invalidCount += invalid
 	}

+	totalProcessed := processedCount + failedCount + invalidCount
+	if totalProcessed == 0 {
+		if err := recordNoProcessableFilesOutcome(ctx, batchService, batchInfo); err != nil {
+			_ = batchService.MarkFailed(ctx, batchInfo.ID)
+			tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
+			return fmt.Errorf("record no-processable-files outcome: %w", err)
+		}
+		failedCount = 1
+		totalProcessed = 1
+	}
+
 	// Update batch progress - convert safely to int32
 	const maxInt32 = 2147483647
 	var processedInt32, failedInt32, invalidInt32 int32
@@ -229,8 +277,6 @@ func processSingleBatch(
 		logger.Error("Failed to update batch progress", "batch_id", batchInfo.ID, "error", err)
 	}

-	// Mark batch as completed if any files were processed successfully
-	totalProcessed := processedCount + failedCount + invalidCount
 	if totalProcessed > 0 {
 		if processedCount > 0 {
 			// Mark as completed if at least one document was processed successfully
@@ -243,6 +289,7 @@ func processSingleBatch(
 				logger.Error("Failed to mark batch as failed", "batch_id", batchInfo.ID, "error", err)
 			}
 		}
+		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
 	}

 	logger.Info("Batch processing completed",
@@ -254,6 +301,15 @@ func processSingleBatch(
 	return nil
 }

+func tryFinalizeBatchFolders(ctx context.Context, logger *slog.Logger, cleanupService *foldercleanup.Service, batchID uuid.UUID) {
+	if cleanupService == nil {
+		return
+	}
+	if _, err := cleanupService.TryFinalizeBatch(ctx, batchID); err != nil {
+		logger.Error("Failed to finalize batch folder cleanup", "batch_id", batchID, "error", err)
+	}
+}
+
 // downloadZipFromS3 downloads a zip file from S3
 func downloadZipFromS3(ctx context.Context, s3Client *s3.Client, bucket, key string) ([]byte, error) {
 	result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
@@ -291,33 +347,23 @@ func uploadToS3(ctx context.Context, s3Client *s3.Client, bucket, key string, da
 func processZipFile(
 	ctx context.Context,
 	logger *slog.Logger,
-	batchService *batch.Service,
+	batchService batchWorkerService,
 	s3Client *s3.Client,
 	bucket, extractPath string,
 	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
 	batchInfo *batch.BatchUploadDetails,
 	file *zip.File,
-) (int, int, int) {
+) (int, int, int, error) {
 	if file.FileInfo().IsDir() {
-		return 0, 0, 0
+		return 0, 0, 0, nil
 	}

-	// Skip system files and hidden files (but allow relative path prefixes like ./file.pdf)
-	filename := file.Name
-	if strings.HasPrefix(filename, "__MACOSX/") {
-		return 0, 0, 0
-	}
-
-	// Skip hidden files (files that start with . but are not relative paths like ./file)
-	if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") {
-		return 0, 0, 0
-	}
-
-	// Security: Clean filename to prevent path traversal attacks while preserving internal structure
-	filename = sanitizeZipPath(filename)
-	if filename == "" {
-		logger.Warn("Skipping file with invalid path", "original", file.Name)
-		return 0, 0, 0
+	filename, skip := batchZipFilename(file.Name)
+	if skip {
+		if filename == "" {
+			logger.Warn("Skipping file with invalid path", "original", file.Name)
+		}
+		return 0, 0, 0, nil
 	}

 	// Extract file content
@@ -326,8 +372,10 @@ func processZipFile(
 		logger.Error("Failed to open file in zip", "filename", filename, "error", err)
 		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
 		errMsg := err.Error()
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil)
-		return 0, 1, 0
+		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil); recordErr != nil {
+			return 0, 0, 0, fmt.Errorf("record failed_open outcome: %w", recordErr)
+		}
+		return 0, 1, 0, nil
 	}

 	fileData, err := io.ReadAll(fileReader)
@@ -336,8 +384,10 @@ func processZipFile(
 		logger.Error("Failed to read file from zip", "filename", filename, "error", err)
 		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
 		errMsg := err.Error()
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil)
-		return 0, 1, 0
+		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil); recordErr != nil {
+			return 0, 0, 0, fmt.Errorf("record failed_read outcome: %w", recordErr)
+		}
+		return 0, 1, 0, nil
 	}

 	// Store extracted file in S3 temporary directory (filename now includes sanitized path)
@@ -346,16 +396,20 @@ func processZipFile(
 		logger.Error("Failed to upload extracted file to S3", "filename", filename, "error", err)
 		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
 		errMsg := err.Error()
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil)
-		return 0, 1, 0
+		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil); recordErr != nil {
+			return 0, 0, 0, fmt.Errorf("record failed_s3_upload outcome: %w", recordErr)
+		}
+		return 0, 1, 0, nil
 	}

 	// Check if file type is supported
 	if !isAcceptedExtension(filename) {
 		logger.Warn("Invalid file type", "filename", filename)
 		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil)
-		return 0, 0, 1
+		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil); recordErr != nil {
+			return 0, 0, 0, fmt.Errorf("record invalid_type outcome: %w", recordErr)
+		}
+		return 0, 0, 1, nil
 	}

 	// Call the document upload handler with batch ID
@@ -363,20 +417,62 @@ func processZipFile(
 		logger.Error("Failed to upload document", "filename", filename, "error", err)
 		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
 		errMsg := err.Error()
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil)
-		return 0, 1, 0
+		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); recordErr != nil {
+			return 0, 0, 0, fmt.Errorf("record failed_upload outcome: %w", recordErr)
+		}
+		return 0, 1, 0, nil
 	}

-	// Check for duplicate by computing file hash and looking up existing document
+	if err := recordSuccessfulBatchOutcome(ctx, batchService, batchInfo, filename, fileData); err != nil {
+		return 0, 0, 0, err
+	}
+
+	logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
+	return 1, 0, 0, nil
+}
+
+func recordNoProcessableFilesOutcome(ctx context.Context, batchService batchWorkerService, batchInfo *batch.BatchUploadDetails) error {
+	filename := batchInfo.OriginalFilename
+	if filename == "" {
+		filename = batchInfo.ID.String() + ".zip"
+	}
+	errMsg := "archive contained no processable files"
+	_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
+	if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); err != nil {
+		return fmt.Errorf("record failed_upload outcome: %w", err)
+	}
+	return nil
+}
+
+func batchZipFilename(filename string) (string, bool) {
+	if strings.HasPrefix(filename, "__MACOSX/") {
+		return filename, true
+	}
+	if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") {
+		return filename, true
+	}
+	cleanFilename := sanitizeZipPath(filename)
+	return cleanFilename, cleanFilename == ""
+}
+
+func recordSuccessfulBatchOutcome(
+	ctx context.Context,
+	batchService batchWorkerService,
+	batchInfo *batch.BatchUploadDetails,
+	filename string,
+	fileData []byte,
+) error {
 	hash := sha256.Sum256(fileData)
 	hashStr := hex.EncodeToString(hash[:])
 	existingID, _ := batchService.CheckDuplicate(ctx, batchInfo.ClientID, hashStr)
 	if existingID != nil {
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID)
-	} else {
-		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil)
+		if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID); err != nil {
+			return fmt.Errorf("record duplicate outcome: %w", err)
+		}
+		return nil
 	}
-
-	logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
-	return 1, 0, 0
+	if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil); err != nil {
+		return fmt.Errorf("record submitted outcome: %w", err)
+	}
+	return nil
 }