batch upload impl
This commit is contained in:
jay brown
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user