fix: keep frozen artifact metadata consistent during merges

docs/development.md

@@ -14,7 +14,7 @@ go test ./...
 
 `AGENTS.md` lists `go mod tidy && go test ./... && ./e2e/run.sh` as the required pre-commit check.
 
-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`.
+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`. Worker log assertions read Docker’s log-driver stream and retain each container until its exit status is collected, so a truncated Compose attach stream cannot hide the last worker messages.
 
 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:
 

docs/rewrite-flow.md

@@ -62,6 +62,8 @@ graph TD
 
 A rewrite producing zero destination parts skips artifact upload and transitions directly to `FINISHED`, with `empty_output` recorded in state. Importing it marks it imported without downloading or attaching parts. Nonempty artifacts still require matching frozen part directories.
 
+Before measuring final output and freezing it, both rewrite and compaction stop destination merge selection and wait for in-flight merges to finish. The drain has a separate five-minute limit (or the caller's shorter deadline) and fails the operation if it cannot establish a stable part set. Merges remain stopped through measurement and FREEZE, so recorded counts describe the archived snapshot, including for aggregate tables whose physical row counts shrink during merges. Merges are restored on success or failure; restoration errors also fail the operation. Ordinary merge settling can still time out and keep partially compacted output, but the final snapshot guard cannot be skipped. This applies to newly produced artifacts; existing mismatched metadata requires separate reconciliation.
+
 ## Local Insert Checkpoints
 
 Workers use up to 20 source-offset chunks per part. `-insert-chunk-min-rows` defaults to 10,000,000; the chunk count is `min(20, max(1, floor(source_rows / minimum)))`. Ranges divide the rows evenly (differing by at most one row), so no final tiny tail is created. Parts below 20 million rows use one insert, 20 million rows use two chunks, and 200 million or more use 20. Set the flag to `0` to disable chunking, or increase it for larger minimum chunks. This row-based default is a starting point, not a guarantee of chunk duration or memory usage.

e2e/run.sh

@@ -19,6 +19,22 @@ export CLICKHOUSE_UID="$(id -u)" CLICKHOUSE_GID="$(id -g)"
 
 cd "$ROOT"
 
+# Read the log driver's stream instead of Compose's attach stream, which can
+# close before the final worker logs arrive when a child process exits.
+run_worker_logged() (
+  local container="partforge-e2e-worker-${BASHPID}-${RANDOM}"
+  local exit_code
+  if [[ "${1:-}" == "--name" ]]; then
+    container="$2"
+    shift 2
+  fi
+  trap 'docker rm -f "$container" >/dev/null 2>&1 || true' EXIT
+  docker compose run --detach --name "$container" "$@" >/dev/null
+  docker logs --follow "$container"
+  exit_code="$(docker wait "$container")"
+  exit "$exit_code"
+)
+
 log_value() {
   local line="$1"
   local key="$2"
@@ -222,7 +238,7 @@ fi
 
 for i in $(seq 1 "$incremental_part_count"); do
   backup_worker_log="$ROOT/.e2e/backup-worker-${i}.log"
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged worker \
     worker \
     -role=inserter \
     -insert-chunk-min-rows=1 \
@@ -357,7 +373,7 @@ fi
 
 for i in $(seq 1 "$part_count"); do
   worker_log="$ROOT/.e2e/worker-${i}.log"
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged worker \
     worker \
     -role=inserter \
     -insert-chunk-min-rows=1 \
@@ -411,7 +427,7 @@ CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm \
   -s3-endpoint=http://localstack:4566 -postgres-url="$POSTGRES_URL"
 
 for i in $(seq 1 "$part_count"); do
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker worker \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged worker worker \
     -role=inserter -insert-chunk-min-rows=1 -merge-max-runtime=1ns -once \
     -s3-endpoint=http://localstack:4566 -postgres-url="$POSTGRES_URL"
 done
@@ -452,7 +468,7 @@ exec /usr/local/bin/s5cmd "$@"
 SH
   chmod +x "$ROOT/.e2e/hold-compact-download"
   active_log="$ROOT/.e2e/finalize-active-worker.log"
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm --name "$compact_container" \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged --name "$compact_container" \
     -v "$ROOT/.e2e:/work/.e2e" worker worker -role=compactor -once \
     -s5cmd-binary=/work/.e2e/hold-compact-download \
     -s3-endpoint=http://localstack:4566 -postgres-url="$POSTGRES_URL" \
@@ -509,7 +525,7 @@ SH
 batched=0
 for i in 1 2; do
   compact_log="$ROOT/.e2e/compact-$i.log"
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged worker \
     worker \
     -role=compactor \
     -s3-endpoint=http://localstack:4566 \
@@ -659,7 +675,7 @@ fi
 
 for i in $(seq 1 "$chain_part_count"); do
   chain_worker_log="$ROOT/.e2e/chain-worker-${i}.log"
-  CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \
+  CLICKHOUSE_DATA_DIR="$DATA_DIR" run_worker_logged worker \
     worker \
     -role=inserter \
     -insert-chunk-min-rows=1 \

internal/rewrite/compactor.go

@@ -88,7 +88,7 @@ type CompactResult struct {
 	Inputs                []CompactInput
 }
 
-func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (CompactResult, error) {
+func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (result CompactResult, err error) {
 	if item.JobID == "" || item.OutputPartID == "" || item.OutputFinishedKey == "" {
 		return CompactResult{}, fmt.Errorf("compact work item is missing job id, output part id, or output finished key")
 	}
@@ -272,6 +272,12 @@ func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (CompactRe
 	}
 	cancelPhase()
 
+	resumeMerges, err := p.pauseDestinationMerges(ctx, target)
+	if err != nil {
+		return CompactResult{}, err
+	}
+	defer func() { err = errors.Join(err, resumeMerges()) }()
+
 	destPartitions, err := p.activePartPartitionStats(ctx, item.DestinationDatabase, item.DestinationTable)
 	if err != nil {
 		return CompactResult{}, fmt.Errorf("measure compact output active part partitions: %w", err)

internal/rewrite/compactor_test.go

@@ -124,6 +124,9 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 
 			var partitionQueries atomic.Uint64
 			var mergeCountQueries atomic.Uint64
+			var mergesPaused atomic.Bool
+			var snapshotDrained atomic.Bool
+			var snapshotMeasured atomic.Bool
 			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				body, err := io.ReadAll(r.Body)
 				if err != nil {
@@ -135,12 +138,26 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 				case strings.HasPrefix(query, "SELECT arrayElement(data_paths"):
 					_, _ = io.WriteString(w, dataPath+"\n")
 				case strings.HasPrefix(query, "SELECT partition_id, count()"):
+					if snapshotDrained.Load() {
+						// The in-flight aggregate merge reduced rows before draining.
+						if !mergesPaused.Load() {
+							t.Error("measured final output with merges enabled")
+						}
+						snapshotMeasured.Store(true)
+						_, _ = io.WriteString(w, "202601\t"+strconv.FormatUint(tt.partsAfterFirstQuery, 10)+"\t12\t120\n")
+						return
+					}
 					parts := uint64(2)
 					if partitionQueries.Add(1) > 1 {
 						parts = tt.partsAfterFirstQuery
 					}
 					_, _ = io.WriteString(w, "202601\t"+strconv.FormatUint(parts, 10)+"\t20\t200\n")
 				case strings.HasPrefix(query, "SELECT count() FROM system.merges"):
+					if mergesPaused.Load() {
+						snapshotDrained.Store(true)
+						_, _ = io.WriteString(w, "0\n")
+						return
+					}
 					mergeCountQueries.Add(1)
 					if tt.stop != "" {
 						if tt.stop == "unexpected" {
@@ -164,16 +181,22 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 				case strings.HasPrefix(query, "SELECT name, path, type FROM system.disks"):
 					_, _ = io.WriteString(w, "default\t"+diskPath+"\tlocal\n")
 				case strings.Contains(query, " FREEZE WITH NAME '"):
+					if !mergesPaused.Load() || !snapshotDrained.Load() || !snapshotMeasured.Load() {
+						t.Error("froze output without stopping and draining destination merges")
+					}
 					marker := " FREEZE WITH NAME '"
 					freezeName := strings.TrimSuffix(query[strings.Index(query, marker)+len(marker):], "'")
 					if err := writeCompactTestPart(filepath.Join(diskPath, "shadow", freezeName, "store", "abc", "def", "202601_1_2_1")); err != nil {
 						http.Error(w, err.Error(), http.StatusInternalServerError)
 					}
+				case strings.HasPrefix(query, "SYSTEM STOP MERGES "):
+					mergesPaused.Store(true)
+					snapshotDrained.Store(false)
+				case strings.HasPrefix(query, "SYSTEM START MERGES "):
+					mergesPaused.Store(false)
 				case query == "SYSTEM FLUSH LOGS",
 					strings.HasPrefix(query, "CREATE DATABASE "),
 					strings.HasPrefix(query, "CREATE TABLE "),
-					strings.HasPrefix(query, "SYSTEM STOP MERGES "),
-					strings.HasPrefix(query, "SYSTEM START MERGES "),
 					strings.HasPrefix(query, "ALTER TABLE "):
 				default:
 					http.Error(w, "unexpected query: "+query, http.StatusInternalServerError)
@@ -231,6 +254,12 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 			if result.Reduced != tt.wantReduced {
 				t.Fatalf("reduced = %t, want %t", result.Reduced, tt.wantReduced)
 			}
+			if result.DestinationStats.Rows != 12 {
+				t.Fatalf("destination rows = %d, want post-drain count 12", result.DestinationStats.Rows)
+			}
+			if mergesPaused.Load() {
+				t.Fatal("destination merges remained stopped after compaction")
+			}
 			if (result.FinishedKey == "finished/compact-1") != tt.wantReduced {
 				t.Fatalf("finished key = %q, want uploaded %t", result.FinishedKey, tt.wantReduced)
 			}

internal/rewrite/processor.go

@@ -643,6 +643,11 @@ func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourceP
 	if _, err := p.waitForDestinationMerges(ctx, m, stageTracker, mergeTarget, "after_insert", false); err != nil {
 		return rewriteResult{}, err
 	}
+	resumeMerges, err := p.pauseDestinationMerges(ctx, mergeTarget)
+	if err != nil {
+		return rewriteResult{}, err
+	}
+	defer func() { err = errors.Join(err, resumeMerges()) }()
 	if err := p.reportStageProgress(ctx, m, stageTracker, stageMeasureDestinationParts); err != nil {
 		return rewriteResult{}, err
 	}

internal/rewrite/snapshot.go

@@ -0,0 +1,47 @@
+package rewrite
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+)
+
+// pauseDestinationMerges makes the active part set stable until measurement and
+// FREEZE finish. STOP MERGES prevents selection, but existing merges must drain.
+func (p Processor) pauseDestinationMerges(ctx context.Context, target mergeWaitTarget) (func() error, error) {
+	table := target.tableSQL()
+	resume := func() error {
+		// Cleanup must still work after the operation's context is cancelled.
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
+		defer cancel()
+		if err := p.ClickHouse.Exec(cleanupCtx, "SYSTEM START MERGES "+table); err != nil {
+			return fmt.Errorf("restore destination merges for %s: %w", table, err)
+		}
+		return nil
+	}
+	if err := p.ClickHouse.Exec(ctx, "SYSTEM STOP MERGES "+table); err != nil {
+		return nil, errors.Join(fmt.Errorf("stop destination merges for snapshot: %w", err), resume())
+	}
+	// MergeTimeout is an optimization budget and may be deliberately tiny.
+	// Snapshot correctness needs its own bounded wait, including after settling
+	// timed out. The caller can still impose a shorter operation deadline.
+	waitCtx, cancel := context.WithTimeout(ctx, DefaultMergeTimeout)
+	defer cancel()
+	interval := p.MergePollInterval
+	if interval <= 0 {
+		interval = defaultMergePollInterval
+	}
+	for {
+		count, err := p.destinationMergeCount(waitCtx, target)
+		if err != nil {
+			return nil, errors.Join(fmt.Errorf("drain destination merges for snapshot: %w", err), resume())
+		}
+		if count == 0 {
+			return resume, nil
+		}
+		if err := sleepOrDone(waitCtx, interval); err != nil {
+			return nil, errors.Join(fmt.Errorf("drain destination merges for snapshot: %w", err), resume())
+		}
+	}
+}

internal/rewrite/snapshot_test.go

@@ -0,0 +1,97 @@
+package rewrite
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/PostHog/partforge/internal/chhttp"
+)
+
+func TestSnapshotWaitsForInFlightMerge(t *testing.T) {
+	var mu sync.Mutex
+	var queries []string
+	counts := 0
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b, _ := io.ReadAll(r.Body)
+		q := string(b)
+		mu.Lock()
+		defer mu.Unlock()
+		queries = append(queries, q)
+		if strings.HasPrefix(q, "SELECT count() FROM system.merges") {
+			counts++
+			if counts == 1 {
+				_, _ = io.WriteString(w, "1\n")
+			} else {
+				_, _ = io.WriteString(w, "0\n")
+			}
+		}
+	}))
+	defer server.Close()
+	ctx, cancel := context.WithCancel(context.Background())
+	p := Processor{ClickHouse: chhttp.Client{URL: server.URL}, MergeTimeout: time.Nanosecond, MergePollInterval: time.Millisecond}
+	resume, err := p.pauseDestinationMerges(ctx, testMergeWaitTarget())
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Cancellation after the snapshot must not prevent restoring merges.
+	cancel()
+	if err := resume(); err != nil {
+		t.Fatal(err)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if len(queries) != 4 || !strings.HasPrefix(queries[0], "SYSTEM STOP MERGES ") ||
+		counts != 2 || !strings.HasPrefix(queries[3], "SYSTEM START MERGES ") {
+		t.Fatalf("snapshot queries = %v", queries)
+	}
+}
+
+func TestSnapshotDrainFailureRestoresMerges(t *testing.T) {
+	for _, failure := range []string{"query", "timeout", "restore"} {
+		t.Run(failure, func(t *testing.T) {
+			var mu sync.Mutex
+			restored := false
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				b, _ := io.ReadAll(r.Body)
+				q := string(b)
+				switch {
+				case strings.HasPrefix(q, "SELECT count() FROM system.merges"):
+					if failure == "query" {
+						http.Error(w, "cannot read merges", http.StatusInternalServerError)
+					} else {
+						_, _ = io.WriteString(w, "1\n")
+					}
+				case strings.HasPrefix(q, "SYSTEM START MERGES "):
+					mu.Lock()
+					restored = true
+					mu.Unlock()
+					if failure == "restore" {
+						http.Error(w, "cannot restore merges", http.StatusInternalServerError)
+					}
+				}
+			}))
+			defer server.Close()
+			p := Processor{ClickHouse: chhttp.Client{URL: server.URL}, MergeTimeout: time.Nanosecond, MergePollInterval: time.Millisecond}
+			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+			defer cancel()
+			resume, err := p.pauseDestinationMerges(ctx, testMergeWaitTarget())
+			if err == nil || resume != nil || !strings.Contains(err.Error(), "drain destination merges") {
+				t.Fatalf("pause = %v, resume nil = %t", err, resume == nil)
+			}
+			if failure == "restore" && !strings.Contains(err.Error(), "restore destination merges") {
+				t.Fatalf("lost restore error: %v", err)
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			if !restored {
+				t.Fatal("did not restore merges after failed snapshot guard")
+			}
+		})
+	}
+}