Harden distributed pipeline

coordinator/internal/storage/postgres/artifact_repo.go

@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
 	a.Kind = domain.ArtifactKind(kind)
 	return &a, nil
 }
+
+func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
+	sql, args, err := psql.Select(artifactColumns...).
+		From("artifacts").
+		Where(sq.Eq{
+			"task_id": taskID,
+			"attempt": attempt,
+			"kind":    string(domain.ArtifactPartialResult),
+		}).
+		ToSql()
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		a    domain.Artifact
+		kind string
+	)
+	err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
+		&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
+		&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
+	if errors.Is(err, pgx.ErrNoRows) {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, fmt.Errorf("find partial result: %w", err)
+	}
+	a.Kind = domain.ArtifactKind(kind)
+	return &a, nil
+}

coordinator/internal/storage/postgres/integration_test.go

@@ -360,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
 		t.Errorf("replay must be idempotent, got %v", err)
 	}
 
-	// A different result artifact for the same task is a genuine conflict.
-	art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
-	other := in
-	other.ResultArtifactID = art2.ID
-	if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
-		t.Errorf("err = %v, want ErrResultConflict", err)
+}
+
+func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
+	pool := testPool(t)
+	ctx := context.Background()
+	job, tasks := seedJob(t, pool, 1)
+	taskID := tasks[0].ID
+	first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
+	second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
+	if err != nil {
+		t.Fatal(err)
+	}
+	attempt := 1
+	second.Attempt = &attempt
+	second.SetContent("other-sha", 5)
+	if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
+		t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
 	}
 }
 

coordinator/internal/storage/postgres/task_repo.go

@@ -337,17 +337,26 @@ SET status           = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
                             ELSE error_message END,
     completed_at     = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
     version          = version + 1
-WHERE status IN ('leased','running') AND lease_expires_at < $1`
+WHERE status IN ('leased','running') AND lease_expires_at < $1
+RETURNING job_id`
 
-func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
-	var affected int64
+func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
+	var affected []uuid.UUID
 	err := withRetry(ctx, func(ctx context.Context) error {
-		tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
+		rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
 		if err != nil {
 			return err
 		}
-		affected = tag.RowsAffected()
-		return nil
+		defer rows.Close()
+		affected = affected[:0]
+		for rows.Next() {
+			var jobID uuid.UUID
+			if err := rows.Scan(&jobID); err != nil {
+				return err
+			}
+			affected = append(affected, jobID)
+		}
+		return rows.Err()
 	})
 	return affected, err
 }

coordinator/internal/storage/postgres/ui_read_repo.go

@@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
 	return tasks, rows.Err()
 }
 
+func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
+	out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
+	if len(jobIDs) == 0 {
+		return out, nil
+	}
+	sql, args, err := psql.Select(taskColumns...).From("tasks").
+		Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
+	if err != nil {
+		return nil, err
+	}
+	rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
+	if err != nil {
+		return nil, fmt.Errorf("list tasks by jobs: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		task, err := scanTask(rows)
+		if err != nil {
+			return nil, err
+		}
+		out[task.JobID] = append(out[task.JobID], *task)
+	}
+	return out, rows.Err()
+}
+
 func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
 	if limit < 1 || limit > 100 {
 		return nil, domain.ErrInvalidInput