jacob-mathison/query-orchestration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in jmathison/v2-upload-batch (pull request #224)

Browse Source 
Implement v2 upload batch cleanup

* Implement v2 upload batch cleanup

* Merge remote-tracking branch 'origin/main' into jmathison/v2-upload-batch

* Address upload batch review feedback

* Raise batch worker coverage

* Fix batch cleanup review issues


Approved-by: Jay Brown
This commit is contained in:
Jacob Mathison
2026-05-08 21:42:46 +00:00
parent 7b820b18c2
commit a080ca59d8
35 changed files with 2587 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/document/batch/foldercleanup/service.go
+194
View File
@@ -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
}
internal/document/batch/foldercleanup/service_test.go
+293
View File
@@ -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)
}