feat(coordinator): upload a dataset and chunk it into shard tasks (CTX-05, part 4)
The coordinator can now ingest a dataset itself, not only accept client-supplied
chunk URIs.
- internal/chunk: a deterministic, generic TSV row splitter — repeats the header
per shard, buffers one shard at a time, rejects header-only input. Unit-tested.
- POST /jobs/upload (multipart): streams the dataset into an input artifact,
splits it into shard artifacts, and creates one shard task per shard, all in
one transaction; blobs are cleaned up if the transaction fails.
- GET /tasks/{id}/input streams a task's input shard back to the worker.
- domain: NewUploadedJob, NewShardTask, Task/Job.InputArtifactID; a shard task's
input is an artifact, not a URI. Claim response nests input:{uri,sha256} per
the contract, with uri = /tasks/{id}/input for shards.
- migration 0005 makes input_uri nullable and adds a has-input check.
- The existing URI-based POST /jobs path is untouched; both coexist.
This commit is contained in:
@@ -28,7 +28,7 @@ var _ usecase.BlobStore = (*FSStore)(nil)
|
||||
// rename is only atomic within one filesystem.
|
||||
func NewFSStore(dir string) (*FSStore, error) {
|
||||
staging := filepath.Join(dir, ".staging")
|
||||
if err := os.MkdirAll(staging, 0o755); err != nil {
|
||||
if err := os.MkdirAll(staging, 0o750); err != nil {
|
||||
return nil, fmt.Errorf("create blob dirs: %w", err)
|
||||
}
|
||||
return &FSStore{dir: dir, staging: staging}, nil
|
||||
@@ -85,7 +85,9 @@ func (s *FSStore) Open(ctx context.Context, key string) (io.ReadCloser, error) {
|
||||
if err := checkKey(key); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f, err := os.Open(filepath.Join(s.dir, key))
|
||||
// checkKey has rejected any traversal, so the joined path stays under s.dir.
|
||||
f, err := os.Open(filepath.Join(s.dir, key)) //nolint:gosec // key validated by checkKey
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -30,7 +30,7 @@ var _ usecase.TaskRepository = (*TaskRepo)(nil)
|
||||
// Every query that returns a task selects exactly this list, in this order —
|
||||
// three hand-written column lists would drift apart within a week.
|
||||
var taskColumns = []string{
|
||||
"id", "job_id", "chunk_index", "workload", "input_uri", "input_sha256",
|
||||
"id", "job_id", "chunk_index", "workload", "input_uri", "input_artifact_id", "input_sha256",
|
||||
"parameters", "status", "attempt", "max_attempts", "lease_owner", "lease_expires_at",
|
||||
"result_artifact_id", "metrics", "error_code", "error_message",
|
||||
"created_at", "started_at", "completed_at", "version",
|
||||
@@ -49,9 +49,12 @@ func scanTask(row pgx.Row) (*domain.Task, error) {
|
||||
var (
|
||||
t domain.Task
|
||||
status string
|
||||
// input_uri is nullable now (uploaded shards have none), so it cannot
|
||||
// scan straight into a string; NULL becomes the empty InputURI.
|
||||
inputURI *string
|
||||
)
|
||||
err := row.Scan(
|
||||
&t.ID, &t.JobID, &t.ChunkIndex, &t.Workload, &t.InputURI, &t.InputSHA256,
|
||||
&t.ID, &t.JobID, &t.ChunkIndex, &t.Workload, &inputURI, &t.InputArtifactID, &t.InputSHA256,
|
||||
&t.Parameters, &status, &t.Attempt, &t.MaxAttempts, &t.LeaseOwner, &t.LeaseExpiresAt,
|
||||
&t.ResultArtifactID, &t.Metrics, &t.ErrorCode, &t.ErrorMessage,
|
||||
&t.CreatedAt, &t.StartedAt, &t.CompletedAt, &t.Version,
|
||||
@@ -59,6 +62,9 @@ func scanTask(row pgx.Row) (*domain.Task, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if inputURI != nil {
|
||||
t.InputURI = *inputURI
|
||||
}
|
||||
t.Status = domain.TaskStatus(status)
|
||||
return &t, nil
|
||||
}
|
||||
@@ -206,10 +212,12 @@ func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error
|
||||
batch := &pgx.Batch{}
|
||||
for _, t := range tasks {
|
||||
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).
|
||||
Columns("id", "job_id", "chunk_index", "workload", "input_uri", "input_artifact_id",
|
||||
"input_sha256", "parameters", "status", "attempt", "max_attempts", "created_at", "version").
|
||||
// input_uri is stored NULL (not "") when empty, so the ck_tasks_has_input
|
||||
// check actually bites: a task with neither a URI nor an artifact fails.
|
||||
Values(t.ID, t.JobID, t.ChunkIndex, t.Workload, nullIfEmpty(t.InputURI), t.InputArtifactID,
|
||||
t.InputSHA256, jsonbOrEmpty(t.Parameters), string(t.Status), t.Attempt, t.MaxAttempts, t.CreatedAt, t.Version).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -84,6 +84,15 @@ func jsonbOrEmpty(m map[string]any) map[string]any {
|
||||
return m
|
||||
}
|
||||
|
||||
// nullIfEmpty maps "" to a SQL NULL, so an absent optional string is stored as
|
||||
// NULL rather than an empty string that would defeat a NOT-NULL-or check.
|
||||
func nullIfEmpty(s string) any {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// conn returns the transaction bound to ctx, or the pool when there is none.
|
||||
func conn(ctx context.Context, pool *pgxpool.Pool) querier {
|
||||
if tx, ok := ctx.Value(txKey{}).(pgx.Tx); ok {
|
||||
|
||||
Reference in New Issue
Block a user