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
@@ -6,14 +6,17 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"os"
+	"strings"
 	"testing"

 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/document/batch"
+	"queryorchestration/internal/document/batch/foldercleanup"
 	"queryorchestration/internal/serviceconfig"
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/test"
@@ -21,10 +24,20 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )

+type outcomeFailingBatchService struct {
+	*batch.Service
+}
+
+func (s *outcomeFailingBatchService) RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string,
+	outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error {
+	return errors.New("forced outcome write failure")
+}
+
 // TestProcessBatchWork tests the main batch processing work function
 func TestProcessBatchWork(t *testing.T) {
 	if testing.Short() {
@@ -110,6 +123,33 @@ func TestProcessBatchWork(t *testing.T) {
 	}
 }

+func TestProcessBatchWorkRunsFolderCleanupWithoutBatches(t *testing.T) {
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	test.CreateDB(t, cfg)
+
+	workerConfig := map[string]any{
+		ConfigKeyBatchService:  batch.New(cfg),
+		ConfigKeyS3Client:      &s3.Client{},
+		ConfigKeyBucket:        "bucket",
+		ConfigKeyUploadHandler: func(context.Context, string, io.Reader, string, *uuid.UUID) error { return nil },
+		ConfigKeyFolderCleanup: foldercleanup.New(cfg),
+	}
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	require.NoError(t, processBatchWork(ctx, logger, workerConfig))
+}
+
+func TestTryFinalizeBatchFoldersHandlesNilAndErrors(t *testing.T) {
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	test.CreateDB(t, cfg)
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+
+	tryFinalizeBatchFolders(ctx, logger, nil, uuid.New())
+	tryFinalizeBatchFolders(ctx, logger, foldercleanup.New(cfg), uuid.New())
+}
+
 // TestProcessSingleBatch tests processing a single batch
 func TestProcessSingleBatch(t *testing.T) {
 	if testing.Short() {
@@ -169,7 +209,7 @@ func TestProcessSingleBatch(t *testing.T) {

 	// Process the batch
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// All three file types (PDF, TXT, DOCX) are now accepted
@@ -253,7 +293,7 @@ func TestProcessSingleBatchWithFailure(t *testing.T) {

 	// Process the batch
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Verify results
@@ -357,6 +397,61 @@ func createTestZIPForWorker(t *testing.T, numFiles int) []byte {
 	return buf.Bytes()
 }

+func firstZipFile(t *testing.T, files map[string][]byte) *zip.File {
+	t.Helper()
+
+	zipFiles := zipFiles(t, files)
+	require.NotEmpty(t, zipFiles)
+	return zipFiles[0]
+}
+
+func zipFiles(t *testing.T, files map[string][]byte) []*zip.File {
+	t.Helper()
+
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+	for name, content := range files {
+		if content == nil && strings.HasSuffix(name, "/") {
+			header := &zip.FileHeader{Name: name}
+			header.SetMode(os.ModeDir | 0o755)
+			_, err := zipWriter.CreateHeader(header)
+			require.NoError(t, err)
+			continue
+		}
+		fileWriter, err := zipWriter.Create(name)
+		require.NoError(t, err)
+		_, err = fileWriter.Write(content)
+		require.NoError(t, err)
+	}
+	require.NoError(t, zipWriter.Close())
+
+	reader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
+	require.NoError(t, err)
+	return reader.File
+}
+
+func zipBytes(t *testing.T, files map[string][]byte) []byte {
+	t.Helper()
+
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+	for name, content := range files {
+		if content == nil && strings.HasSuffix(name, "/") {
+			header := &zip.FileHeader{Name: name}
+			header.SetMode(os.ModeDir | 0o755)
+			_, err := zipWriter.CreateHeader(header)
+			require.NoError(t, err)
+			continue
+		}
+		fileWriter, err := zipWriter.Create(name)
+		require.NoError(t, err)
+		_, err = fileWriter.Write(content)
+		require.NoError(t, err)
+	}
+	require.NoError(t, zipWriter.Close())
+	return buf.Bytes()
+}
+
 // Helper function to create a ZIP with mixed content types
 func createMixedContentZIP(t *testing.T) []byte {
 	var buf bytes.Buffer
@@ -539,7 +634,7 @@ func TestProcessBatchWithNestedFolders(t *testing.T) {

 	// Process the batch
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Verify batch status
@@ -677,7 +772,7 @@ func TestProcessBatchWithPathTraversal(t *testing.T) {

 	// Process the batch
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Verify that path traversal attempts were blocked
@@ -757,6 +852,334 @@ func TestProcessZipFile_RecordsOutcomes(t *testing.T) {
 	}
 }

+func TestProcessZipFile_RecordsFailedS3UploadOutcome(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	clientID := "test_zip_failed_s3"
+	test.CreateTestClient(t, cfg, clientID, "Test ZIP Failed S3 Client")
+
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok)
+
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "failed-s3.zip", 1, cfg.GetBucket(), "archive.zip", 1)
+	require.NoError(t, err)
+	file := firstZipFile(t, map[string][]byte{
+		"document.pdf": []byte("%%PDF-1.4\n%%Test PDF\n%%%%EOF"),
+	})
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	processed, failed, invalid, err := processZipFile(
+		ctx,
+		logger,
+		batchService,
+		s3Client,
+		"missing-bucket",
+		"extract/path",
+		func(context.Context, string, io.Reader, string, *uuid.UUID) error {
+			t.Fatal("upload handler should not be called when extracted S3 upload fails")
+			return nil
+		},
+		&batch.BatchUploadDetails{
+			BatchUploadSummary: batch.BatchUploadSummary{ID: batchID, ClientID: clientID},
+		},
+		file,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, 0, processed)
+	assert.Equal(t, 1, failed)
+	assert.Equal(t, 0, invalid)
+
+	outcomes, err := batchService.ListOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	require.Len(t, outcomes, 1)
+	assert.Equal(t, string(repository.BatchOutcomeStatusFailedS3Upload), outcomes[0].Outcome)
+	assert.Equal(t, "document.pdf", outcomes[0].Filename)
+	require.NotNil(t, outcomes[0].ErrorDetail)
+	assert.Contains(t, *outcomes[0].ErrorDetail, "failed to put object to S3")
+}
+
+func TestProcessZipFile_RecordsInvalidTypeOutcome(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	clientID := "test_zip_invalid_type"
+	test.CreateTestClient(t, cfg, clientID, "Test ZIP Invalid Type Client")
+
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok)
+
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "invalid-type.zip", 1, cfg.GetBucket(), "archive.zip", 1)
+	require.NoError(t, err)
+	file := firstZipFile(t, map[string][]byte{
+		"unsupported.exe": []byte("not a supported document"),
+	})
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	processed, failed, invalid, err := processZipFile(
+		ctx,
+		logger,
+		batchService,
+		s3Client,
+		cfg.GetBucket(),
+		"extract/path",
+		func(context.Context, string, io.Reader, string, *uuid.UUID) error {
+			t.Fatal("upload handler should not be called for invalid file types")
+			return nil
+		},
+		&batch.BatchUploadDetails{
+			BatchUploadSummary: batch.BatchUploadSummary{ID: batchID, ClientID: clientID},
+		},
+		file,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, 0, processed)
+	assert.Equal(t, 0, failed)
+	assert.Equal(t, 1, invalid)
+
+	outcomes, err := batchService.ListOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	require.Len(t, outcomes, 1)
+	assert.Equal(t, string(repository.BatchOutcomeStatusInvalidType), outcomes[0].Outcome)
+	assert.Equal(t, "unsupported.exe", outcomes[0].Filename)
+}
+
+func TestProcessZipFile_SkipsDirectoriesAndHiddenEntries(t *testing.T) {
+	ctx := t.Context()
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	files := zipFiles(t, map[string][]byte{
+		"folder/":               nil,
+		"__MACOSX/document.pdf": []byte("metadata"),
+		".DS_Store":             []byte("metadata"),
+		"../escape.pdf":         []byte("blocked"),
+	})
+
+	for _, file := range files {
+		processed, failed, invalid, err := processZipFile(
+			ctx,
+			logger,
+			nil,
+			nil,
+			"",
+			"",
+			func(context.Context, string, io.Reader, string, *uuid.UUID) error {
+				t.Fatal("upload handler should not be called for skipped entries")
+				return nil
+			},
+			&batch.BatchUploadDetails{},
+			file,
+		)
+		require.NoError(t, err)
+		assert.Equal(t, 0, processed)
+		assert.Equal(t, 0, failed)
+		assert.Equal(t, 0, invalid)
+	}
+}
+
+func TestOutcomeWriteFailureDoesNotCompleteBatch(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	clientID := "test_outcome_write_failure"
+	test.CreateTestClient(t, cfg, clientID, "Test Outcome Write Failure Client")
+
+	batchService := batch.New(cfg)
+	failingBatchService := &outcomeFailingBatchService{Service: batchService}
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok)
+	bucket := cfg.GetBucket()
+
+	zipContent := createTestZIPForWorker(t, 1)
+	archiveKey := fmt.Sprintf("test/%s/outcome-failure.zip", clientID)
+	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(archiveKey),
+		Body:   bytes.NewReader(zipContent),
+	})
+	require.NoError(t, err)
+
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "outcome-failure.zip", 0, bucket, archiveKey, int64(len(zipContent)))
+	require.NoError(t, err)
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "outcome-failure-folder",
+		Parentid:  nil,
+		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,
+	}))
+
+	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
+		_, err := io.ReadAll(docData)
+		return err
+	}
+	batchDetails := &batch.BatchUploadDetails{
+		BatchUploadSummary: batch.BatchUploadSummary{
+			ID:               batchID,
+			ClientID:         clientID,
+			OriginalFilename: "outcome-failure.zip",
+			Status:           "processing",
+		},
+		ArchiveKey:    archiveKey,
+		ArchiveBucket: bucket,
+	}
+
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	cleanupService := foldercleanup.New(cfg)
+	err = processSingleBatch(ctx, logger, failingBatchService, s3Client, bucket, uploadHandler, batchDetails, cleanupService)
+	require.Error(t, err)
+
+	batchInfo, err := cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{
+		ID:       batchID,
+		ClientID: clientID,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, repository.BatchStatusFailed, batchInfo.Status)
+	assert.Zero(t, batchInfo.TotalDocuments, "failed outcome writes must not count toward batch totals")
+
+	outcomeCount, err := cfg.GetDBQueries().CountBatchOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	assert.Zero(t, outcomeCount)
+
+	cleanupBatch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID)
+	require.NoError(t, err)
+	assert.False(t, cleanupBatch.FolderCleanupCompletedAt.Valid)
+	_, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	require.NoError(t, err)
+}
+
+func TestAllSkippedZipRecordsFailedOutcomeAndFinalizes(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	clientID := "test_all_skipped_zip"
+	test.CreateTestClient(t, cfg, clientID, "Test All Skipped ZIP Client")
+
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok)
+	bucket := cfg.GetBucket()
+
+	zipContent := zipBytes(t, map[string][]byte{
+		"folder/":               nil,
+		"__MACOSX/document.pdf": []byte("metadata"),
+		".DS_Store":             []byte("metadata"),
+		"../escape.pdf":         []byte("blocked"),
+	})
+	archiveKey := fmt.Sprintf("test/%s/all-skipped.zip", clientID)
+	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(archiveKey),
+		Body:   bytes.NewReader(zipContent),
+	})
+	require.NoError(t, err)
+
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "all-skipped.zip", 0, bucket, archiveKey, int64(len(zipContent)))
+	require.NoError(t, err)
+	folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{
+		Path:      "/all-skipped",
+		Parentid:  nil,
+		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,
+	}))
+
+	batchDetails := &batch.BatchUploadDetails{
+		BatchUploadSummary: batch.BatchUploadSummary{
+			ID:               batchID,
+			ClientID:         clientID,
+			OriginalFilename: "all-skipped.zip",
+			Status:           "processing",
+		},
+		ArchiveKey:    archiveKey,
+		ArchiveBucket: bucket,
+	}
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket,
+		func(context.Context, string, io.Reader, string, *uuid.UUID) error {
+			t.Fatal("upload handler should not be called for skipped-only archives")
+			return nil
+		},
+		batchDetails,
+		foldercleanup.New(cfg),
+	)
+	require.NoError(t, err)
+
+	batchInfo, err := cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{
+		ID:       batchID,
+		ClientID: clientID,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, repository.BatchStatusFailed, batchInfo.Status)
+	assert.Equal(t, int32(1), batchInfo.TotalDocuments)
+	assert.Equal(t, int32(1), batchInfo.FailedDocuments)
+
+	outcomes, err := batchService.ListOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	require.Len(t, outcomes, 1)
+	assert.Equal(t, "all-skipped.zip", outcomes[0].Filename)
+	assert.Equal(t, string(repository.BatchOutcomeStatusFailedUpload), outcomes[0].Outcome)
+	require.NotNil(t, outcomes[0].ErrorDetail)
+	assert.Contains(t, *outcomes[0].ErrorDetail, "archive contained no processable files")
+
+	unprocessed, err := cfg.GetDBQueries().GetUnprocessedBatches(ctx)
+	require.NoError(t, err)
+	for _, batch := range unprocessed {
+		assert.NotEqual(t, batchID, batch.ID)
+	}
+
+	cleanupBatch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID)
+	require.NoError(t, err)
+	assert.True(t, cleanupBatch.FolderCleanupCompletedAt.Valid)
+	_, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID)
+	assert.True(t, errors.Is(err, pgx.ErrNoRows))
+}
+
 // TestProcessZipFile_DetectsDuplicates verifies that when a file in the batch
 // has the same content hash as an existing document in the system, it is
 // recorded as a duplicate with the existing document's ID.
@@ -999,7 +1422,7 @@ func TestProcessBatchWithDuplicateFilenames(t *testing.T) {

 	// Process the batch
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Verify all files were uploaded with correct paths
@@ -1066,7 +1489,7 @@ func TestProcessBatchWithAllSupportedFileTypes(t *testing.T) {
 	}

 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Verify batch completed successfully
@@ -1166,7 +1589,7 @@ func TestProcessBatchWithMixedValidInvalidFileTypes(t *testing.T) {
 	}

 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil)
 	require.NoError(t, err)

 	// Batch should still complete (not fail) because some files succeeded