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

Merged in feature/batch-status (pull request #215)

Browse Source 
Implement and test the batch status feature

* working
This commit is contained in:
Jay Brown
2026-03-12 18:53:42 +00:00
parent 8a4029058b
commit aafe7d5b5f
39 changed files with 4252 additions and 516 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/server/api/batch_worker.go
+22
View File
@@ -4,12 +4,15 @@ import (
"archive/zip"
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"log/slog"
"path"
"strings"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document/batch"
"github.com/aws/aws-sdk-go-v2/aws"
@@ -292,6 +295,8 @@ func processZipFile(
if err != nil {
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
}
@@ -300,6 +305,8 @@ func processZipFile(
if err != nil {
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
}
@@ -308,6 +315,8 @@ func processZipFile(
if err := uploadToS3(ctx, s3Client, bucket, extractedKey, fileData); err != nil {
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
}
@@ -315,6 +324,7 @@ func processZipFile(
if !strings.HasSuffix(strings.ToLower(filename), ".pdf") {
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
}
@@ -322,9 +332,21 @@ func processZipFile(
if err := uploadHandler(ctx, batchInfo.ClientID, bytes.NewReader(fileData), filename, &batchInfo.ID); err != nil {
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
}
// Check for duplicate by computing file hash and looking up existing document
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)
}
logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
return 1, 0, 0
}
internal/server/api/batch_worker_test.go
+217
View File
@@ -4,6 +4,8 @@ import (
"archive/zip"
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"log/slog"
@@ -668,6 +670,221 @@ func TestProcessBatchWithPathTraversal(t *testing.T) {
assert.Equal(t, 3, len(uploadedDocs), "Should have processed exactly 3 safe PDFs")
}
// TestProcessZipFile_RecordsOutcomes tests that processZipFile records per-file
// outcomes for each extracted file.
func TestProcessZipFile_RecordsOutcomes(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_outcomes_record"
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Record Client")
batchService := batch.New(cfg)
s3ClientInterface := cfg.GetStoreClient()
s3Client, ok := s3ClientInterface.(*s3.Client)
require.True(t, ok)
bucket := cfg.GetBucket()
// Create ZIP with 2 PDFs
zipContent := createTestZIPForWorker(t, 2)
archiveKey := fmt.Sprintf("test/%s/outcomes.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, "outcomes.zip", 2, bucket, archiveKey, int64(len(zipContent)))
require.NoError(t, err)
uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
_, _ = io.ReadAll(docData)
return nil
}
workerConfig := map[string]any{
ConfigKeyBatchService: batchService,
ConfigKeyS3Client: s3Client,
ConfigKeyBucket: bucket,
ConfigKeyUploadHandler: uploadHandler,
}
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
err = processBatchWork(ctx, logger, workerConfig)
require.NoError(t, err)
// Verify outcomes were recorded for each file
outcomes, err := batchService.ListOutcomes(ctx, batchID)
require.NoError(t, err)
assert.Len(t, outcomes, 2)
for _, o := range outcomes {
assert.Equal(t, "submitted", o.Outcome, "each PDF should be recorded as submitted")
assert.Contains(t, o.Filename, ".pdf")
}
}
// 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.
func TestProcessZipFile_DetectsDuplicates(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_outcomes_dup"
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Dup Client")
batchService := batch.New(cfg)
s3ClientInterface := cfg.GetStoreClient()
s3Client, ok := s3ClientInterface.(*s3.Client)
require.True(t, ok)
bucket := cfg.GetBucket()
// Pre-create an existing document with a known hash
knownContent := []byte("%%PDF-1.4\n%%Already existing content\n%%%%EOF")
knownHash := sha256Hash(knownContent)
existingDocID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: knownHash,
})
require.NoError(t, err)
// Create a ZIP with one file that has the same content as the existing doc
var buf bytes.Buffer
zipWriter := zip.NewWriter(&buf)
w, err := zipWriter.Create("duplicate_of_existing.pdf")
require.NoError(t, err)
_, err = w.Write(knownContent)
require.NoError(t, err)
require.NoError(t, zipWriter.Close())
zipContent := buf.Bytes()
archiveKey := fmt.Sprintf("test/%s/dup.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, "dup.zip", 1, bucket, archiveKey, int64(len(zipContent)))
require.NoError(t, err)
uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
_, _ = io.ReadAll(docData)
return nil
}
workerConfig := map[string]any{
ConfigKeyBatchService: batchService,
ConfigKeyS3Client: s3Client,
ConfigKeyBucket: bucket,
ConfigKeyUploadHandler: uploadHandler,
}
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
err = processBatchWork(ctx, logger, workerConfig)
require.NoError(t, err)
// Verify outcome is "duplicate" with the existing document's ID
outcomes, err := batchService.ListOutcomes(ctx, batchID)
require.NoError(t, err)
require.Len(t, outcomes, 1)
assert.Equal(t, "duplicate", outcomes[0].Outcome)
require.NotNil(t, outcomes[0].DocumentID)
assert.Equal(t, existingDocID, *outcomes[0].DocumentID)
}
// TestProcessZipFile_MixedOutcomes tests a ZIP with PDF + non-PDF + paths to
// verify all outcome types (submitted, invalid_type) are correctly recorded.
func TestProcessZipFile_MixedOutcomes(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_outcomes_mixed"
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Mixed Client")
batchService := batch.New(cfg)
s3ClientInterface := cfg.GetStoreClient()
s3Client, ok := s3ClientInterface.(*s3.Client)
require.True(t, ok)
bucket := cfg.GetBucket()
// Create mixed ZIP
zipContent := createMixedContentZIP(t)
archiveKey := fmt.Sprintf("test/%s/mixed_outcomes.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, "mixed_outcomes.zip", 3, bucket, archiveKey, int64(len(zipContent)))
require.NoError(t, err)
uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
_, _ = io.ReadAll(docData)
return nil
}
workerConfig := map[string]any{
ConfigKeyBatchService: batchService,
ConfigKeyS3Client: s3Client,
ConfigKeyBucket: bucket,
ConfigKeyUploadHandler: uploadHandler,
}
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
err = processBatchWork(ctx, logger, workerConfig)
require.NoError(t, err)
// Verify outcomes
outcomes, err := batchService.ListOutcomes(ctx, batchID)
require.NoError(t, err)
assert.Len(t, outcomes, 3) // 1 PDF + 1 txt + 1 docx
outcomeMap := map[string]string{}
for _, o := range outcomes {
outcomeMap[o.Filename] = o.Outcome
}
assert.Equal(t, "submitted", outcomeMap["document.pdf"])
assert.Equal(t, "invalid_type", outcomeMap["notes.txt"])
assert.Equal(t, "invalid_type", outcomeMap["report.docx"])
}
// sha256Hash computes the hex-encoded SHA-256 hash of data.
func sha256Hash(data []byte) string {
h := sha256.Sum256(data)
return hex.EncodeToString(h[:])
}
// TestProcessBatchWithDuplicateFilenames tests handling of duplicate filenames in different folders
func TestProcessBatchWithDuplicateFilenames(t *testing.T) {
if testing.Short() {