diff --git a/api/docCleanRunner/runner.go b/api/docCleanRunner/runner.go index 34fa200f..730583d3 100644 --- a/api/docCleanRunner/runner.go +++ b/api/docCleanRunner/runner.go @@ -5,6 +5,7 @@ import ( "log/slog" "queryorchestration/internal/database/repository" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentclean "queryorchestration/internal/document/clean" @@ -21,6 +22,7 @@ const Name = "docCleanRunner" type Services struct { Clean *documentclean.Service Outcome *outcome.Reporter + Cleanup *foldercleanup.Service } // Runner processes document clean messages from the SQS queue. @@ -63,10 +65,18 @@ func (s Runner) Process(ctx context.Context, body Body) bool { outcomeName = repository.BatchOutcomeStatusCleanFailed errorDetail = result.FailReason } - if updateErr := s.svc.Outcome.Update(ctx, body.DocumentID, outcomeName, errorDetail); updateErr != nil { + batchIDs, updateErr := s.svc.Outcome.Update(ctx, body.DocumentID, outcomeName, errorDetail) + if updateErr != nil { slog.Error("failed to update batch outcome", "document_id", body.DocumentID, "outcome", outcomeName, "error", updateErr) } + if s.svc.Cleanup != nil { + for _, batchID := range batchIDs { + if _, err := s.svc.Cleanup.TryFinalizeBatch(ctx, batchID); err != nil { + slog.Error("failed to finalize batch folder cleanup", "batch_id", batchID, "error", err) + } + } + } return true // validation pass or fail -> acknowledge (don't requeue bad PDFs) } diff --git a/api/docInitRunner/runner.go b/api/docInitRunner/runner.go index 672a0610..e76df453 100644 --- a/api/docInitRunner/runner.go +++ b/api/docInitRunner/runner.go @@ -7,6 +7,7 @@ import ( "log/slog" "queryorchestration/internal/database/repository" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentinit "queryorchestration/internal/document/init" "queryorchestration/internal/serviceconfig/objectstore" @@ -26,6 +27,7 @@ type Services struct { Document *documentinit.Service Queries *repository.Queries Outcome *outcome.Reporter + Cleanup *foldercleanup.Service } // Runner processes document init messages from the SQS queue. @@ -108,6 +110,13 @@ func (s Runner) Process(ctx context.Context, body Body) bool { "document_id", result.ID, "outcome", outcomeName, "error", resolveErr) // Non-fatal: document was created successfully } + if outcomeName == repository.BatchOutcomeStatusInitDuplicate { + if s.svc.Cleanup != nil && result.BatchID != nil { + if _, err := s.svc.Cleanup.TryFinalizeBatch(ctx, *result.BatchID); err != nil { + slog.Error("failed to finalize batch folder cleanup", "batch_id", *result.BatchID, "error", err) + } + } + } return true } diff --git a/api/docSyncRunner/runner.go b/api/docSyncRunner/runner.go index b1be0006..28c00cbd 100644 --- a/api/docSyncRunner/runner.go +++ b/api/docSyncRunner/runner.go @@ -5,6 +5,7 @@ import ( "log/slog" "queryorchestration/internal/database/repository" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentsync "queryorchestration/internal/document/sync" @@ -21,6 +22,7 @@ const Name = "docSyncRunner" type Services struct { Document *documentsync.Service Outcome *outcome.Reporter + Cleanup *foldercleanup.Service } // Runner processes document sync messages from the SQS queue. @@ -59,10 +61,18 @@ func (s Runner) Process(ctx context.Context, body Body) bool { if result.Skipped { outcomeName = repository.BatchOutcomeStatusSyncSkipped } - if updateErr := s.svc.Outcome.Update(ctx, body.DocumentID, outcomeName, nil); updateErr != nil { + batchIDs, updateErr := s.svc.Outcome.Update(ctx, body.DocumentID, outcomeName, nil) + if updateErr != nil { slog.Error("failed to update batch outcome", "document_id", body.DocumentID, "outcome", outcomeName, "error", updateErr) } + if outcomeName == repository.BatchOutcomeStatusSyncSkipped && s.svc.Cleanup != nil { + for _, batchID := range batchIDs { + if _, err := s.svc.Cleanup.TryFinalizeBatch(ctx, batchID); err != nil { + slog.Error("failed to finalize batch folder cleanup", "batch_id", batchID, "error", err) + } + } + } return true } diff --git a/cmd/docCleanRunner/main.go b/cmd/docCleanRunner/main.go index d1b84c54..1cc79c2b 100644 --- a/cmd/docCleanRunner/main.go +++ b/cmd/docCleanRunner/main.go @@ -12,6 +12,7 @@ import ( "os" doccleanrunner "queryorchestration/api/docCleanRunner" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentclean "queryorchestration/internal/document/clean" "queryorchestration/internal/server/runner" @@ -42,6 +43,7 @@ func main() { return doccleanrunner.New(&doccleanrunner.Services{ Clean: clean, Outcome: outcome.New(cfg.GetDBQueries()), + Cleanup: foldercleanup.New(cfg), }) } diff --git a/cmd/docInitRunner/main.go b/cmd/docInitRunner/main.go index c61de5bb..9a46baa7 100644 --- a/cmd/docInitRunner/main.go +++ b/cmd/docInitRunner/main.go @@ -21,6 +21,7 @@ import ( "os" docinitrunner "queryorchestration/api/docInitRunner" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentinit "queryorchestration/internal/document/init" "queryorchestration/internal/server/runner" @@ -52,6 +53,7 @@ func main() { Document: docinit, Queries: cfg.GetDBQueries(), Outcome: outcome.New(cfg.GetDBQueries()), + Cleanup: foldercleanup.New(cfg), }) } diff --git a/cmd/docSyncRunner/main.go b/cmd/docSyncRunner/main.go index a796db97..4d6924bb 100644 --- a/cmd/docSyncRunner/main.go +++ b/cmd/docSyncRunner/main.go @@ -17,6 +17,7 @@ import ( docsyncrunner "queryorchestration/api/docSyncRunner" "queryorchestration/internal/client" "queryorchestration/internal/document" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/document/batch/outcome" documentsync "queryorchestration/internal/document/sync" "queryorchestration/internal/server/runner" @@ -50,6 +51,7 @@ func main() { return docsyncrunner.New(&docsyncrunner.Services{ Document: docsync, Outcome: outcome.New(cfg.GetDBQueries()), + Cleanup: foldercleanup.New(cfg), }) } diff --git a/deployments/compose.local.yaml b/deployments/compose.local.yaml index 217be4f6..1b6f5889 100644 --- a/deployments/compose.local.yaml +++ b/deployments/compose.local.yaml @@ -300,6 +300,12 @@ services: COGNITO_REGION: ${COGNITO_REGION} PERMIT_IO_PDP_URL: http://permit_pdp:7000 PERMIT_IO_TENANT: ${PERMIT_IO_TENANT} + LOCAL_DEV_CLIENT_ID: ${LOCAL_DEV_CLIENT_ID} + LOCAL_DEV_CLIENT_NAME: ${LOCAL_DEV_CLIENT_NAME} + LOCAL_DEV_EULA_VERSION: ${LOCAL_DEV_EULA_VERSION} + LOCAL_DEV_EULA_TITLE: ${LOCAL_DEV_EULA_TITLE} + LOCAL_DEV_EULA_AUTO_ACCEPT_SUBJECT: ${LOCAL_DEV_EULA_AUTO_ACCEPT_SUBJECT} + LOCAL_DEV_EULA_AUTO_ACCEPT_EMAIL: ${LOCAL_DEV_EULA_AUTO_ACCEPT_EMAIL} BUCKET: ${BUCKET_IN} CHATBOT_SERVICE_URL: ${CHATBOT_SERVICE_URL} CHATBOT_API_KEY: ${CHATBOT_API_KEY} diff --git a/docs/ai.generated/04-data-architecture.md b/docs/ai.generated/04-data-architecture.md index bf4018b4..cd3b07f7 100644 --- a/docs/ai.generated/04-data-architecture.md +++ b/docs/ai.generated/04-data-architecture.md @@ -28,6 +28,7 @@ erDiagram batch_uploads ||--o{ documents : "contains" batch_uploads ||--o{ documentUploads : "tracks uploads" batch_uploads ||--o{ batch_document_outcomes : "tracks outcomes" + batch_uploads ||--o{ batch_folder_candidates : "tracks cleanup candidates" folders ||--o{ documents : "contains" folders ||--o{ folders : "has children" @@ -487,6 +488,7 @@ CREATE INDEX idx_batch_uploads_status ON batch_uploads(status); - Temporal tracking with creation and completion timestamps - Storage metadata tracking (archive bucket/key and file size) - Storage consistency constraint ensures bucket and key are both set or both null +- `folder_cleanup_completed_at` records when backend folder cleanup has finalized for the batch - Indexed by client_id and status for efficient queries - Note: Uses snake_case naming convention (unlike other tables) @@ -530,6 +532,29 @@ CREATE INDEX idx_batch_doc_outcomes_document_id ON batch_document_outcomes(docum - Indexed by batch_id for efficient batch detail lookups and by document_id for pipeline updates - Added in migration 124 +### Batch Folder Candidates (`batch_folder_candidates`) +Tracks folders touched by legacy ZIP batch extraction so the backend can remove empty auto-created folders after all file outcomes are terminal. + +```sql +CREATE TABLE batch_folder_candidates ( + batch_id uuid NOT NULL, + folder_id uuid NOT NULL, + path text NOT NULL, + existed_before boolean NOT NULL, + created_at timestamp NOT NULL DEFAULT NOW(), + PRIMARY KEY (batch_id, folder_id), + FOREIGN KEY (batch_id) REFERENCES batch_uploads(id) ON DELETE CASCADE, + FOREIGN KEY (folder_id) REFERENCES folders(id) ON DELETE CASCADE +); +``` + +**Batch Folder Cleanup Features**: +- Records each non-root folder segment touched by a batch upload +- Preserves `existed_before=false` once any batch creates the folder, enabling idempotent retries +- Cleanup finalizes only after terminal batch outcome evidence exists; zero-outcome failed batches stay cleanup-pending +- Safe deletion never removes `/`, pre-existing folders, renamed folders, folders with documents/uploads/children/chatbot scopes, or folders still needed by another pending batch +- Added in migration 132 + ### Folders (`folders`) Virtual folder hierarchy for document organization. Folders are database-only entities with no direct relationship to S3 storage locations. diff --git a/docs/ai.generated/07-configuration-management.md b/docs/ai.generated/07-configuration-management.md index 30bdf9ee..47ec30fb 100644 --- a/docs/ai.generated/07-configuration-management.md +++ b/docs/ai.generated/07-configuration-management.md @@ -91,6 +91,14 @@ Auth path customization: - `COGNITO_LOGOUT_PATH` (default `/logout`) - `HEALTH_PATH` (default `/health`) +Local development bootstrap, active only when `DISABLE_AUTH=true`: +- `LOCAL_DEV_CLIENT_ID` (optional) - Creates this client at queryAPI startup when it does not already exist. +- `LOCAL_DEV_CLIENT_NAME` (optional) - Display name for `LOCAL_DEV_CLIENT_ID`; defaults to the client ID. +- `LOCAL_DEV_EULA_VERSION` (optional) - Creates or updates this local EULA version and marks it current. +- `LOCAL_DEV_EULA_TITLE` (optional) - Title for the local EULA; defaults to `AArete DoczyAI Terms of Use`. +- `LOCAL_DEV_EULA_AUTO_ACCEPT_SUBJECT` (optional) - Cognito subject/user key to mark as accepted for the local EULA. +- `LOCAL_DEV_EULA_AUTO_ACCEPT_EMAIL` (optional) - Email for the auto-accepted user; defaults to the subject value. + ## Rate Limiting From `internal/serviceconfig/ratelimit/config.go`: @@ -149,4 +157,3 @@ reference, including the grace-window floor. - Some deployment files still include legacy query/text environment variables because deprecated compatibility services remain in compose definitions. - For API/runtime behavior, prefer `serviceAPIs/queryAPI.yaml` plus `internal/serviceconfig/*` as source of truth. - diff --git a/docs/ai.generated/08-operations-guide.md b/docs/ai.generated/08-operations-guide.md index 73421435..f1d29667 100644 --- a/docs/ai.generated/08-operations-guide.md +++ b/docs/ai.generated/08-operations-guide.md @@ -446,6 +446,33 @@ kubectl get pods -l app=query-orchestration task health:check:all ``` +### Batch Folder Cleanup Operations + +Legacy ZIP batch uploads record auto-created folder candidates in `batch_folder_candidates`. Cleanup finalization is run by the batch worker after terminal batch status changes and by the QueryAPI background worker through `FinalizeReadyBatches`. + +Operational checks: + +```sql +-- Batches waiting for folder cleanup +SELECT id, client_id, status, total_documents, folder_cleanup_completed_at +FROM batch_uploads +WHERE folder_cleanup_completed_at IS NULL + AND id IN (SELECT batch_id FROM batch_folder_candidates); + +-- Candidate folders for one batch +SELECT batch_id, folder_id, path, existed_before +FROM batch_folder_candidates +WHERE batch_id = ''; + +-- Outcome evidence used by cleanup readiness +SELECT outcome, COUNT(*) +FROM batch_document_outcomes +WHERE batch_id = '' +GROUP BY outcome; +``` + +Cleanup should not complete for failed batches with `total_documents = 0` and no outcome rows. That state usually means outcome persistence failed before the worker could record a terminal per-file result. All-skipped archives are handled separately with a synthetic `failed_upload` outcome so the batch can fail visibly and cleanup can remove empty auto-created folders. + ## Capacity Planning ### Resource Requirements diff --git a/docs/ai.generated/batch-outcome-tracking-guide.md b/docs/ai.generated/batch-outcome-tracking-guide.md index 8cf10d61..5a0ff225 100644 --- a/docs/ai.generated/batch-outcome-tracking-guide.md +++ b/docs/ai.generated/batch-outcome-tracking-guide.md @@ -148,7 +148,7 @@ Each terminal outcome indicates a specific action for the caller: - **`failed_open`** / **`failed_read`** -- The file inside the ZIP was damaged or unreadable. Re-create the ZIP archive and resubmit. -- **`failed_s3_upload`** / **`failed_upload`** -- An infrastructure error occurred during processing. Retry the entire batch. +- **`failed_s3_upload`** / **`failed_upload`** -- An infrastructure error occurred during processing. `failed_upload` is also used for ZIP archives that contain no processable files after directories, hidden metadata, and unsafe paths are skipped. - **`sync_skipped`** -- The client's sync configuration prevents processing. Contact an administrator. @@ -181,3 +181,11 @@ ZIP extracted ``` Non-terminal outcomes (`submitted`, `init_complete`, `sync_complete`) indicate the document is still being processed. Poll the batch details endpoint to see updates as files progress through each stage. + +## 7. Folder Cleanup Coupling + +Legacy ZIP batch uploads still create folders during extraction. For batch uploads, the backend records touched folder segments in `batch_folder_candidates` and finalizes cleanup when the batch has terminal outcome evidence. + +Accepted terminal outcomes (`clean_passed`, `sync_skipped`) keep their folder assignments. Rejected outcomes (`duplicate`, `init_duplicate`, `invalid_type`, `failed_open`, `failed_read`, `failed_s3_upload`, `failed_upload`, `clean_failed`) detach current-batch documents/uploads from auto-created candidate folders. Empty auto-created folders are then deleted only when they are not root, did not pre-exist, were not renamed, have no documents/uploads/children/chatbot scopes, and are not pending cleanup for another batch. + +If outcome persistence fails before any outcome row exists, cleanup remains pending instead of marking `folder_cleanup_completed_at`. This keeps failed zero-outcome batches visible for retry/diagnosis and prevents cleanup from closing before rejected artifacts can be detached. diff --git a/docs/ai.generated/service.details.md b/docs/ai.generated/service.details.md index f9459fe8..a9b34512 100644 --- a/docs/ai.generated/service.details.md +++ b/docs/ai.generated/service.details.md @@ -89,6 +89,25 @@ This document describes current service behavior in the processing pipeline. **Output**: `{document_id}` to `DOCSYNC` for each client document. +--- + +### 6. batch folder cleanup +**Purpose**: Remove empty folders auto-created by legacy ZIP batch extraction when all file outcomes are terminal. + +**Code locations**: +- `internal/document/batch/foldercleanup/` +- `internal/server/api/batch_worker.go` +- `internal/database/queries/batch.sql` + +**Behavior**: +- Locks the batch row before finalizing cleanup +- Requires terminal batch status plus persisted outcome evidence +- Detaches rejected current-batch documents and upload rows from candidate folders +- Deletes only safe empty auto-created folders +- Retries ready cleanup from the QueryAPI background worker + +**Output**: updates `batch_uploads.folder_cleanup_completed_at`; may null rejected `documents.folderId` and `documentUploads.folder_id`; may delete safe candidate folders. + ## Deprecated Compatibility Services These services are retained for deployment compatibility and are currently health-only/no-op: @@ -105,6 +124,7 @@ S3 Upload Event -> docInitRunner -> docSyncRunner -> docCleanRunner + -> batch folder cleanup (for legacy ZIP batches) ``` ```text @@ -113,4 +133,3 @@ Client Sync Trigger -> docSyncRunner -> docCleanRunner ``` - diff --git a/docs/ai.generated/software-architecture.md b/docs/ai.generated/software-architecture.md index 525aa7e8..66890bb2 100644 --- a/docs/ai.generated/software-architecture.md +++ b/docs/ai.generated/software-architecture.md @@ -368,6 +368,7 @@ s.cfg.GetQueueClient() | `models.go` | Generated Go types from database schema | | `client.sql.go` | Generated client query methods | | `document.sql.go` | Generated document query methods | +| `batch.sql.go` | Generated batch, outcome, and folder cleanup query methods | | etc. | One file per SQL source file | **SQL Query Definition** (`internal/database/queries/client.sql`): @@ -406,9 +407,13 @@ q.WithTx(tx) **Examples**: - `00000000000003_clients.up.sql` - Create clients table - `00000000000109_create_folders_table.up.sql` - Add folders feature +- `00000000000124_batch_document_outcomes.up.sql` - Track per-file batch outcomes +- `00000000000132_batch_folder_candidates.up.sql` - Track batch folder cleanup candidates **Migration Tool**: golang-migrate +Batch cleanup queries live in `internal/database/queries/batch.sql` with the rest of the batch upload workflow. Keep raw production SQL there so SQLC regenerates typed methods for readiness checks, rejected artifact detachment, cleanup completion, and safe candidate folder deletion. + --- ### 9. Configuration Layer (`internal/serviceconfig/`) diff --git a/internal/database/migrations/00000000000132_batch_folder_candidates.down.sql b/internal/database/migrations/00000000000132_batch_folder_candidates.down.sql new file mode 100644 index 00000000..5fc4e806 --- /dev/null +++ b/internal/database/migrations/00000000000132_batch_folder_candidates.down.sql @@ -0,0 +1,8 @@ +DROP INDEX IF EXISTS idx_batch_uploads_folder_cleanup_pending; + +ALTER TABLE batch_uploads +DROP COLUMN IF EXISTS folder_cleanup_completed_at; + +DROP INDEX IF EXISTS idx_batch_folder_candidates_folder_id; +DROP INDEX IF EXISTS idx_batch_folder_candidates_batch_id; +DROP TABLE IF EXISTS batch_folder_candidates; diff --git a/internal/database/migrations/00000000000132_batch_folder_candidates.up.sql b/internal/database/migrations/00000000000132_batch_folder_candidates.up.sql new file mode 100644 index 00000000..7682c4ea --- /dev/null +++ b/internal/database/migrations/00000000000132_batch_folder_candidates.up.sql @@ -0,0 +1,18 @@ +CREATE TABLE batch_folder_candidates ( + batch_id uuid NOT NULL REFERENCES batch_uploads(id) ON DELETE CASCADE, + folder_id uuid NOT NULL REFERENCES folders(id) ON DELETE CASCADE, + path text NOT NULL, + existed_before boolean NOT NULL, + created_at timestamp NOT NULL DEFAULT NOW(), + PRIMARY KEY (batch_id, folder_id) +); + +CREATE INDEX idx_batch_folder_candidates_batch_id ON batch_folder_candidates(batch_id); +CREATE INDEX idx_batch_folder_candidates_folder_id ON batch_folder_candidates(folder_id); + +ALTER TABLE batch_uploads +ADD COLUMN folder_cleanup_completed_at timestamp; + +CREATE INDEX idx_batch_uploads_folder_cleanup_pending +ON batch_uploads(created_at) +WHERE folder_cleanup_completed_at IS NULL; diff --git a/internal/database/queries/batch.sql b/internal/database/queries/batch.sql index 5cf0c8b8..d5fa3f9d 100644 --- a/internal/database/queries/batch.sql +++ b/internal/database/queries/batch.sql @@ -55,6 +55,11 @@ SET status = $2::batch_status, completed_at = CASE WHEN $2::batch_status IN ('completed', 'failed', 'cancelled') THEN NOW() ELSE NULL END WHERE id = $1; +-- name: SetBatchCompletedAt :exec +UPDATE batch_uploads +SET completed_at = @completed_at +WHERE id = @id; + -- name: AddFailedFilename :exec UPDATE batch_uploads SET failed_filenames = failed_filenames || jsonb_build_array($2::text) @@ -89,7 +94,7 @@ SET outcome = EXCLUDED.outcome, document_id = EXCLUDED.document_id, updated_at = NOW(); --- name: UpdateBatchDocumentOutcomeByDocumentID :exec +-- name: UpdateBatchDocumentOutcomeByDocumentID :many -- Updates the outcome for a document progressing through the pipeline. -- Only updates rows in non-terminal pipeline states to prevent overwriting -- terminal outcomes (e.g. a 'duplicate' row from a different batch that @@ -99,7 +104,8 @@ SET outcome = $2::batch_outcome_status, error_detail = $3, updated_at = NOW() WHERE document_id = $1 - AND outcome IN ('submitted', 'init_complete', 'sync_complete'); + AND outcome IN ('submitted', 'init_complete', 'sync_complete') +RETURNING batch_id; -- name: ResolveBatchDocumentOutcome :exec -- Called by docInitRunner to set the document_id on a previously-submitted @@ -128,4 +134,145 @@ SELECT FROM batch_document_outcomes bdo LEFT JOIN currentCleanEntries cce ON cce.documentId = bdo.document_id WHERE bdo.batch_id = $1 -ORDER BY bdo.created_at ASC; \ No newline at end of file +ORDER BY bdo.created_at ASC; + +-- name: RecordBatchFolderCandidate :exec +-- Records a non-root folder touched by a batch upload. Once a batch observes a +-- folder as newly created, retries must preserve that fact. +INSERT INTO batch_folder_candidates (batch_id, folder_id, path, existed_before) +VALUES (@batch_id, @folder_id, @path, @existed_before) +ON CONFLICT (batch_id, folder_id) DO UPDATE +SET path = EXCLUDED.path, + existed_before = batch_folder_candidates.existed_before AND EXCLUDED.existed_before; + +-- name: ListBatchFolderCandidates :many +SELECT batch_id, folder_id, path, existed_before, created_at +FROM batch_folder_candidates +WHERE batch_id = @batch_id +ORDER BY path; + +-- name: LockBatchUploadForFolderCleanup :one +SELECT id, status, total_documents, created_at, completed_at, folder_cleanup_completed_at +FROM batch_uploads +WHERE id = @batch_id +FOR UPDATE; + +-- name: CountBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = @batch_id; + +-- name: CountNonTerminalBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = @batch_id + AND outcome IN ('submitted', 'init_complete', 'sync_complete'); + +-- name: CountAcceptedBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = @batch_id + AND outcome IN ('clean_passed', 'sync_skipped'); + +-- name: ListFolderCleanupReadyBatchIDs :many +-- @sqlc-vet-disable +SELECT bu.id +FROM batch_uploads bu +WHERE bu.folder_cleanup_completed_at IS NULL + AND bu.status IN ('completed', 'failed', 'cancelled') + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = bu.id + ) + AND ( + ( + bu.total_documents > 0 + AND + (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) = bu.total_documents + AND NOT EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = bu.id + AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete') + ) + ) + OR ( + bu.status IN ('failed', 'cancelled') + AND (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) > 0 + AND bu.completed_at IS NOT NULL + AND bu.completed_at <= NOW() - (sqlc.arg(stale_after_seconds)::int * INTERVAL '1 second') + ) + ) +ORDER BY bu.created_at ASC +LIMIT sqlc.arg(limit_count); + +-- name: DetachRejectedBatchDocumentsFromCandidateFolders :execrows +-- @sqlc-vet-disable +UPDATE documents d +SET folderId = NULL +WHERE d.batch_id = @batch_id + AND d.folderId IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = @batch_id + AND bfc.folder_id = d.folderId + ) + AND EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = @batch_id + AND bdo.document_id = d.id + AND ( + bdo.outcome IN ( + 'duplicate', + 'init_duplicate', + 'invalid_type', + 'failed_open', + 'failed_read', + 'failed_s3_upload', + 'failed_upload', + 'clean_failed' + ) + OR (@include_nonterminal::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) + ) + ); + +-- name: DetachRejectedBatchDocumentUploadsFromCandidateFolders :execrows +-- @sqlc-vet-disable +UPDATE documentUploads du +SET folder_id = NULL +WHERE du.batch_id = @batch_id + AND du.folder_id IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = @batch_id + AND bfc.folder_id = du.folder_id + ) + AND EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = @batch_id + AND bdo.filename = du.filename + AND ( + bdo.outcome IN ( + 'duplicate', + 'init_duplicate', + 'invalid_type', + 'failed_open', + 'failed_read', + 'failed_s3_upload', + 'failed_upload', + 'clean_failed' + ) + OR (@include_nonterminal::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) + ) + ); + +-- name: MarkBatchFolderCleanupCompleted :exec +UPDATE batch_uploads +SET folder_cleanup_completed_at = NOW() +WHERE id = @batch_id + AND folder_cleanup_completed_at IS NULL; diff --git a/internal/database/queries/eula.sql b/internal/database/queries/eula.sql index 7a188f5c..a6e2b8ff 100644 --- a/internal/database/queries/eula.sql +++ b/internal/database/queries/eula.sql @@ -108,3 +108,30 @@ FROM eulaAgreements WHERE eulaVersionId = $1 ORDER BY agreedAt DESC LIMIT $2 OFFSET $3; + +-- name: ActivateLocalDevEulaVersion :one +-- Activates or updates the local development EULA version and makes it current. +WITH cleared AS ( + UPDATE eulaVersions + SET isCurrent = FALSE + WHERE isCurrent = TRUE +), +upserted AS ( + INSERT INTO eulaVersions (version, title, content, effectiveDate, createdBy, isCurrent, activatedAt, activatedBy) + VALUES (@version, @title, @content, NOW(), 'local-dev', TRUE, NOW(), 'local-dev') + ON CONFLICT (version) DO UPDATE + SET title = EXCLUDED.title, + content = EXCLUDED.content, + effectiveDate = EXCLUDED.effectiveDate, + isCurrent = TRUE, + activatedAt = NOW(), + activatedBy = EXCLUDED.activatedBy + RETURNING id +) +SELECT id FROM upserted; + +-- name: CreateLocalDevEulaAgreement :exec +-- Records an idempotent local development EULA agreement. +INSERT INTO eulaAgreements (cognitoSubjectId, userEmail, eulaVersionId, agreedFromIp) +VALUES (@cognito_subject_id, @user_email, @eula_version_id, '127.0.0.1') +ON CONFLICT (cognitoSubjectId, eulaVersionId) DO NOTHING; diff --git a/internal/database/queries/folders.sql b/internal/database/queries/folders.sql index 860e4e54..ba30fd6b 100644 --- a/internal/database/queries/folders.sql +++ b/internal/database/queries/folders.sql @@ -88,6 +88,15 @@ VALUES ($1, $2, $3, $4) ON CONFLICT (clientId, path) DO UPDATE SET path = EXCLUDED.path RETURNING *; +-- name: CreateFolderIfAbsent :one +-- Insert a folder only when it does not already exist. +-- Batch uploads use this to distinguish newly-created folder paths from +-- pre-existing paths without mutating existing folders. +INSERT INTO folders (path, parentId, clientId, createdBy) +VALUES (@path, @parent_id, @client_id, @created_by) +ON CONFLICT (clientId, path) DO NOTHING +RETURNING *; + -- name: GetDocumentsByFolderEnriched :many -- Gets documents in a folder with all enriched metadata fields SELECT @@ -180,3 +189,35 @@ WITH RECURSIVE folder_tree AS ( JOIN folder_tree ft ON f.parentId = ft.id ) DELETE FROM folders WHERE id IN (SELECT id FROM folder_tree); + +-- name: ListSafeAutoCreatedFolderCandidates :many +-- @sqlc-vet-disable +-- Candidate folders can be deleted only after every batch that touched that +-- folder has finished folder cleanup. The path equality check protects renamed +-- folders from deletion. +SELECT DISTINCT f.id, f.path +FROM folders f +JOIN batch_folder_candidates bfc ON bfc.folder_id = f.id +WHERE bfc.existed_before = false + AND bfc.path = f.path + AND f.path <> '/' + AND NOT EXISTS ( + SELECT 1 + FROM batch_folder_candidates other_bfc + JOIN batch_uploads bu ON bu.id = other_bfc.batch_id + WHERE other_bfc.folder_id = f.id + AND bu.folder_cleanup_completed_at IS NULL + ) +ORDER BY f.path DESC +LIMIT sqlc.arg(limit_count); + +-- name: DeleteSafeAutoCreatedFolder :execrows +-- @sqlc-vet-disable +-- Delete one auto-created folder only when it is currently empty and no longer +-- referenced by documents, uploads, or child folders. +DELETE FROM folders f +WHERE f.id = @folder_id + AND f.path <> '/' + AND NOT EXISTS (SELECT 1 FROM documents d WHERE d.folderId = f.id) + AND NOT EXISTS (SELECT 1 FROM documentUploads du WHERE du.folder_id = f.id) + AND NOT EXISTS (SELECT 1 FROM folders child WHERE child.parentId = f.id); diff --git a/internal/database/repository/batch.sql.go b/internal/database/repository/batch.sql.go index b28b960a..1c745848 100644 --- a/internal/database/repository/batch.sql.go +++ b/internal/database/repository/batch.sql.go @@ -33,6 +33,44 @@ func (q *Queries) AddFailedFilename(ctx context.Context, arg *AddFailedFilenameP return err } +const countAcceptedBatchOutcomes = `-- name: CountAcceptedBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = $1 + AND outcome IN ('clean_passed', 'sync_skipped') +` + +// CountAcceptedBatchOutcomes +// +// SELECT COUNT(*)::int +// FROM batch_document_outcomes +// WHERE batch_id = $1 +// AND outcome IN ('clean_passed', 'sync_skipped') +func (q *Queries) CountAcceptedBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) { + row := q.db.QueryRow(ctx, countAcceptedBatchOutcomes, batchID) + var column_1 int32 + err := row.Scan(&column_1) + return column_1, err +} + +const countBatchOutcomes = `-- name: CountBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = $1 +` + +// CountBatchOutcomes +// +// SELECT COUNT(*)::int +// FROM batch_document_outcomes +// WHERE batch_id = $1 +func (q *Queries) CountBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) { + row := q.db.QueryRow(ctx, countBatchOutcomes, batchID) + var column_1 int32 + err := row.Scan(&column_1) + return column_1, err +} + const countDocumentsByBatchId = `-- name: CountDocumentsByBatchId :one SELECT COUNT(*) as count FROM documents @@ -51,6 +89,26 @@ func (q *Queries) CountDocumentsByBatchId(ctx context.Context, batchID *uuid.UUI return count, err } +const countNonTerminalBatchOutcomes = `-- name: CountNonTerminalBatchOutcomes :one +SELECT COUNT(*)::int +FROM batch_document_outcomes +WHERE batch_id = $1 + AND outcome IN ('submitted', 'init_complete', 'sync_complete') +` + +// CountNonTerminalBatchOutcomes +// +// SELECT COUNT(*)::int +// FROM batch_document_outcomes +// WHERE batch_id = $1 +// AND outcome IN ('submitted', 'init_complete', 'sync_complete') +func (q *Queries) CountNonTerminalBatchOutcomes(ctx context.Context, batchID uuid.UUID) (int32, error) { + row := q.db.QueryRow(ctx, countNonTerminalBatchOutcomes, batchID) + var column_1 int32 + err := row.Scan(&column_1) + return column_1, err +} + const createBatchUpload = `-- name: CreateBatchUpload :one INSERT INTO batch_uploads (client_id, original_filename, total_documents) VALUES ($1, $2, $3) RETURNING id @@ -128,6 +186,158 @@ func (q *Queries) CreateBatchUploadWithStorage(ctx context.Context, arg *CreateB return &i, err } +const detachRejectedBatchDocumentUploadsFromCandidateFolders = `-- name: DetachRejectedBatchDocumentUploadsFromCandidateFolders :execrows +UPDATE documentUploads du +SET folder_id = NULL +WHERE du.batch_id = $1 + AND du.folder_id IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = $1 + AND bfc.folder_id = du.folder_id + ) + AND EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = $1 + AND bdo.filename = du.filename + AND ( + bdo.outcome IN ( + 'duplicate', + 'init_duplicate', + 'invalid_type', + 'failed_open', + 'failed_read', + 'failed_s3_upload', + 'failed_upload', + 'clean_failed' + ) + OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) + ) + ) +` + +type DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams struct { + BatchID *uuid.UUID `db:"batch_id"` + IncludeNonterminal bool `db:"include_nonterminal"` +} + +// @sqlc-vet-disable +// +// UPDATE documentUploads du +// SET folder_id = NULL +// WHERE du.batch_id = $1 +// AND du.folder_id IS NOT NULL +// AND EXISTS ( +// SELECT 1 +// FROM batch_folder_candidates bfc +// WHERE bfc.batch_id = $1 +// AND bfc.folder_id = du.folder_id +// ) +// AND EXISTS ( +// SELECT 1 +// FROM batch_document_outcomes bdo +// WHERE bdo.batch_id = $1 +// AND bdo.filename = du.filename +// AND ( +// bdo.outcome IN ( +// 'duplicate', +// 'init_duplicate', +// 'invalid_type', +// 'failed_open', +// 'failed_read', +// 'failed_s3_upload', +// 'failed_upload', +// 'clean_failed' +// ) +// OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) +// ) +// ) +func (q *Queries) DetachRejectedBatchDocumentUploadsFromCandidateFolders(ctx context.Context, arg *DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams) (int64, error) { + result, err := q.db.Exec(ctx, detachRejectedBatchDocumentUploadsFromCandidateFolders, arg.BatchID, arg.IncludeNonterminal) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const detachRejectedBatchDocumentsFromCandidateFolders = `-- name: DetachRejectedBatchDocumentsFromCandidateFolders :execrows +UPDATE documents d +SET folderId = NULL +WHERE d.batch_id = $1 + AND d.folderId IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = $1 + AND bfc.folder_id = d.folderId + ) + AND EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = $1 + AND bdo.document_id = d.id + AND ( + bdo.outcome IN ( + 'duplicate', + 'init_duplicate', + 'invalid_type', + 'failed_open', + 'failed_read', + 'failed_s3_upload', + 'failed_upload', + 'clean_failed' + ) + OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) + ) + ) +` + +type DetachRejectedBatchDocumentsFromCandidateFoldersParams struct { + BatchID *uuid.UUID `db:"batch_id"` + IncludeNonterminal bool `db:"include_nonterminal"` +} + +// @sqlc-vet-disable +// +// UPDATE documents d +// SET folderId = NULL +// WHERE d.batch_id = $1 +// AND d.folderId IS NOT NULL +// AND EXISTS ( +// SELECT 1 +// FROM batch_folder_candidates bfc +// WHERE bfc.batch_id = $1 +// AND bfc.folder_id = d.folderId +// ) +// AND EXISTS ( +// SELECT 1 +// FROM batch_document_outcomes bdo +// WHERE bdo.batch_id = $1 +// AND bdo.document_id = d.id +// AND ( +// bdo.outcome IN ( +// 'duplicate', +// 'init_duplicate', +// 'invalid_type', +// 'failed_open', +// 'failed_read', +// 'failed_s3_upload', +// 'failed_upload', +// 'clean_failed' +// ) +// OR ($2::boolean AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete')) +// ) +// ) +func (q *Queries) DetachRejectedBatchDocumentsFromCandidateFolders(ctx context.Context, arg *DetachRejectedBatchDocumentsFromCandidateFoldersParams) (int64, error) { + result, err := q.db.Exec(ctx, detachRejectedBatchDocumentsFromCandidateFolders, arg.BatchID, arg.IncludeNonterminal) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const getBatchUpload = `-- name: GetBatchUpload :one SELECT id, client_id, original_filename, total_documents, processed_documents, failed_documents, invalid_type_documents, status, progress_percent, @@ -459,6 +669,45 @@ func (q *Queries) ListBatchDocumentOutcomes(ctx context.Context, batchID uuid.UU return items, nil } +const listBatchFolderCandidates = `-- name: ListBatchFolderCandidates :many +SELECT batch_id, folder_id, path, existed_before, created_at +FROM batch_folder_candidates +WHERE batch_id = $1 +ORDER BY path +` + +// ListBatchFolderCandidates +// +// SELECT batch_id, folder_id, path, existed_before, created_at +// FROM batch_folder_candidates +// WHERE batch_id = $1 +// ORDER BY path +func (q *Queries) ListBatchFolderCandidates(ctx context.Context, batchID uuid.UUID) ([]*BatchFolderCandidate, error) { + rows, err := q.db.Query(ctx, listBatchFolderCandidates, batchID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []*BatchFolderCandidate{} + for rows.Next() { + var i BatchFolderCandidate + if err := rows.Scan( + &i.BatchID, + &i.FolderID, + &i.Path, + &i.ExistedBefore, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const listBatchUploads = `-- name: ListBatchUploads :many SELECT id, client_id, original_filename, total_documents, processed_documents, failed_documents, invalid_type_documents, status, progress_percent, @@ -530,6 +779,183 @@ func (q *Queries) ListBatchUploads(ctx context.Context, arg *ListBatchUploadsPar return items, nil } +const listFolderCleanupReadyBatchIDs = `-- name: ListFolderCleanupReadyBatchIDs :many +SELECT bu.id +FROM batch_uploads bu +WHERE bu.folder_cleanup_completed_at IS NULL + AND bu.status IN ('completed', 'failed', 'cancelled') + AND EXISTS ( + SELECT 1 + FROM batch_folder_candidates bfc + WHERE bfc.batch_id = bu.id + ) + AND ( + ( + bu.total_documents > 0 + AND + (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) = bu.total_documents + AND NOT EXISTS ( + SELECT 1 + FROM batch_document_outcomes bdo + WHERE bdo.batch_id = bu.id + AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete') + ) + ) + OR ( + bu.status IN ('failed', 'cancelled') + AND (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) > 0 + AND bu.completed_at IS NOT NULL + AND bu.completed_at <= NOW() - ($1::int * INTERVAL '1 second') + ) + ) +ORDER BY bu.created_at ASC +LIMIT $2 +` + +type ListFolderCleanupReadyBatchIDsParams struct { + StaleAfterSeconds int32 `db:"stale_after_seconds"` + LimitCount int64 `db:"limit_count"` +} + +// @sqlc-vet-disable +// +// SELECT bu.id +// FROM batch_uploads bu +// WHERE bu.folder_cleanup_completed_at IS NULL +// AND bu.status IN ('completed', 'failed', 'cancelled') +// AND EXISTS ( +// SELECT 1 +// FROM batch_folder_candidates bfc +// WHERE bfc.batch_id = bu.id +// ) +// AND ( +// ( +// bu.total_documents > 0 +// AND +// (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) = bu.total_documents +// AND NOT EXISTS ( +// SELECT 1 +// FROM batch_document_outcomes bdo +// WHERE bdo.batch_id = bu.id +// AND bdo.outcome IN ('submitted', 'init_complete', 'sync_complete') +// ) +// ) +// OR ( +// bu.status IN ('failed', 'cancelled') +// AND (SELECT COUNT(*) FROM batch_document_outcomes bdo WHERE bdo.batch_id = bu.id) > 0 +// AND bu.completed_at IS NOT NULL +// AND bu.completed_at <= NOW() - ($1::int * INTERVAL '1 second') +// ) +// ) +// ORDER BY bu.created_at ASC +// LIMIT $2 +func (q *Queries) ListFolderCleanupReadyBatchIDs(ctx context.Context, arg *ListFolderCleanupReadyBatchIDsParams) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, listFolderCleanupReadyBatchIDs, arg.StaleAfterSeconds, arg.LimitCount) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const lockBatchUploadForFolderCleanup = `-- name: LockBatchUploadForFolderCleanup :one +SELECT id, status, total_documents, created_at, completed_at, folder_cleanup_completed_at +FROM batch_uploads +WHERE id = $1 +FOR UPDATE +` + +type LockBatchUploadForFolderCleanupRow struct { + ID uuid.UUID `db:"id"` + Status BatchStatus `db:"status"` + TotalDocuments int32 `db:"total_documents"` + CreatedAt pgtype.Timestamp `db:"created_at"` + CompletedAt pgtype.Timestamp `db:"completed_at"` + FolderCleanupCompletedAt pgtype.Timestamp `db:"folder_cleanup_completed_at"` +} + +// LockBatchUploadForFolderCleanup +// +// SELECT id, status, total_documents, created_at, completed_at, folder_cleanup_completed_at +// FROM batch_uploads +// WHERE id = $1 +// FOR UPDATE +func (q *Queries) LockBatchUploadForFolderCleanup(ctx context.Context, batchID uuid.UUID) (*LockBatchUploadForFolderCleanupRow, error) { + row := q.db.QueryRow(ctx, lockBatchUploadForFolderCleanup, batchID) + var i LockBatchUploadForFolderCleanupRow + err := row.Scan( + &i.ID, + &i.Status, + &i.TotalDocuments, + &i.CreatedAt, + &i.CompletedAt, + &i.FolderCleanupCompletedAt, + ) + return &i, err +} + +const markBatchFolderCleanupCompleted = `-- name: MarkBatchFolderCleanupCompleted :exec +UPDATE batch_uploads +SET folder_cleanup_completed_at = NOW() +WHERE id = $1 + AND folder_cleanup_completed_at IS NULL +` + +// MarkBatchFolderCleanupCompleted +// +// UPDATE batch_uploads +// SET folder_cleanup_completed_at = NOW() +// WHERE id = $1 +// AND folder_cleanup_completed_at IS NULL +func (q *Queries) MarkBatchFolderCleanupCompleted(ctx context.Context, batchID uuid.UUID) error { + _, err := q.db.Exec(ctx, markBatchFolderCleanupCompleted, batchID) + return err +} + +const recordBatchFolderCandidate = `-- name: RecordBatchFolderCandidate :exec +INSERT INTO batch_folder_candidates (batch_id, folder_id, path, existed_before) +VALUES ($1, $2, $3, $4) +ON CONFLICT (batch_id, folder_id) DO UPDATE +SET path = EXCLUDED.path, + existed_before = batch_folder_candidates.existed_before AND EXCLUDED.existed_before +` + +type RecordBatchFolderCandidateParams struct { + BatchID uuid.UUID `db:"batch_id"` + FolderID uuid.UUID `db:"folder_id"` + Path string `db:"path"` + ExistedBefore bool `db:"existed_before"` +} + +// Records a non-root folder touched by a batch upload. Once a batch observes a +// folder as newly created, retries must preserve that fact. +// +// INSERT INTO batch_folder_candidates (batch_id, folder_id, path, existed_before) +// VALUES ($1, $2, $3, $4) +// ON CONFLICT (batch_id, folder_id) DO UPDATE +// SET path = EXCLUDED.path, +// existed_before = batch_folder_candidates.existed_before AND EXCLUDED.existed_before +func (q *Queries) RecordBatchFolderCandidate(ctx context.Context, arg *RecordBatchFolderCandidateParams) error { + _, err := q.db.Exec(ctx, recordBatchFolderCandidate, + arg.BatchID, + arg.FolderID, + arg.Path, + arg.ExistedBefore, + ) + return err +} + const resolveBatchDocumentOutcome = `-- name: ResolveBatchDocumentOutcome :exec UPDATE batch_document_outcomes SET document_id = $3, @@ -568,6 +994,27 @@ func (q *Queries) ResolveBatchDocumentOutcome(ctx context.Context, arg *ResolveB return err } +const setBatchCompletedAt = `-- name: SetBatchCompletedAt :exec +UPDATE batch_uploads +SET completed_at = $1 +WHERE id = $2 +` + +type SetBatchCompletedAtParams struct { + CompletedAt pgtype.Timestamp `db:"completed_at"` + ID uuid.UUID `db:"id"` +} + +// SetBatchCompletedAt +// +// UPDATE batch_uploads +// SET completed_at = $1 +// WHERE id = $2 +func (q *Queries) SetBatchCompletedAt(ctx context.Context, arg *SetBatchCompletedAtParams) error { + _, err := q.db.Exec(ctx, setBatchCompletedAt, arg.CompletedAt, arg.ID) + return err +} + const setBatchTotalDocuments = `-- name: SetBatchTotalDocuments :exec UPDATE batch_uploads SET total_documents = $2 WHERE id = $1 ` @@ -585,13 +1032,14 @@ func (q *Queries) SetBatchTotalDocuments(ctx context.Context, arg *SetBatchTotal return err } -const updateBatchDocumentOutcomeByDocumentID = `-- name: UpdateBatchDocumentOutcomeByDocumentID :exec +const updateBatchDocumentOutcomeByDocumentID = `-- name: UpdateBatchDocumentOutcomeByDocumentID :many UPDATE batch_document_outcomes SET outcome = $2::batch_outcome_status, error_detail = $3, updated_at = NOW() WHERE document_id = $1 AND outcome IN ('submitted', 'init_complete', 'sync_complete') +RETURNING batch_id ` type UpdateBatchDocumentOutcomeByDocumentIDParams struct { @@ -611,9 +1059,25 @@ type UpdateBatchDocumentOutcomeByDocumentIDParams struct { // updated_at = NOW() // WHERE document_id = $1 // AND outcome IN ('submitted', 'init_complete', 'sync_complete') -func (q *Queries) UpdateBatchDocumentOutcomeByDocumentID(ctx context.Context, arg *UpdateBatchDocumentOutcomeByDocumentIDParams) error { - _, err := q.db.Exec(ctx, updateBatchDocumentOutcomeByDocumentID, arg.DocumentID, arg.Column2, arg.ErrorDetail) - return err +// RETURNING batch_id +func (q *Queries) UpdateBatchDocumentOutcomeByDocumentID(ctx context.Context, arg *UpdateBatchDocumentOutcomeByDocumentIDParams) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, updateBatchDocumentOutcomeByDocumentID, arg.DocumentID, arg.Column2, arg.ErrorDetail) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var batch_id uuid.UUID + if err := rows.Scan(&batch_id); err != nil { + return nil, err + } + items = append(items, batch_id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const updateBatchProgress = `-- name: UpdateBatchProgress :exec diff --git a/internal/database/repository/eula.sql.go b/internal/database/repository/eula.sql.go index 1b302f1e..eb430ef5 100644 --- a/internal/database/repository/eula.sql.go +++ b/internal/database/repository/eula.sql.go @@ -12,6 +12,60 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) +const activateLocalDevEulaVersion = `-- name: ActivateLocalDevEulaVersion :one +WITH cleared AS ( + UPDATE eulaVersions + SET isCurrent = FALSE + WHERE isCurrent = TRUE +), +upserted AS ( + INSERT INTO eulaVersions (version, title, content, effectiveDate, createdBy, isCurrent, activatedAt, activatedBy) + VALUES ($1, $2, $3, NOW(), 'local-dev', TRUE, NOW(), 'local-dev') + ON CONFLICT (version) DO UPDATE + SET title = EXCLUDED.title, + content = EXCLUDED.content, + effectiveDate = EXCLUDED.effectiveDate, + isCurrent = TRUE, + activatedAt = NOW(), + activatedBy = EXCLUDED.activatedBy + RETURNING id +) +SELECT id FROM upserted +` + +type ActivateLocalDevEulaVersionParams struct { + Version string `db:"version"` + Title string `db:"title"` + Content string `db:"content"` +} + +// Activates or updates the local development EULA version and makes it current. +// +// WITH cleared AS ( +// UPDATE eulaVersions +// SET isCurrent = FALSE +// WHERE isCurrent = TRUE +// ), +// upserted AS ( +// INSERT INTO eulaVersions (version, title, content, effectiveDate, createdBy, isCurrent, activatedAt, activatedBy) +// VALUES ($1, $2, $3, NOW(), 'local-dev', TRUE, NOW(), 'local-dev') +// ON CONFLICT (version) DO UPDATE +// SET title = EXCLUDED.title, +// content = EXCLUDED.content, +// effectiveDate = EXCLUDED.effectiveDate, +// isCurrent = TRUE, +// activatedAt = NOW(), +// activatedBy = EXCLUDED.activatedBy +// RETURNING id +// ) +// SELECT id FROM upserted +func (q *Queries) ActivateLocalDevEulaVersion(ctx context.Context, arg *ActivateLocalDevEulaVersionParams) (uuid.UUID, error) { + row := q.db.QueryRow(ctx, activateLocalDevEulaVersion, arg.Version, arg.Title, arg.Content) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + const clearCurrentEulaVersions = `-- name: ClearCurrentEulaVersions :exec UPDATE eulaVersions SET isCurrent = FALSE @@ -167,6 +221,28 @@ func (q *Queries) CreateEulaVersion(ctx context.Context, arg *CreateEulaVersionP return &i, err } +const createLocalDevEulaAgreement = `-- name: CreateLocalDevEulaAgreement :exec +INSERT INTO eulaAgreements (cognitoSubjectId, userEmail, eulaVersionId, agreedFromIp) +VALUES ($1, $2, $3, '127.0.0.1') +ON CONFLICT (cognitoSubjectId, eulaVersionId) DO NOTHING +` + +type CreateLocalDevEulaAgreementParams struct { + CognitoSubjectID string `db:"cognito_subject_id"` + UserEmail string `db:"user_email"` + EulaVersionID uuid.UUID `db:"eula_version_id"` +} + +// Records an idempotent local development EULA agreement. +// +// INSERT INTO eulaAgreements (cognitoSubjectId, userEmail, eulaVersionId, agreedFromIp) +// VALUES ($1, $2, $3, '127.0.0.1') +// ON CONFLICT (cognitoSubjectId, eulaVersionId) DO NOTHING +func (q *Queries) CreateLocalDevEulaAgreement(ctx context.Context, arg *CreateLocalDevEulaAgreementParams) error { + _, err := q.db.Exec(ctx, createLocalDevEulaAgreement, arg.CognitoSubjectID, arg.UserEmail, arg.EulaVersionID) + return err +} + const getCurrentEulaVersion = `-- name: GetCurrentEulaVersion :one SELECT id, version, title, content, effectiveDate, createdAt, createdBy, isCurrent, activatedAt, activatedBy FROM eulaVersions diff --git a/internal/database/repository/folders.sql.go b/internal/database/repository/folders.sql.go index 0b2bf60c..05ace8a4 100644 --- a/internal/database/repository/folders.sql.go +++ b/internal/database/repository/folders.sql.go @@ -91,6 +91,47 @@ func (q *Queries) CreateFolder(ctx context.Context, arg *CreateFolderParams) (*F return &i, err } +const createFolderIfAbsent = `-- name: CreateFolderIfAbsent :one +INSERT INTO folders (path, parentId, clientId, createdBy) +VALUES ($1, $2, $3, $4) +ON CONFLICT (clientId, path) DO NOTHING +RETURNING id, path, parentid, clientid, createdat, createdby +` + +type CreateFolderIfAbsentParams struct { + Path string `db:"path"` + ParentID *uuid.UUID `db:"parent_id"` + ClientID string `db:"client_id"` + CreatedBy string `db:"created_by"` +} + +// Insert a folder only when it does not already exist. +// Batch uploads use this to distinguish newly-created folder paths from +// pre-existing paths without mutating existing folders. +// +// INSERT INTO folders (path, parentId, clientId, createdBy) +// VALUES ($1, $2, $3, $4) +// ON CONFLICT (clientId, path) DO NOTHING +// RETURNING id, path, parentid, clientid, createdat, createdby +func (q *Queries) CreateFolderIfAbsent(ctx context.Context, arg *CreateFolderIfAbsentParams) (*Folder, error) { + row := q.db.QueryRow(ctx, createFolderIfAbsent, + arg.Path, + arg.ParentID, + arg.ClientID, + arg.CreatedBy, + ) + var i Folder + err := row.Scan( + &i.ID, + &i.Path, + &i.Parentid, + &i.Clientid, + &i.Createdat, + &i.Createdby, + ) + return &i, err +} + const deleteDocumentUploadsInFolderTree = `-- name: DeleteDocumentUploadsInFolderTree :exec WITH RECURSIVE folder_tree AS ( SELECT id FROM folders WHERE id = $1 @@ -116,6 +157,33 @@ func (q *Queries) DeleteDocumentUploadsInFolderTree(ctx context.Context, folderI return err } +const deleteSafeAutoCreatedFolder = `-- name: DeleteSafeAutoCreatedFolder :execrows +DELETE FROM folders f +WHERE f.id = $1 + AND f.path <> '/' + AND NOT EXISTS (SELECT 1 FROM documents d WHERE d.folderId = f.id) + AND NOT EXISTS (SELECT 1 FROM documentUploads du WHERE du.folder_id = f.id) + AND NOT EXISTS (SELECT 1 FROM folders child WHERE child.parentId = f.id) +` + +// @sqlc-vet-disable +// Delete one auto-created folder only when it is currently empty and no longer +// referenced by documents, uploads, or child folders. +// +// DELETE FROM folders f +// WHERE f.id = $1 +// AND f.path <> '/' +// AND NOT EXISTS (SELECT 1 FROM documents d WHERE d.folderId = f.id) +// AND NOT EXISTS (SELECT 1 FROM documentUploads du WHERE du.folder_id = f.id) +// AND NOT EXISTS (SELECT 1 FROM folders child WHERE child.parentId = f.id) +func (q *Queries) DeleteSafeAutoCreatedFolder(ctx context.Context, folderID uuid.UUID) (int64, error) { + result, err := q.db.Exec(ctx, deleteSafeAutoCreatedFolder, folderID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const getClientRootFolder = `-- name: GetClientRootFolder :one SELECT id, path, parentid, clientid, createdat, createdby FROM folders WHERE clientId = $1 AND path = '/' @@ -743,6 +811,69 @@ func (q *Queries) HardDeleteFolderTree(ctx context.Context, folderID *uuid.UUID) return result.RowsAffected(), nil } +const listSafeAutoCreatedFolderCandidates = `-- name: ListSafeAutoCreatedFolderCandidates :many +SELECT DISTINCT f.id, f.path +FROM folders f +JOIN batch_folder_candidates bfc ON bfc.folder_id = f.id +WHERE bfc.existed_before = false + AND bfc.path = f.path + AND f.path <> '/' + AND NOT EXISTS ( + SELECT 1 + FROM batch_folder_candidates other_bfc + JOIN batch_uploads bu ON bu.id = other_bfc.batch_id + WHERE other_bfc.folder_id = f.id + AND bu.folder_cleanup_completed_at IS NULL + ) +ORDER BY f.path DESC +LIMIT $1 +` + +type ListSafeAutoCreatedFolderCandidatesRow struct { + ID uuid.UUID `db:"id"` + Path string `db:"path"` +} + +// @sqlc-vet-disable +// Candidate folders can be deleted only after every batch that touched that +// folder has finished folder cleanup. The path equality check protects renamed +// folders from deletion. +// +// SELECT DISTINCT f.id, f.path +// FROM folders f +// JOIN batch_folder_candidates bfc ON bfc.folder_id = f.id +// WHERE bfc.existed_before = false +// AND bfc.path = f.path +// AND f.path <> '/' +// AND NOT EXISTS ( +// SELECT 1 +// FROM batch_folder_candidates other_bfc +// JOIN batch_uploads bu ON bu.id = other_bfc.batch_id +// WHERE other_bfc.folder_id = f.id +// AND bu.folder_cleanup_completed_at IS NULL +// ) +// ORDER BY f.path DESC +// LIMIT $1 +func (q *Queries) ListSafeAutoCreatedFolderCandidates(ctx context.Context, limitCount int64) ([]*ListSafeAutoCreatedFolderCandidatesRow, error) { + rows, err := q.db.Query(ctx, listSafeAutoCreatedFolderCandidates, limitCount) + if err != nil { + return nil, err + } + defer rows.Close() + items := []*ListSafeAutoCreatedFolderCandidatesRow{} + for rows.Next() { + var i ListSafeAutoCreatedFolderCandidatesRow + if err := rows.Scan(&i.ID, &i.Path); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const renameFolder = `-- name: RenameFolder :exec UPDATE folders SET path = $2 diff --git a/internal/database/repository/models.go b/internal/database/repository/models.go index afe82820..6b59e9ef 100644 --- a/internal/database/repository/models.go +++ b/internal/database/repository/models.go @@ -348,22 +348,31 @@ type BatchDocumentOutcome struct { CreatedAt pgtype.Timestamp `db:"created_at"` } +type BatchFolderCandidate struct { + BatchID uuid.UUID `db:"batch_id"` + FolderID uuid.UUID `db:"folder_id"` + Path string `db:"path"` + ExistedBefore bool `db:"existed_before"` + CreatedAt pgtype.Timestamp `db:"created_at"` +} + type BatchUpload 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"` - ArchiveBucket *string `db:"archive_bucket"` - ArchiveKey *string `db:"archive_key"` - FileSizeBytes *int64 `db:"file_size_bytes"` + 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"` + ArchiveBucket *string `db:"archive_bucket"` + ArchiveKey *string `db:"archive_key"` + FileSizeBytes *int64 `db:"file_size_bytes"` + FolderCleanupCompletedAt pgtype.Timestamp `db:"folder_cleanup_completed_at"` } type BotSession struct { diff --git a/internal/document/batch/foldercleanup/service.go b/internal/document/batch/foldercleanup/service.go new file mode 100644 index 00000000..0e6a7793 --- /dev/null +++ b/internal/document/batch/foldercleanup/service.go @@ -0,0 +1,194 @@ +package foldercleanup + +import ( + "context" + "fmt" + "time" + + "queryorchestration/internal/database/repository" + "queryorchestration/internal/serviceconfig" + + "github.com/google/uuid" +) + +const ( + DefaultReadyBatchLimit int64 = 100 + DefaultFolderDeleteLimit int64 = 500 + defaultStaleAfter = 15 * time.Minute +) + +type Service struct { + cfg serviceconfig.ConfigProvider +} + +type FinalizeResult struct { + Ready bool + AcceptedDocuments int32 + RejectedDocuments int32 + DetachedDocuments int64 + DetachedUploads int64 + DeletedFolders int64 +} + +func New(cfg serviceconfig.ConfigProvider) *Service { + return &Service{cfg: cfg} +} + +func (s *Service) TryFinalizeBatch(ctx context.Context, batchID uuid.UUID) (*FinalizeResult, error) { + result := &FinalizeResult{} + + err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error { + batch, err := q.LockBatchUploadForFolderCleanup(ctx, batchID) + if err != nil { + return err + } + if batch.FolderCleanupCompletedAt.Valid { + result.Ready = true + return nil + } + + outcomeCount, err := q.CountBatchOutcomes(ctx, batchID) + if err != nil { + return err + } + nonTerminalCount, err := q.CountNonTerminalBatchOutcomes(ctx, batchID) + if err != nil { + return err + } + acceptedCount, err := q.CountAcceptedBatchOutcomes(ctx, batchID) + if err != nil { + return err + } + + staleTerminal := isStaleTerminalBatch(batch, defaultStaleAfter) && outcomeCount > 0 + if !isBatchReadyForFolderCleanup(batch, outcomeCount, nonTerminalCount, staleTerminal) { + result.Ready = false + return nil + } + + includeNonTerminal := staleTerminal && nonTerminalCount > 0 + result.Ready = true + result.AcceptedDocuments = acceptedCount + result.RejectedDocuments = outcomeCount - acceptedCount + + detachedDocuments, err := q.DetachRejectedBatchDocumentsFromCandidateFolders(ctx, &repository.DetachRejectedBatchDocumentsFromCandidateFoldersParams{ + BatchID: &batchID, + IncludeNonterminal: includeNonTerminal, + }) + if err != nil { + return err + } + result.DetachedDocuments = detachedDocuments + + detachedUploads, err := q.DetachRejectedBatchDocumentUploadsFromCandidateFolders(ctx, &repository.DetachRejectedBatchDocumentUploadsFromCandidateFoldersParams{ + BatchID: &batchID, + IncludeNonterminal: includeNonTerminal, + }) + if err != nil { + return err + } + result.DetachedUploads = detachedUploads + + return q.MarkBatchFolderCleanupCompleted(ctx, batchID) + }) + if err != nil { + return nil, err + } + if !result.Ready { + return result, nil + } + + deleted, err := s.DeleteSafeAutoCreatedFolders(ctx, DefaultFolderDeleteLimit) + if err != nil { + return nil, err + } + result.DeletedFolders = deleted + return result, nil +} + +func (s *Service) FinalizeReadyBatches(ctx context.Context, limit int64) (int, error) { + if limit <= 0 { + limit = DefaultReadyBatchLimit + } + batchIDs, err := s.cfg.GetDBQueries().ListFolderCleanupReadyBatchIDs(ctx, &repository.ListFolderCleanupReadyBatchIDsParams{ + StaleAfterSeconds: int32(defaultStaleAfter.Seconds()), + LimitCount: limit, + }) + if err != nil { + return 0, err + } + + finalized := 0 + for _, batchID := range batchIDs { + result, err := s.TryFinalizeBatch(ctx, batchID) + if err != nil { + return finalized, fmt.Errorf("finalize batch %s: %w", batchID, err) + } + if result.Ready { + finalized++ + } + } + return finalized, nil +} + +func (s *Service) DeleteSafeAutoCreatedFolders(ctx context.Context, limit int64) (int64, error) { + if limit <= 0 { + limit = DefaultFolderDeleteLimit + } + + var deleted int64 + for deleted < limit { + candidates, err := s.cfg.GetDBQueries().ListSafeAutoCreatedFolderCandidates(ctx, limit-deleted) + if err != nil { + return deleted, err + } + if len(candidates) == 0 { + return deleted, nil + } + + deletedThisPass := int64(0) + for _, candidate := range candidates { + rows, err := s.cfg.GetDBQueries().DeleteSafeAutoCreatedFolder(ctx, candidate.ID) + if err != nil { + return deleted, err + } + deleted += rows + deletedThisPass += rows + if deleted >= limit { + break + } + } + if deletedThisPass == 0 { + return deleted, nil + } + } + return deleted, nil +} + +func isTerminalStatus(status repository.BatchStatus) bool { + return status == repository.BatchStatusCompleted || + status == repository.BatchStatusFailed || + status == repository.BatchStatusCancelled +} + +func isStaleTerminalBatch(batch *repository.LockBatchUploadForFolderCleanupRow, staleAfter time.Duration) bool { + if batch.Status != repository.BatchStatusFailed && batch.Status != repository.BatchStatusCancelled { + return false + } + if !batch.CompletedAt.Valid { + return false + } + return time.Since(batch.CompletedAt.Time) >= staleAfter +} + +func isBatchReadyForFolderCleanup( + batch *repository.LockBatchUploadForFolderCleanupRow, + outcomeCount, nonTerminalCount int32, + staleTerminal bool, +) bool { + if !isTerminalStatus(batch.Status) { + return false + } + allOutcomesTerminal := batch.TotalDocuments > 0 && outcomeCount == batch.TotalDocuments && nonTerminalCount == 0 + return allOutcomesTerminal || staleTerminal +} diff --git a/internal/document/batch/foldercleanup/service_test.go b/internal/document/batch/foldercleanup/service_test.go new file mode 100644 index 00000000..8f4c9cba --- /dev/null +++ b/internal/document/batch/foldercleanup/service_test.go @@ -0,0 +1,293 @@ +package foldercleanup_test + +import ( + "errors" + "testing" + "time" + + "queryorchestration/internal/database/repository" + "queryorchestration/internal/document/batch/foldercleanup" + "queryorchestration/internal/serviceconfig" + "queryorchestration/internal/test" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type TestConfig struct { + serviceconfig.BaseConfig +} + +func createClientRoot(t *testing.T, cfg *TestConfig, clientID string) *repository.Folder { + t.Helper() + ctx := t.Context() + err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ + Clientid: clientID, + Name: "Test Client", + }) + require.NoError(t, err) + root, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/", + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + return root +} + +func createBatch(t *testing.T, cfg *TestConfig, clientID string, total int32) uuid.UUID { + t.Helper() + batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ + ClientID: clientID, + OriginalFilename: "batch.zip", + TotalDocuments: total, + }) + require.NoError(t, err) + return batchID +} + +func TestFinalizeDeletesNewRejectedFolders(t *testing.T) { + cfg := &TestConfig{} + test.CreateDB(t, cfg) + ctx := t.Context() + + clientID := "cleanup_rejected_" + uuid.New().String()[:8] + root := createClientRoot(t, cfg, clientID) + batchID := createBatch(t, cfg, clientID, 1) + + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/auto", + Parentid: &root.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + }) + require.NoError(t, err) + + filename := "duplicate.pdf" + err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{ + ID: uuid.New(), + Clientid: clientID, + Bucket: "bucket", + Key: "key", + Part: 1, + Createdat: pgtype.Timestamp{Valid: true}, + Filename: &filename, + BatchID: &batchID, + FolderID: &folder.ID, + }) + require.NoError(t, err) + err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ + BatchID: batchID, + Filename: filename, + Column3: repository.BatchOutcomeStatusDuplicate, + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ + ID: batchID, + Column2: repository.BatchStatusFailed, + })) + + result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) + require.NoError(t, err) + require.True(t, result.Ready) + assert.Equal(t, int64(1), result.DetachedUploads) + assert.Equal(t, int64(1), result.DeletedFolders) + + upload, err := cfg.GetDBQueries().GetDocumentUploadByKey(ctx, &repository.GetDocumentUploadByKeyParams{ + Bucket: "bucket", + Key: "key", + }) + require.NoError(t, err) + assert.Nil(t, upload.FolderID) + + _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + assert.True(t, errors.Is(err, pgx.ErrNoRows)) +} + +func TestFinalizePreservesExistingFolders(t *testing.T) { + cfg := &TestConfig{} + test.CreateDB(t, cfg) + ctx := t.Context() + + clientID := "cleanup_existing_" + uuid.New().String()[:8] + root := createClientRoot(t, cfg, clientID) + batchID := createBatch(t, cfg, clientID, 1) + + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/existing", + Parentid: &root.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: true, + }) + require.NoError(t, err) + err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ + BatchID: batchID, + Filename: "duplicate.pdf", + Column3: repository.BatchOutcomeStatusDuplicate, + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ + ID: batchID, + Column2: repository.BatchStatusFailed, + })) + + result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) + require.NoError(t, err) + require.True(t, result.Ready) + assert.Zero(t, result.DeletedFolders) + + persisted, err := cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + require.NoError(t, err) + assert.Equal(t, "/existing", persisted.Path) +} + +func TestFinalizeReadyBatchesRetriesReadyCleanup(t *testing.T) { + cfg := &TestConfig{} + test.CreateDB(t, cfg) + ctx := t.Context() + + clientID := "cleanup_retry_" + uuid.New().String()[:8] + root := createClientRoot(t, cfg, clientID) + batchID := createBatch(t, cfg, clientID, 1) + + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/retry", + Parentid: &root.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + })) + require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ + BatchID: batchID, + Filename: "invalid.exe", + Column3: repository.BatchOutcomeStatusInvalidType, + })) + require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ + ID: batchID, + Column2: repository.BatchStatusFailed, + })) + + service := foldercleanup.New(cfg) + finalized, err := service.FinalizeReadyBatches(ctx, 10) + require.NoError(t, err) + assert.Equal(t, 1, finalized) + + batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) + require.NoError(t, err) + assert.True(t, batch.FolderCleanupCompletedAt.Valid) + _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + assert.True(t, errors.Is(err, pgx.ErrNoRows)) + + finalized, err = service.FinalizeReadyBatches(ctx, 10) + require.NoError(t, err) + assert.Zero(t, finalized) +} + +func TestFinalizeDeclinesFailedBatchWithoutOutcomes(t *testing.T) { + cfg := &TestConfig{} + test.CreateDB(t, cfg) + ctx := t.Context() + + clientID := "cleanup_zero_" + uuid.New().String()[:8] + root := createClientRoot(t, cfg, clientID) + batchID := createBatch(t, cfg, clientID, 0) + + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/zero", + Parentid: &root.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + })) + require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ + ID: batchID, + Column2: repository.BatchStatusFailed, + })) + + service := foldercleanup.New(cfg) + result, err := service.TryFinalizeBatch(ctx, batchID) + require.NoError(t, err) + require.False(t, result.Ready) + + batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) + require.NoError(t, err) + assert.False(t, batch.FolderCleanupCompletedAt.Valid) + _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + require.NoError(t, err) + + finalized, err := service.FinalizeReadyBatches(ctx, 10) + require.NoError(t, err) + assert.Zero(t, finalized) +} + +func TestFinalizeStaleFailedBatchWithOutcomeEvidence(t *testing.T) { + cfg := &TestConfig{} + test.CreateDB(t, cfg) + ctx := t.Context() + + clientID := "cleanup_stale_" + uuid.New().String()[:8] + root := createClientRoot(t, cfg, clientID) + batchID := createBatch(t, cfg, clientID, 2) + + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/stale", + Parentid: &root.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + })) + require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ + BatchID: batchID, + Filename: "stuck.pdf", + Column3: repository.BatchOutcomeStatusSubmitted, + })) + require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ + ID: batchID, + Column2: repository.BatchStatusFailed, + })) + require.NoError(t, cfg.GetDBQueries().SetBatchCompletedAt(ctx, &repository.SetBatchCompletedAtParams{ + ID: batchID, + CompletedAt: pgtype.Timestamp{Time: time.Now().Add(-20 * time.Minute), Valid: true}, + })) + + result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) + require.NoError(t, err) + require.True(t, result.Ready) + assert.Equal(t, int64(1), result.DeletedFolders) + assert.Equal(t, int32(1), result.RejectedDocuments) +} diff --git a/internal/document/batch/outcome/reporter.go b/internal/document/batch/outcome/reporter.go index af1d1a43..c501827a 100644 --- a/internal/document/batch/outcome/reporter.go +++ b/internal/document/batch/outcome/reporter.go @@ -78,7 +78,7 @@ func (r *Reporter) Resolve(ctx context.Context, batchID *uuid.UUID, filename str // If no outcome row exists for this document_id (non-batch document), this // affects 0 rows and is a no-op. func (r *Reporter) Update(ctx context.Context, documentID uuid.UUID, - outcomeStatus repository.BatchOutcomeStatus, errorDetail *string) error { + outcomeStatus repository.BatchOutcomeStatus, errorDetail *string) ([]uuid.UUID, error) { return r.queries.UpdateBatchDocumentOutcomeByDocumentID(ctx, &repository.UpdateBatchDocumentOutcomeByDocumentIDParams{ DocumentID: &documentID, Column2: outcomeStatus, diff --git a/internal/document/batch/outcome/reporter_test.go b/internal/document/batch/outcome/reporter_test.go index 6ea19cd6..5d6665a0 100644 --- a/internal/document/batch/outcome/reporter_test.go +++ b/internal/document/batch/outcome/reporter_test.go @@ -125,8 +125,9 @@ func TestUpdate_UpdatesOutcomeByDocumentID(t *testing.T) { require.NoError(t, err) // Now update to sync_complete - err = reporter.Update(ctx, docID, repository.BatchOutcomeStatusSyncComplete, nil) + batchIDs, err := reporter.Update(ctx, docID, repository.BatchOutcomeStatusSyncComplete, nil) require.NoError(t, err) + require.Equal(t, []uuid.UUID{batchID}, batchIDs) // Verify the outcome was advanced rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID) @@ -145,6 +146,7 @@ func TestUpdate_NoOpForNonBatchDocument(t *testing.T) { reporter := outcome.New(cfg.GetDBQueries()) // Call Update with a random document_id that has no outcome row - err := reporter.Update(ctx, uuid.New(), repository.BatchOutcomeStatusSyncComplete, nil) + batchIDs, err := reporter.Update(ctx, uuid.New(), repository.BatchOutcomeStatusSyncComplete, nil) require.NoError(t, err) + require.Empty(t, batchIDs) } diff --git a/internal/document/upload/get.go b/internal/document/upload/get.go index c95ca907..4071402a 100644 --- a/internal/document/upload/get.go +++ b/internal/document/upload/get.go @@ -2,6 +2,7 @@ package documentupload import ( "context" + "errors" "fmt" "io" "mime" @@ -14,6 +15,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) @@ -80,7 +82,7 @@ func parsePath(path string) (folderPath string, filename string) { // Returns the UUID of the leaf folder (/contracts/2025/Q1 in this example). // If folderPath is empty, returns the root folder ID (documents without paths go to root). // The root folder "/" must exist (created when the client was created). -func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queries, clientID string, folderPath string, createdBy string) (uuid.UUID, error) { +func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queries, clientID string, folderPath string, createdBy string, batchID *uuid.UUID) (uuid.UUID, error) { // Get the root folder for this client (must exist - created when client was created) rootFolder, err := q.GetFolderByPath(ctx, &repository.GetFolderByPathParams{ Clientid: clientID, @@ -119,16 +121,20 @@ func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queri currentPath = currentPath + "/" + part } - // Upsert the folder (create if not exists, return existing if exists) - folder, err := q.UpsertFolder(ctx, &repository.UpsertFolderParams{ - Path: currentPath, - Parentid: parentID, - Clientid: clientID, - Createdby: createdBy, - }) + folder, existedBefore, err := s.createOrGetFolder(ctx, q, clientID, currentPath, parentID, createdBy, batchID) if err != nil { return uuid.Nil, err } + if batchID != nil { + if err := q.RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: *batchID, + FolderID: folder.ID, + Path: currentPath, + ExistedBefore: existedBefore, + }); err != nil { + return uuid.Nil, err + } + } // This folder becomes the parent for the next level lastFolderID = folder.ID @@ -138,6 +144,48 @@ func (s *Service) createFolderHierarchy(ctx context.Context, q *repository.Queri return lastFolderID, nil } +func (s *Service) createOrGetFolder( + ctx context.Context, + q *repository.Queries, + clientID string, + currentPath string, + parentID *uuid.UUID, + createdBy string, + batchID *uuid.UUID, +) (*repository.Folder, bool, error) { + if batchID == nil { + folder, err := q.UpsertFolder(ctx, &repository.UpsertFolderParams{ + Path: currentPath, + Parentid: parentID, + Clientid: clientID, + Createdby: createdBy, + }) + return folder, true, err + } + + folder, err := q.CreateFolderIfAbsent(ctx, &repository.CreateFolderIfAbsentParams{ + Path: currentPath, + ParentID: parentID, + ClientID: clientID, + CreatedBy: createdBy, + }) + if err == nil { + return folder, false, nil + } + if !errors.Is(err, pgx.ErrNoRows) { + return nil, false, err + } + + folder, err = q.GetFolderByPath(ctx, &repository.GetFolderByPathParams{ + Clientid: clientID, + Path: currentPath, + }) + if err != nil { + return nil, false, err + } + return folder, true, nil +} + func (s *Service) Upload(ctx context.Context, file File) error { _, err := s.UploadWithResult(ctx, file) return err @@ -197,7 +245,7 @@ func (s *Service) UploadWithResult(ctx context.Context, file File) (UploadResult // Create folder hierarchy if folder path is provided // Use system email as createdBy since this is an automated upload process - folderID, err := s.createFolderHierarchy(ctx, q, file.ClientID, folderPath, "system@doczy.local") + folderID, err := s.createFolderHierarchy(ctx, q, file.ClientID, folderPath, "system@doczy.local", file.BatchID) if err != nil { return err } diff --git a/internal/document/upload/get_test.go b/internal/document/upload/get_test.go index bfc9cf27..0709f32f 100644 --- a/internal/document/upload/get_test.go +++ b/internal/document/upload/get_test.go @@ -14,6 +14,7 @@ import ( "queryorchestration/internal/test" "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -144,6 +145,99 @@ func TestUploadWithPath(t *testing.T) { assert.Len(t, os.Contents, 1) } +func TestUploadWithBatchRecordsFolderCandidates(t *testing.T) { + t.Parallel() + cfg := &Config{} + test.CreateDB(t, cfg) + acfg := test.CreateAWSContainer(t, cfg) + + test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) + test.CreateBucket(t, cfg) + + clientID := "client_batch_candidates" + createTestClient(t, cfg, clientID, "client") + + rootFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{ + Clientid: clientID, + Path: "/", + }) + require.NoError(t, err) + parentFolder, err := cfg.GetDBQueries().CreateFolder(t.Context(), &repository.CreateFolderParams{ + Path: "/existing", + Parentid: &rootFolder.ID, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + + batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ + ClientID: clientID, + OriginalFilename: "batch.zip", + TotalDocuments: 1, + }) + require.NoError(t, err) + + svc := documentupload.New(cfg) + _, err = svc.UploadWithResult(t.Context(), documentupload.File{ + ClientID: clientID, + Content: strings.NewReader("batch content"), + BatchID: &batchID, + Filename: "agreement.pdf", + Path: "existing/new-child/agreement.pdf", + }) + require.NoError(t, err) + + candidates, err := cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID) + require.NoError(t, err) + require.Len(t, candidates, 2) + assert.Equal(t, "/existing", candidates[0].Path) + assert.Equal(t, parentFolder.ID, candidates[0].FolderID) + assert.True(t, candidates[0].ExistedBefore) + assert.Equal(t, "/existing/new-child", candidates[1].Path) + assert.False(t, candidates[1].ExistedBefore) + + _, err = svc.UploadWithResult(t.Context(), documentupload.File{ + ClientID: clientID, + Content: strings.NewReader("batch content retry"), + BatchID: &batchID, + Filename: "agreement-2.pdf", + Path: "existing/new-child/agreement-2.pdf", + }) + require.NoError(t, err) + + candidates, err = cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID) + require.NoError(t, err) + require.Len(t, candidates, 2) + assert.False(t, candidates[1].ExistedBefore, "new folder candidates stay marked as newly created across retries") +} + +func TestUploadWithPathDoesNotRecordCandidatesForNonBatch(t *testing.T) { + t.Parallel() + cfg := &Config{} + test.CreateDB(t, cfg) + acfg := test.CreateAWSContainer(t, cfg) + + test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) + test.CreateBucket(t, cfg) + + clientID := "client_non_batch_candidates" + createTestClient(t, cfg, clientID, "client") + + batchID := uuid.New() + svc := documentupload.New(cfg) + _, err := svc.UploadWithResult(t.Context(), documentupload.File{ + ClientID: clientID, + Content: strings.NewReader("non batch content"), + Filename: "agreement.pdf", + Path: "contracts/agreement.pdf", + }) + require.NoError(t, err) + + candidates, err := cfg.GetDBQueries().ListBatchFolderCandidates(t.Context(), batchID) + require.NoError(t, err) + assert.Empty(t, candidates) +} + func TestUploadWithPathNoFolder(t *testing.T) { t.Parallel() cfg := &Config{} diff --git a/internal/server/api/batch_worker.go b/internal/server/api/batch_worker.go index 54c078df..bceb43da 100644 --- a/internal/server/api/batch_worker.go +++ b/internal/server/api/batch_worker.go @@ -14,18 +14,30 @@ import ( "queryorchestration/internal/database/repository" "queryorchestration/internal/document/batch" + "queryorchestration/internal/document/batch/foldercleanup" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" ) +type batchWorkerService interface { + AddFailedFilename(ctx context.Context, batchID uuid.UUID, filename string) error + CheckDuplicate(ctx context.Context, clientID string, hash string) (*uuid.UUID, error) + MarkCompleted(ctx context.Context, batchID uuid.UUID) error + MarkFailed(ctx context.Context, batchID uuid.UUID) error + RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string, outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error + SetTotalDocuments(ctx context.Context, batchID uuid.UUID, total int32) error + UpdateProgress(ctx context.Context, clientID string, batchID uuid.UUID, processed int32, failed int32, invalid int32) error +} + // Config key constants for batch worker const ( ConfigKeyBatchService = "batchService" ConfigKeyS3Client = "s3Client" ConfigKeyBucket = "bucket" ConfigKeyUploadHandler = "uploadHandler" + ConfigKeyFolderCleanup = "folderCleanupService" ) // acceptedExtensions defines the file extensions allowed in batch uploads. @@ -115,6 +127,15 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin return fmt.Errorf("upload handler not found in config") } + cleanupService, _ := config[ConfigKeyFolderCleanup].(*foldercleanup.Service) + if cleanupService != nil { + if finalized, err := cleanupService.FinalizeReadyBatches(ctx, foldercleanup.DefaultReadyBatchLimit); err != nil { + logger.Error("Failed to finalize ready batch folders", "error", err) + } else if finalized > 0 { + logger.Info("Finalized ready batch folders", "count", finalized) + } + } + // Query unprocessed batches batches, err := batchService.ListUnprocessed(ctx) if err != nil { @@ -136,7 +157,7 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin continue } - if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails); err != nil { + if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, cleanupService); err != nil { logger.Error("Failed to process batch", "batch_id", batchDetails.ID, "client_id", batchDetails.ClientID, @@ -146,6 +167,14 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin } } + if cleanupService != nil { + if deleted, err := cleanupService.DeleteSafeAutoCreatedFolders(ctx, foldercleanup.DefaultFolderDeleteLimit); err != nil { + logger.Error("Failed to delete safe auto-created folders", "error", err) + } else if deleted > 0 { + logger.Info("Deleted safe auto-created folders", "count", deleted) + } + } + return nil } @@ -153,11 +182,12 @@ func processBatchWork(ctx context.Context, logger *slog.Logger, config map[strin func processSingleBatch( ctx context.Context, logger *slog.Logger, - batchService *batch.Service, + batchService batchWorkerService, s3Client *s3.Client, bucket string, uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error, batchInfo *batch.BatchUploadDetails, + cleanupService *foldercleanup.Service, ) error { logger.Info("Processing batch", "batch_id", batchInfo.ID, @@ -171,6 +201,7 @@ func processSingleBatch( if err != nil { // Mark batch as failed if we can't download the zip _ = batchService.MarkFailed(ctx, batchInfo.ID) + tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID) return fmt.Errorf("failed to download zip from S3: %w", err) } @@ -188,17 +219,34 @@ func processSingleBatch( zipReader, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData))) if err != nil { _ = batchService.MarkFailed(ctx, batchInfo.ID) + tryFinalizeBatchFolders(ctx, logger, cleanupService, 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) + processed, failed, invalid, err := processZipFile(ctx, logger, batchService, s3Client, bucket, extractPath, uploadHandler, batchInfo, file) + if err != nil { + _ = batchService.MarkFailed(ctx, batchInfo.ID) + tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID) + return fmt.Errorf("failed to process zip file %q: %w", file.Name, err) + } processedCount += processed failedCount += failed invalidCount += invalid } + totalProcessed := processedCount + failedCount + invalidCount + if totalProcessed == 0 { + if err := recordNoProcessableFilesOutcome(ctx, batchService, batchInfo); err != nil { + _ = batchService.MarkFailed(ctx, batchInfo.ID) + tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID) + return fmt.Errorf("record no-processable-files outcome: %w", err) + } + failedCount = 1 + totalProcessed = 1 + } + // Update batch progress - convert safely to int32 const maxInt32 = 2147483647 var processedInt32, failedInt32, invalidInt32 int32 @@ -229,8 +277,6 @@ func processSingleBatch( 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 @@ -243,6 +289,7 @@ func processSingleBatch( logger.Error("Failed to mark batch as failed", "batch_id", batchInfo.ID, "error", err) } } + tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID) } logger.Info("Batch processing completed", @@ -254,6 +301,15 @@ func processSingleBatch( return nil } +func tryFinalizeBatchFolders(ctx context.Context, logger *slog.Logger, cleanupService *foldercleanup.Service, batchID uuid.UUID) { + if cleanupService == nil { + return + } + if _, err := cleanupService.TryFinalizeBatch(ctx, batchID); err != nil { + logger.Error("Failed to finalize batch folder cleanup", "batch_id", batchID, "error", err) + } +} + // 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{ @@ -291,33 +347,23 @@ func uploadToS3(ctx context.Context, s3Client *s3.Client, bucket, key string, da func processZipFile( ctx context.Context, logger *slog.Logger, - batchService *batch.Service, + batchService batchWorkerService, 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) { +) (int, int, int, error) { if file.FileInfo().IsDir() { - return 0, 0, 0 + return 0, 0, 0, nil } - // Skip system files and hidden files (but allow relative path prefixes like ./file.pdf) - filename := file.Name - if strings.HasPrefix(filename, "__MACOSX/") { - return 0, 0, 0 - } - - // Skip hidden files (files that start with . but are not relative paths like ./file) - if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") { - return 0, 0, 0 - } - - // Security: Clean filename to prevent path traversal attacks while preserving internal structure - filename = sanitizeZipPath(filename) - if filename == "" { - logger.Warn("Skipping file with invalid path", "original", file.Name) - return 0, 0, 0 + filename, skip := batchZipFilename(file.Name) + if skip { + if filename == "" { + logger.Warn("Skipping file with invalid path", "original", file.Name) + } + return 0, 0, 0, nil } // Extract file content @@ -326,8 +372,10 @@ func processZipFile( logger.Error("Failed to open file in zip", "filename", filename, "error", err) _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) errMsg := err.Error() - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil) - return 0, 1, 0 + if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil); recordErr != nil { + return 0, 0, 0, fmt.Errorf("record failed_open outcome: %w", recordErr) + } + return 0, 1, 0, nil } fileData, err := io.ReadAll(fileReader) @@ -336,8 +384,10 @@ func processZipFile( logger.Error("Failed to read file from zip", "filename", filename, "error", err) _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) errMsg := err.Error() - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil) - return 0, 1, 0 + if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil); recordErr != nil { + return 0, 0, 0, fmt.Errorf("record failed_read outcome: %w", recordErr) + } + return 0, 1, 0, nil } // Store extracted file in S3 temporary directory (filename now includes sanitized path) @@ -346,16 +396,20 @@ func processZipFile( logger.Error("Failed to upload extracted file to S3", "filename", filename, "error", err) _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) errMsg := err.Error() - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil) - return 0, 1, 0 + if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil); recordErr != nil { + return 0, 0, 0, fmt.Errorf("record failed_s3_upload outcome: %w", recordErr) + } + return 0, 1, 0, nil } // Check if file type is supported if !isAcceptedExtension(filename) { logger.Warn("Invalid file type", "filename", filename) _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil) - return 0, 0, 1 + if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil); recordErr != nil { + return 0, 0, 0, fmt.Errorf("record invalid_type outcome: %w", recordErr) + } + return 0, 0, 1, nil } // Call the document upload handler with batch ID @@ -363,20 +417,62 @@ func processZipFile( logger.Error("Failed to upload document", "filename", filename, "error", err) _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) errMsg := err.Error() - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil) - return 0, 1, 0 + if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); recordErr != nil { + return 0, 0, 0, fmt.Errorf("record failed_upload outcome: %w", recordErr) + } + return 0, 1, 0, nil } - // Check for duplicate by computing file hash and looking up existing document + if err := recordSuccessfulBatchOutcome(ctx, batchService, batchInfo, filename, fileData); err != nil { + return 0, 0, 0, err + } + + logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID) + return 1, 0, 0, nil +} + +func recordNoProcessableFilesOutcome(ctx context.Context, batchService batchWorkerService, batchInfo *batch.BatchUploadDetails) error { + filename := batchInfo.OriginalFilename + if filename == "" { + filename = batchInfo.ID.String() + ".zip" + } + errMsg := "archive contained no processable files" + _ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename) + if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); err != nil { + return fmt.Errorf("record failed_upload outcome: %w", err) + } + return nil +} + +func batchZipFilename(filename string) (string, bool) { + if strings.HasPrefix(filename, "__MACOSX/") { + return filename, true + } + if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") { + return filename, true + } + cleanFilename := sanitizeZipPath(filename) + return cleanFilename, cleanFilename == "" +} + +func recordSuccessfulBatchOutcome( + ctx context.Context, + batchService batchWorkerService, + batchInfo *batch.BatchUploadDetails, + filename string, + fileData []byte, +) error { hash := sha256.Sum256(fileData) hashStr := hex.EncodeToString(hash[:]) existingID, _ := batchService.CheckDuplicate(ctx, batchInfo.ClientID, hashStr) if existingID != nil { - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID) - } else { - _ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil) + if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID); err != nil { + return fmt.Errorf("record duplicate outcome: %w", err) + } + return nil } - - logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID) - return 1, 0, 0 + if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil); err != nil { + return fmt.Errorf("record submitted outcome: %w", err) + } + return nil } diff --git a/internal/server/api/batch_worker_test.go b/internal/server/api/batch_worker_test.go index 656374db..d3fa4e8a 100644 --- a/internal/server/api/batch_worker_test.go +++ b/internal/server/api/batch_worker_test.go @@ -6,14 +6,17 @@ import ( "context" "crypto/sha256" "encoding/hex" + "errors" "fmt" "io" "log/slog" "os" + "strings" "testing" "queryorchestration/internal/database/repository" "queryorchestration/internal/document/batch" + "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/test" @@ -21,10 +24,20 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +type outcomeFailingBatchService struct { + *batch.Service +} + +func (s *outcomeFailingBatchService) RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string, + outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error { + return errors.New("forced outcome write failure") +} + // TestProcessBatchWork tests the main batch processing work function func TestProcessBatchWork(t *testing.T) { if testing.Short() { @@ -110,6 +123,33 @@ func TestProcessBatchWork(t *testing.T) { } } +func TestProcessBatchWorkRunsFolderCleanupWithoutBatches(t *testing.T) { + ctx := t.Context() + cfg := &BaseConfig{} + test.CreateDB(t, cfg) + + workerConfig := map[string]any{ + ConfigKeyBatchService: batch.New(cfg), + ConfigKeyS3Client: &s3.Client{}, + ConfigKeyBucket: "bucket", + ConfigKeyUploadHandler: func(context.Context, string, io.Reader, string, *uuid.UUID) error { return nil }, + ConfigKeyFolderCleanup: foldercleanup.New(cfg), + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + require.NoError(t, processBatchWork(ctx, logger, workerConfig)) +} + +func TestTryFinalizeBatchFoldersHandlesNilAndErrors(t *testing.T) { + ctx := t.Context() + cfg := &BaseConfig{} + test.CreateDB(t, cfg) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + + tryFinalizeBatchFolders(ctx, logger, nil, uuid.New()) + tryFinalizeBatchFolders(ctx, logger, foldercleanup.New(cfg), uuid.New()) +} + // TestProcessSingleBatch tests processing a single batch func TestProcessSingleBatch(t *testing.T) { if testing.Short() { @@ -169,7 +209,7 @@ func TestProcessSingleBatch(t *testing.T) { // Process the batch logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // All three file types (PDF, TXT, DOCX) are now accepted @@ -253,7 +293,7 @@ func TestProcessSingleBatchWithFailure(t *testing.T) { // Process the batch logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Verify results @@ -357,6 +397,61 @@ func createTestZIPForWorker(t *testing.T, numFiles int) []byte { return buf.Bytes() } +func firstZipFile(t *testing.T, files map[string][]byte) *zip.File { + t.Helper() + + zipFiles := zipFiles(t, files) + require.NotEmpty(t, zipFiles) + return zipFiles[0] +} + +func zipFiles(t *testing.T, files map[string][]byte) []*zip.File { + t.Helper() + + var buf bytes.Buffer + zipWriter := zip.NewWriter(&buf) + for name, content := range files { + if content == nil && strings.HasSuffix(name, "/") { + header := &zip.FileHeader{Name: name} + header.SetMode(os.ModeDir | 0o755) + _, err := zipWriter.CreateHeader(header) + require.NoError(t, err) + continue + } + fileWriter, err := zipWriter.Create(name) + require.NoError(t, err) + _, err = fileWriter.Write(content) + require.NoError(t, err) + } + require.NoError(t, zipWriter.Close()) + + reader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len())) + require.NoError(t, err) + return reader.File +} + +func zipBytes(t *testing.T, files map[string][]byte) []byte { + t.Helper() + + var buf bytes.Buffer + zipWriter := zip.NewWriter(&buf) + for name, content := range files { + if content == nil && strings.HasSuffix(name, "/") { + header := &zip.FileHeader{Name: name} + header.SetMode(os.ModeDir | 0o755) + _, err := zipWriter.CreateHeader(header) + require.NoError(t, err) + continue + } + fileWriter, err := zipWriter.Create(name) + require.NoError(t, err) + _, err = fileWriter.Write(content) + require.NoError(t, err) + } + require.NoError(t, zipWriter.Close()) + return buf.Bytes() +} + // Helper function to create a ZIP with mixed content types func createMixedContentZIP(t *testing.T) []byte { var buf bytes.Buffer @@ -539,7 +634,7 @@ func TestProcessBatchWithNestedFolders(t *testing.T) { // Process the batch logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Verify batch status @@ -677,7 +772,7 @@ func TestProcessBatchWithPathTraversal(t *testing.T) { // Process the batch logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Verify that path traversal attempts were blocked @@ -757,6 +852,334 @@ func TestProcessZipFile_RecordsOutcomes(t *testing.T) { } } +func TestProcessZipFile_RecordsFailedS3UploadOutcome(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + ctx := t.Context() + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + test.CreateAWSResources(t, cfg) + + clientID := "test_zip_failed_s3" + test.CreateTestClient(t, cfg, clientID, "Test ZIP Failed S3 Client") + + batchService := batch.New(cfg) + s3ClientInterface := cfg.GetStoreClient() + s3Client, ok := s3ClientInterface.(*s3.Client) + require.True(t, ok) + + batchID, err := batchService.CreateWithStorage(ctx, clientID, "failed-s3.zip", 1, cfg.GetBucket(), "archive.zip", 1) + require.NoError(t, err) + file := firstZipFile(t, map[string][]byte{ + "document.pdf": []byte("%%PDF-1.4\n%%Test PDF\n%%%%EOF"), + }) + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + processed, failed, invalid, err := processZipFile( + ctx, + logger, + batchService, + s3Client, + "missing-bucket", + "extract/path", + func(context.Context, string, io.Reader, string, *uuid.UUID) error { + t.Fatal("upload handler should not be called when extracted S3 upload fails") + return nil + }, + &batch.BatchUploadDetails{ + BatchUploadSummary: batch.BatchUploadSummary{ID: batchID, ClientID: clientID}, + }, + file, + ) + require.NoError(t, err) + assert.Equal(t, 0, processed) + assert.Equal(t, 1, failed) + assert.Equal(t, 0, invalid) + + outcomes, err := batchService.ListOutcomes(ctx, batchID) + require.NoError(t, err) + require.Len(t, outcomes, 1) + assert.Equal(t, string(repository.BatchOutcomeStatusFailedS3Upload), outcomes[0].Outcome) + assert.Equal(t, "document.pdf", outcomes[0].Filename) + require.NotNil(t, outcomes[0].ErrorDetail) + assert.Contains(t, *outcomes[0].ErrorDetail, "failed to put object to S3") +} + +func TestProcessZipFile_RecordsInvalidTypeOutcome(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + ctx := t.Context() + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + test.CreateAWSResources(t, cfg) + + clientID := "test_zip_invalid_type" + test.CreateTestClient(t, cfg, clientID, "Test ZIP Invalid Type Client") + + batchService := batch.New(cfg) + s3ClientInterface := cfg.GetStoreClient() + s3Client, ok := s3ClientInterface.(*s3.Client) + require.True(t, ok) + + batchID, err := batchService.CreateWithStorage(ctx, clientID, "invalid-type.zip", 1, cfg.GetBucket(), "archive.zip", 1) + require.NoError(t, err) + file := firstZipFile(t, map[string][]byte{ + "unsupported.exe": []byte("not a supported document"), + }) + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + processed, failed, invalid, err := processZipFile( + ctx, + logger, + batchService, + s3Client, + cfg.GetBucket(), + "extract/path", + func(context.Context, string, io.Reader, string, *uuid.UUID) error { + t.Fatal("upload handler should not be called for invalid file types") + return nil + }, + &batch.BatchUploadDetails{ + BatchUploadSummary: batch.BatchUploadSummary{ID: batchID, ClientID: clientID}, + }, + file, + ) + require.NoError(t, err) + assert.Equal(t, 0, processed) + assert.Equal(t, 0, failed) + assert.Equal(t, 1, invalid) + + outcomes, err := batchService.ListOutcomes(ctx, batchID) + require.NoError(t, err) + require.Len(t, outcomes, 1) + assert.Equal(t, string(repository.BatchOutcomeStatusInvalidType), outcomes[0].Outcome) + assert.Equal(t, "unsupported.exe", outcomes[0].Filename) +} + +func TestProcessZipFile_SkipsDirectoriesAndHiddenEntries(t *testing.T) { + ctx := t.Context() + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + files := zipFiles(t, map[string][]byte{ + "folder/": nil, + "__MACOSX/document.pdf": []byte("metadata"), + ".DS_Store": []byte("metadata"), + "../escape.pdf": []byte("blocked"), + }) + + for _, file := range files { + processed, failed, invalid, err := processZipFile( + ctx, + logger, + nil, + nil, + "", + "", + func(context.Context, string, io.Reader, string, *uuid.UUID) error { + t.Fatal("upload handler should not be called for skipped entries") + return nil + }, + &batch.BatchUploadDetails{}, + file, + ) + require.NoError(t, err) + assert.Equal(t, 0, processed) + assert.Equal(t, 0, failed) + assert.Equal(t, 0, invalid) + } +} + +func TestOutcomeWriteFailureDoesNotCompleteBatch(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + ctx := t.Context() + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + test.CreateAWSResources(t, cfg) + + clientID := "test_outcome_write_failure" + test.CreateTestClient(t, cfg, clientID, "Test Outcome Write Failure Client") + + batchService := batch.New(cfg) + failingBatchService := &outcomeFailingBatchService{Service: batchService} + s3ClientInterface := cfg.GetStoreClient() + s3Client, ok := s3ClientInterface.(*s3.Client) + require.True(t, ok) + bucket := cfg.GetBucket() + + zipContent := createTestZIPForWorker(t, 1) + archiveKey := fmt.Sprintf("test/%s/outcome-failure.zip", clientID) + _, err := s3Client.PutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(archiveKey), + Body: bytes.NewReader(zipContent), + }) + require.NoError(t, err) + + batchID, err := batchService.CreateWithStorage(ctx, clientID, "outcome-failure.zip", 0, bucket, archiveKey, int64(len(zipContent))) + require.NoError(t, err) + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "outcome-failure-folder", + Parentid: nil, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + })) + + uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error { + _, err := io.ReadAll(docData) + return err + } + batchDetails := &batch.BatchUploadDetails{ + BatchUploadSummary: batch.BatchUploadSummary{ + ID: batchID, + ClientID: clientID, + OriginalFilename: "outcome-failure.zip", + Status: "processing", + }, + ArchiveKey: archiveKey, + ArchiveBucket: bucket, + } + + logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) + cleanupService := foldercleanup.New(cfg) + err = processSingleBatch(ctx, logger, failingBatchService, s3Client, bucket, uploadHandler, batchDetails, cleanupService) + require.Error(t, err) + + batchInfo, err := cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{ + ID: batchID, + ClientID: clientID, + }) + require.NoError(t, err) + assert.Equal(t, repository.BatchStatusFailed, batchInfo.Status) + assert.Zero(t, batchInfo.TotalDocuments, "failed outcome writes must not count toward batch totals") + + outcomeCount, err := cfg.GetDBQueries().CountBatchOutcomes(ctx, batchID) + require.NoError(t, err) + assert.Zero(t, outcomeCount) + + cleanupBatch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) + require.NoError(t, err) + assert.False(t, cleanupBatch.FolderCleanupCompletedAt.Valid) + _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + require.NoError(t, err) +} + +func TestAllSkippedZipRecordsFailedOutcomeAndFinalizes(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + ctx := t.Context() + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + test.CreateAWSResources(t, cfg) + + clientID := "test_all_skipped_zip" + test.CreateTestClient(t, cfg, clientID, "Test All Skipped ZIP Client") + + batchService := batch.New(cfg) + s3ClientInterface := cfg.GetStoreClient() + s3Client, ok := s3ClientInterface.(*s3.Client) + require.True(t, ok) + bucket := cfg.GetBucket() + + zipContent := zipBytes(t, map[string][]byte{ + "folder/": nil, + "__MACOSX/document.pdf": []byte("metadata"), + ".DS_Store": []byte("metadata"), + "../escape.pdf": []byte("blocked"), + }) + archiveKey := fmt.Sprintf("test/%s/all-skipped.zip", clientID) + _, err := s3Client.PutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(archiveKey), + Body: bytes.NewReader(zipContent), + }) + require.NoError(t, err) + + batchID, err := batchService.CreateWithStorage(ctx, clientID, "all-skipped.zip", 0, bucket, archiveKey, int64(len(zipContent))) + require.NoError(t, err) + folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ + Path: "/all-skipped", + Parentid: nil, + Clientid: clientID, + Createdby: "test", + }) + require.NoError(t, err) + require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ + BatchID: batchID, + FolderID: folder.ID, + Path: folder.Path, + ExistedBefore: false, + })) + + batchDetails := &batch.BatchUploadDetails{ + BatchUploadSummary: batch.BatchUploadSummary{ + ID: batchID, + ClientID: clientID, + OriginalFilename: "all-skipped.zip", + Status: "processing", + }, + ArchiveKey: archiveKey, + ArchiveBucket: bucket, + } + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, + func(context.Context, string, io.Reader, string, *uuid.UUID) error { + t.Fatal("upload handler should not be called for skipped-only archives") + return nil + }, + batchDetails, + foldercleanup.New(cfg), + ) + require.NoError(t, err) + + batchInfo, err := cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{ + ID: batchID, + ClientID: clientID, + }) + require.NoError(t, err) + assert.Equal(t, repository.BatchStatusFailed, batchInfo.Status) + assert.Equal(t, int32(1), batchInfo.TotalDocuments) + assert.Equal(t, int32(1), batchInfo.FailedDocuments) + + outcomes, err := batchService.ListOutcomes(ctx, batchID) + require.NoError(t, err) + require.Len(t, outcomes, 1) + assert.Equal(t, "all-skipped.zip", outcomes[0].Filename) + assert.Equal(t, string(repository.BatchOutcomeStatusFailedUpload), outcomes[0].Outcome) + require.NotNil(t, outcomes[0].ErrorDetail) + assert.Contains(t, *outcomes[0].ErrorDetail, "archive contained no processable files") + + unprocessed, err := cfg.GetDBQueries().GetUnprocessedBatches(ctx) + require.NoError(t, err) + for _, batch := range unprocessed { + assert.NotEqual(t, batchID, batch.ID) + } + + cleanupBatch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) + require.NoError(t, err) + assert.True(t, cleanupBatch.FolderCleanupCompletedAt.Valid) + _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) + assert.True(t, errors.Is(err, pgx.ErrNoRows)) +} + // TestProcessZipFile_DetectsDuplicates verifies that when a file in the batch // has the same content hash as an existing document in the system, it is // recorded as a duplicate with the existing document's ID. @@ -999,7 +1422,7 @@ func TestProcessBatchWithDuplicateFilenames(t *testing.T) { // Process the batch logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Verify all files were uploaded with correct paths @@ -1066,7 +1489,7 @@ func TestProcessBatchWithAllSupportedFileTypes(t *testing.T) { } logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Verify batch completed successfully @@ -1166,7 +1589,7 @@ func TestProcessBatchWithMixedValidInvalidFileTypes(t *testing.T) { } logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) - err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails) + err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, nil) require.NoError(t, err) // Batch should still complete (not fail) because some files succeeded diff --git a/internal/server/api/listener.go b/internal/server/api/listener.go index dc26de97..871199fa 100644 --- a/internal/server/api/listener.go +++ b/internal/server/api/listener.go @@ -16,6 +16,7 @@ import ( "queryorchestration/internal/backgroundtask" "queryorchestration/internal/cognitoauth" "queryorchestration/internal/document/batch" + "queryorchestration/internal/document/batch/foldercleanup" documentupload "queryorchestration/internal/document/upload" awsc "queryorchestration/internal/serviceconfig/aws" "queryorchestration/internal/serviceconfig/build" @@ -224,8 +225,16 @@ func New(ctx context.Context, cfg Config) (*Server, error) { return nil, err } + if err := ensureLocalDevClient(ctx, cfg); err != nil { + return nil, err + } + if err := ensureLocalDevEula(ctx, cfg); err != nil { + return nil, err + } + // Initialize background task runner for batch processing batchService := batch.New(cfg) + cleanupService := foldercleanup.New(cfg) // Create upload handler function that calls the document upload service uploadHandler := createDocumentUploadHandler(cfg) @@ -236,6 +245,7 @@ func New(ctx context.Context, cfg Config) (*Server, error) { ConfigKeyS3Client: cfg.GetStoreClient(), ConfigKeyBucket: cfg.GetBucket(), ConfigKeyUploadHandler: uploadHandler, + ConfigKeyFolderCleanup: cleanupService, } runner, err := backgroundtask.Initialize( diff --git a/internal/server/api/local_dev_bootstrap.go b/internal/server/api/local_dev_bootstrap.go new file mode 100644 index 00000000..b764153b --- /dev/null +++ b/internal/server/api/local_dev_bootstrap.go @@ -0,0 +1,127 @@ +package api + +import ( + "context" + "errors" + "log/slog" + "os" + "strings" + + "queryorchestration/internal/client" + "queryorchestration/internal/database/repository" + "queryorchestration/internal/serviceconfig" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +const ( + localDevClientIDEnv = "LOCAL_DEV_CLIENT_ID" + localDevClientNameEnv = "LOCAL_DEV_CLIENT_NAME" + localDevEulaVersionEnv = "LOCAL_DEV_EULA_VERSION" + localDevEulaTitleEnv = "LOCAL_DEV_EULA_TITLE" + localDevEulaAutoAcceptSubjectEnv = "LOCAL_DEV_EULA_AUTO_ACCEPT_SUBJECT" + localDevEulaAutoAcceptEmailEnv = "LOCAL_DEV_EULA_AUTO_ACCEPT_EMAIL" +) + +const localDevEulaContent = `

Acceptance of Terms

By accessing and using AArete DoczyAI in this local development environment, you agree to use the application only for authorized testing and evaluation.

Authorized Use

You are responsible for ensuring uploaded documents and generated outputs are appropriate for local development use and comply with applicable internal policies.

Data Handling

This local instance is intended for development workflows. Do not upload production, confidential, or regulated data unless the environment has been approved for that data.

No Warranty

This local development instance is provided as-is for testing. Features, data, and outputs may change during development.

` + +type localDevClientConfig interface { + serviceconfig.ConfigProvider + GetLogger() *slog.Logger +} + +func ensureLocalDevClient(ctx context.Context, cfg localDevClientConfig) error { + if strings.ToLower(os.Getenv("DISABLE_AUTH")) != "true" { + return nil + } + + clientID := strings.TrimSpace(os.Getenv(localDevClientIDEnv)) + if clientID == "" { + return nil + } + + clientName := strings.TrimSpace(os.Getenv(localDevClientNameEnv)) + if clientName == "" { + clientName = clientID + } + + _, err := cfg.GetDBQueries().GetClient(ctx, clientID) + if err == nil { + return nil + } + if !errors.Is(err, pgx.ErrNoRows) { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "42P01" { + localDevBootstrapLogger(cfg).Warn("Skipping local dev client bootstrap because client tables are unavailable") + return nil + } + return err + } + + _, err = client.New(cfg).Create(ctx, client.CreateParams{ + ID: clientID, + Name: clientName, + }) + if err != nil { + return err + } + + localDevBootstrapLogger(cfg).Info("Created local dev client", slog.String("client_id", clientID)) + return nil +} + +func ensureLocalDevEula(ctx context.Context, cfg localDevClientConfig) error { + if strings.ToLower(os.Getenv("DISABLE_AUTH")) != "true" { + return nil + } + + version := strings.TrimSpace(os.Getenv(localDevEulaVersionEnv)) + if version == "" { + return nil + } + + title := strings.TrimSpace(os.Getenv(localDevEulaTitleEnv)) + if title == "" { + title = "AArete DoczyAI Terms of Use" + } + + versionID, err := cfg.GetDBQueries().ActivateLocalDevEulaVersion(ctx, &repository.ActivateLocalDevEulaVersionParams{ + Version: version, + Title: title, + Content: localDevEulaContent, + }) + if err != nil { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "42P01" { + localDevBootstrapLogger(cfg).Warn("Skipping local dev EULA bootstrap because EULA tables are unavailable") + return nil + } + return err + } + + subject := strings.TrimSpace(os.Getenv(localDevEulaAutoAcceptSubjectEnv)) + email := strings.TrimSpace(os.Getenv(localDevEulaAutoAcceptEmailEnv)) + if subject != "" { + if email == "" { + email = subject + } + if err := cfg.GetDBQueries().CreateLocalDevEulaAgreement(ctx, &repository.CreateLocalDevEulaAgreementParams{ + CognitoSubjectID: subject, + UserEmail: email, + EulaVersionID: versionID, + }); err != nil { + return err + } + } + + localDevBootstrapLogger(cfg).Info("Activated local dev EULA", slog.String("version", version)) + return nil +} + +func localDevBootstrapLogger(cfg localDevClientConfig) *slog.Logger { + if logger := cfg.GetLogger(); logger != nil { + return logger + } + return slog.Default() +} diff --git a/internal/server/api/local_dev_bootstrap_test.go b/internal/server/api/local_dev_bootstrap_test.go new file mode 100644 index 00000000..5894ae1a --- /dev/null +++ b/internal/server/api/local_dev_bootstrap_test.go @@ -0,0 +1,120 @@ +package api + +import ( + "testing" + + "queryorchestration/internal/database/repository" + "queryorchestration/internal/serviceconfig" + "queryorchestration/internal/test" + + "github.com/stretchr/testify/require" +) + +func TestEnsureLocalDevClientCreatesConfiguredClient(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + t.Setenv("DISABLE_AUTH", "true") + t.Setenv(localDevClientIDEnv, "mathisonprojects-local") + t.Setenv(localDevClientNameEnv, "Mathison Projects Local") + + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + + err := ensureLocalDevClient(t.Context(), cfg) + require.NoError(t, err) + + dbClient, err := cfg.GetDBQueries().GetClient(t.Context(), "mathisonprojects-local") + require.NoError(t, err) + require.Equal(t, "Mathison Projects Local", dbClient.Name) + + rootFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{ + Clientid: "mathisonprojects-local", + Path: "/", + }) + require.NoError(t, err) + require.Equal(t, "/", rootFolder.Path) + + require.NoError(t, ensureLocalDevClient(t.Context(), cfg)) +} + +func TestEnsureLocalDevClientSkipsWhenUnconfigured(t *testing.T) { + t.Setenv("DISABLE_AUTH", "true") + t.Setenv(localDevClientIDEnv, "") + + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + + require.NoError(t, ensureLocalDevClient(t.Context(), cfg)) +} + +func TestEnsureLocalDevEula(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + t.Setenv("DISABLE_AUTH", "true") + t.Setenv(localDevEulaVersionEnv, "local-2026-05-07") + t.Setenv(localDevEulaTitleEnv, "AArete DoczyAI Terms of Use") + t.Setenv(localDevEulaAutoAcceptSubjectEnv, "jacob+2@mathisonprojects.com") + t.Setenv(localDevEulaAutoAcceptEmailEnv, "jacob+2@mathisonprojects.com") + + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + + require.NoError(t, ensureLocalDevEula(t.Context(), cfg)) + + current, err := cfg.GetDBQueries().GetCurrentEulaVersion(t.Context()) + require.NoError(t, err) + require.Equal(t, "local-2026-05-07", current.Version) + require.Equal(t, "AArete DoczyAI Terms of Use", current.Title) + + agreement, err := cfg.GetDBQueries().GetUserCurrentEulaAgreement(t.Context(), "jacob+2@mathisonprojects.com") + require.NoError(t, err) + require.Equal(t, current.ID, agreement.Eulaversionid) + + require.NoError(t, ensureLocalDevEula(t.Context(), cfg)) +} + +func TestEnsureLocalDevEulaUsesDefaults(t *testing.T) { + if testing.Short() { + t.SkipNow() + } + + t.Setenv("DISABLE_AUTH", "true") + t.Setenv(localDevEulaVersionEnv, "local-defaults") + t.Setenv(localDevEulaTitleEnv, "") + t.Setenv(localDevEulaAutoAcceptSubjectEnv, "local-subject") + t.Setenv(localDevEulaAutoAcceptEmailEnv, "") + + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + test.CreateDB(t, cfg) + + require.NoError(t, ensureLocalDevEula(t.Context(), cfg)) + + current, err := cfg.GetDBQueries().GetCurrentEulaVersion(t.Context()) + require.NoError(t, err) + require.Equal(t, "local-defaults", current.Version) + require.Equal(t, "AArete DoczyAI Terms of Use", current.Title) + + agreement, err := cfg.GetDBQueries().GetUserCurrentEulaAgreement(t.Context(), "local-subject") + require.NoError(t, err) + require.Equal(t, "local-subject", agreement.Useremail) +} + +func TestEnsureLocalDevEulaSkipsWhenDisabledOrUnconfigured(t *testing.T) { + cfg := &BaseConfig{} + _ = serviceconfig.InitializeConfig(cfg) + + t.Setenv("DISABLE_AUTH", "false") + t.Setenv(localDevEulaVersionEnv, "local-disabled") + require.NoError(t, ensureLocalDevEula(t.Context(), cfg)) + + t.Setenv("DISABLE_AUTH", "true") + t.Setenv(localDevEulaVersionEnv, "") + require.NoError(t, ensureLocalDevEula(t.Context(), cfg)) +} diff --git a/internal/usermanagement/permitio.go b/internal/usermanagement/permitio.go index 2edcb3b9..3ad9570a 100644 --- a/internal/usermanagement/permitio.go +++ b/internal/usermanagement/permitio.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "net/http" + "net/url" "strings" ) @@ -32,7 +33,7 @@ func CreatePermitUser(client *http.Client, config *PermitConfig, cognitoUser *Co return fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/users", config.BaseURL, projectID, envID) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/users", config.BaseURL, projectID, envID) attrs := map[string]interface{}{ "cognito_username": cognitoUser.Username, @@ -55,7 +56,7 @@ func CreatePermitUser(client *http.Client, config *PermitConfig, cognitoUser *Co return fmt.Errorf("error marshaling user data: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + req, err := http.NewRequest("POST", endpoint, bytes.NewBuffer(jsonData)) if err != nil { return fmt.Errorf("error creating request: %w", err) } @@ -121,9 +122,9 @@ func GetPermitUser(client *http.Client, config *PermitConfig, userKey string) (* return nil, fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/users/%s", config.BaseURL, projectID, envID, userKey) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/users/%s", config.BaseURL, projectID, envID, userKey) - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } @@ -172,9 +173,9 @@ func DeletePermitUser(client *http.Client, config *PermitConfig, userKey string) return fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/users/%s", config.BaseURL, projectID, envID, userKey) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/users/%s", config.BaseURL, projectID, envID, userKey) - req, err := http.NewRequest("DELETE", url, nil) + req, err := http.NewRequest("DELETE", endpoint, nil) if err != nil { return fmt.Errorf("error creating request: %w", err) } @@ -226,9 +227,9 @@ func FindPermitUserByEmail(client *http.Client, config *PermitConfig, email stri } // Search for user by email - url := fmt.Sprintf("%s/v2/facts/%s/%s/users?search=%s", config.BaseURL, projectID, envID, email) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/users?search=%s", config.BaseURL, projectID, envID, email) - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return "", fmt.Errorf("error creating request: %w", err) } @@ -328,7 +329,7 @@ func AssignRole(client *http.Client, config *PermitConfig, userID, role string) return fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments", config.BaseURL, projectID, envID) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments", config.BaseURL, projectID, envID) tenant := config.DefaultTenant if tenant == "" { @@ -346,7 +347,7 @@ func AssignRole(client *http.Client, config *PermitConfig, userID, role string) return fmt.Errorf("error marshaling role assignment: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + req, err := http.NewRequest("POST", endpoint, bytes.NewBuffer(jsonData)) if err != nil { return fmt.Errorf("error creating request: %w", err) } @@ -402,7 +403,7 @@ func UnassignRole(client *http.Client, config *PermitConfig, userID, role string return fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments", config.BaseURL, projectID, envID) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments", config.BaseURL, projectID, envID) tenant := config.DefaultTenant if tenant == "" { @@ -420,7 +421,7 @@ func UnassignRole(client *http.Client, config *PermitConfig, userID, role string return fmt.Errorf("error marshaling role assignment: %w", err) } - req, err := http.NewRequest("DELETE", url, bytes.NewBuffer(jsonData)) + req, err := http.NewRequest("DELETE", endpoint, bytes.NewBuffer(jsonData)) if err != nil { return fmt.Errorf("error creating request: %w", err) } @@ -476,10 +477,10 @@ func GetPermitUserRoles(client *http.Client, config *PermitConfig, userKey strin return nil, fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments?user=%s", - config.BaseURL, projectID, envID, userKey) + endpoint := fmt.Sprintf("%s/v2/facts/%s/%s/role_assignments?user=%s", + config.BaseURL, projectID, envID, url.QueryEscape(userKey)) - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } @@ -559,10 +560,10 @@ func GetExistingRoles(client *http.Client, config *PermitConfig) ([]string, erro } // Construct API URL for roles endpoint - url := fmt.Sprintf("%s/v2/schema/%s/%s/roles", config.BaseURL, projectID, envID) + endpoint := fmt.Sprintf("%s/v2/schema/%s/%s/roles", config.BaseURL, projectID, envID) // Create HTTP request - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } @@ -714,9 +715,9 @@ type RoleDetails struct { // - *APIKeyScope: The scope containing organization, project, and environment IDs // - error: Error if the API call fails or response cannot be parsed func GetAPIKeyScope(client *http.Client, apiKey, baseURL string) (*APIKeyScope, error) { - url := fmt.Sprintf("%s/v2/api-key/scope", baseURL) + endpoint := fmt.Sprintf("%s/v2/api-key/scope", baseURL) - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } @@ -767,10 +768,10 @@ func GetRoleDetails(client *http.Client, config *PermitConfig, roleKey string) ( return nil, fmt.Errorf("failed to get environment ID: %w", err) } - url := fmt.Sprintf("%s/v2/schema/%s/%s/roles/%s", + endpoint := fmt.Sprintf("%s/v2/schema/%s/%s/roles/%s", config.BaseURL, projectID, envID, roleKey) - req, err := http.NewRequest("GET", url, nil) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } diff --git a/internal/usermanagement/permitio_test.go b/internal/usermanagement/permitio_test.go index 75cfc088..7d245dad 100644 --- a/internal/usermanagement/permitio_test.go +++ b/internal/usermanagement/permitio_test.go @@ -71,6 +71,24 @@ func TestGetAPIKeyScope_Success(t *testing.T) { assert.Equal(t, "env-789", scope.EnvironmentID) } +func TestGetPermitUserRolesEscapesUserKey(t *testing.T) { + server, config := createMockPermitServer(t, map[string]http.HandlerFunc{ + "/v2/facts/test-project/test-env/role_assignments": func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "jacob+2@mathisonprojects.com", r.URL.Query().Get("user")) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, err := w.Write([]byte(`[{"role":"super_admin","user":"jacob+2@mathisonprojects.com","tenant":"default"}]`)) + require.NoError(t, err) + }, + }) + defer server.Close() + + roles, err := GetPermitUserRoles(server.Client(), config, "jacob+2@mathisonprojects.com") + + require.NoError(t, err) + assert.Equal(t, []string{"super_admin"}, roles) +} + // TestGetAPIKeyScope_OrgLevelKey tests API key with only organization scope func TestGetAPIKeyScope_OrgLevelKey(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {