diff --git a/coordinator/go.mod b/coordinator/go.mod index e424a5e..4fc8bcc 100644 --- a/coordinator/go.mod +++ b/coordinator/go.mod @@ -3,6 +3,7 @@ module github.com/emil28092005/SciMesh/coordinator go 1.22 require ( + github.com/Masterminds/squirrel v1.5.4 github.com/cenkalti/backoff/v4 v4.3.0 github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.6.0 @@ -14,6 +15,8 @@ require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect github.com/jackc/puddle/v2 v2.2.1 // indirect + github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect + github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect golang.org/x/crypto v0.17.0 // indirect golang.org/x/sync v0.1.0 // indirect golang.org/x/text v0.14.0 // indirect diff --git a/coordinator/go.sum b/coordinator/go.sum index 7e95bc7..d9c880a 100644 --- a/coordinator/go.sum +++ b/coordinator/go.sum @@ -1,3 +1,5 @@ +github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8afzqM= +github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10= github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8= github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -15,9 +17,14 @@ github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= +github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 h1:SOEGU9fKiNWd/HOJuq6+3iTQz8KNCLtVX6idSoTLdUw= +github.com/lann/builder v0.0.0-20180802200727-47ae307949d0/go.mod h1:dXGbAdH5GtBTC4WfIxhKZfyBF/HBFgRZSWwZ9g/He9o= +github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 h1:P6pPBnrTSX3DEVR4fDembhRWSsG5rVo6hYhAB/ADZrk= +github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0/go.mod h1:vmVJ0l/dxyfGW6FmdpVm2joNMFikkuWg0EoCKLGUMNw= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= diff --git a/coordinator/internal/storage/postgres/builder.go b/coordinator/internal/storage/postgres/builder.go new file mode 100644 index 0000000..0025347 --- /dev/null +++ b/coordinator/internal/storage/postgres/builder.go @@ -0,0 +1,11 @@ +package postgres + +import sq "github.com/Masterminds/squirrel" + +// psql is the shared statement builder, fixed to PostgreSQL $N placeholders so +// no call site repeats PlaceholderFormat(sq.Dollar). +// +// Not everything goes through it. Two genuinely set-based statements stay as +// raw SQL — claimNext (a FOR UPDATE SKIP LOCKED CTE) and expireLeases (CASE +// logic in the SET) — because a builder would obscure them, not clarify them. +var psql = sq.StatementBuilder.PlaceholderFormat(sq.Dollar) diff --git a/coordinator/internal/storage/postgres/job_repo.go b/coordinator/internal/storage/postgres/job_repo.go index 1945a07..f2a7e2e 100644 --- a/coordinator/internal/storage/postgres/job_repo.go +++ b/coordinator/internal/storage/postgres/job_repo.go @@ -5,6 +5,7 @@ import ( "errors" "time" + sq "github.com/Masterminds/squirrel" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -24,28 +25,36 @@ func NewJobRepo(pool *pgxpool.Pool) *JobRepo { var _ usecase.JobRepository = (*JobRepo)(nil) -const jobColumns = `id, workload, input_uri, parameters, status, created_at, completed_at` - -const insertJobSQL = ` -INSERT INTO jobs (id, workload, input_uri, parameters, status, created_at) -VALUES ($1, $2, $3, $4, $5, $6)` +var jobColumns = []string{"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at"} // Insert runs inside the caller's transaction, alongside the job's tasks — that // is what makes "all tasks or none" hold. func (r *JobRepo) Insert(ctx context.Context, j *domain.Job) error { - _, err := conn(ctx, r.pool).Exec(ctx, insertJobSQL, - j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt) + sql, args, err := psql.Insert("jobs"). + Columns("id", "workload", "input_uri", "parameters", "status", "created_at"). + Values(j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt). + ToSql() + if err != nil { + return err + } + _, err = conn(ctx, r.pool).Exec(ctx, sql, args...) return err } -const getJobSQL = `SELECT ` + jobColumns + ` FROM jobs WHERE id = $1` - func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) { + sql, args, err := psql.Select(jobColumns...). + From("jobs"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + var ( j domain.Job status string ) - err := conn(ctx, r.pool).QueryRow(ctx, getJobSQL, id).Scan( + err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( &j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt) if errors.Is(err, pgx.ErrNoRows) { return nil, domain.ErrJobNotFound @@ -57,12 +66,21 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) { return &j, nil } -const updateJobStatusSQL = `UPDATE jobs SET status = $2, completed_at = $3 WHERE id = $1` - func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error { - tag, err := conn(ctx, r.pool).Exec(ctx, updateJobStatusSQL, id, string(status), completedAt) + sql, args, err := psql.Update("jobs"). + SetMap(map[string]any{ + "status": string(status), + "completed_at": completedAt, + }). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return err + } + + tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...) if err != nil { return err } diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go index 4c0abe1..3b26038 100644 --- a/coordinator/internal/storage/postgres/task_repo.go +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -3,8 +3,10 @@ package postgres import ( "context" "errors" + "strings" "time" + sq "github.com/Masterminds/squirrel" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -27,10 +29,16 @@ var _ usecase.TaskRepository = (*TaskRepo)(nil) // taskColumns is the single source of truth for the shape scanTask expects. // Every query that returns a task selects exactly this list, in this order — // three hand-written column lists would drift apart within a week. -const taskColumns = `id, job_id, chunk_index, workload, input_uri, input_sha256, - parameters, status, attempt, max_attempts, lease_owner, lease_expires_at, - result_uri, result_sha256, metrics, error_code, error_message, - created_at, started_at, completed_at, version` +var taskColumns = []string{ + "id", "job_id", "chunk_index", "workload", "input_uri", "input_sha256", + "parameters", "status", "attempt", "max_attempts", "lease_owner", "lease_expires_at", + "result_uri", "result_sha256", "metrics", "error_code", "error_message", + "created_at", "started_at", "completed_at", "version", +} + +// taskColumnList is the same set as a comma string, for the raw claim query's +// RETURNING clause, which the builder does not touch. +var taskColumnList = strings.Join(taskColumns, ", ") // scanTask maps one row onto an entity. // @@ -57,14 +65,13 @@ func scanTask(row pgx.Row) (*domain.Task, error) { // claimNextSQL leases one task in a single statement. // -// FOR UPDATE SKIP LOCKED is what makes concurrent coordinators safe: each -// process locks a different candidate row instead of queueing on the same one, -// so no task is ever handed to two workers and no claim blocks behind another. -// Splitting this into SELECT + UPDATE would reintroduce exactly that race. -// -// The CTE column is aliased to cid so the RETURNING list below can use bare -// column names without colliding with the candidate relation. -const claimNextSQL = ` +// Left as raw SQL on purpose: it is a data-modifying CTE with FOR UPDATE SKIP +// LOCKED, which no query builder expresses — and which is the whole point. +// SKIP LOCKED is what makes concurrent coordinators safe: each process locks a +// different candidate row instead of queueing on the same one, so no task is +// ever handed to two workers and no claim blocks behind another. Splitting this +// into SELECT + UPDATE would reintroduce exactly that race. +var claimNextSQL = ` WITH candidate AS ( SELECT id AS cid FROM tasks @@ -84,7 +91,7 @@ SET status = 'leased', version = version + 1 FROM candidate WHERE tasks.id = candidate.cid -RETURNING ` + taskColumns +RETURNING ` + taskColumnList // ClaimNext atomically leases the next eligible task. func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domain.Task, error) { @@ -113,12 +120,18 @@ func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domai return task, nil } -const getForUpdateSQL = `SELECT ` + taskColumns + ` FROM tasks WHERE id = $1 FOR UPDATE` - // GetForUpdate reads a task and holds its row lock until the caller's // transaction ends, so read-modify-write use cases cannot interleave. func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) { - t, err := scanTask(conn(ctx, r.pool).QueryRow(ctx, getForUpdateSQL, id)) + sql, args, err := psql.Select(taskColumns...). + From("tasks"). + Where(sq.Eq{"id": id}). + Suffix("FOR UPDATE"). + ToSql() + if err != nil { + return nil, err + } + t, err := scanTask(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) if errors.Is(err, pgx.ErrNoRows) { return nil, domain.ErrTaskNotFound } @@ -128,32 +141,32 @@ func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task return t, nil } -// updateTaskSQL writes the mutated entity back under optimistic concurrency. -// -// The entity has already incremented its Version in memory, so the new value -// goes into SET while the guard in WHERE compares against the previous one. -const updateTaskSQL = ` -UPDATE tasks -SET status = $2, - attempt = $3, - lease_owner = $4, - lease_expires_at = $5, - result_uri = $6, - result_sha256 = $7, - metrics = $8, - error_code = $9, - error_message = $10, - started_at = $11, - completed_at = $12, - version = $13 -WHERE id = $1 AND version = $13 - 1` - +// Update writes the mutated entity back under optimistic concurrency. The entity +// has already incremented its Version in memory, so the new value goes into SET +// while the WHERE guard matches against the previous one (Version-1). func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { - tag, err := conn(ctx, r.pool).Exec(ctx, updateTaskSQL, - t.ID, string(t.Status), t.Attempt, t.LeaseOwner, t.LeaseExpiresAt, - t.ResultURI, t.ResultSHA256, t.Metrics, t.ErrorCode, t.ErrorMessage, - t.StartedAt, t.CompletedAt, t.Version, - ) + sql, args, err := psql.Update("tasks"). + SetMap(map[string]any{ + "status": string(t.Status), + "attempt": t.Attempt, + "lease_owner": t.LeaseOwner, + "lease_expires_at": t.LeaseExpiresAt, + "result_uri": t.ResultURI, + "result_sha256": t.ResultSHA256, + "metrics": t.Metrics, + "error_code": t.ErrorCode, + "error_message": t.ErrorMessage, + "started_at": t.StartedAt, + "completed_at": t.CompletedAt, + "version": t.Version, + }). + Where(sq.Eq{"id": t.ID, "version": t.Version - 1}). + ToSql() + if err != nil { + return err + } + + tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...) if err != nil { return err } @@ -165,11 +178,6 @@ func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { return nil } -const insertTaskSQL = ` -INSERT INTO tasks (id, job_id, chunk_index, workload, input_uri, input_sha256, - parameters, status, attempt, max_attempts, created_at, version) -VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)` - // InsertBatch writes every task in one round trip. It runs inside the caller's // transaction, which is what makes "all tasks or none" hold. func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error { @@ -179,10 +187,16 @@ func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error batch := &pgx.Batch{} for _, t := range tasks { - batch.Queue(insertTaskSQL, - t.ID, t.JobID, t.ChunkIndex, t.Workload, t.InputURI, t.InputSHA256, - jsonbOrEmpty(t.Parameters), string(t.Status), t.Attempt, t.MaxAttempts, t.CreatedAt, t.Version, - ) + sql, args, err := psql.Insert("tasks"). + Columns("id", "job_id", "chunk_index", "workload", "input_uri", "input_sha256", + "parameters", "status", "attempt", "max_attempts", "created_at", "version"). + Values(t.ID, t.JobID, t.ChunkIndex, t.Workload, t.InputURI, t.InputSHA256, + jsonbOrEmpty(t.Parameters), string(t.Status), t.Attempt, t.MaxAttempts, t.CreatedAt, t.Version). + ToSql() + if err != nil { + return err + } + batch.Queue(sql, args...) } results := conn(ctx, r.pool).SendBatch(ctx, batch) @@ -195,17 +209,20 @@ func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error return results.Close() } -const listCompletedSQL = ` -SELECT ` + taskColumns + ` -FROM tasks -WHERE job_id = $1 AND status = 'completed' -ORDER BY chunk_index` - // ListCompleted returns results in chunk order, which the stitcher relies on: // a non-deterministic order would make the merged output depend on which worker // happened to finish first. func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) { - rows, err := conn(ctx, r.pool).Query(ctx, listCompletedSQL, jobID) + sql, args, err := psql.Select(taskColumns...). + From("tasks"). + Where(sq.Eq{"job_id": jobID, "status": "completed"}). + OrderBy("chunk_index"). + ToSql() + if err != nil { + return nil, err + } + + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) if err != nil { return nil, err } @@ -222,10 +239,17 @@ func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domai return tasks, rows.Err() } -const countByStatusSQL = `SELECT status, count(*) FROM tasks WHERE job_id = $1 GROUP BY status` - func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) { - rows, err := conn(ctx, r.pool).Query(ctx, countByStatusSQL, jobID) + sql, args, err := psql.Select("status", "count(*)"). + From("tasks"). + Where(sq.Eq{"job_id": jobID}). + GroupBy("status"). + ToSql() + if err != nil { + return nil, err + } + + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) if err != nil { return nil, err } @@ -248,10 +272,11 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma // expireLeasesSQL applies the lease-expiry rule set-based, mirroring // domain.Task.ExpireLease: requeue while attempts remain, otherwise fail. // -// It is one statement rather than a load-decide-save loop because several -// coordinators run it concurrently; an atomic UPDATE makes the duplicate work -// harmless — the loser simply updates zero rows. -const expireLeasesSQL = ` +// Left as raw SQL: the branching lives in CASE expressions inside the SET, which +// a builder cannot express more clearly than this. It is one statement rather +// than a load-decide-save loop because several coordinators run it concurrently; +// an atomic UPDATE makes the duplicate work harmless — the loser updates zero rows. +var expireLeasesSQL = ` UPDATE tasks SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_status ELSE 'failed'::task_status END, diff --git a/coordinator/internal/storage/postgres/worker_repo.go b/coordinator/internal/storage/postgres/worker_repo.go index 68f6aaf..034a97f 100644 --- a/coordinator/internal/storage/postgres/worker_repo.go +++ b/coordinator/internal/storage/postgres/worker_repo.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" + sq "github.com/Masterminds/squirrel" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -21,27 +22,34 @@ func NewWorkerRepo(pool *pgxpool.Pool) *WorkerRepo { return &WorkerRepo{pool: pool} } -const workerColumns = `id, name, capabilities, status, last_heartbeat_at, created_at, updated_at` - -const insertWorkerSQL = ` -INSERT INTO workers (` + workerColumns + `) -VALUES ($1, $2, $3, $4, $5, $6, $7)` +var workerColumns = []string{"id", "name", "capabilities", "status", "last_heartbeat_at", "created_at", "updated_at"} func (r *WorkerRepo) Insert(ctx context.Context, w *domain.Worker) error { - // capabilities is a jsonb column; pgx marshals the []string to a JSON array. - _, err := conn(ctx, r.pool).Exec(ctx, insertWorkerSQL, - w.ID, w.Name, w.Capabilities, string(w.Status), - w.LastHeartbeatAt, w.CreatedAt, w.UpdatedAt) + sql, args, err := psql.Insert("workers"). + Columns(workerColumns...). + // capabilities is a jsonb column; pgx marshals the []string to a JSON array. + Values(w.ID, w.Name, w.Capabilities, string(w.Status), + w.LastHeartbeatAt, w.CreatedAt, w.UpdatedAt). + ToSql() if err != nil { + return err + } + if _, err := conn(ctx, r.pool).Exec(ctx, sql, args...); err != nil { return fmt.Errorf("insert worker: %w", err) } return nil } -const getWorkerSQL = `SELECT ` + workerColumns + ` FROM workers WHERE id = $1` - func (r *WorkerRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Worker, error) { - w, err := scanWorker(conn(ctx, r.pool).QueryRow(ctx, getWorkerSQL, id)) + sql, args, err := psql.Select(workerColumns...). + From("workers"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + + w, err := scanWorker(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) if errors.Is(err, pgx.ErrNoRows) { return nil, domain.ErrWorkerNotFound }