batch upload impl

internal/document/batch/service.go

@@ -32,6 +32,9 @@ type BatchUploadSummary struct {
 type BatchUploadDetails struct {
 	BatchUploadSummary
 	FailedFilenames []string `json:"failed_filenames"`
+	ArchiveBucket   string   `json:"archive_bucket,omitempty"`
+	ArchiveKey      string   `json:"archive_key,omitempty"`
+	FileSizeBytes   int64    `json:"file_size_bytes,omitempty"`
 }
 
 // Service handles batch upload operations
@@ -55,9 +58,27 @@ func (s *Service) Create(ctx context.Context, clientID string, filename string,
 	})
 }
 
+// CreateWithStorage creates a batch record with S3 storage metadata
+func (s *Service) CreateWithStorage(ctx context.Context, clientID string, filename string, totalDocs int32, bucket string, key string, sizeBytes int64) (uuid.UUID, error) {
+	result, err := s.cfg.GetDBQueries().CreateBatchUploadWithStorage(ctx, &repository.CreateBatchUploadWithStorageParams{
+		ClientID:         clientID,
+		OriginalFilename: filename,
+		TotalDocuments:   totalDocs,
+		Column4:          repository.BatchStatusProcessing,
+		ArchiveBucket:    &bucket,
+		ArchiveKey:       &key,
+		FileSizeBytes:    &sizeBytes,
+	})
+	if err != nil {
+		return uuid.Nil, fmt.Errorf("failed to create batch upload record: %w", err)
+	}
+
+	return result.ID, nil
+}
+
 // Get retrieves batch upload details
 func (s *Service) Get(ctx context.Context, clientID string, batchID uuid.UUID) (*BatchUploadDetails, error) {
-	batch, err := s.cfg.GetDBQueries().GetBatchUpload(ctx, &repository.GetBatchUploadParams{
+	batch, err := s.cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{
 		ID:       batchID,
 		ClientID: clientID,
 	})
@@ -65,7 +86,7 @@ func (s *Service) Get(ctx context.Context, clientID string, batchID uuid.UUID) (
 		return nil, err
 	}
 
-	return s.convertToDetails(batch), nil
+	return s.convertToDetailsWithStorage(batch), nil
 }
 
 // List retrieves all batch uploads for a client with pagination
@@ -153,39 +174,6 @@ func (s *Service) AddFailedFilename(ctx context.Context, batchID uuid.UUID, file
 	})
 }
 
-// convertToDetails converts repository model to service details model
-func (s *Service) convertToDetails(batch *repository.BatchUpload) *BatchUploadDetails {
-	details := &BatchUploadDetails{
-		BatchUploadSummary: BatchUploadSummary{
-			ID:                   batch.ID,
-			ClientID:             batch.ClientID,
-			OriginalFilename:     batch.OriginalFilename,
-			Status:               string(batch.Status),
-			TotalDocuments:       batch.TotalDocuments,
-			ProcessedDocuments:   batch.ProcessedDocuments,
-			FailedDocuments:      batch.FailedDocuments,
-			InvalidTypeDocuments: batch.InvalidTypeDocuments,
-			ProgressPercent:      batch.ProgressPercent,
-			CreatedAt:            batch.CreatedAt.Time,
-		},
-		FailedFilenames: make([]string, 0),
-	}
-
-	if batch.CompletedAt.Valid {
-		details.CompletedAt = &batch.CompletedAt.Time
-	}
-
-	// Parse failed filenames from JSONB
-	if len(batch.FailedFilenames) > 0 {
-		var filenames []string
-		if err := json.Unmarshal(batch.FailedFilenames, &filenames); err == nil {
-			details.FailedFilenames = filenames
-		}
-	}
-
-	return details
-}
-
 // convertToSummary converts repository list row to service summary model
 func (s *Service) convertToSummary(batch *repository.ListBatchUploadsRow) *BatchUploadSummary {
 	summary := &BatchUploadSummary{
@@ -216,3 +204,53 @@ func (s *Service) pgTimestampToTime(ts pgtype.Timestamp) time.Time {
 	}
 	return time.Time{}
 }
+
+// convertToDetailsWithStorage converts repository model with storage to service details model
+func (s *Service) convertToDetailsWithStorage(batch *repository.GetBatchUploadWithStorageRow) *BatchUploadDetails {
+	details := &BatchUploadDetails{
+		BatchUploadSummary: BatchUploadSummary{
+			ID:                   batch.ID,
+			ClientID:             batch.ClientID,
+			OriginalFilename:     batch.OriginalFilename,
+			Status:               string(batch.Status),
+			TotalDocuments:       batch.TotalDocuments,
+			ProcessedDocuments:   batch.ProcessedDocuments,
+			FailedDocuments:      batch.FailedDocuments,
+			InvalidTypeDocuments: batch.InvalidTypeDocuments,
+			ProgressPercent:      0, // Calculate if needed
+			CreatedAt:            batch.CreatedAt.Time,
+		},
+		FailedFilenames: make([]string, 0),
+	}
+
+	if batch.CompletedAt.Valid {
+		details.CompletedAt = &batch.CompletedAt.Time
+	}
+
+	// Add storage fields
+	if batch.ArchiveBucket != nil {
+		details.ArchiveBucket = *batch.ArchiveBucket
+	}
+	if batch.ArchiveKey != nil {
+		details.ArchiveKey = *batch.ArchiveKey
+	}
+	if batch.FileSizeBytes != nil {
+		details.FileSizeBytes = *batch.FileSizeBytes
+	}
+
+	// Parse failed filenames from JSONB
+	if len(batch.FailedFilenames) > 0 {
+		var filenames []string
+		if err := json.Unmarshal(batch.FailedFilenames, &filenames); err == nil {
+			details.FailedFilenames = filenames
+		}
+	}
+
+	// Calculate progress percent
+	if batch.TotalDocuments > 0 {
+		processed := batch.ProcessedDocuments + batch.FailedDocuments + batch.InvalidTypeDocuments
+		details.ProgressPercent = int32(float64(processed) / float64(batch.TotalDocuments) * 100)
+	}
+
+	return details
+}

internal/document/batch/service_test.go

@@ -1,19 +1,27 @@
 package batch_test
 
 import (
+	"archive/zip"
+	"bytes"
+	"fmt"
+	"io"
 	"testing"
 
 	"queryorchestration/internal/database/repository"
 	documentbatch "queryorchestration/internal/document/batch"
+	"queryorchestration/internal/document/batch/storage"
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/test"
 
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 type TestConfig struct {
 	serviceconfig.BaseConfig
+	objectstore.ObjectStoreConfig
 }
 
 func TestUpdateProgress(t *testing.T) {
@@ -183,3 +191,232 @@ func TestAddFailedFilename(t *testing.T) {
 	require.NoError(t, err)
 	assert.Contains(t, details.FailedFilenames, "failed_file.pdf")
 }
+
+func TestCreateWithStorage_Success(t *testing.T) {
+	// Setup: Create database testcontainer
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	service := documentbatch.New(cfg)
+
+	// Create test client
+	clientID := "test_client_storage"
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client Storage",
+	})
+	require.NoError(t, err)
+
+	// Test: Create batch with storage metadata
+	batchID, err := service.CreateWithStorage(
+		t.Context(),
+		clientID,
+		"test.zip",
+		0, // totalDocs unknown in Part 1
+		"test-bucket",
+		"batches/test-client/2025/01/01/test.zip",
+		1024*1024, // 1MB
+	)
+
+	// Verify: Batch created successfully
+	require.NoError(t, err)
+	assert.NotEmpty(t, batchID)
+
+	// Verify: Batch can be retrieved with storage metadata
+	batch, err := service.Get(t.Context(), clientID, batchID)
+	require.NoError(t, err)
+	assert.Equal(t, "test.zip", batch.OriginalFilename)
+	assert.Equal(t, "processing", batch.Status)
+	assert.Equal(t, "test-bucket", batch.ArchiveBucket)
+	assert.Equal(t, "batches/test-client/2025/01/01/test.zip", batch.ArchiveKey)
+	assert.Equal(t, int64(1024*1024), batch.FileSizeBytes)
+}
+
+func TestCreateWithStorage_DatabaseFailure(t *testing.T) {
+	// Setup: Create database
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	_ = documentbatch.New(cfg) // service unused in skipped test
+
+	// Create test client
+	clientID := "test_client_storage_fail"
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client Storage Fail",
+	})
+	require.NoError(t, err)
+
+	// Note: Skipping database connection close test as Pool interface doesn't expose Close method
+	// This test would require mocking the database layer to simulate connection failures
+	t.Skip("Database connection failure testing requires mocking")
+}
+
+// TestCreate_Success tests the Create function with a valid ZIP file uploaded to S3
+func TestCreate_Success(t *testing.T) {
+	// Setup: Create database and AWS resources with testcontainers
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create batch service and storage service
+	batchService := documentbatch.New(cfg)
+	storageService := storage.New(cfg)
+
+	// Create test client
+	clientID := "test_client_create"
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client Create",
+	})
+	require.NoError(t, err)
+
+	// Create test ZIP content with 3 PDFs
+	zipContent := createTestZIPContent(t, 3)
+
+	// First upload the ZIP to S3 (simulating what would happen in the full flow)
+	storageResult, err := storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)
+	require.NoError(t, err)
+
+	// Test: Use the Create function to create a batch record
+	batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 3)
+
+	// Verify: Batch created successfully
+	require.NoError(t, err)
+	assert.NotEmpty(t, batchID)
+
+	// Verify: Batch can be retrieved
+	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
+		ID:       batchID,
+		ClientID: clientID,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "test.zip", batch.OriginalFilename)
+	assert.Equal(t, int32(3), batch.TotalDocuments)
+	assert.Equal(t, clientID, batch.ClientID)
+	assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
+
+	// Verify: File still exists in S3
+	s3Client := cfg.GetStoreClient()
+	_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
+		Bucket: &storageResult.Bucket,
+		Key:    &storageResult.Key,
+	})
+	require.NoError(t, err, "ZIP file should exist in S3")
+}
+
+// TestCreate_InvalidZIP tests Create function with invalid ZIP data
+func TestCreate_InvalidZIP(t *testing.T) {
+	// Setup: Create database and AWS resources
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	batchService := documentbatch.New(cfg)
+	storageService := storage.New(cfg)
+
+	// Create test client
+	clientID := "test_client_invalid_zip"
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client Invalid ZIP",
+	})
+	require.NoError(t, err)
+
+	// Create invalid ZIP content (just plain text, not a ZIP)
+	invalidContent := bytes.NewReader([]byte("This is not a valid ZIP file"))
+
+	// Test: Try to validate before storing (storage service should validate)
+	err = storageService.ValidateZIPFile("notazip.txt", int64(invalidContent.Len()))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "invalid file type")
+
+	// Even if we bypass validation and try with .zip extension, the content is invalid
+	// The Create function itself doesn't validate ZIP content, it just creates a DB record
+	// So we test that Create works but processing would fail later
+	batchID, err := batchService.Create(t.Context(), clientID, "fake.zip", 1)
+
+	// The Create function should succeed (it only creates DB record)
+	require.NoError(t, err)
+	assert.NotEmpty(t, batchID)
+
+	// Verify the batch exists but would fail during actual processing
+	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
+		ID:       batchID,
+		ClientID: clientID,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "fake.zip", batch.OriginalFilename)
+}
+
+// TestCreate_S3UploadFailure tests Create function when S3 upload fails
+func TestCreate_S3UploadFailure(t *testing.T) {
+	// Setup: Create database and AWS resources
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	batchService := documentbatch.New(cfg)
+	storageService := storage.New(cfg)
+
+	// Create test client
+	clientID := "test_client_s3_fail"
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Client S3 Fail",
+	})
+	require.NoError(t, err)
+
+	// Create valid ZIP content
+	zipContent := createTestZIPContent(t, 2)
+
+	// Temporarily break S3 configuration to simulate upload failure
+	// Save original bucket and set invalid one
+	originalBucket := cfg.GetBucket()
+	cfg.SetBucket("non-existent-bucket-that-should-not-exist")
+
+	// Test: Try to upload to S3 with invalid bucket
+	_, err = storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)
+
+	// Verify: S3 upload should fail
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to upload ZIP to S3")
+
+	// Restore original bucket for cleanup
+	cfg.SetBucket(originalBucket)
+
+	// The Create function itself is independent of S3, so it would still work
+	// This demonstrates separation of concerns
+	batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 2)
+	require.NoError(t, err)
+	assert.NotEmpty(t, batchID)
+
+	// Verify batch exists in DB even though S3 upload failed
+	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
+		ID:       batchID,
+		ClientID: clientID,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "test.zip", batch.OriginalFilename)
+	assert.Equal(t, int32(2), batch.TotalDocuments)
+}
+
+// Helper function to create test ZIP content with specified number of PDF files
+func createTestZIPContent(t *testing.T, numFiles int) io.Reader {
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	for i := 1; i <= numFiles; i++ {
+		filename := fmt.Sprintf("document_%d.pdf", i)
+		fileWriter, err := zipWriter.Create(filename)
+		require.NoError(t, err)
+
+		// Write minimal PDF content
+		pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
+		_, err = fileWriter.Write([]byte(pdfContent))
+		require.NoError(t, err)
+	}
+
+	err := zipWriter.Close()
+	require.NoError(t, err)
+
+	return bytes.NewReader(buf.Bytes())
+}

internal/document/batch/storage/service.go

@@ -0,0 +1,84 @@
+package storage
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"path/filepath"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/google/uuid"
+
+	"queryorchestration/internal/serviceconfig/objectstore"
+)
+
+type Service struct {
+	cfg objectstore.ConfigProvider
+}
+
+type StorageResult struct {
+	Bucket    string
+	Key       string
+	SizeBytes int64
+}
+
+func New(cfg objectstore.ConfigProvider) *Service {
+	return &Service{cfg: cfg}
+}
+
+// StoreZIPFile uploads ZIP file to S3 and returns storage metadata
+func (s *Service) StoreZIPFile(ctx context.Context, clientID string, filename string, content io.Reader) (*StorageResult, error) {
+	// Create unique S3 key with timestamp and UUID
+	timestamp := time.Now().Format("2006/01/02/15")
+	batchID := uuid.New()
+	key := fmt.Sprintf("batches/%s/%s/%s-%s", clientID, timestamp, batchID.String(), filename)
+
+	// Get bucket from existing configuration
+	bucketName := s.cfg.GetBucket()
+
+	// Read content to get size (required for S3 PutObject)
+	buf := &bytes.Buffer{}
+	size, err := io.Copy(buf, content)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read file content: %w", err)
+	}
+
+	// Upload to S3 using the store client
+	_, err = s.cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
+		Bucket: &bucketName,
+		Key:    &key,
+		Body:   bytes.NewReader(buf.Bytes()),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to upload ZIP to S3: %w", err)
+	}
+
+	return &StorageResult{
+		Bucket:    bucketName,
+		Key:       key,
+		SizeBytes: size,
+	}, nil
+}
+
+// ValidateZIPFile performs basic ZIP file validation
+func (s *Service) ValidateZIPFile(filename string, sizeBytes int64) error {
+	// Check file extension
+	ext := filepath.Ext(filename)
+	if ext != ".zip" {
+		return fmt.Errorf("invalid file type: expected .zip, got %s", ext)
+	}
+
+	// Check file size (100MB limit for Part 1)
+	const maxSizeBytes = 100 * 1024 * 1024 // 100MB
+	if sizeBytes > maxSizeBytes {
+		return fmt.Errorf("file too large: %d bytes exceeds limit of %d bytes", sizeBytes, maxSizeBytes)
+	}
+
+	if sizeBytes == 0 {
+		return fmt.Errorf("empty file not allowed")
+	}
+
+	return nil
+}

internal/document/batch/storage/service_test.go

@@ -0,0 +1,165 @@
+package storage
+
+import (
+	"archive/zip"
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
+	"queryorchestration/internal/test"
+)
+
+// TestConfig provides a test configuration that implements both service config and AWS interfaces
+type TestConfig struct {
+	serviceconfig.BaseConfig
+	objectstore.ObjectStoreConfig
+}
+
+func TestStoreZIPFile_Success(t *testing.T) {
+	// Setup: Create testcontainer with localstack
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create storage service
+	service := New(cfg)
+
+	// Create test ZIP content with 3 PDFs
+	zipContent := createTestZIPContent(t)
+
+	// Test: Store ZIP file
+	result, err := service.StoreZIPFile(t.Context(), "test-client", "test.zip", zipContent)
+
+	// Verify: Storage result
+	require.NoError(t, err)
+	assert.NotEmpty(t, result.Bucket)
+	assert.NotEmpty(t, result.Key)
+	assert.Greater(t, result.SizeBytes, int64(0))
+	assert.Contains(t, result.Key, "test-client")
+	assert.Contains(t, result.Key, "test.zip")
+
+	// Verify: File exists in S3 (using actual S3 client)
+	s3Client := cfg.GetStoreClient()
+	_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
+		Bucket: &result.Bucket,
+		Key:    &result.Key,
+	})
+	require.NoError(t, err, "ZIP file should exist in S3")
+}
+
+func TestStoreZIPFile_LargeFile(t *testing.T) {
+	// Setup: Create testcontainer with localstack
+	cfg := &TestConfig{}
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create storage service
+	service := New(cfg)
+
+	// Create a large test ZIP content (but still under 100MB limit)
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	// Create multiple PDFs to make a larger ZIP
+	for i := 0; i < 10; i++ {
+		filename := fmt.Sprintf("document_%d.pdf", i)
+		fileWriter, err := zipWriter.Create(filename)
+		require.NoError(t, err)
+
+		// Write a larger PDF content
+		pdfContent := fmt.Sprintf("%%PDF-1.4\n")
+		// Add some bulk content
+		for j := 0; j < 1000; j++ {
+			pdfContent += fmt.Sprintf("%%Page %d Line %d content\n", i, j)
+		}
+		pdfContent += "%%EOF"
+		_, err = fileWriter.Write([]byte(pdfContent))
+		require.NoError(t, err)
+	}
+
+	err := zipWriter.Close()
+	require.NoError(t, err)
+
+	// Test: Store larger ZIP file
+	result, err := service.StoreZIPFile(t.Context(), "test-client", "large-batch.zip", bytes.NewReader(buf.Bytes()))
+
+	// Verify: Storage result
+	require.NoError(t, err)
+	assert.NotEmpty(t, result.Bucket)
+	assert.NotEmpty(t, result.Key)
+	assert.Greater(t, result.SizeBytes, int64(0)) // Should have some size
+}
+
+func TestValidateZIPFile_Success(t *testing.T) {
+	service := New(nil) // No config needed for validation
+
+	// Test: Valid ZIP file
+	err := service.ValidateZIPFile("test.zip", 1024*1024) // 1MB
+
+	// Verify: No error
+	assert.NoError(t, err)
+}
+
+func TestValidateZIPFile_InvalidExtension(t *testing.T) {
+	service := New(nil)
+
+	// Test: Invalid file extension
+	err := service.ValidateZIPFile("test.pdf", 1024*1024)
+
+	// Verify: Error returned
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "invalid file type")
+}
+
+func TestValidateZIPFile_TooLarge(t *testing.T) {
+	service := New(nil)
+
+	// Test: File too large (101MB)
+	err := service.ValidateZIPFile("test.zip", 101*1024*1024)
+
+	// Verify: Error returned
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "file too large")
+}
+
+func TestValidateZIPFile_Empty(t *testing.T) {
+	service := New(nil)
+
+	// Test: Empty file
+	err := service.ValidateZIPFile("test.zip", 0)
+
+	// Verify: Error returned
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "empty file not allowed")
+}
+
+// Helper function to create test ZIP content with 3 PDF files
+func createTestZIPContent(t *testing.T) *bytes.Reader {
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	// Create 3 PDF files in the ZIP
+	pdfFiles := []string{"test1.pdf", "test2.pdf", "test3.pdf"}
+
+	for _, filename := range pdfFiles {
+		fileWriter, err := zipWriter.Create(filename)
+		require.NoError(t, err)
+
+		// Write minimal PDF content (PDF magic header)
+		// Each file has slightly different content for identification
+		pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
+		_, err = fileWriter.Write([]byte(pdfContent))
+		require.NoError(t, err)
+	}
+
+	err := zipWriter.Close()
+	require.NoError(t, err)
+
+	return bytes.NewReader(buf.Bytes())
+}