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

schema

Browse Source
This commit is contained in:
jay brown
2025-08-04 10:54:22 -07:00
parent 582e5b77c0
commit 56fb9b4ab6
20 changed files with 1026 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/database/migrations/00000000000103_batch_uploads.down.sql
+13
View File
@@ -0,0 +1,13 @@
-- Drop indexes
DROP INDEX IF EXISTS idx_documents_batch_id;
DROP INDEX IF EXISTS idx_batch_uploads_status;
DROP INDEX IF EXISTS idx_batch_uploads_client_id;
-- Remove batch_id from documents table
ALTER TABLE documents DROP COLUMN IF EXISTS batch_id;
-- Drop batch uploads table
DROP TABLE IF EXISTS batch_uploads;
-- Drop batch status ENUM type
DROP TYPE IF EXISTS batch_status;
internal/database/migrations/00000000000103_batch_uploads.up.sql
+44
View File
@@ -0,0 +1,44 @@
-- Create batch status ENUM type
CREATE TYPE batch_status AS ENUM (
'processing', -- Currently processing
'completed', -- All documents processed (may have individual failures)
'failed', -- Batch-level failure (ZIP corrupt, etc.)
'cancelled' -- Cancelled by user
);
-- Create batch uploads table
CREATE TABLE batch_uploads (
id uuid PRIMARY KEY DEFAULT uuid_generate_v7(),
client_id varchar(255) NOT NULL,
original_filename text NOT NULL,
-- Essential batch metrics
total_documents integer NOT NULL,
processed_documents integer NOT NULL DEFAULT 0,
failed_documents integer NOT NULL DEFAULT 0,
invalid_type_documents integer NOT NULL DEFAULT 0,
-- Status and progress
status batch_status NOT NULL DEFAULT 'processing',
progress_percent integer NOT NULL DEFAULT 0,
-- Failed document filenames (JSON array)
failed_filenames jsonb DEFAULT '[]'::jsonb,
-- Timestamps
created_at timestamp NOT NULL DEFAULT NOW(),
completed_at timestamp,
FOREIGN KEY (client_id) REFERENCES clients(clientId)
);
-- Create indexes for efficient queries
CREATE INDEX idx_batch_uploads_client_id ON batch_uploads(client_id);
CREATE INDEX idx_batch_uploads_status ON batch_uploads(status);
-- Add batch_id to documents table
ALTER TABLE documents ADD COLUMN batch_id uuid NULL;
ALTER TABLE documents ADD FOREIGN KEY (batch_id) REFERENCES batch_uploads(id);
-- Create index on batch_id for efficient batch document queries
CREATE INDEX idx_documents_batch_id ON documents(batch_id);
internal/database/queries/batch.sql
+54
View File
@@ -0,0 +1,54 @@
-- name: CreateBatchUpload :one
INSERT INTO batch_uploads (client_id, original_filename, total_documents)
VALUES ($1, $2, $3) RETURNING id;
-- name: GetBatchUpload :one
SELECT id, client_id, original_filename, total_documents, processed_documents,
failed_documents, invalid_type_documents, status, progress_percent,
failed_filenames, created_at, completed_at
FROM batch_uploads
WHERE id = $1 AND client_id = $2;
-- name: ListBatchUploads :many
SELECT id, client_id, original_filename, total_documents, processed_documents,
failed_documents, invalid_type_documents, status, progress_percent,
created_at, completed_at
FROM batch_uploads
WHERE client_id = $1
ORDER BY created_at DESC
LIMIT $2 OFFSET $3;
-- name: UpdateBatchProgress :exec
UPDATE batch_uploads
SET processed_documents = $2,
failed_documents = $3,
invalid_type_documents = $4,
progress_percent = $5
WHERE id = $1;
-- name: UpdateBatchStatus :exec
UPDATE batch_uploads
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: AddFailedFilename :exec
UPDATE batch_uploads
SET failed_filenames = failed_filenames || jsonb_build_array($2::text)
WHERE id = $1;
-- name: CancelBatchUpload :exec
UPDATE batch_uploads
SET status = 'cancelled',
completed_at = NOW()
WHERE id = $1 AND client_id = $2 AND status = 'processing';
-- name: GetDocumentsByBatchId :many
SELECT id, clientId, hash
FROM documents
WHERE batch_id = $1;
-- name: CountDocumentsByBatchId :one
SELECT COUNT(*) as count
FROM documents
WHERE batch_id = $1;
internal/database/queries/document.sql
+1 -1
View File
@@ -31,7 +31,7 @@ GROUP BY d.id, d.clientId, d.hash;
SELECT id, hash from documents where clientId = @clientId;
-- name: CreateDocument :one
INSERT INTO documents (clientId, hash) VALUES ($1, $2) RETURNING id;
INSERT INTO documents (clientId, hash, batch_id) VALUES ($1, $2, $3) RETURNING id;
-- name: AddDocumentEntry :exec
INSERT INTO documentEntries (documentId, bucket, key) VALUES ($1, $2, $3);
internal/database/repository/batch.sql.go
+304
View File
@@ -0,0 +1,304 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.27.0
// source: batch.sql
package repository
import (
"context"
"github.com/google/uuid"
"github.com/jackc/pgx/v5/pgtype"
)
const addFailedFilename = `-- name: AddFailedFilename :exec
UPDATE batch_uploads
SET failed_filenames = failed_filenames || jsonb_build_array($2::text)
WHERE id = $1
`
type AddFailedFilenameParams struct {
ID uuid.UUID `db:"id"`
Column2 string `db:"column_2"`
}
// AddFailedFilename
//
// UPDATE batch_uploads
// SET failed_filenames = failed_filenames || jsonb_build_array($2::text)
// WHERE id = $1
func (q *Queries) AddFailedFilename(ctx context.Context, arg *AddFailedFilenameParams) error {
_, err := q.db.Exec(ctx, addFailedFilename, arg.ID, arg.Column2)
return err
}
const cancelBatchUpload = `-- name: CancelBatchUpload :exec
UPDATE batch_uploads
SET status = 'cancelled',
completed_at = NOW()
WHERE id = $1 AND client_id = $2 AND status = 'processing'
`
type CancelBatchUploadParams struct {
ID uuid.UUID `db:"id"`
ClientID string `db:"client_id"`
}
// CancelBatchUpload
//
// UPDATE batch_uploads
// SET status = 'cancelled',
// completed_at = NOW()
// WHERE id = $1 AND client_id = $2 AND status = 'processing'
func (q *Queries) CancelBatchUpload(ctx context.Context, arg *CancelBatchUploadParams) error {
_, err := q.db.Exec(ctx, cancelBatchUpload, arg.ID, arg.ClientID)
return err
}
const countDocumentsByBatchId = `-- name: CountDocumentsByBatchId :one
SELECT COUNT(*) as count
FROM documents
WHERE batch_id = $1
`
// CountDocumentsByBatchId
//
// SELECT COUNT(*) as count
// FROM documents
// WHERE batch_id = $1
func (q *Queries) CountDocumentsByBatchId(ctx context.Context, batchID *uuid.UUID) (int64, error) {
row := q.db.QueryRow(ctx, countDocumentsByBatchId, batchID)
var count int64
err := row.Scan(&count)
return count, err
}
const createBatchUpload = `-- name: CreateBatchUpload :one
INSERT INTO batch_uploads (client_id, original_filename, total_documents)
VALUES ($1, $2, $3) RETURNING id
`
type CreateBatchUploadParams struct {
ClientID string `db:"client_id"`
OriginalFilename string `db:"original_filename"`
TotalDocuments int32 `db:"total_documents"`
}
// CreateBatchUpload
//
// INSERT INTO batch_uploads (client_id, original_filename, total_documents)
// VALUES ($1, $2, $3) RETURNING id
func (q *Queries) CreateBatchUpload(ctx context.Context, arg *CreateBatchUploadParams) (uuid.UUID, error) {
row := q.db.QueryRow(ctx, createBatchUpload, arg.ClientID, arg.OriginalFilename, arg.TotalDocuments)
var id uuid.UUID
err := row.Scan(&id)
return id, err
}
const getBatchUpload = `-- name: GetBatchUpload :one
SELECT id, client_id, original_filename, total_documents, processed_documents,
failed_documents, invalid_type_documents, status, progress_percent,
failed_filenames, created_at, completed_at
FROM batch_uploads
WHERE id = $1 AND client_id = $2
`
type GetBatchUploadParams struct {
ID uuid.UUID `db:"id"`
ClientID string `db:"client_id"`
}
// GetBatchUpload
//
// SELECT id, client_id, original_filename, total_documents, processed_documents,
// failed_documents, invalid_type_documents, status, progress_percent,
// failed_filenames, created_at, completed_at
// FROM batch_uploads
// WHERE id = $1 AND client_id = $2
func (q *Queries) GetBatchUpload(ctx context.Context, arg *GetBatchUploadParams) (*BatchUpload, error) {
row := q.db.QueryRow(ctx, getBatchUpload, arg.ID, arg.ClientID)
var i BatchUpload
err := row.Scan(
&i.ID,
&i.ClientID,
&i.OriginalFilename,
&i.TotalDocuments,
&i.ProcessedDocuments,
&i.FailedDocuments,
&i.InvalidTypeDocuments,
&i.Status,
&i.ProgressPercent,
&i.FailedFilenames,
&i.CreatedAt,
&i.CompletedAt,
)
return &i, err
}
const getDocumentsByBatchId = `-- name: GetDocumentsByBatchId :many
SELECT id, clientId, hash
FROM documents
WHERE batch_id = $1
`
type GetDocumentsByBatchIdRow struct {
ID uuid.UUID `db:"id"`
Clientid string `db:"clientid"`
Hash string `db:"hash"`
}
// GetDocumentsByBatchId
//
// SELECT id, clientId, hash
// FROM documents
// WHERE batch_id = $1
func (q *Queries) GetDocumentsByBatchId(ctx context.Context, batchID *uuid.UUID) ([]*GetDocumentsByBatchIdRow, error) {
rows, err := q.db.Query(ctx, getDocumentsByBatchId, batchID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []*GetDocumentsByBatchIdRow{}
for rows.Next() {
var i GetDocumentsByBatchIdRow
if err := rows.Scan(&i.ID, &i.Clientid, &i.Hash); 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,
created_at, completed_at
FROM batch_uploads
WHERE client_id = $1
ORDER BY created_at DESC
LIMIT $2 OFFSET $3
`
type ListBatchUploadsParams struct {
ClientID string `db:"client_id"`
Limit int64 `db:"limit"`
Offset int64 `db:"offset"`
}
type ListBatchUploadsRow 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"`
CreatedAt pgtype.Timestamp `db:"created_at"`
CompletedAt pgtype.Timestamp `db:"completed_at"`
}
// ListBatchUploads
//
// SELECT id, client_id, original_filename, total_documents, processed_documents,
// failed_documents, invalid_type_documents, status, progress_percent,
// created_at, completed_at
// FROM batch_uploads
// WHERE client_id = $1
// ORDER BY created_at DESC
// LIMIT $2 OFFSET $3
func (q *Queries) ListBatchUploads(ctx context.Context, arg *ListBatchUploadsParams) ([]*ListBatchUploadsRow, error) {
rows, err := q.db.Query(ctx, listBatchUploads, arg.ClientID, arg.Limit, arg.Offset)
if err != nil {
return nil, err
}
defer rows.Close()
items := []*ListBatchUploadsRow{}
for rows.Next() {
var i ListBatchUploadsRow
if err := rows.Scan(
&i.ID,
&i.ClientID,
&i.OriginalFilename,
&i.TotalDocuments,
&i.ProcessedDocuments,
&i.FailedDocuments,
&i.InvalidTypeDocuments,
&i.Status,
&i.ProgressPercent,
&i.CreatedAt,
&i.CompletedAt,
); err != nil {
return nil, err
}
items = append(items, &i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const updateBatchProgress = `-- name: UpdateBatchProgress :exec
UPDATE batch_uploads
SET processed_documents = $2,
failed_documents = $3,
invalid_type_documents = $4,
progress_percent = $5
WHERE id = $1
`
type UpdateBatchProgressParams struct {
ID uuid.UUID `db:"id"`
ProcessedDocuments int32 `db:"processed_documents"`
FailedDocuments int32 `db:"failed_documents"`
InvalidTypeDocuments int32 `db:"invalid_type_documents"`
ProgressPercent int32 `db:"progress_percent"`
}
// UpdateBatchProgress
//
// UPDATE batch_uploads
// SET processed_documents = $2,
// failed_documents = $3,
// invalid_type_documents = $4,
// progress_percent = $5
// WHERE id = $1
func (q *Queries) UpdateBatchProgress(ctx context.Context, arg *UpdateBatchProgressParams) error {
_, err := q.db.Exec(ctx, updateBatchProgress,
arg.ID,
arg.ProcessedDocuments,
arg.FailedDocuments,
arg.InvalidTypeDocuments,
arg.ProgressPercent,
)
return err
}
const updateBatchStatus = `-- name: UpdateBatchStatus :exec
UPDATE batch_uploads
SET status = $2::batch_status,
completed_at = CASE WHEN $2::batch_status IN ('completed', 'failed', 'cancelled') THEN NOW() ELSE NULL END
WHERE id = $1
`
type UpdateBatchStatusParams struct {
ID uuid.UUID `db:"id"`
Column2 BatchStatus `db:"column_2"`
}
// UpdateBatchStatus
//
// UPDATE batch_uploads
// SET status = $2::batch_status,
// completed_at = CASE WHEN $2::batch_status IN ('completed', 'failed', 'cancelled') THEN NOW() ELSE NULL END
// WHERE id = $1
func (q *Queries) UpdateBatchStatus(ctx context.Context, arg *UpdateBatchStatusParams) error {
_, err := q.db.Exec(ctx, updateBatchStatus, arg.ID, arg.Column2)
return err
}
internal/database/repository/batch_test.go
+475
View File
@@ -0,0 +1,475 @@
package repository_test
import (
"fmt"
"testing"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/test"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestBatchUpload(t *testing.T) {
if testing.Short() {
t.SkipNow()
}
t.Run("CreateBatchUpload", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client first
clientID := fmt.Sprintf("TEST_CLIENT_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
params := &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "test_batch.zip",
TotalDocuments: 10,
}
batchID, err := queries.CreateBatchUpload(ctx, params)
require.NoError(t, err)
assert.NotEqual(t, uuid.Nil, batchID)
})
t.Run("GetBatchUpload", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_GET_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Get",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "test_get.zip",
TotalDocuments: 5,
})
require.NoError(t, err)
// Get batch upload
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
assert.Equal(t, batchID, batch.ID)
assert.Equal(t, clientID, batch.ClientID)
assert.Equal(t, "test_get.zip", batch.OriginalFilename)
assert.Equal(t, int32(5), batch.TotalDocuments)
assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
assert.Equal(t, int32(0), batch.ProcessedDocuments)
assert.Equal(t, int32(0), batch.FailedDocuments)
assert.NotNil(t, batch.CreatedAt)
assert.False(t, batch.CompletedAt.Valid) // pgtype.Timestamp has Valid field
})
t.Run("ListBatchUploads", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_LIST_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client List",
Clientid: clientID,
})
require.NoError(t, err)
// Create multiple batch uploads
for i := 0; i < 3; i++ {
totalDocs := int32(1) + int32(i) //nolint:gosec // safe conversion for test data
_, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: fmt.Sprintf("batch%d.zip", i),
TotalDocuments: totalDocs,
})
require.NoError(t, err)
}
// List batch uploads
batches, err := queries.ListBatchUploads(ctx, &repository.ListBatchUploadsParams{
ClientID: clientID,
Limit: 10,
Offset: 0,
})
require.NoError(t, err)
assert.Len(t, batches, 3)
})
t.Run("UpdateBatchProgress", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_PROGRESS_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Progress",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "progress.zip",
TotalDocuments: 100,
})
require.NoError(t, err)
// Update progress
err = queries.UpdateBatchProgress(ctx, &repository.UpdateBatchProgressParams{
ID: batchID,
ProcessedDocuments: 42,
FailedDocuments: 2,
InvalidTypeDocuments: 1,
ProgressPercent: 45,
})
require.NoError(t, err)
// Verify update
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
assert.Equal(t, int32(42), batch.ProcessedDocuments)
assert.Equal(t, int32(2), batch.FailedDocuments)
assert.Equal(t, int32(1), batch.InvalidTypeDocuments)
assert.Equal(t, int32(45), batch.ProgressPercent)
})
t.Run("UpdateBatchStatus", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_STATUS_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Status",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "status.zip",
TotalDocuments: 10,
})
require.NoError(t, err)
// Update status to completed
err = queries.UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
ID: batchID,
Column2: repository.BatchStatusCompleted,
})
require.NoError(t, err)
// Verify update
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
assert.Equal(t, repository.BatchStatusCompleted, batch.Status)
assert.True(t, batch.CompletedAt.Valid)
assert.True(t, batch.CompletedAt.Time.After(batch.CreatedAt.Time))
})
t.Run("AddFailedFilename", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_FAILED_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Failed",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "failed.zip",
TotalDocuments: 10,
})
require.NoError(t, err)
// Add failed filenames
err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{
ID: batchID,
Column2: "document1.pdf",
})
require.NoError(t, err)
err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{
ID: batchID,
Column2: "document2.pdf",
})
require.NoError(t, err)
// Verify failed filenames were added
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
// Verify failed filenames were added as JSON array
assert.NotNil(t, batch.FailedFilenames)
// The failed filenames should contain our test filenames
failedFilenamesStr := string(batch.FailedFilenames)
assert.Contains(t, failedFilenamesStr, "document1.pdf")
assert.Contains(t, failedFilenamesStr, "document2.pdf")
})
t.Run("CancelBatchUpload", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_CANCEL_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Cancel",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "cancel.zip",
TotalDocuments: 10,
})
require.NoError(t, err)
// Cancel batch upload
err = queries.CancelBatchUpload(ctx, &repository.CancelBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
// Verify cancellation
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
ID: batchID,
ClientID: clientID,
})
require.NoError(t, err)
assert.Equal(t, repository.BatchStatusCancelled, batch.Status)
assert.True(t, batch.CompletedAt.Valid)
})
t.Run("GetDocumentsByBatchId", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_DOCS_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Docs",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "docs.zip",
TotalDocuments: 3,
})
require.NoError(t, err)
// Create documents with batch ID
var docIDs []uuid.UUID
for i := 0; i < 3; i++ {
docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: fmt.Sprintf("hash%d", i),
BatchID: &batchID,
})
require.NoError(t, err)
docIDs = append(docIDs, docID)
}
// Get documents by batch ID
docs, err := queries.GetDocumentsByBatchId(ctx, &batchID)
require.NoError(t, err)
assert.Len(t, docs, 3)
// Verify all documents are returned
for _, doc := range docs {
assert.Contains(t, docIDs, doc.ID)
assert.Equal(t, clientID, doc.Clientid)
}
})
t.Run("CountDocumentsByBatchId", func(t *testing.T) {
t.Parallel()
ctx := t.Context()
cfg := &serviceconfig.BaseConfig{}
test.CreateDB(t, cfg)
queries := cfg.GetDBQueries()
// Create a client
clientID := fmt.Sprintf("TEST_CLIENT_COUNT_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "Test Client Count",
Clientid: clientID,
})
require.NoError(t, err)
// Create batch upload
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
ClientID: clientID,
OriginalFilename: "count.zip",
TotalDocuments: 5,
})
require.NoError(t, err)
// Create documents with batch ID
for i := 0; i < 4; i++ {
_, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: fmt.Sprintf("hash%d", i),
BatchID: &batchID,
})
require.NoError(t, err)
}
// Count documents by batch ID
count, err := queries.CountDocumentsByBatchId(ctx, &batchID)
require.NoError(t, err)
assert.Equal(t, int64(4), count)
})
}
func TestBatchStatusEnum(t *testing.T) {
t.Run("Scan method", func(t *testing.T) {
// Test valid scan with string
var status repository.BatchStatus
err := status.Scan("processing")
require.NoError(t, err)
assert.Equal(t, repository.BatchStatusProcessing, status)
// Test valid scan with []byte
var statusBytes repository.BatchStatus
err = statusBytes.Scan([]byte("completed"))
require.NoError(t, err)
assert.Equal(t, repository.BatchStatusCompleted, statusBytes)
// Test all enum values
testCases := []struct {
str string
expected repository.BatchStatus
}{
{"processing", repository.BatchStatusProcessing},
{"completed", repository.BatchStatusCompleted},
{"failed", repository.BatchStatusFailed},
{"cancelled", repository.BatchStatusCancelled}, //nolint:misspell // matches database enum value
}
for _, tc := range testCases {
var s repository.BatchStatus
err := s.Scan(tc.str)
require.NoError(t, err)
assert.Equal(t, tc.expected, s)
}
// Test scan with unsupported type
var invalidStatus repository.BatchStatus
err = invalidStatus.Scan(123) // int is not supported
assert.Error(t, err)
assert.Contains(t, err.Error(), "unsupported scan type")
})
t.Run("NullBatchStatus", func(t *testing.T) {
// Test scan with nil
var nullStatus repository.NullBatchStatus
err := nullStatus.Scan(nil)
require.NoError(t, err)
assert.False(t, nullStatus.Valid)
assert.Equal(t, repository.BatchStatus(""), nullStatus.BatchStatus)
// Test scan with valid value
err = nullStatus.Scan("processing")
require.NoError(t, err)
assert.True(t, nullStatus.Valid)
assert.Equal(t, repository.BatchStatusProcessing, nullStatus.BatchStatus)
// Test Value method when Valid is false
nullStatus.Valid = false // reset to false
val, err := nullStatus.Value()
require.NoError(t, err)
assert.Nil(t, val)
// Test Value method when Valid is true
nullStatus.Valid = true
nullStatus.BatchStatus = repository.BatchStatusCompleted
val, err = nullStatus.Value()
require.NoError(t, err)
assert.Equal(t, "completed", val)
})
t.Run("BatchStatus Valid method", func(t *testing.T) {
// Test valid enum values
assert.True(t, repository.BatchStatusProcessing.Valid())
assert.True(t, repository.BatchStatusCompleted.Valid())
assert.True(t, repository.BatchStatusFailed.Valid())
assert.True(t, repository.BatchStatusCancelled.Valid()) //nolint:misspell // matches database enum value
// Test invalid enum value
invalidStatus := repository.BatchStatus("invalid")
assert.False(t, invalidStatus.Valid())
})
}
internal/database/repository/clean_test.go
+1
View File
@@ -34,6 +34,7 @@ func TestClean(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
internal/database/repository/document.sql.go
+14 -7
View File
@@ -59,19 +59,20 @@ func (q *Queries) AddDocumentUpload(ctx context.Context, arg *AddDocumentUploadP
}
const createDocument = `-- name: CreateDocument :one
INSERT INTO documents (clientId, hash) VALUES ($1, $2) RETURNING id
INSERT INTO documents (clientId, hash, batch_id) VALUES ($1, $2, $3) RETURNING id
`
type CreateDocumentParams struct {
Clientid string `db:"clientid"`
Hash string `db:"hash"`
Clientid string `db:"clientid"`
Hash string `db:"hash"`
BatchID *uuid.UUID `db:"batch_id"`
}
// CreateDocument
//
// INSERT INTO documents (clientId, hash) VALUES ($1, $2) RETURNING id
// INSERT INTO documents (clientId, hash, batch_id) VALUES ($1, $2, $3) RETURNING id
func (q *Queries) CreateDocument(ctx context.Context, arg *CreateDocumentParams) (uuid.UUID, error) {
row := q.db.QueryRow(ctx, createDocument, arg.Clientid, arg.Hash)
row := q.db.QueryRow(ctx, createDocument, arg.Clientid, arg.Hash, arg.BatchID)
var id uuid.UUID
err := row.Scan(&id)
return id, err
@@ -192,12 +193,18 @@ const getDocumentSummary = `-- name: GetDocumentSummary :one
SELECT id, clientId, hash FROM documents WHERE id = $1
`
type GetDocumentSummaryRow struct {
ID uuid.UUID `db:"id"`
Clientid string `db:"clientid"`
Hash string `db:"hash"`
}
// GetDocumentSummary
//
// SELECT id, clientId, hash FROM documents WHERE id = $1
func (q *Queries) GetDocumentSummary(ctx context.Context, id uuid.UUID) (*Document, error) {
func (q *Queries) GetDocumentSummary(ctx context.Context, id uuid.UUID) (*GetDocumentSummaryRow, error) {
row := q.db.QueryRow(ctx, getDocumentSummary, id)
var i Document
var i GetDocumentSummaryRow
err := row.Scan(&i.ID, &i.Clientid, &i.Hash)
return &i, err
}
internal/database/repository/document_test.go
+23 -11
View File
@@ -1,12 +1,14 @@
package repository_test
import (
"fmt"
"testing"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/test"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -24,7 +26,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -44,7 +46,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -55,6 +57,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -74,7 +77,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -85,6 +88,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -92,6 +96,7 @@ func TestDocument(t *testing.T) {
documentTwoID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_hash_two",
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, documentTwoID)
@@ -109,7 +114,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -120,11 +125,12 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
clientTwoId := "EXAMPLE TWO"
clientTwoId := fmt.Sprintf("EXAMPLE_TWO_%s", uuid.New().String()[:8])
err = queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_name_two",
Clientid: clientTwoId,
@@ -135,6 +141,7 @@ func TestDocument(t *testing.T) {
idTwo, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientTwoId,
Hash: hashTwo,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, idTwo)
@@ -155,7 +162,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -166,13 +173,14 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
doc, err := queries.GetDocumentSummary(ctx, id)
require.NoError(t, err)
assert.EqualExportedValues(t, &repository.Document{
assert.EqualExportedValues(t, &repository.GetDocumentSummaryRow{
ID: id,
Clientid: clientId,
Hash: hash,
@@ -187,7 +195,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -198,6 +206,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -220,7 +229,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -231,6 +240,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -251,7 +261,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -262,6 +272,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -292,7 +303,7 @@ func TestDocument(t *testing.T) {
queries := cfg.GetDBQueries()
clientId := "EXAMPLE"
clientId := fmt.Sprintf("EXAMPLE_%s", uuid.New().String()[:8])
err := queries.CreateClient(ctx, &repository.CreateClientParams{
Name: "example_client",
Clientid: clientId,
@@ -303,6 +314,7 @@ func TestDocument(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
internal/database/repository/models.go
+74 -3
View File
@@ -12,6 +12,61 @@ import (
"github.com/jackc/pgx/v5/pgtype"
)
type BatchStatus string
const (
BatchStatusProcessing BatchStatus = "processing"
BatchStatusCompleted BatchStatus = "completed"
BatchStatusFailed BatchStatus = "failed"
BatchStatusCancelled BatchStatus = "cancelled"
)
func (e *BatchStatus) Scan(src interface{}) error {
switch s := src.(type) {
case []byte:
*e = BatchStatus(s)
case string:
*e = BatchStatus(s)
default:
return fmt.Errorf("unsupported scan type for BatchStatus: %T", src)
}
return nil
}
type NullBatchStatus struct {
BatchStatus BatchStatus
Valid bool // Valid is true if BatchStatus is not NULL
}
// Scan implements the Scanner interface.
func (ns *NullBatchStatus) Scan(value interface{}) error {
if value == nil {
ns.BatchStatus, ns.Valid = "", false
return nil
}
ns.Valid = true
return ns.BatchStatus.Scan(value)
}
// Value implements the driver Valuer interface.
func (ns NullBatchStatus) Value() (driver.Value, error) {
if !ns.Valid {
return nil, nil
}
return string(ns.BatchStatus), nil
}
func (e BatchStatus) Valid() bool {
switch e {
case BatchStatusProcessing,
BatchStatusCompleted,
BatchStatusFailed,
BatchStatusCancelled:
return true
}
return false
}
type Cleanfailtype string
const (
@@ -177,6 +232,21 @@ func (e Querytype) Valid() bool {
return false
}
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"`
}
type Client struct {
Clientid string `db:"clientid"`
Name string `db:"name"`
@@ -292,9 +362,10 @@ type Currenttextentry struct {
}
type Document struct {
ID uuid.UUID `db:"id"`
Clientid string `db:"clientid"`
Hash string `db:"hash"`
ID uuid.UUID `db:"id"`
Clientid string `db:"clientid"`
Hash string `db:"hash"`
BatchID *uuid.UUID `db:"batch_id"`
}
type Documentclean struct {
internal/database/repository/query_test.go
+2
View File
@@ -249,6 +249,7 @@ func TestQueryDependencyTree(t *testing.T) {
docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: "sample",
BatchID: nil,
})
require.NoError(t, err)
@@ -604,6 +605,7 @@ func BenchmarkListQueryDirectDependentsByDocId(b *testing.B) {
docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientID,
Hash: "sample",
BatchID: nil,
})
require.NoError(b, err)
internal/database/repository/result_test.go
+3
View File
@@ -65,11 +65,13 @@ func TestResultValues(t *testing.T) {
documentID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_hash",
BatchID: nil,
})
require.NoError(t, err)
documentTwoID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_hash_two",
BatchID: nil,
})
require.NoError(t, err)
@@ -519,6 +521,7 @@ func BenchmarkListUnsynced(b *testing.B) {
documentID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_hash",
BatchID: nil,
})
require.NoError(b, err)
internal/database/repository/sync_test.go
+4
View File
@@ -45,6 +45,7 @@ func TestListClientDocumentIDs(t *testing.T) {
docOne, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: id,
Hash: "example_hash",
BatchID: nil,
})
require.NoError(t, err)
@@ -66,6 +67,7 @@ func TestListClientDocumentIDs(t *testing.T) {
docTwo, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: id,
Hash: "example_hash_two",
BatchID: nil,
})
require.NoError(t, err)
@@ -268,6 +270,7 @@ func TestClientSync(t *testing.T) {
documentID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_noclean",
BatchID: nil,
})
require.NoError(t, err)
@@ -1003,6 +1006,7 @@ func createDocumentWithCollector(t testing.TB, queries *repository.Queries, clie
documentID, err := queries.CreateDocument(t.Context(), &repository.CreateDocumentParams{
Clientid: clientId,
Hash: "example_hash",
BatchID: nil,
})
require.NoError(t, err)
internal/database/repository/text_test.go
+3
View File
@@ -37,6 +37,7 @@ func TestTextExtraction(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -144,6 +145,7 @@ func TestTextTextractPart(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)
@@ -273,6 +275,7 @@ func TestTextOutPart(t *testing.T) {
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
Clientid: clientId,
Hash: hash,
BatchID: nil,
})
require.NoError(t, err)
assert.NotEmpty(t, id)