Bind result artifacts to lease attempts

coordinator/internal/storage/postgres/artifact_repo.go

@@ -26,14 +26,14 @@ func NewArtifactRepo(pool *pgxpool.Pool) *ArtifactRepo {
 var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil)
 
 var artifactColumns = []string{
-	"id", "job_id", "task_id", "kind", "filename", "storage_key",
+	"id", "job_id", "task_id", "attempt", "kind", "filename", "storage_key",
 	"content_type", "size_bytes", "sha256", "created_at",
 }
 
 func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error {
 	sql, args, err := psql.Insert("artifacts").
 		Columns(artifactColumns...).
-		Values(a.ID, a.JobID, a.TaskID, string(a.Kind), a.Filename, a.StorageKey,
+		Values(a.ID, a.JobID, a.TaskID, a.Attempt, string(a.Kind), a.Filename, a.StorageKey,
 			a.ContentType, a.SizeBytes, a.SHA256, a.CreatedAt).
 		ToSql()
 	if err != nil {
@@ -59,7 +59,7 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
 		kind string
 	)
 	err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
-		&a.ID, &a.JobID, &a.TaskID, &kind, &a.Filename, &a.StorageKey,
+		&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, domain.ErrArtifactNotFound

coordinator/internal/storage/postgres/integration_test.go

@@ -259,6 +259,8 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) {
 				return err
 			}
 			art.SetContent(fmt.Sprintf("rsha-%d", task.ChunkIndex), 1)
+			attempt := 1
+			art.Attempt = &attempt
 			if err := artifacts.Insert(ctx, art); err != nil {
 				return err
 			}
@@ -269,6 +271,7 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) {
 			}
 			owner := "worker-1"
 			fresh.Status = domain.TaskLeased
+			fresh.Attempt = attempt
 			fresh.LeaseOwner = &owner
 			expires := now.Add(time.Minute)
 			fresh.LeaseExpiresAt = &expires
@@ -350,6 +353,10 @@ func seedArtifact(t *testing.T, pool *pgxpool.Pool, jobID uuid.UUID, taskID *uui
 		t.Fatalf("build artifact: %v", err)
 	}
 	art.SetContent(fmt.Sprintf("sha-%s", art.ID), 3)
+	if kind == domain.ArtifactPartialResult {
+		attempt := 1
+		art.Attempt = &attempt
+	}
 	if err := NewArtifactRepo(pool).Insert(context.Background(), art); err != nil {
 		t.Fatalf("insert artifact: %v", err)
 	}
@@ -446,6 +453,32 @@ func TestArtifactRepoRoundTrip(t *testing.T) {
 	}
 }
 
+func TestPartialResultArtifactRoundTripsAttempt(t *testing.T) {
+	pool := testPool(t)
+	ctx := context.Background()
+	job, tasks := seedJob(t, pool, 1)
+	taskID := tasks[0].ID
+	art, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult,
+		"result.csv", "text/csv", time.Now().UTC())
+	if err != nil {
+		t.Fatal(err)
+	}
+	attempt := 2
+	art.Attempt = &attempt
+	art.SetContent("sha", 3)
+	repo := NewArtifactRepo(pool)
+	if err := repo.Insert(ctx, art); err != nil {
+		t.Fatalf("insert: %v", err)
+	}
+	got, err := repo.Get(ctx, art.ID)
+	if err != nil {
+		t.Fatalf("get: %v", err)
+	}
+	if got.Attempt == nil || *got.Attempt != attempt {
+		t.Fatalf("attempt = %v, want %d", got.Attempt, attempt)
+	}
+}
+
 // A shard task stores its input as an artifact and no URI: this exercises the
 // nullable input_uri column, the input_artifact_id round-trip, and the
 // ck_tasks_has_input check that requires one or the other.