From 7014fa790be242682c080ae75fb0ff4f9159c0f1 Mon Sep 17 00:00:00 2001 From: jay brown Date: Wed, 6 Aug 2025 15:06:18 -0700 Subject: [PATCH] batch upload impl --- api/queryAPI/README.md | 119 +++++++++ api/queryAPI/authHandlers.go | 12 +- api/queryAPI/controllers.go | 19 +- api/queryAPI/documents.go | 87 ++++--- api/queryAPI/documents_test.go | 243 +++++++++++++++++- api/queryAPI/test_helpers.go | 37 ++- internal/database/README.md | 73 +++++- ...0000104_add_batch_storage_columns.down.sql | 5 + ...000000104_add_batch_storage_columns.up.sql | 11 + internal/database/queries/batch.sql | 20 ++ internal/database/repository/batch.sql.go | 135 +++++++++- internal/database/repository/models.go | 3 + internal/document/batch/service.go | 108 +++++--- internal/document/batch/service_test.go | 237 +++++++++++++++++ internal/document/batch/storage/service.go | 84 ++++++ .../document/batch/storage/service_test.go | 165 ++++++++++++ internal/test/aws.go | 11 + 17 files changed, 1272 insertions(+), 97 deletions(-) create mode 100644 api/queryAPI/README.md create mode 100644 internal/database/migrations/00000000000104_add_batch_storage_columns.down.sql create mode 100644 internal/database/migrations/00000000000104_add_batch_storage_columns.up.sql create mode 100644 internal/document/batch/storage/service.go create mode 100644 internal/document/batch/storage/service_test.go diff --git a/api/queryAPI/README.md b/api/queryAPI/README.md new file mode 100644 index 00000000..9d78e13f --- /dev/null +++ b/api/queryAPI/README.md @@ -0,0 +1,119 @@ +# queryAPI Package + +This directory contains the REST API implementation for the Query Orchestration platform. The code here serves as the HTTP interface layer between external clients and the internal business logic. + +## Architecture Overview + +### Controllers Struct (controllers.go) + +The `Controllers` struct is the central architectural component that implements the **Service Layer/Facade pattern** for the entire REST API. + +#### Purpose +- **OpenAPI Contract Implementation**: Implements the `ServerInterface` interface generated from `serviceAPIs/queryAPI.yaml` +- **Service Aggregation**: Consolidates 11 different business services into a single, cohesive API layer +- **Dependency Injection Container**: Provides clean separation between HTTP handlers and business logic + +#### Structure +```go +type Controllers struct { + svc *Services // Aggregated business services + cfg ConfigProvider // Combined auth & object storage configuration +} +``` + +#### Pattern Benefits +1. **Single Responsibility**: Each handler method focuses solely on HTTP↔Service translation +2. **Testability**: Services can be easily mocked for unit testing +3. **Modularity**: New services can be added without changing the controller pattern +4. **Type Safety**: Ensures compile-time checking against OpenAPI specification + +## Key Architectural Components + +### 1. Generated Code (api.gen.go) +- **DO NOT EDIT**: This file is auto-generated from `serviceAPIs/queryAPI.yaml` +- Contains all request/response types and the `ServerInterface` interface +- Regenerated via `task generate` command +- Ensures API contract consistency + +### 2. Handler Files +Each file implements specific domain endpoints: +- `client.go` - Client management operations +- `collector.go` - Query collector configuration +- `documents.go` - Document upload and management +- `export.go` - Export operations +- `query.go` - Query execution and management +- `status.go` - Health check and status endpoints +- `authHandlers.go` - Authentication flow handlers (login, logout, callbacks) + +### 3. Service Dependencies +The Controllers struct aggregates these internal services: +- **Export Service** - Manages export processes +- **Collector Service** - Handles query collections +- **Query Service** - Query execution logic +- **Client Service** - Client entity management +- **Document Service** - Document processing +- **DocumentUpload Service** - S3 upload handling +- **DocumentBatch Service** - Batch document operations + +### 4. Authentication Integration +- Uses AWS Cognito for JWT-based authentication +- Permit.io integration for RBAC (Role-Based Access Control) +- Authentication middleware validates tokens before reaching handlers +- `ConfigProvider` interface provides auth configuration access + +## Handler Pattern + +All handlers follow a consistent pattern: + +```go +func (s *Controllers) HandlerName(ctx echo.Context, params...) error { + // 1. Bind/validate request data + req := RequestType{} + if err := ctx.Bind(&req); err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err) + } + + // 2. Call appropriate service method + result, err := s.svc.ServiceName.Method(ctx.Request().Context(), ...) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, "Error message") + } + + // 3. Return formatted response + return ctx.JSON(http.StatusOK, ResponseType{...}) +} +``` + +## Testing Structure + +- `*_test.go` files provide unit tests for each handler group +- `test_helpers.go` contains shared test utilities and fixtures +- Tests use mocked services to validate handler logic in isolation +- Integration tests exist at the service layer level + +## Important Notes for Maintenance + +1. **OpenAPI-First Development**: Always update `serviceAPIs/queryAPI.yaml` first, then regenerate code +2. **Service Layer Separation**: Business logic belongs in `internal/` services, not in handlers +3. **Error Handling**: Use Echo's `HTTPError` for consistent error responses +4. **Context Propagation**: Always pass `ctx.Request().Context()` to service methods for proper cancellation +5. **No Direct Database Access**: All data operations must go through service layer +6. **Consistent Status Codes**: Follow REST conventions (201 for creation, 200 for success, 400 for client errors) + +## Development Workflow + +1. Modify OpenAPI spec in `serviceAPIs/queryAPI.yaml` +2. Run `task generate` to update generated code +3. Implement/update handler methods in Controllers +4. Add/update tests for new functionality +5. Run `task lint` to ensure code quality +6. Run `task test:functional` to validate changes + +## Configuration + +The `ConfigProvider` interface combines: +- Authentication configuration (Cognito settings, JWT validation) +- Object storage configuration (S3 bucket settings) +- Base service configuration (ports, logging, database) + +This unified configuration is injected into Controllers at startup, providing all necessary runtime settings. \ No newline at end of file diff --git a/api/queryAPI/authHandlers.go b/api/queryAPI/authHandlers.go index 3e0f8586..a22fdb1d 100644 --- a/api/queryAPI/authHandlers.go +++ b/api/queryAPI/authHandlers.go @@ -28,20 +28,20 @@ func (s *Controllers) Logout(ctx echo.Context) error { // the middleware will have already removed the auth cookie // Construct the Cognito logout URL with proper parameters - logoutURL := s.authConfig.GetAuthLogoutURL() - clientID := s.authConfig.GetAuthClientID() + logoutURL := s.cfg.GetAuthLogoutURL() + clientID := s.cfg.GetAuthClientID() // The logout_uri should redirect back to our home page // Use the same base URL that was used for redirect URI setup baseURL := ctx.Scheme() + "://" + ctx.Request().Host - logoutRedirectURI := baseURL + s.authConfig.GetAuthHomePath() + logoutRedirectURI := baseURL + s.cfg.GetAuthHomePath() // Build the full Cognito logout URL with query parameters parsedURL, err := url.Parse(logoutURL) if err != nil { - s.authConfig.GetAuthLogger().Error("Failed to parse logout URL", "error", err) + s.cfg.GetAuthLogger().Error("Failed to parse logout URL", "error", err) // Fallback to simple redirect to home - return ctx.Redirect(http.StatusFound, s.authConfig.GetAuthHomePath()) + return ctx.Redirect(http.StatusFound, s.cfg.GetAuthHomePath()) } query := parsedURL.Query() @@ -51,7 +51,7 @@ func (s *Controllers) Logout(ctx echo.Context) error { cognitoLogoutURL := parsedURL.String() - s.authConfig.GetAuthLogger().Debug("Redirecting to Cognito logout", + s.cfg.GetAuthLogger().Debug("Redirecting to Cognito logout", "logout_url", cognitoLogoutURL, "client_id", clientID, "logout_uri", logoutRedirectURI) diff --git a/api/queryAPI/controllers.go b/api/queryAPI/controllers.go index bbe5f035..bc6bba4c 100644 --- a/api/queryAPI/controllers.go +++ b/api/queryAPI/controllers.go @@ -13,6 +13,7 @@ import ( querytest "queryorchestration/internal/query/test" queryupdate "queryorchestration/internal/query/update" "queryorchestration/internal/serviceconfig/auth" + "queryorchestration/internal/serviceconfig/objectstore" ) const Name = "queryAPI" @@ -31,14 +32,20 @@ type Services struct { DocumentBatch *documentbatch.Service } -type Controllers struct { - svc *Services - authConfig auth.ConfigProvider +// ConfigProvider combines auth and objectstore interfaces for the Controllers +type ConfigProvider interface { + auth.ConfigProvider + objectstore.ConfigProvider } -func NewControllers(services *Services, authConfig auth.ConfigProvider) *Controllers { +type Controllers struct { + svc *Services + cfg ConfigProvider +} + +func NewControllers(services *Services, cfg ConfigProvider) *Controllers { return &Controllers{ - svc: services, - authConfig: authConfig, + svc: services, + cfg: cfg, } } diff --git a/api/queryAPI/documents.go b/api/queryAPI/documents.go index f3c32a8e..d0be4bf4 100644 --- a/api/queryAPI/documents.go +++ b/api/queryAPI/documents.go @@ -5,6 +5,7 @@ import ( "log/slog" "net/http" + "queryorchestration/internal/document/batch/storage" documentupload "queryorchestration/internal/document/upload" "github.com/google/uuid" @@ -128,48 +129,70 @@ func (s *Controllers) ListDocumentBatches(ctx echo.Context, clientId ClientID, p // UploadDocumentBatch handles ZIP archive upload containing multiple PDFs func (s *Controllers) UploadDocumentBatch(ctx echo.Context, clientId ClientID) error { - // Get the multipart form - form, err := ctx.MultipartForm() + // Parse multipart form to get ZIP file + file, err := ctx.FormFile("archive") if err != nil { - slog.Error("unable to get form data", "error", err) - return echo.NewHTTPError(http.StatusBadRequest, "unable to get form data") + slog.Error("failed to get uploaded file", "error", err, "client_id", clientId) + return echo.NewHTTPError(http.StatusBadRequest, "file upload required") } - // Get the archive file - files := form.File["archive"] - if len(files) == 0 { - return echo.NewHTTPError(http.StatusBadRequest, "missing archive file") - } else if len(files) > 1 { - return echo.NewHTTPError(http.StatusBadRequest, "too many files") - } - file := files[0] - - // Skip ZIP validation for tests - in production this would check file content - // if file.Header.Get("Content-Type") != "application/zip" { - // return echo.NewHTTPError(http.StatusBadRequest, "file must be a ZIP archive") - // } - - // For now, we'll create a batch record with placeholder total documents - // In the real implementation, we would: - // 1. Store the ZIP file to S3 - // 2. Start an async process to extract and process PDFs - // 3. Return immediately with 202 Accepted - - // Create batch record - batchID, err := s.svc.DocumentBatch.Create(ctx.Request().Context(), clientId, file.Filename, 0) + // Open uploaded file + src, err := file.Open() if err != nil { - slog.Error("unable to create batch", "error", err) - return echo.NewHTTPError(http.StatusInternalServerError, "unable to create batch") + slog.Error("failed to open uploaded file", "error", err, "client_id", clientId) + return echo.NewHTTPError(http.StatusInternalServerError, "failed to process uploaded file") + } + defer src.Close() + + // Initialize storage service with objectstore config + storageService := storage.New(s.cfg) + + // Validate ZIP file + err = storageService.ValidateZIPFile(file.Filename, file.Size) + if err != nil { + slog.Error("ZIP validation failed", "filename", file.Filename, "size", file.Size, "error", err, "client_id", clientId) + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid ZIP file: %v", err)) } - // Build response + // Store ZIP file to S3 + storageResult, err := storageService.StoreZIPFile(ctx.Request().Context(), string(clientId), file.Filename, src) + if err != nil { + slog.Error("failed to store ZIP file", "error", err, "client_id", clientId) + return echo.NewHTTPError(http.StatusInternalServerError, "failed to store ZIP file") + } + + // Create batch record with storage metadata + // Note: totalDocs is unknown at this point, will be determined during extraction in later parts + batchID, err := s.svc.DocumentBatch.CreateWithStorage( + ctx.Request().Context(), + string(clientId), + file.Filename, + 0, // totalDocs unknown in Part 1 + storageResult.Bucket, + storageResult.Key, + storageResult.SizeBytes, + ) + if err != nil { + slog.Error("failed to create batch record", "error", err, "client_id", clientId) + return echo.NewHTTPError(http.StatusInternalServerError, "failed to create batch record") + } + + slog.Info("ZIP batch upload completed", + "batch_id", batchID, + "filename", file.Filename, + "size_bytes", storageResult.SizeBytes, + "s3_key", storageResult.Key, + "client_id", clientId, + ) + + // Return 202 Accepted with batch ID and status URL + statusUrl := fmt.Sprintf("/clients/%s/documents/batches/%s", clientId, batchID) response := BatchUploadResponse{ BatchId: BatchID(batchID), - Status: BatchStatusProcessing, - StatusUrl: fmt.Sprintf("/client/%s/document/batch/%s", clientId, batchID.String()), + Status: "processing", + StatusUrl: statusUrl, } - // Return 202 Accepted for async processing return ctx.JSON(http.StatusAccepted, response) } diff --git a/api/queryAPI/documents_test.go b/api/queryAPI/documents_test.go index 8c385de2..9560d0d1 100644 --- a/api/queryAPI/documents_test.go +++ b/api/queryAPI/documents_test.go @@ -1,19 +1,27 @@ package queryapi_test import ( + "archive/zip" "bytes" + "encoding/json" + "fmt" + "io" "mime/multipart" "net/http" + "net/http/httptest" "path/filepath" "testing" + "time" - "queryorchestration/internal/database/repository" - "queryorchestration/internal/test" - - queryapi "queryorchestration/api/queryAPI" - + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/google/uuid" + "github.com/labstack/echo/v4" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + queryapi "queryorchestration/api/queryAPI" + "queryorchestration/internal/database/repository" + "queryorchestration/internal/test" ) func TestUploadDocument(t *testing.T) { @@ -175,6 +183,7 @@ func TestUploadDocumentBatch(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) + test.CreateAWSResources(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) @@ -283,3 +292,227 @@ func TestCancelDocumentBatch(t *testing.T) { require.NoError(t, err) assert.Equal(t, repository.BatchStatusCancelled, batch.Status) } + +func TestUploadDocumentBatch_Part1_Success(t *testing.T) { + // Setup: Create testcontainers (database + localstack) + cfg := &ControllerConfig{} + test.CreateDB(t, cfg) + initializeTestConfig(t, cfg) + test.CreateAWSResources(t, cfg) + + svc := createControllerServices(cfg) + cons := queryapi.NewControllers(svc, cfg) + + // Create test client with unique ID + clientID := fmt.Sprintf("test_client_part1_%d", time.Now().UnixNano()) + err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ + Clientid: clientID, + Name: "Test Client Part1", + }) + require.NoError(t, err) + + // Create test ZIP file with 3 PDFs (test1.pdf, test2.pdf, test3.pdf) + zipContent := createTestZIPFile(t, 3) + + // Create multipart form request + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + + part, err := writer.CreateFormFile("archive", "test-batch.zip") + require.NoError(t, err) + + _, err = io.Copy(part, zipContent) + require.NoError(t, err) + + err = writer.Close() + require.NoError(t, err) + + // Test: Upload ZIP batch + e := echo.New() + rec := httptest.NewRecorder() + ctx := e.NewContext( + httptest.NewRequest(http.MethodPost, "/", body), + rec, + ) + ctx.Request().Header.Set("Content-Type", writer.FormDataContentType()) + + err = cons.UploadDocumentBatch(ctx, clientID) + require.NoError(t, err) + + // Verify: HTTP response + assert.Equal(t, http.StatusAccepted, rec.Code) + + var response queryapi.BatchUploadResponse + err = json.Unmarshal(rec.Body.Bytes(), &response) + require.NoError(t, err) + + assert.NotEmpty(t, response.BatchId) + assert.Equal(t, queryapi.BatchStatusProcessing, response.Status) + assert.NotEmpty(t, response.StatusUrl) + + // Verify: Batch exists in database with storage metadata + batch, err := cfg.GetDBQueries().GetBatchUploadWithStorage(t.Context(), &repository.GetBatchUploadWithStorageParams{ + ID: uuid.UUID(response.BatchId), + ClientID: clientID, + }) + require.NoError(t, err) + + assert.Equal(t, "test-batch.zip", batch.OriginalFilename) + assert.Equal(t, repository.BatchStatusProcessing, batch.Status) + assert.NotNil(t, batch.ArchiveBucket) + assert.NotEmpty(t, *batch.ArchiveBucket) + assert.NotNil(t, batch.ArchiveKey) + assert.NotEmpty(t, *batch.ArchiveKey) + assert.Contains(t, *batch.ArchiveKey, clientID) + assert.Contains(t, *batch.ArchiveKey, "test-batch.zip") + assert.NotNil(t, batch.FileSizeBytes) + assert.Greater(t, *batch.FileSizeBytes, int64(0)) + + // Verify: ZIP file exists in S3 + s3Client := cfg.GetStoreClient() + _, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{ + Bucket: batch.ArchiveBucket, + Key: batch.ArchiveKey, + }) + require.NoError(t, err, "ZIP file should exist in S3") +} + +func TestUploadDocumentBatch_Part1_InvalidFile(t *testing.T) { + cfg := &ControllerConfig{} + test.CreateDB(t, cfg) + initializeTestConfig(t, cfg) + + svc := createControllerServices(cfg) + cons := queryapi.NewControllers(svc, cfg) + + // Create test client with unique ID + clientID := fmt.Sprintf("test_client_invalid_%d", time.Now().UnixNano()) + err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ + Clientid: clientID, + Name: "Test Client Invalid", + }) + require.NoError(t, err) + + // Create invalid file (not ZIP) + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + + part, err := writer.CreateFormFile("archive", "test.pdf") + require.NoError(t, err) + + _, err = part.Write([]byte("not a zip file")) + require.NoError(t, err) + + err = writer.Close() + require.NoError(t, err) + + // Test: Upload invalid file + ctx := echo.New().NewContext( + httptest.NewRequest(http.MethodPost, "/", body), + httptest.NewRecorder(), + ) + ctx.Request().Header.Set("Content-Type", writer.FormDataContentType()) + + err = cons.UploadDocumentBatch(ctx, clientID) + + // Verify: Error response + require.Error(t, err) + httpErr, ok := err.(*echo.HTTPError) + require.True(t, ok) + assert.Equal(t, http.StatusBadRequest, httpErr.Code) + assert.Contains(t, httpErr.Message, "invalid file type") +} + +func TestUploadDocumentBatch_Part1_TooLarge(t *testing.T) { + cfg := &ControllerConfig{} + test.CreateDB(t, cfg) + initializeTestConfig(t, cfg) + + svc := createControllerServices(cfg) + cons := queryapi.NewControllers(svc, cfg) + + // Create test client with unique ID + clientID := fmt.Sprintf("test_client_toolarge_%d", time.Now().UnixNano()) + err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ + Clientid: clientID, + Name: "Test Client TooLarge", + }) + require.NoError(t, err) + + // Create oversized ZIP file + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + + // Set a fake Content-Length header to simulate large file + part, err := writer.CreateFormFile("archive", "huge.zip") + require.NoError(t, err) + + // Write minimal content but set file size in header + _, err = part.Write([]byte("PK")) // ZIP file header + require.NoError(t, err) + + err = writer.Close() + require.NoError(t, err) + + // Test: Upload with modified request to simulate large file + req := httptest.NewRequest(http.MethodPost, "/", body) + req.Header.Set("Content-Type", writer.FormDataContentType()) + + // Override the file size in the form data + ctx := echo.New().NewContext(req, httptest.NewRecorder()) + + // Manually create form with large file size + form, _ := ctx.MultipartForm() + if len(form.File["archive"]) > 0 { + // Simulate 101MB file + form.File["archive"][0].Size = 101 * 1024 * 1024 + } + + err = cons.UploadDocumentBatch(ctx, clientID) + + // Verify: Error response + require.Error(t, err) + httpErr, ok := err.(*echo.HTTPError) + require.True(t, ok) + assert.Equal(t, http.StatusBadRequest, httpErr.Code) + assert.Contains(t, httpErr.Message, "file too large") +} + +// Helper function to create test ZIP with multiple PDFs +func createTestZIPFile(t *testing.T, pdfCount int) *bytes.Reader { + var buf bytes.Buffer + zipWriter := zip.NewWriter(&buf) + + // Create standard test PDF files + standardFiles := []string{"test1.pdf", "test2.pdf", "test3.pdf"} + + // If requested count is <= 3, use standard names + if pdfCount <= 3 { + for i := 0; i < pdfCount; i++ { + fileWriter, err := zipWriter.Create(standardFiles[i]) + require.NoError(t, err) + + // Write minimal PDF content + pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", standardFiles[i]) + _, err = fileWriter.Write([]byte(pdfContent)) + require.NoError(t, err) + } + } else { + // For more than 3 files, use numbered names + for i := 0; i < pdfCount; 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%%Document %d content\n%%%%EOF", i) + _, err = fileWriter.Write([]byte(pdfContent)) + require.NoError(t, err) + } + } + + err := zipWriter.Close() + require.NoError(t, err) + + return bytes.NewReader(buf.Bytes()) +} diff --git a/api/queryAPI/test_helpers.go b/api/queryAPI/test_helpers.go index 5d1413c2..17f31eea 100644 --- a/api/queryAPI/test_helpers.go +++ b/api/queryAPI/test_helpers.go @@ -1,6 +1,13 @@ package queryapi -import "log/slog" +import ( + "context" + "io" + "log/slog" + + "queryorchestration/internal/serviceconfig/aws" + "queryorchestration/internal/serviceconfig/objectstore" +) // mockAuthConfig provides a minimal implementation of auth.ConfigProvider for testing type mockAuthConfig struct{} @@ -47,6 +54,34 @@ func (m *mockAuthConfig) InitializeAuthConfig(string, *slog.Logger) error { retu func (m *mockAuthConfig) GetNoJWTValidation() bool { return false } func (m *mockAuthConfig) SetNoJWTValidation(bool) {} +// Objectstore interface methods +func (m *mockAuthConfig) GetAWSKeyID() string { return "test-key" } +func (m *mockAuthConfig) SetAWSKeyID(string) {} +func (m *mockAuthConfig) GetAWSSecretKey() string { return "test-secret" } +func (m *mockAuthConfig) SetAWSSecretKey(string) {} +func (m *mockAuthConfig) GetAWSSessionToken() string { return "" } +func (m *mockAuthConfig) SetAWSSessionToken(string) {} +func (m *mockAuthConfig) GetAWSRegion() string { return "us-east-1" } +func (m *mockAuthConfig) SetAWSRegion(string) {} +func (m *mockAuthConfig) GetAWSEndpoint() string { return "http://localhost:4566" } +func (m *mockAuthConfig) SetAWSEndpoint(string) {} +func (m *mockAuthConfig) SetStoreClient(context.Context) error { return nil } +func (m *mockAuthConfig) SetStoreClientFromS3Cfg(objectstore.S3Client) {} +func (m *mockAuthConfig) SetStoreClientAndPingByName(context.Context, string) error { return nil } +func (m *mockAuthConfig) GetStoreClient() objectstore.S3Client { return nil } +func (m *mockAuthConfig) PingStoreByName(context.Context, string) error { return nil } +func (m *mockAuthConfig) SetS3Endpoint(string) {} +func (m *mockAuthConfig) GetS3Endpoint() string { return "http://localhost:4566" } +func (m *mockAuthConfig) SetS3UsePathStyle(bool) {} +func (m *mockAuthConfig) CalculateETag(context.Context, io.Reader) (string, error) { + return "test-etag", nil +} +func (m *mockAuthConfig) GetDirectoryPart(uint16, int64) uint16 { return 0 } +func (m *mockAuthConfig) GetBucket() string { return "test-bucket" } +func (m *mockAuthConfig) SetBucket(string) {} +func (m *mockAuthConfig) GetAWSProfile() aws.Profile { return aws.Profile("default") } +func (m *mockAuthConfig) SetAWSProfile(aws.Profile) {} + // NewTestControllers creates a Controllers instance for testing with a mock auth config func NewTestControllers(services *Services) *Controllers { return NewControllers(services, &mockAuthConfig{}) diff --git a/internal/database/README.md b/internal/database/README.md index a89adb94..33c3ad02 100644 --- a/internal/database/README.md +++ b/internal/database/README.md @@ -6,6 +6,46 @@ Package database provides the data access layer for the query orchestration plat ## Architecture +### Service Layer Architecture + +The database access follows a clean, layered architecture pattern: + +1. **Repository Layer** (`repository/`) + - SQLC-generated code providing type-safe database queries + - `Queries` struct contains all database operations + - Generated from SQL files in `queries/` directory + - Accessed throughout codebase via `ConfigProvider.GetDBQueries()` + +2. **Service Layer** (`internal/*/service.go`) + - Each domain has its own service containing business logic + - Services use the repository layer for all database access + - Examples: `client/service.go`, `document/service.go`, `query/service.go` + +3. **Configuration Layer** (`internal/serviceconfig/database/`) + - Manages database connections and connection pooling + - Provides transaction support via `ExecuteDBTransaction()` + - Handles migration execution on startup + +4. **Database Access Pattern** + ``` + API Controllers → Service Layer (business logic) + ↓ + ConfigProvider.GetDBQueries() + ↓ + Repository Layer (SQLC-generated) + ↓ + PostgreSQL Database + ``` + +### Architecture Principles + +- **No Direct SQL in Business Logic**: All SQL queries are defined in `.sql` files +- **Type Safety**: SQLC generates Go types from SQL schemas +- **Clean Separation**: HTTP handlers → Services → Repository → Database +- **Transaction Support**: Coordinated through ConfigProvider interface +- **Code Generation**: Repository code is generated, not hand-written +- **Single Source of Truth**: Database schema defined in migrations + ### Core Components #### Database Schema @@ -16,6 +56,7 @@ The database schema supports document processing and query orchestration with th - **Collectors**: Query collection specifications per client - **Results**: Processed query outputs linked to documents - **Text Extractions**: OCR results from document processing +- **Batch Uploads**: Batch document upload tracking with progress monitoring #### Migration System - Uses `golang-migrate` for schema version management @@ -53,6 +94,7 @@ The database schema supports document processing and query orchestration with th - **results**: Query execution results per document - **collectors**: Client-specific query collections - **collectorVersions**: Collector version history +- **batch_uploads**: Batch document upload tracking with progress and failure metrics ### Views @@ -83,21 +125,20 @@ The database schema supports document processing and query orchestration with th // Create database configuration dbConfig := &serviceconfig.DBConfig{ DBHost: "localhost", - DBPort: "5432", + DBPort: 5432, DBName: "querydb", DBUser: "dbuser", DBPassword: "dbpass", } -// Get connection pool -pool, err := dbConfig.GetDBPool(ctx) +// Initialize connection pool +err := dbConfig.SetDBPoolConfig() if err != nil { return err } -defer pool.Close() -// Create repository queries instance -queries := repository.New(pool) +// Get repository queries instance +queries := dbConfig.GetDBQueries() ``` ### Transaction Management @@ -168,6 +209,15 @@ err := migrator.Up() ## SQLC Integration +### Configuration +SQLC is configured via `sqlc.yml` in the project root with: +- PostgreSQL engine with pgx/v5 driver +- Queries sourced from `internal/database/queries/` +- Schema from `internal/database/migrations/` +- Generated code output to `internal/database/repository/` +- Custom rules enforcing query performance and safety +- UUID type mapping to `github.com/google/uuid` + ### Query Definition Format ```sql @@ -187,12 +237,15 @@ SET status = $2, updated = CURRENT_TIMESTAMP WHERE id = $1; ``` -### Generated Code +### Generated Code Features SQLC generates: -- Type-safe query functions +- Type-safe query functions with parameter structs - Struct definitions matching table schemas -- Null handling with sql.NullString, sql.NullTime, etc. -- Parameter validation +- Null handling with pointers for nullable fields +- UUID type support via google/uuid package +- SQL comments included in generated code +- Enum validation methods +- Empty slice initialization for array returns ## Testing diff --git a/internal/database/migrations/00000000000104_add_batch_storage_columns.down.sql b/internal/database/migrations/00000000000104_add_batch_storage_columns.down.sql new file mode 100644 index 00000000..8126eccd --- /dev/null +++ b/internal/database/migrations/00000000000104_add_batch_storage_columns.down.sql @@ -0,0 +1,5 @@ +-- Remove storage columns and constraint +ALTER TABLE batch_uploads DROP CONSTRAINT IF EXISTS batch_storage_required; +ALTER TABLE batch_uploads DROP COLUMN IF EXISTS file_size_bytes; +ALTER TABLE batch_uploads DROP COLUMN IF EXISTS archive_key; +ALTER TABLE batch_uploads DROP COLUMN IF EXISTS archive_bucket; \ No newline at end of file diff --git a/internal/database/migrations/00000000000104_add_batch_storage_columns.up.sql b/internal/database/migrations/00000000000104_add_batch_storage_columns.up.sql new file mode 100644 index 00000000..02acef97 --- /dev/null +++ b/internal/database/migrations/00000000000104_add_batch_storage_columns.up.sql @@ -0,0 +1,11 @@ +-- Add storage metadata columns to batch_uploads table +ALTER TABLE batch_uploads ADD COLUMN archive_bucket text; +ALTER TABLE batch_uploads ADD COLUMN archive_key text; +ALTER TABLE batch_uploads ADD COLUMN file_size_bytes bigint; + +-- Add constraint to ensure storage metadata is consistent (if any storage field is set, bucket and key must both be set) +ALTER TABLE batch_uploads ADD CONSTRAINT batch_storage_consistent +CHECK ( + (archive_bucket IS NULL AND archive_key IS NULL AND file_size_bytes IS NULL) OR + (archive_bucket IS NOT NULL AND archive_key IS NOT NULL) +); \ No newline at end of file diff --git a/internal/database/queries/batch.sql b/internal/database/queries/batch.sql index d0c36f59..0e035697 100644 --- a/internal/database/queries/batch.sql +++ b/internal/database/queries/batch.sql @@ -2,6 +2,18 @@ INSERT INTO batch_uploads (client_id, original_filename, total_documents) VALUES ($1, $2, $3) RETURNING id; +-- name: CreateBatchUploadWithStorage :one +INSERT INTO batch_uploads ( + client_id, + original_filename, + total_documents, + status, + archive_bucket, + archive_key, + file_size_bytes +) VALUES ($1, $2, $3, $4::batch_status, $5, $6, $7) +RETURNING id, created_at; + -- name: GetBatchUpload :one SELECT id, client_id, original_filename, total_documents, processed_documents, failed_documents, invalid_type_documents, status, progress_percent, @@ -9,6 +21,14 @@ SELECT id, client_id, original_filename, total_documents, processed_documents, FROM batch_uploads WHERE id = $1 AND client_id = $2; +-- name: GetBatchUploadWithStorage :one +SELECT id, client_id, original_filename, total_documents, + processed_documents, failed_documents, invalid_type_documents, + status, archive_bucket, archive_key, file_size_bytes, + failed_filenames, created_at, completed_at +FROM batch_uploads +WHERE id = $1 AND client_id = $2; + -- name: ListBatchUploads :many SELECT id, client_id, original_filename, total_documents, processed_documents, failed_documents, invalid_type_documents, status, progress_percent, diff --git a/internal/database/repository/batch.sql.go b/internal/database/repository/batch.sql.go index e1355308..90a2ef87 100644 --- a/internal/database/repository/batch.sql.go +++ b/internal/database/repository/batch.sql.go @@ -96,6 +96,61 @@ func (q *Queries) CreateBatchUpload(ctx context.Context, arg *CreateBatchUploadP return id, err } +const createBatchUploadWithStorage = `-- name: CreateBatchUploadWithStorage :one +INSERT INTO batch_uploads ( + client_id, + original_filename, + total_documents, + status, + archive_bucket, + archive_key, + file_size_bytes +) VALUES ($1, $2, $3, $4::batch_status, $5, $6, $7) +RETURNING id, created_at +` + +type CreateBatchUploadWithStorageParams struct { + ClientID string `db:"client_id"` + OriginalFilename string `db:"original_filename"` + TotalDocuments int32 `db:"total_documents"` + Column4 BatchStatus `db:"column_4"` + ArchiveBucket *string `db:"archive_bucket"` + ArchiveKey *string `db:"archive_key"` + FileSizeBytes *int64 `db:"file_size_bytes"` +} + +type CreateBatchUploadWithStorageRow struct { + ID uuid.UUID `db:"id"` + CreatedAt pgtype.Timestamp `db:"created_at"` +} + +// CreateBatchUploadWithStorage +// +// INSERT INTO batch_uploads ( +// client_id, +// original_filename, +// total_documents, +// status, +// archive_bucket, +// archive_key, +// file_size_bytes +// ) VALUES ($1, $2, $3, $4::batch_status, $5, $6, $7) +// RETURNING id, created_at +func (q *Queries) CreateBatchUploadWithStorage(ctx context.Context, arg *CreateBatchUploadWithStorageParams) (*CreateBatchUploadWithStorageRow, error) { + row := q.db.QueryRow(ctx, createBatchUploadWithStorage, + arg.ClientID, + arg.OriginalFilename, + arg.TotalDocuments, + arg.Column4, + arg.ArchiveBucket, + arg.ArchiveKey, + arg.FileSizeBytes, + ) + var i CreateBatchUploadWithStorageRow + err := row.Scan(&i.ID, &i.CreatedAt) + return &i, err +} + const getBatchUpload = `-- name: GetBatchUpload :one SELECT id, client_id, original_filename, total_documents, processed_documents, failed_documents, invalid_type_documents, status, progress_percent, @@ -109,6 +164,21 @@ type GetBatchUploadParams struct { ClientID string `db:"client_id"` } +type GetBatchUploadRow struct { + ID uuid.UUID `db:"id"` + ClientID string `db:"client_id"` + OriginalFilename string `db:"original_filename"` + TotalDocuments int32 `db:"total_documents"` + ProcessedDocuments int32 `db:"processed_documents"` + FailedDocuments int32 `db:"failed_documents"` + InvalidTypeDocuments int32 `db:"invalid_type_documents"` + Status BatchStatus `db:"status"` + ProgressPercent int32 `db:"progress_percent"` + FailedFilenames []byte `db:"failed_filenames"` + CreatedAt pgtype.Timestamp `db:"created_at"` + CompletedAt pgtype.Timestamp `db:"completed_at"` +} + // GetBatchUpload // // SELECT id, client_id, original_filename, total_documents, processed_documents, @@ -116,9 +186,9 @@ type GetBatchUploadParams struct { // failed_filenames, created_at, completed_at // FROM batch_uploads // WHERE id = $1 AND client_id = $2 -func (q *Queries) GetBatchUpload(ctx context.Context, arg *GetBatchUploadParams) (*BatchUpload, error) { +func (q *Queries) GetBatchUpload(ctx context.Context, arg *GetBatchUploadParams) (*GetBatchUploadRow, error) { row := q.db.QueryRow(ctx, getBatchUpload, arg.ID, arg.ClientID) - var i BatchUpload + var i GetBatchUploadRow err := row.Scan( &i.ID, &i.ClientID, @@ -136,6 +206,67 @@ func (q *Queries) GetBatchUpload(ctx context.Context, arg *GetBatchUploadParams) return &i, err } +const getBatchUploadWithStorage = `-- name: GetBatchUploadWithStorage :one +SELECT id, client_id, original_filename, total_documents, + processed_documents, failed_documents, invalid_type_documents, + status, archive_bucket, archive_key, file_size_bytes, + failed_filenames, created_at, completed_at +FROM batch_uploads +WHERE id = $1 AND client_id = $2 +` + +type GetBatchUploadWithStorageParams struct { + ID uuid.UUID `db:"id"` + ClientID string `db:"client_id"` +} + +type GetBatchUploadWithStorageRow struct { + ID uuid.UUID `db:"id"` + ClientID string `db:"client_id"` + OriginalFilename string `db:"original_filename"` + TotalDocuments int32 `db:"total_documents"` + ProcessedDocuments int32 `db:"processed_documents"` + FailedDocuments int32 `db:"failed_documents"` + InvalidTypeDocuments int32 `db:"invalid_type_documents"` + Status BatchStatus `db:"status"` + ArchiveBucket *string `db:"archive_bucket"` + ArchiveKey *string `db:"archive_key"` + FileSizeBytes *int64 `db:"file_size_bytes"` + FailedFilenames []byte `db:"failed_filenames"` + CreatedAt pgtype.Timestamp `db:"created_at"` + CompletedAt pgtype.Timestamp `db:"completed_at"` +} + +// GetBatchUploadWithStorage +// +// SELECT id, client_id, original_filename, total_documents, +// processed_documents, failed_documents, invalid_type_documents, +// status, archive_bucket, archive_key, file_size_bytes, +// failed_filenames, created_at, completed_at +// FROM batch_uploads +// WHERE id = $1 AND client_id = $2 +func (q *Queries) GetBatchUploadWithStorage(ctx context.Context, arg *GetBatchUploadWithStorageParams) (*GetBatchUploadWithStorageRow, error) { + row := q.db.QueryRow(ctx, getBatchUploadWithStorage, arg.ID, arg.ClientID) + var i GetBatchUploadWithStorageRow + err := row.Scan( + &i.ID, + &i.ClientID, + &i.OriginalFilename, + &i.TotalDocuments, + &i.ProcessedDocuments, + &i.FailedDocuments, + &i.InvalidTypeDocuments, + &i.Status, + &i.ArchiveBucket, + &i.ArchiveKey, + &i.FileSizeBytes, + &i.FailedFilenames, + &i.CreatedAt, + &i.CompletedAt, + ) + return &i, err +} + const getDocumentsByBatchId = `-- name: GetDocumentsByBatchId :many SELECT id, clientId, hash FROM documents diff --git a/internal/database/repository/models.go b/internal/database/repository/models.go index 2e9580a4..e9d71fdb 100644 --- a/internal/database/repository/models.go +++ b/internal/database/repository/models.go @@ -245,6 +245,9 @@ type BatchUpload struct { FailedFilenames []byte `db:"failed_filenames"` CreatedAt pgtype.Timestamp `db:"created_at"` CompletedAt pgtype.Timestamp `db:"completed_at"` + ArchiveBucket *string `db:"archive_bucket"` + ArchiveKey *string `db:"archive_key"` + FileSizeBytes *int64 `db:"file_size_bytes"` } type Client struct { diff --git a/internal/document/batch/service.go b/internal/document/batch/service.go index 14bddf79..72f4b514 100644 --- a/internal/document/batch/service.go +++ b/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 +} diff --git a/internal/document/batch/service_test.go b/internal/document/batch/service_test.go index bd96207b..ee18fc72 100644 --- a/internal/document/batch/service_test.go +++ b/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()) +} diff --git a/internal/document/batch/storage/service.go b/internal/document/batch/storage/service.go new file mode 100644 index 00000000..77dca942 --- /dev/null +++ b/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 +} diff --git a/internal/document/batch/storage/service_test.go b/internal/document/batch/storage/service_test.go new file mode 100644 index 00000000..0fa2c698 --- /dev/null +++ b/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()) +} diff --git a/internal/test/aws.go b/internal/test/aws.go index 3dd74fd8..1abd68c9 100644 --- a/internal/test/aws.go +++ b/internal/test/aws.go @@ -96,3 +96,14 @@ func CreateAWSContainer(t testing.TB, cfg AWSConfigProvider) *AWSContainerConfig ExternalEndpoint: awsAddress, } } + +// CreateAWSResources sets up AWS resources for testing +func CreateAWSResources(t testing.TB, cfg AWSConfigProvider) *AWSContainerConfig { + awsConfig := CreateAWSContainer(t, cfg) + + // Set up the S3 client and bucket + SetStoreClient(t, t.Context(), cfg, awsConfig.ExternalEndpoint) + CreateBucket(t, cfg) + + return awsConfig +}