Add state migrations and reduce Postgres scheduling overhead

README.md

@@ -33,7 +33,15 @@ on source node         parts to S3;          local ClickHouse;         destinati
 
 ## Getting started
 
-Two SQL files define your migration; everything else is mechanical. Write them, then run four commands.
+Initialize or upgrade the Postgres state schema before running uploads or workers:
+
+```sh
+partforge migrate -postgres-url="$POSTGRES_URL"
+```
+
+For an existing deployment, stop workers while migrating. See [Postgres migrations](docs/postgres.md#migrate-before-starting-workers).
+
+Two SQL files define your ClickHouse rewrite. Write them, then run the pipeline commands below.
 
 ### 1. Destination schema — `dest.sql`
 

cmd/partforge/main.go

@@ -71,6 +71,7 @@ type commandHelp struct {
 }
 
 var commandHelps = []commandHelp{
+	{Name: "migrate", Usage: "[flags]", Summary: "Apply pending Postgres state schema migrations.", Details: "Run once before starting workers. Adopts existing state tables. Stop workers while migrating; schema changes and backfills run in one transaction."},
 	{
 		Name:    "upload-backup",
 		Usage:   "[flags]",
@@ -266,6 +267,8 @@ func run() error {
 	defer stop()
 
 	switch os.Args[1] {
+	case "migrate":
+		return runMigrate(ctx, os.Args[2:])
 	case "upload-backup":
 		return runUploadBackup(ctx, os.Args[2:])
 	case "upload-freeze":
@@ -1725,6 +1728,7 @@ func runWorker(ctx context.Context, args []string) error {
 			if roleSettings.Compact {
 				didCompactWork, err := runWorkerCompaction(ctx, workerCompactionConfig{
 					StateStore:                    stateStore,
+					Once:                          *once,
 					WorkerID:                      resolvedWorkerID,
 					WorkDir:                       *workDir,
 					ClickHouseURL:                 *clickHouseURL,
@@ -2018,6 +2022,7 @@ func createWorkerRunDirs(workDir string) (workerRunDirs, error) {
 }
 
 type workerCompactionConfig struct {
+	Once                          bool
 	StateStore                    *state.Store
 	WorkerID                      string
 	WorkDir                       string
@@ -2068,38 +2073,17 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool,
 			return false, err
 		}
 	}
-	if cfg.CompactLeaseStaleAfter > 0 {
-		now := time.Now().UTC()
-		released, err := cfg.StateStore.ReleaseStaleCompactingParts(ctx, now, cfg.CompactLeaseStaleAfter)
-		if err != nil {
-			if ctx.Err() != nil {
-				slog.Info("worker shutdown requested while releasing stale compact work", "stage", "shutdown")
-				return false, nil
-			}
-			return false, err
-		}
-		if released > 0 {
-			slog.Warn("released stale compacting parts", "stage", "release_stale_compact", "worker_id", cfg.WorkerID, "released", released, "stale_after", cfg.CompactLeaseStaleAfter)
-		}
+	finalized, err := cfg.StateStore.MaintainCompaction(ctx, cfg.CompactWindow, cfg.CompactLeaseStaleAfter, time.Now().UTC(), cfg.Once)
+	if err != nil {
+		return false, err
 	}
-	finalization := compactFinalizationResult{}
-	if cfg.CompactWindow > 0 {
-		var err error
-		finalization, err = finalizeCompactReadyJobs(ctx, cfg.StateStore, cfg.CompactWindow, time.Now().UTC())
-		if err != nil {
-			return false, err
-		}
-		if finalization.Finalized > 0 {
-			slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalization.Finalized)
-			return true, nil
-		}
-		if len(finalization.ExpiredJobIDs) > 0 {
-			slog.Info("skipping compact claims for jobs past compact window", "stage", "claim_compact", "worker_id", cfg.WorkerID, "jobs", len(finalization.ExpiredJobIDs))
-		}
+	if finalized > 0 {
+		slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalized)
+		return true, nil
 	}
 	slog.Info("claiming compact-ready batch", "stage", "claim_compact", "worker_id", cfg.WorkerID)
 	batch, err := cfg.StateStore.ClaimNextCompactBatch(ctx, cfg.WorkerID, time.Now().UTC(), state.CompactClaimOptions{
-		ExcludedJobIDs: finalization.ExpiredJobIDs,
+		CompactWindow: cfg.CompactWindow,
 	})
 	if err != nil {
 		if ctx.Err() != nil {
@@ -2109,14 +2093,6 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool,
 		return false, err
 	}
 	if batch == nil {
-		finalization, err := finalizeCompactReadyJobs(ctx, cfg.StateStore, cfg.CompactWindow, time.Now().UTC())
-		if err != nil {
-			return false, err
-		}
-		if finalization.Finalized > 0 {
-			slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalization.Finalized)
-			return true, nil
-		}
 		return false, nil
 	}
 	currentBatch := func() state.CompactBatch {
@@ -2174,11 +2150,7 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool,
 		DestinationSchema:   batch.Parts[0].DestinationSchema,
 		Inputs:              compactInputs(batch.Parts),
 	}
-	jobParts, err := cfg.StateStore.ListJobParts(ctx, batch.JobID)
-	if err != nil {
-		return true, markCompactBatchFailed(fmt.Errorf("list job parts for compact deadline: %w", err))
-	}
-	compactDeadline, err := compactBatchDeadline(jobParts, cfg.CompactWindow, time.Now().UTC())
+	compactDeadline, err := cfg.StateStore.CompactDeadline(ctx, batch.JobID, cfg.CompactWindow)
 	if err != nil {
 		return true, markCompactBatchFailed(err)
 	}
@@ -2191,13 +2163,13 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool,
 		}
 		slog.Info("compact window expired before compact batch started; released", "stage", "compact_window_expired", "job_id", batch.JobID, "output_part_id", outputPartID)
 		finalizeCtx, finalizeCancel := workerStateUpdateContext()
-		finalization, finalizeErr := finalizeCompactReadyJob(finalizeCtx, cfg.StateStore, batch.JobID, cfg.CompactWindow, time.Now().UTC())
+		finalized, finalizeErr := cfg.StateStore.FinalizeCompactReadyJob(finalizeCtx, batch.JobID, cfg.CompactWindow, time.Now().UTC())
 		finalizeCancel()
 		if finalizeErr != nil {
 			return true, finalizeErr
 		}
-		if finalization.Finalized > 0 {
-			slog.Info("finalized compact-ready artifacts after compact window expiration", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalization.Finalized)
+		if finalized > 0 {
+			slog.Info("finalized compact-ready artifacts after compact window expiration", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalized)
 		}
 		return true, nil
 	}
@@ -2274,14 +2246,14 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool,
 			return true, nil
 		}
 		finalizeCtx, finalizeCancel := workerStateUpdateContext()
-		finalization, finalizeErr := finalizeCompactReadyJob(finalizeCtx, cfg.StateStore, batch.JobID, cfg.CompactWindow, time.Now().UTC())
+		finalized, finalizeErr := cfg.StateStore.FinalizeCompactReadyJob(finalizeCtx, batch.JobID, cfg.CompactWindow, time.Now().UTC())
 		finalizeCancel()
 		if finalizeErr != nil {
 			cleanupCompactNow()
 			return true, finalizeErr
 		}
-		if finalization.Finalized > 0 {
-			slog.Info("finalized compact-ready artifacts after no-reduction compaction", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalization.Finalized)
+		if finalized > 0 {
+			slog.Info("finalized compact-ready artifacts after no-reduction compaction", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalized)
 		}
 		cleanupCompactNow()
 		return true, nil
@@ -2591,59 +2563,6 @@ func compactClaimSplayMax(compactWindow time.Duration) time.Duration {
 	return 250 * time.Millisecond
 }
 
-type compactFinalizationResult struct {
-	Finalized     int
-	ExpiredJobIDs map[string]struct{}
-}
-
-func finalizeCompactReadyJobs(ctx context.Context, store *state.Store, compactWindow time.Duration, now time.Time) (compactFinalizationResult, error) {
-	result := compactFinalizationResult{ExpiredJobIDs: map[string]struct{}{}}
-	jobIDs, err := store.ListJobIDsByStatus(ctx, state.StatusCompactReady)
-	if err != nil {
-		return result, err
-	}
-	for _, jobID := range jobIDs {
-		jobResult, err := finalizeCompactReadyJob(ctx, store, jobID, compactWindow, now)
-		if err != nil {
-			return result, err
-		}
-		result.Finalized += jobResult.Finalized
-		for expiredJobID := range jobResult.ExpiredJobIDs {
-			result.ExpiredJobIDs[expiredJobID] = struct{}{}
-		}
-	}
-	return result, nil
-}
-
-func finalizeCompactReadyJob(ctx context.Context, store *state.Store, jobID string, compactWindow time.Duration, now time.Time) (compactFinalizationResult, error) {
-	result := compactFinalizationResult{ExpiredJobIDs: map[string]struct{}{}}
-	parts, err := store.ListJobParts(ctx, jobID)
-	if err != nil {
-		return result, err
-	}
-	expired, err := compactWindowExpired(parts, compactWindow, now)
-	if err != nil {
-		return result, err
-	}
-	if expired {
-		result.ExpiredJobIDs[jobID] = struct{}{}
-	}
-	compactReady, ok, err := finalizableCompactReadyParts(parts, compactWindow, now)
-	if err != nil {
-		return result, err
-	}
-	if !ok {
-		return result, nil
-	}
-	for _, part := range compactReady {
-		if err := store.MarkCompactReadyFinished(ctx, part, now); err != nil {
-			return result, err
-		}
-		result.Finalized++
-	}
-	return result, nil
-}
-
 func finalizableCompactReadyParts(parts []state.Part, compactWindow time.Duration, now time.Time) ([]state.Part, bool, error) {
 	compactReady := compactReadyParts(parts)
 	if len(compactReady) == 0 {
@@ -2862,6 +2781,27 @@ func runImportFinished(ctx context.Context, args []string) error {
 	})
 }
 
+func runMigrate(ctx context.Context, args []string) error {
+	fs := newCommandFlagSet("migrate")
+	configPath := fs.String("config", defaultConfigPath, "JSON config file path; CLI flags override config values")
+	stateTable := fs.String("state-table", defaultStateTable, "Postgres table used for PartForge state")
+	region := fs.String("aws-region", "", "AWS region for Postgres IAM auth")
+	postgresURL := fs.String("postgres-url", "", "Postgres state store connection URL")
+	postgresIAMAuth := fs.Bool("postgres-iam-auth", false, "use AWS IAM authentication for the Postgres state store")
+	if err := parseFlags(fs, args); err != nil {
+		return err
+	}
+	if err := applyConfigDefaults(fs, *configPath, "migrate"); err != nil {
+		return err
+	}
+	applied, err := state.Migrate(ctx, state.Config{Region: *region, Endpoint: *postgresURL, IAMAuth: *postgresIAMAuth, Table: *stateTable})
+	if err != nil {
+		return err
+	}
+	slog.Info("state schema is up to date", "migrations_applied", applied, "state_table", *stateTable)
+	return nil
+}
+
 func runListJobs(ctx context.Context, args []string) error {
 	fs := newCommandFlagSet("list-jobs")
 	var (

docs/deployment.md

@@ -4,6 +4,8 @@ The worker image is published on every push to `main` to `ghcr.io/<owner>/partfo
 
 The image is a single Ubuntu container with `clickhouse-server`, `clickhouse-client`, `s5cmd`, and the Go binary. Its entrypoint is the binary and the default command is `worker`. It **runs as root** (so it can write its work directory on root-owned host mounts) and starts a local `clickhouse server` child process for each claimed part.
 
+Before starting a new worker release, stop the existing workers and run `partforge migrate` once with that release's image and the same Postgres/state-table settings. The migration backfills and indexes the existing table under table locks; restart workers after it succeeds. See [the migration procedure](postgres.md#migrate-before-starting-workers).
+
 ## Recommended: workers on ECS with an IAM task role
 
 Run the workers as an ECS service and give the task an **IAM role** scoped to the S3 bucket and the RDS/Aurora PostgreSQL database user. This is the recommended setup:

docs/development.md

@@ -16,6 +16,15 @@ go test ./...
 
 The e2e script stands up LocalStack, Postgres, and a ClickHouse container, builds the worker image, and runs the full pipeline against `e2e/sql/`, diffing the result against `e2e/expected.tsv`. It builds the image each run; set `PARTFORGE_E2E_SKIP_BUILD=1` to reuse an existing `partforge-worker:latest`.
 
+PostgreSQL integration checks exercise fresh and hand-created schemas, migration rollback and concurrent migration runs, concurrent work claims, ownership, maintenance, and 20,000-row query plans. With local compose Postgres running:
+
+```sh
+PARTFORGE_TEST_POSTGRES_URL='postgres://partforge:partforge@localhost:15432/partforge?sslmode=disable' \
+  go test ./internal/state -run Postgres -count=1
+```
+
+These tests create and remove isolated schemas. They skip when the connection variable is unset. The e2e script also runs `migrate` twice before uploading, checking both initial setup and a no-op rerun.
+
 ## Build
 
 ```sh

docs/postgres.md

@@ -20,30 +20,47 @@ For local compose runs:
 
 Default table name is `partforge_state`; override with `-state-table` or `state_table`.
 
-## Schema
+## Migrate before starting workers
 
-The app creates the table and indexes on startup:
-
-```sql
-CREATE TABLE IF NOT EXISTS partforge_state (
-    job_id text NOT NULL,
-    part_id text NOT NULL,
-    status text NOT NULL,
-    worker_id text NOT NULL DEFAULT '',
-    created_at text NOT NULL,
-    updated_at text NOT NULL,
-    data jsonb NOT NULL,
-    PRIMARY KEY (job_id, part_id)
-);
-
-CREATE INDEX IF NOT EXISTS partforge_state_status_idx
-    ON partforge_state (status, created_at, job_id, part_id);
-
-CREATE INDEX IF NOT EXISTS partforge_state_job_status_idx
-    ON partforge_state (job_id, status, part_id);
+```sh
+partforge migrate -postgres-url="$POSTGRES_URL"
 ```
 
-The scalar columns support claims and job/status scans. The full part record is stored in `data`.
+Use the same `-config`, `-state-table`, `-postgres-iam-auth`, and `-aws-region` settings as the workers. A schema-qualified state table is supported. The database role running migrations must own the existing state table and have permission to create tables and indexes in its schema.
+
+For an existing deployment:
+
+1. Stop workers and pause uploads and state-changing admin commands.
+2. Run `partforge migrate` with the new binary against the existing state table.
+3. Start the new workers after the command succeeds.
+
+Migrations backfill columns and build indexes in one transaction, holding table locks. Plan a maintenance window; this is not an online migration command. Existing rows, statuses, progress, and lineage are retained. The first migration adopts the original hand-created table with `CREATE TABLE IF NOT EXISTS`; do not recreate or empty it.
+
+The command records numbered versions in `<state-table>_migrations`. An advisory transaction lock serializes concurrent migration commands, and a failure rolls back the pending changes and their version records together. Re-running a successful migration command is a no-op. Unexpected versions or gaps fail explicitly. Workers check the version at startup and tell you to run `partforge migrate` if an upgrade is needed; they no longer execute schema DDL.
+
+Migrations are compiled into the binary in [internal/state/migrations.go](../internal/state/migrations.go). Append new SQL migrations to the list; never edit, reorder, or remove an already released migration. There is no automatic downgrade command.
+
+## Schema and scheduling
+
+The original primary key `(job_id, part_id)`, state columns, and full `data` JSON record remain. Additional scalar columns support scheduling without indexing the mutable JSON itself:
+
+| Columns | Purpose |
+|---|---|
+| `source_artifact_bytes` | Ordered partial index for largest-first READY claims |
+| `compact_bytes`, `compact_eligible` | Ordered partial index for eligible compact claims |
+| `compact_normalized` | Identify artifacts containing one physical part |
+| `compact_stale_at` | Indexed stale-compaction lookup, preserving the existing earlier-of-heartbeat-and-claim timeout |
+| `original_compact_ready_at` | Indexed job compact deadline lookup |
+
+Application writes update these projections in the same SQL statement as the JSON or status change. No triggers are installed. Use PartForge commands for state changes; direct SQL must maintain the corresponding derived columns too. The full Part record remains in `data`; no job metadata or progress table is required.
+
+Compactors claim the largest eligible unlocked artifact. Job, destination, and explicit partition filters still apply; partitions with an active compactor are no longer deprioritized.
+
+A single `<state-table>_maintenance` record reserves compaction maintenance once every ten seconds across the worker fleet. `worker -once` requests an immediate pass but still skips a pass already in progress. Only the reservation holder scans job summaries and expires stale work. The reservation commits with the work, and active maintenance is skipped by other workers without waiting. Workers sharing a state table should use the same compact-window and stale-timeout settings. New finalizable output can wait up to the next maintenance pass; a long maintenance run or database contention can extend that delay.
+
+Rewrite progress combines live query counters and stage timing in one periodic heartbeat. Rewrite and compact progress updates use conditional SQL patches that check worker ownership. A failed rewrite heartbeat cancels processing and surfaces the error.
+
+Connection pool configuration is unchanged.
 
 ## IAM Auth
 

docs/setup.md

@@ -14,6 +14,9 @@ Requirements, configuration, and how to run the four stages by hand. For the hig
 
 ```sh
 docker compose up -d localstack postgres
+docker compose build worker
+docker compose run --rm worker migrate \
+  -postgres-url='postgres://partforge:partforge@postgres:5432/partforge?sslmode=disable'
 ```
 
 This creates the `partforge` S3 bucket in LocalStack and starts a local Postgres database. Point commands at them with:

e2e/run.sh

@@ -114,6 +114,9 @@ for _ in $(seq 1 60); do
 done
 docker compose exec -T postgres pg_isready -U partforge -d partforge >/dev/null
 
+CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker migrate -postgres-url="$POSTGRES_URL"
+CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker migrate -postgres-url="$POSTGRES_URL"
+
 docker compose exec -T clickhouse clickhouse-client --multiquery < e2e/sql/setup_and_freeze.sql
 
 docker compose exec -T clickhouse clickhouse-client --query \