query-orchestration/internal/server/api/batch_worker.go

package api

import (
	"archive/zip"
	"bytes"
	"context"
	"fmt"
	"io"
	"log/slog"
	"path"
	"strings"

	"queryorchestration/internal/document/batch"

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

// Config key constants for batch worker
const (
	ConfigKeyBatchService  = "batchService"
	ConfigKeyS3Client      = "s3Client"
	ConfigKeyBucket        = "bucket"
	ConfigKeyUploadHandler = "uploadHandler"
)

// processBatchWork is the background worker function that processes batch uploads
func processBatchWork(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	// Extract the batch service from config
	batchService, ok := config[ConfigKeyBatchService].(*batch.Service)
	if !ok {
		return fmt.Errorf("batch service not found in config")
	}

	// Extract the S3 client from config
	s3Client, ok := config[ConfigKeyS3Client].(*s3.Client)
	if !ok {
		return fmt.Errorf("S3 client not found in config")
	}

	// Extract the bucket name from config
	bucket, ok := config[ConfigKeyBucket].(string)
	if !ok {
		return fmt.Errorf("bucket not found in config")
	}

	// Extract the document upload handler from config
	uploadHandler, ok := config[ConfigKeyUploadHandler].(func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error)
	if !ok {
		return fmt.Errorf("upload handler not found in config")
	}

	// Query unprocessed batches
	batches, err := batchService.ListUnprocessed(ctx)
	if err != nil {
		logger.Error("Failed to query unprocessed batches", "error", err)
		return err
	}

	logger.Info("Background task executed", "unprocessed_batches", len(batches))

	// Process each batch
	for _, batchSummary := range batches {
		// Get full batch details with archive information
		batchDetails, err := batchService.Get(ctx, batchSummary.ClientID, batchSummary.ID)
		if err != nil {
			logger.Error("Failed to get batch details",
				"batch_id", batchSummary.ID,
				"client_id", batchSummary.ClientID,
				"error", err)
			continue
		}

		if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails); err != nil {
			logger.Error("Failed to process batch",
				"batch_id", batchDetails.ID,
				"client_id", batchDetails.ClientID,
				"error", err)
			// Continue processing other batches even if one fails
			continue
		}
	}

	return nil
}

// processSingleBatch processes a single batch upload
func processSingleBatch(
	ctx context.Context,
	logger *slog.Logger,
	batchService *batch.Service,
	s3Client *s3.Client,
	bucket string,
	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
	batchInfo *batch.BatchUploadDetails,
) error {
	logger.Info("Processing batch",
		"batch_id", batchInfo.ID,
		"client_id", batchInfo.ClientID,
		"filename", batchInfo.OriginalFilename)

	// Batch is already in 'processing' status by default when created

	// Download the zip file from S3
	zipData, err := downloadZipFromS3(ctx, s3Client, bucket, batchInfo.ArchiveKey)
	if err != nil {
		// Mark batch as failed if we can't download the zip
		_ = batchService.MarkFailed(ctx, batchInfo.ID)
		return fmt.Errorf("failed to download zip from S3: %w", err)
	}

	// Extract and process files from the zip
	processedCount := 0
	failedCount := 0
	invalidCount := 0

	// Create a temporary S3 directory for extracted files
	tempDir := fmt.Sprintf("%s_burst", batchInfo.ID.String())
	basePath := path.Dir(batchInfo.ArchiveKey)
	extractPath := path.Join(basePath, tempDir)

	// Open the zip archive
	zipReader, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
	if err != nil {
		_ = batchService.MarkFailed(ctx, 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)
		processedCount += processed
		failedCount += failed
		invalidCount += invalid
	}

	// Update batch progress - convert safely to int32
	const maxInt32 = 2147483647
	var processedInt32, failedInt32, invalidInt32 int32
	if processedCount <= maxInt32 {
		processedInt32 = int32(processedCount) // #nosec G115 - checked bounds above
	} else {
		processedInt32 = maxInt32 // Max int32 value
	}
	if failedCount <= maxInt32 {
		failedInt32 = int32(failedCount) // #nosec G115 - checked bounds above
	} else {
		failedInt32 = maxInt32 // Max int32 value
	}
	if invalidCount <= maxInt32 {
		invalidInt32 = int32(invalidCount) // #nosec G115 - checked bounds above
	} else {
		invalidInt32 = maxInt32 // Max int32 value
	}

	if err := batchService.UpdateProgress(ctx, batchInfo.ClientID, batchInfo.ID,
		processedInt32, failedInt32, invalidInt32); err != nil {
		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
			if err := batchService.MarkCompleted(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as completed", "batch_id", batchInfo.ID, "error", err)
			}
		} else {
			// Mark as failed only if no documents were processed successfully
			if err := batchService.MarkFailed(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as failed", "batch_id", batchInfo.ID, "error", err)
			}
		}
	}

	logger.Info("Batch processing completed",
		"batch_id", batchInfo.ID,
		"processed", processedCount,
		"failed", failedCount,
		"invalid", invalidCount)

	return nil
}

// 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{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to get object from S3: %w", err)
	}
	defer result.Body.Close()

	data, err := io.ReadAll(result.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read S3 object body: %w", err)
	}

	return data, nil
}

// uploadToS3 uploads data to S3
func uploadToS3(ctx context.Context, s3Client *s3.Client, bucket, key string, data []byte) error {
	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(data),
	})
	if err != nil {
		return fmt.Errorf("failed to put object to S3: %w", err)
	}
	return nil
}

// processZipFile processes a single file from the zip archive
// Returns (processedCount, failedCount, invalidCount)
func processZipFile(
	ctx context.Context,
	logger *slog.Logger,
	batchService *batch.Service,
	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) {
	if file.FileInfo().IsDir() {
		return 0, 0, 0
	}

	// Skip system files and sanitize filename
	filename := file.Name
	if strings.HasPrefix(filename, "__MACOSX/") || strings.HasPrefix(filename, ".") {
		return 0, 0, 0
	}

	// Security: Clean filename to prevent path traversal attacks
	filename = path.Base(filename)

	// Extract file content
	fileReader, err := file.Open()
	if err != nil {
		logger.Error("Failed to open file in zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		return 0, 1, 0
	}

	fileData, err := io.ReadAll(fileReader)
	fileReader.Close()
	if err != nil {
		logger.Error("Failed to read file from zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		return 0, 1, 0
	}

	// Store extracted file in S3 temporary directory (filename is already sanitized with path.Base)
	extractedKey := extractPath + "/" + filename
	if err := uploadToS3(ctx, s3Client, bucket, extractedKey, fileData); err != nil {
		logger.Error("Failed to upload extracted file to S3", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		return 0, 1, 0
	}

	// Check if file type is valid (PDF only for now)
	if !strings.HasSuffix(strings.ToLower(filename), ".pdf") {
		logger.Warn("Invalid file type", "filename", filename)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		return 0, 0, 1
	}

	// Call the document upload handler with batch ID
	if err := uploadHandler(ctx, batchInfo.ClientID, bytes.NewReader(fileData), filename, &batchInfo.ID); err != nil {
		logger.Error("Failed to upload document", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		return 0, 1, 0
	}

	logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
	return 1, 0, 0
}