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/upload/get.go
+57 -9
View File
@@ -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
}
internal/document/upload/get_test.go
+94
View File
@@ -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{}