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:
Efremenko Arhip
2026-07-23 16:34:04 +03:00
parent 4a092d2e4e
commit c3243a6b7e
19 changed files with 741 additions and 42 deletions
+4 -2
View File
@@ -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 {