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
api/docInitRunner/runner.go
+37 -2
View File
@@ -7,6 +7,7 @@ import (
"log/slog"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document/batch/outcome"
documentinit "queryorchestration/internal/document/init"
"queryorchestration/internal/serviceconfig/objectstore"
@@ -15,27 +16,45 @@ import (
const Name = "docInitRunner"
// Services holds the dependencies for the docInitRunner.
//
// Fields:
// - Document: service for creating/deduplicating documents
// - Queries: SQLC-generated database queries
// - Outcome: reporter for writing batch pipeline outcomes
type Services struct {
Document *documentinit.Service
Queries *repository.Queries
Outcome *outcome.Reporter
}
// Runner processes document init messages from the SQS queue.
type Runner struct {
svc *Services
}
// New creates a new docInitRunner.
//
// Parameters:
// - svc: the runner's service dependencies
//
// Returns:
// - Runner: the runner instance
func New(svc *Services) Runner {
return Runner{
svc: svc,
}
}
// Body is the SQS message body for docInitRunner.
type Body struct {
Bucket string `json:"bucket" validate:"required"`
Key string `json:"key" validate:"required"`
Hash string `json:"hash" validate:"required"`
}
// Process handles a single document init message. Returns true to acknowledge
// the message (success), or false to requeue (infrastructure error).
func (s Runner) Process(ctx context.Context, body Body) bool {
key, err := objectstore.ParseBucketKey(body.Key)
if err != nil {
@@ -61,7 +80,7 @@ func (s Runner) Process(ctx context.Context, body Body) bool {
folderID = upload.FolderID
}
id, err := s.svc.Document.Create(ctx, &documentinit.Create{
result, err := s.svc.Document.Create(ctx, &documentinit.Create{
Key: key,
Bucket: body.Bucket,
Hash: body.Hash,
@@ -72,7 +91,23 @@ func (s Runner) Process(ctx context.Context, body Body) bool {
if err != nil {
return false
}
slog.Debug("created document", "id", id, "filename", filename, "folder_id", folderID)
slog.Debug("created document", "id", result.ID, "filename", filename, "folder_id", folderID)
// Report outcome for batch documents
outcomeName := repository.BatchOutcomeStatusInitComplete
if result.IsDuplicate {
outcomeName = repository.BatchOutcomeStatusInitDuplicate
}
var resolveFilename string
if result.Filename != nil {
resolveFilename = *result.Filename
}
if resolveErr := s.svc.Outcome.Resolve(ctx, result.BatchID, resolveFilename,
result.ID, outcomeName, nil); resolveErr != nil {
slog.Error("failed to resolve batch outcome",
"document_id", result.ID, "outcome", outcomeName, "error", resolveErr)
// Non-fatal: document was created successfully
}
return true
}
api/docInitRunner/runner_test.go
+241
View File
@@ -8,6 +8,7 @@ import (
docinitrunner "queryorchestration/api/docInitRunner"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document/batch/outcome"
documentinit "queryorchestration/internal/document/init"
"queryorchestration/internal/server/runner"
"queryorchestration/internal/serviceconfig/objectstore"
@@ -51,6 +52,7 @@ func TestDocInitRunner(t *testing.T) {
runnerInstance := docinitrunner.New(&docinitrunner.Services{
Document: documentinit.New(cfg),
Queries: cfg.GetDBQueries(),
Outcome: outcome.New(cfg.GetDBQueries()),
})
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
@@ -101,3 +103,242 @@ func TestDocInitRunner(t *testing.T) {
test.AssertMessageBody(t, cfg, cfg.GetDocumentSyncURL(), regexp.MustCompile(`{"id":".+"}`))
}
// setupDocInitTest creates shared infrastructure (DB, S3, SQS) and returns
// the config, runner, and a cleanup-style helper.
func setupDocInitTest(t *testing.T) *DocInitConfig {
t.Helper()
cfg := &DocInitConfig{}
test.CreateDB(t, cfg)
acfg := test.CreateAWSContainer(t, cfg)
test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
cfg.DocumentSyncURL = test.CreateQueue(t, cfg, test.DocSyncRunnerName)
test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
test.CreateBucket(t, cfg)
return cfg
}
// uploadTestFileToS3 uploads a test file and returns the BucketKey.
// If batchID is non-nil, it is embedded in the key so ParseBucketKey can extract it.
func uploadTestFileToS3(t *testing.T, cfg *DocInitConfig, clientID string, batchID *uuid.UUID) objectstore.BucketKey {
t.Helper()
location := objectstore.BucketKey{
Location: objectstore.Import,
ClientID: clientID,
CreatedAt: time.Now().UTC(),
EntityID: uuid.New(),
BatchID: batchID,
}
testContent := []byte("test doc content " + uuid.New().String())
_, err := cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{
Bucket: aws.String(cfg.GetBucket()),
Key: aws.String(location.String()),
Body: bytes.NewReader(testContent),
})
require.NoError(t, err)
return location
}
// TestDocInitRunner_ReportsInitComplete processes a batch document and verifies
// the outcome row is updated to init_complete with the assigned document_id.
func TestDocInitRunner_ReportsInitComplete(t *testing.T) {
cfg := setupDocInitTest(t)
ctx := t.Context()
clientID := "batch_init_complete_" + uuid.New().String()[:8]
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
Clientid: clientID,
Name: "Batch Init Complete",
})
require.NoError(t, err)
// Create batch
batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "test.zip",
TotalDocuments: 1,
})
require.NoError(t, err)
filename := "invoice.pdf"
// Insert submitted outcome
err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
BatchID: batchID,
Filename: filename,
Column3: repository.BatchOutcomeStatusSubmitted,
})
require.NoError(t, err)
// Upload test file and create upload record with batch ID embedded in key
location := uploadTestFileToS3(t, cfg, clientID, &batchID)
uploadID := uuid.New()
err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
ID: uploadID,
Clientid: clientID,
Bucket: cfg.GetBucket(),
Key: location.String(),
Part: 1,
Createdat: pgtype.Timestamp{
Valid: true,
Time: time.Now().UTC(),
},
Filename: &filename,
BatchID: &batchID,
})
require.NoError(t, err)
runnerInstance := docinitrunner.New(&docinitrunner.Services{
Document: documentinit.New(cfg),
Queries: cfg.GetDBQueries(),
Outcome: outcome.New(cfg.GetDBQueries()),
})
doc := docinitrunner.Body{
Bucket: cfg.GetBucket(),
Key: location.String(),
Hash: "unique_hash_init_complete_" + uuid.New().String()[:8],
}
assert.True(t, runnerInstance.Process(ctx, doc))
// Verify outcome was updated to init_complete with document_id set
rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
require.NoError(t, err)
require.Len(t, rows, 1)
assert.Equal(t, repository.BatchOutcomeStatusInitComplete, rows[0].Outcome)
assert.NotNil(t, rows[0].DocumentID, "document_id should be set after init")
}
// TestDocInitRunner_ReportsInitDuplicate processes a batch document whose hash
// matches an existing document and verifies the outcome is init_duplicate.
func TestDocInitRunner_ReportsInitDuplicate(t *testing.T) {
cfg := setupDocInitTest(t)
ctx := t.Context()
clientID := "batch_init_dup_" + uuid.New().String()[:8]
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
Clientid: clientID,
Name: "Batch Init Duplicate",
})
require.NoError(t, err)
// Create an existing document with a known hash
existingHash := "existing_hash_" + uuid.New().String()[:8]
_, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: existingHash,
})
require.NoError(t, err)
// Create batch
batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "test.zip",
TotalDocuments: 1,
})
require.NoError(t, err)
filename := "duplicate.pdf"
// Insert submitted outcome
err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
BatchID: batchID,
Filename: filename,
Column3: repository.BatchOutcomeStatusSubmitted,
})
require.NoError(t, err)
// Upload test file and create upload record with batch ID in key
location := uploadTestFileToS3(t, cfg, clientID, &batchID)
uploadID := uuid.New()
err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
ID: uploadID,
Clientid: clientID,
Bucket: cfg.GetBucket(),
Key: location.String(),
Part: 1,
Createdat: pgtype.Timestamp{
Valid: true,
Time: time.Now().UTC(),
},
Filename: &filename,
BatchID: &batchID,
})
require.NoError(t, err)
runnerInstance := docinitrunner.New(&docinitrunner.Services{
Document: documentinit.New(cfg),
Queries: cfg.GetDBQueries(),
Outcome: outcome.New(cfg.GetDBQueries()),
})
// Use the same hash as the existing document to trigger duplicate
doc := docinitrunner.Body{
Bucket: cfg.GetBucket(),
Key: location.String(),
Hash: existingHash,
}
assert.True(t, runnerInstance.Process(ctx, doc))
// Verify outcome was updated to init_duplicate
rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
require.NoError(t, err)
require.Len(t, rows, 1)
assert.Equal(t, repository.BatchOutcomeStatusInitDuplicate, rows[0].Outcome)
assert.NotNil(t, rows[0].DocumentID, "document_id should be set for duplicate")
}
// TestDocInitRunner_NoOpForNonBatchDocument verifies that processing a non-batch
// document does not cause outcome errors (the Resolve call is a no-op for nil batchID).
func TestDocInitRunner_NoOpForNonBatchDocument(t *testing.T) {
cfg := setupDocInitTest(t)
ctx := t.Context()
clientID := "nonbatch_init_" + uuid.New().String()[:8]
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
Clientid: clientID,
Name: "NonBatch Init",
})
require.NoError(t, err)
location := uploadTestFileToS3(t, cfg, clientID, nil)
uploadID := uuid.New()
fname := "regular.pdf"
err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
ID: uploadID,
Clientid: clientID,
Bucket: cfg.GetBucket(),
Key: location.String(),
Part: 1,
Createdat: pgtype.Timestamp{
Valid: true,
Time: time.Now().UTC(),
},
Filename: &fname,
BatchID: nil, // Non-batch
})
require.NoError(t, err)
runnerInstance := docinitrunner.New(&docinitrunner.Services{
Document: documentinit.New(cfg),
Queries: cfg.GetDBQueries(),
Outcome: outcome.New(cfg.GetDBQueries()),
})
doc := docinitrunner.Body{
Bucket: cfg.GetBucket(),
Key: location.String(),
Hash: "nonbatch_hash_" + uuid.New().String()[:8],
}
// Should process successfully without any outcome errors
assert.True(t, runnerInstance.Process(ctx, doc))
}