From 3835cc7f1a663b87da3b8c224c469ca38ae30524 Mon Sep 17 00:00:00 2001 From: Rory Shanks Date: Mon, 12 Oct 2026 09:01:03 +0200 Subject: [PATCH] Add upload-chain to start a job from a previous job's finished output New `partforge upload-chain -source-job-id=` registers each finished artifact of a fully finished job as a READY source part of a new job. Only manifests are written; the new manifest field `source_finished_key` points at the parent's finished prefix, and workers download and attach every part tarball in it before running the new insert-select. The new job's source table and schema are the parent job's destination table and schema. Multi-part chained sources fall back to a single unchunked insert because _part_offset chunk ranges only partition one part. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 + cmd/partforge/main.go | 251 +++++++++++++++++++++++++++++ cmd/partforge/main_test.go | 72 +++++++++ docs/operations.md | 2 +- docs/setup.md | 17 ++ e2e/expected-chain.tsv | 4 + e2e/run.sh | 68 ++++++++ e2e/sql/chain-destination.sql | 10 ++ e2e/sql/chain-insert.sql | 7 + internal/manifest/manifest.go | 11 ++ internal/manifest/manifest_test.go | 33 ++++ internal/rewrite/processor.go | 53 +++++- 12 files changed, 521 insertions(+), 9 deletions(-) create mode 100644 e2e/expected-chain.tsv create mode 100644 e2e/sql/chain-destination.sql create mode 100644 e2e/sql/chain-insert.sql diff --git a/README.md b/README.md index 5c745db..07a8058 100644 --- a/README.md +++ b/README.md @@ -128,6 +128,8 @@ The path must be the exact full or incremental backup prefix. Incrementals are r For multiple shards with the same destination schema and insert-select, run the first `upload-freeze` with the SQL files and later shards with `-copy-sql-from-job=`. To test another destination schema against the same uploaded source parts, run `upload-freeze -copy-parts-from-job=` with the new SQL files. +To run a second rewrite over the output of a finished job, run `upload-chain -source-job-id=` with new SQL files. The new job reads the previous job's finished artifacts directly, and its insert-select reads `FROM` the previous job's destination table. See [docs/setup.md](docs/setup.md#chaining-from-a-finished-job). + Part state lifecycle (tracked in Postgres, so a job is resumable): ``` diff --git a/cmd/partforge/main.go b/cmd/partforge/main.go index 0d28f0e..7d1d56e 100644 --- a/cmd/partforge/main.go +++ b/cmd/partforge/main.go @@ -90,6 +90,16 @@ Required: -backup, -database, -table, -bucket, and either both -destination-sche Details: `Run this on a host that can read the source ClickHouse disk paths from system.disks. The command writes a manifest into every frozen part, uploads each raw source part directory, and records the source schema plus destination SQL for workers. With -copy-parts-from-job, it skips local disk scanning and reuses that job's uploaded source parts. Required: -bucket, either both -destination-schema-file and -insert-select-file or -copy-sql-from-job, and either -copy-parts-from-job or all of -database, -table, and -freeze.`, + }, + { + Name: "upload-chain", + Usage: "[flags]", + Summary: "Start a new job whose source is the finished output of a previous job, then create READY work rows in Postgres.", + Details: `Every finished artifact of -source-job-id becomes one READY source part of the new job. Only manifests are written; workers download the previous job's finished part tarballs directly. The new job's source table is the previous job's destination table, so -insert-select-file must read FROM that table. + +The previous job must be fully finished: every non-superseded row must be FINISHED, IMPORTING, or IMPORTED. Keep its finished S3 artifacts until the new job has finished rewriting. + +Required: -source-job-id, -bucket, and either both -destination-schema-file and -insert-select-file or -copy-sql-from-job.`, }, { Name: "worker", @@ -282,6 +292,8 @@ func run() error { return runUploadBackup(ctx, os.Args[2:]) case "upload-freeze": return runUploadFreeze(ctx, os.Args[2:]) + case "upload-chain": + return runUploadChain(ctx, os.Args[2:]) case "worker": return runWorker(ctx, os.Args[2:]) case "import-finished": @@ -1284,6 +1296,245 @@ func registerSourcePartsFromJob(ctx context.Context, store *state.Store, copier return len(sourceParts), nil } +func runUploadChain(ctx context.Context, args []string) error { + fs := newCommandFlagSet("upload-chain") + var ( + configPath = fs.String("config", defaultConfigPath, "JSON config file path; CLI flags override config values") + sourceJobID = fs.String("source-job-id", "", "finished job whose output becomes the new job's source") + destinationSchemaFile = fs.String("destination-schema-file", "", "path to SQL file containing the full destination CREATE TABLE") + insertSelectFile = fs.String("insert-select-file", "", "path to SQL file containing INSERT INTO destination SELECT ... FROM ") + copySQLFromJob = fs.String("copy-sql-from-job", "", "existing job id to copy destination schema and insert-select SQL from") + jobID = fs.String("job-id", "", "job id to store in manifests and Postgres; empty generates one") + jobName = fs.String("job-name", "", "readable job name shown by list-jobs") + clickHouseBuild = fs.String("clickhouse-build", "", "worker ClickHouse build for this job: 26.6 or 26.9-posthog; empty uses the source job's build") + bucket = fs.String("bucket", "", "S3 bucket for source manifests and finished part artifacts; must match the source job's bucket") + prefix = fs.String("prefix", "partforge", "S3 key prefix under the bucket") + stateTable = fs.String("state-table", defaultStateTable, "Postgres table used for PartForge state") + region = fs.String("aws-region", "", "AWS region for Postgres IAM auth; empty resolves from AWS config, then us-east-1") + s3Endpoint = fs.String("s3-endpoint", "", "custom S3 endpoint URL, for example LocalStack") + s5cmdBinary = fs.String("s5cmd-binary", "s5cmd", "path to the s5cmd binary used for S3 transfers") + uploadConcurrency = fs.Int("upload-concurrency", 0, "number of source manifests to upload concurrently; <=0 uses detected CPU count") + 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, "upload-chain"); err != nil { + return err + } + parentJobID := strings.TrimSpace(*sourceJobID) + if parentJobID == "" || *bucket == "" { + return errors.New("source-job-id and bucket are required") + } + if err := validateUploadFreezeSQLInputs(*destinationSchemaFile, *insertSelectFile, *copySQLFromJob); err != nil { + return err + } + resolvedUploadConcurrency, err := resolveUploadConcurrency(*uploadConcurrency) + if err != nil { + return err + } + + startedAt := time.Now() + slog.Info("upload-chain started", "stage", "start", "source_job_id", parentJobID, "bucket", *bucket, "prefix", *prefix) + stateStore, err := state.New(ctx, state.Config{Region: *region, Endpoint: *postgresURL, IAMAuth: *postgresIAMAuth, Table: *stateTable}) + if err != nil { + return err + } + getStateStore := func() (*state.Store, error) { return stateStore, nil } + copier := s3copy.Copier{Binary: *s5cmdBinary, Endpoint: *s3Endpoint} + + parentParts, err := stateStore.ListJobParts(ctx, parentJobID) + if err != nil { + return err + } + chainParts, err := selectChainSourceParts(parentJobID, parentParts, *bucket) + if err != nil { + return err + } + parentSQL, err := copySQLBundleFromJob(ctx, stateStore, copier, parentJobID) + if err != nil { + return fmt.Errorf("read source job SQL: %w", err) + } + sourceSchema := strings.TrimSpace(parentSQL.DestinationSchema) + source, err := destinationTableRefFromSchema(sourceSchema) + if err != nil { + return fmt.Errorf("source job destination schema: %w", err) + } + if _, err := ddl.NormalizeCreateTable(sourceSchema); err != nil { + return fmt.Errorf("source job destination schema is not supported by worker: %w", err) + } + + destinationSchema, insertSelect, err := readUploadFreezeSQL(ctx, *destinationSchemaFile, *insertSelectFile, *copySQLFromJob, getStateStore, copier) + if err != nil { + return err + } + destination, err := destinationTableRefFromSchema(destinationSchema) + if err != nil { + return err + } + if destination.Database == source.Database && destination.Table == source.Table { + return errors.New("source and destination table names must differ inside the worker; the source is the source job's destination table") + } + + buildValue := strings.TrimSpace(*clickHouseBuild) + if buildValue == "" { + buildValue = chainParts[0].ClickHouseBuild + } + resolvedBuild, err := manifest.ResolveClickHouseBuild(buildValue) + if err != nil { + return err + } + + resolvedJobID := strings.TrimSpace(*jobID) + if resolvedJobID == "" { + resolvedJobID, err = manifest.NewJobID() + if err != nil { + return fmt.Errorf("generate job id: %w", err) + } + } + if resolvedJobID == parentJobID { + return errors.New("job-id must differ from source-job-id") + } + slog.Info("registering chained source parts", "stage", "register_chain_parts", "job_id", resolvedJobID, "source_job_id", parentJobID, "source_table", chhttp.TableSQL(source.Database, source.Table), "destination_table", chhttp.TableSQL(destination.Database, destination.Table), "parts", len(chainParts), "clickhouse_build", resolvedBuild) + + params := chainPartParams{ + ClickHouseBuild: resolvedBuild, + JobID: resolvedJobID, + JobName: strings.TrimSpace(*jobName), + SourceJobID: parentJobID, + Source: source, + Dest: destination, + SourceSchema: sourceSchema, + DestinationSchema: destinationSchema, + InsertSelect: insertSelect, + Bucket: *bucket, + Prefix: *prefix, + StateStore: stateStore, + Copier: copier, + } + completed := 0 + err = processTasksInParallel(ctx, chainParts, min(resolvedUploadConcurrency, len(chainParts)), func(ctx context.Context, _ int, part state.Part) (string, error) { + return registerChainPart(ctx, part, params) + }, func(partID string) { + completed++ + slog.Info("chained source part progress", "stage", "register_chain_parts", "job_id", resolvedJobID, "completed_parts", completed, "parts_total", len(chainParts), "part_id", partID) + }, "chain part registration") + if err != nil { + return err + } + slog.Info("upload-chain complete", "stage", "complete", "job_id", resolvedJobID, "source_job_id", parentJobID, "parts", completed, "elapsed", time.Since(startedAt)) + fmt.Println(resolvedJobID) + return nil +} + +// selectChainSourceParts returns the finished artifacts of a job that a chained +// job can read. Partially finished jobs are rejected so no rows are dropped. +func selectChainSourceParts(jobID string, parts []state.Part, bucket string) ([]state.Part, error) { + if len(parts) == 0 { + return nil, fmt.Errorf("source job %s has no state rows", jobID) + } + unfinished := map[state.Status]int{} + var selected []state.Part + for _, part := range parts { + switch part.Status { + case state.StatusSuperseded: + continue + case state.StatusFinished, state.StatusImporting, state.StatusImported: + default: + unfinished[part.Status]++ + continue + } + if part.EmptyOutput { + continue + } + if part.Bucket != bucket { + return nil, fmt.Errorf("source job part %s/%s is in bucket %q, expected %q", part.JobID, part.PartID, part.Bucket, bucket) + } + if strings.TrimSpace(part.FinishedKey) == "" { + return nil, fmt.Errorf("source job part %s/%s has no finished key", part.JobID, part.PartID) + } + selected = append(selected, part) + } + if len(unfinished) > 0 { + var counts []string + for _, status := range statusOrder() { + if count := unfinished[status]; count > 0 { + counts = append(counts, fmt.Sprintf("%s=%d", status, count)) + } + } + return nil, fmt.Errorf("source job %s is not finished (%s); finish compaction and rewrites before chaining", jobID, strings.Join(counts, ", ")) + } + if len(selected) == 0 { + return nil, fmt.Errorf("source job %s has no finished artifacts with output", jobID) + } + sort.Slice(selected, func(i, j int) bool { return selected[i].PartID < selected[j].PartID }) + return selected, nil +} + +type chainPartParams struct { + ClickHouseBuild string + JobID string + JobName string + SourceJobID string + Source manifest.TableRef + Dest manifest.TableRef + SourceSchema string + DestinationSchema string + InsertSelect string + Bucket string + Prefix string + StateStore *state.Store + Copier s3copy.Copier +} + +func chainPartManifest(parent state.Part, params chainPartParams, createdAt time.Time) manifest.Manifest { + partID := manifest.DerivePartID("chain", parent.FinishedKey, parent.PartID, params.SourceSchema, params.DestinationSchema, params.InsertSelect) + return manifest.Manifest{ + ClickHouseBuild: params.ClickHouseBuild, + + Version: manifest.Version, + JobID: params.JobID, + PartID: partID, + Freeze: params.SourceJobID, + Source: params.Source, + Dest: params.Dest, + Part: manifest.SourcePart{Disk: "chain", Name: parent.PartID, RelativePath: parent.FinishedKey}, + SQL: manifest.SQLBundle{SourceSchema: params.SourceSchema, DestinationSchema: params.DestinationSchema, InsertSelect: params.InsertSelect}, + S3: manifest.S3Refs{ + Bucket: params.Bucket, + SourceKey: manifest.SourcePartPrefix(params.Prefix, params.JobID, partID), + FinishedKey: manifest.FinishedPartPrefix(params.Prefix, params.JobID, partID), + SourceFinishedKey: parent.FinishedKey, + }, + CreatedAt: createdAt, + } +} + +func registerChainPart(ctx context.Context, parent state.Part, params chainPartParams) (string, error) { + createdAt := time.Now().UTC() + m := chainPartManifest(parent, params, createdAt) + dir, err := os.MkdirTemp("", "partforge-chain-part-*") + if err != nil { + return "", err + } + defer os.RemoveAll(dir) + if err := artifact.WriteManifest(dir, m); err != nil { + return "", err + } + if err := params.Copier.UploadDir(ctx, dir, m.S3.Bucket, m.S3.SourceKey); err != nil { + return "", fmt.Errorf("upload chained source manifest for %s/%s: %w", parent.JobID, parent.PartID, err) + } + part := state.NewPart(m.JobID, m.PartID, m.S3.Bucket, m.S3.SourceKey, m.S3.FinishedKey, createdAt) + part.JobName = params.JobName + part.ClickHouseBuild = params.ClickHouseBuild + part.SourceArtifactBytes = parent.DestinationActivePartBytes + if err := params.StateStore.CreatePart(ctx, part); err != nil { + return "", fmt.Errorf("create chained source part state for %s/%s from %s/%s: %w", part.JobID, part.PartID, parent.JobID, parent.PartID, err) + } + return part.PartID, nil +} + type uploadPartFunc func(context.Context, int, uploadPartTask) (uploadPartResult, error) func uploadPartsInParallel(ctx context.Context, tasks []uploadPartTask, concurrency int, upload uploadPartFunc, onResult func(uploadPartResult)) error { diff --git a/cmd/partforge/main_test.go b/cmd/partforge/main_test.go index 9c0ae36..c10e4de 100644 --- a/cmd/partforge/main_test.go +++ b/cmd/partforge/main_test.go @@ -2770,3 +2770,75 @@ func TestSummarizeJobCompactWaitsForBatchableSiblings(t *testing.T) { } } } + +func TestSelectChainSourceParts(t *testing.T) { + finished := func(id string, status state.Status) state.Part { + return state.Part{JobID: "job-1", PartID: id, Status: status, Bucket: "partforge", FinishedKey: "p/jobs/job-1/finished/" + id} + } + empty := finished("part-empty", state.StatusFinished) + empty.EmptyOutput = true + parts := []state.Part{ + finished("part-b", state.StatusImported), + finished("part-a", state.StatusFinished), + finished("part-old", state.StatusSuperseded), + finished("compact-c", state.StatusImporting), + empty, + } + got, err := selectChainSourceParts("job-1", parts, "partforge") + if err != nil { + t.Fatal(err) + } + var ids []string + for _, part := range got { + ids = append(ids, part.PartID) + } + if strings.Join(ids, ",") != "compact-c,part-a,part-b" { + t.Fatalf("selected = %v", ids) + } + + unfinished := append(parts, finished("part-r", state.StatusCompactReady), finished("part-f", state.StatusFailed)) + if _, err := selectChainSourceParts("job-1", unfinished, "partforge"); err == nil || !strings.Contains(err.Error(), "COMPACT_READY=1") || !strings.Contains(err.Error(), "FAILED=1") { + t.Fatalf("expected unfinished error with counts, got %v", err) + } + if _, err := selectChainSourceParts("job-1", parts, "other"); err == nil { + t.Fatal("expected bucket mismatch error") + } + if _, err := selectChainSourceParts("job-1", []state.Part{empty}, "partforge"); err == nil { + t.Fatal("expected error when no artifacts have output") + } + if _, err := selectChainSourceParts("job-1", nil, "partforge"); err == nil { + t.Fatal("expected error for missing job") + } +} + +func TestChainPartManifest(t *testing.T) { + parent := state.Part{JobID: "job-1", PartID: "compact-1", Bucket: "partforge", FinishedKey: "p/jobs/job-1/finished/compact-1"} + params := chainPartParams{ + ClickHouseBuild: "26.6", + JobID: "job-2", + SourceJobID: "job-1", + Source: manifest.TableRef{Database: "dst", Table: "events_new"}, + Dest: manifest.TableRef{Database: "dst", Table: "events_v3"}, + SourceSchema: "CREATE TABLE dst.events_new (id UInt64) ENGINE = MergeTree ORDER BY id", + DestinationSchema: "CREATE TABLE dst.events_v3 (id UInt64) ENGINE = MergeTree ORDER BY id", + InsertSelect: "INSERT INTO dst.events_v3 SELECT id FROM dst.events_new", + Bucket: "partforge", + Prefix: "p", + } + m := chainPartManifest(parent, params, time.Unix(0, 0).UTC()) + if err := m.Validate(); err != nil { + t.Fatal(err) + } + if m.S3.SourceFinishedKey != parent.FinishedKey { + t.Fatalf("source finished key = %q", m.S3.SourceFinishedKey) + } + if m.S3.SourceKey != "p/jobs/job-2/source/"+m.PartID || m.S3.FinishedKey != "p/jobs/job-2/finished/"+m.PartID { + t.Fatalf("unexpected keys: %+v", m.S3) + } + if m.Source != params.Source || m.Dest != params.Dest { + t.Fatalf("unexpected tables: %+v -> %+v", m.Source, m.Dest) + } + if again := chainPartManifest(parent, params, time.Now()); again.PartID != m.PartID { + t.Fatal("chained part id is not stable") + } +} diff --git a/docs/operations.md b/docs/operations.md index fa357ba..ade3423 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -125,7 +125,7 @@ Notes: - To finish only waiting artifacts manually, use `set-part-state -job-id=JOB_ID -status=COMPACT_READY -to-status=FINISHED -force`. Replace `-status=COMPACT_READY` with repeated `-part-id=PART_ID` to select individual artifacts. - `retry-failed` moves failed rewrite parts back to `READY` and failed import parts back to `FINISHED` (so `import-finished` retries the import stage without re-running the worker). Any move back to `READY` clears persisted rewrite progress and metrics. - `reset-job` and `reset-compaction` validate compaction lineage (`compact_input_part_ids` / `superseded_by`) and refuse to run if any part has started import. -- `-delete-s3` variants derive the exact S3 target from the job's recorded rows and reject glob metacharacters before deleting. For jobs created with `upload-freeze -copy-parts-from-job`, borrowed source prefixes are not deleted; jobs that own referenced source parts are blocked from deletion while those references exist. +- `-delete-s3` variants derive the exact S3 target from the job's recorded rows and reject glob metacharacters before deleting. For jobs created with `upload-freeze -copy-parts-from-job`, borrowed source prefixes are not deleted; jobs that own referenced source parts are blocked from deletion while those references exist. Jobs used as `upload-chain -source-job-id` are not protected: keep their finished artifacts until the chained job has finished rewriting. ## Shutdown behavior diff --git a/docs/setup.md b/docs/setup.md index 7639a0f..7a4333a 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -195,6 +195,23 @@ Add `-include-partitions=202401,202402` to include only those exact partition ID `-backup` must identify the exact prefix containing `.backup`, `metadata/`, and `data/`. For an incremental, the command follows `base_backup` recursively and validates each `base_backup_uuid`; S3 base locators must have the form `S3('s3://bucket/prefix')`. It resolves ClickHouse's checksum-based file reuse, including files made from a base prefix plus an incremental suffix, and registers `READY` rows. By default it materializes each logical part into PartForge's normal source layout. With `-zero-copy`, it stores ordered source pointers and workers download files directly from all referenced backup layers; finished and compacted artifacts still go to the PartForge bucket. Every referenced backup must remain available until the job finishes. Lightweight and encrypted backups are unsupported. `-copy-sql-from-job` is supported in place of the two SQL files. +#### Chaining from a finished job + +To run another rewrite over the output of a previous job, use `upload-chain`. Each finished artifact of the source job becomes one `READY` source part of the new job: + +```sh +partforge upload-chain \ + -source-job-id= \ + -destination-schema-file=dest_v3.sql \ + -insert-select-file=insert_v3.sql \ + -bucket=partforge \ + -job-name="events migration 002" +``` + +The new job's source table is the source job's destination table, and its source schema is the source job's destination schema. So `insert_v3.sql` reads `FROM dst_db.events_new`, and the new destination table must have a different name. The command writes only manifests: workers download the source job's finished part tarballs directly, attach every part in the artifact, and run the new insert-select. An artifact with more than one physical part runs as one unchunked insert. + +The source job must be fully finished: every non-`SUPERSEDED` row must be `FINISHED`, `IMPORTING`, or `IMPORTED`. Finalize compaction first if needed. Artifacts with empty output are skipped. The source job must use the same bucket. `-clickhouse-build` defaults to the source job's build. Keep the source job's finished S3 artifacts (do not run `delete-job -delete-s3`, `reset-job`, or `reset-compaction` on it) until the chained job has finished rewriting; PartForge does not block those commands. `-copy-sql-from-job` is supported in place of the two SQL files. + ### 4. worker The worker claims the largest `READY` source artifact first, starts a local ClickHouse, downloads and attaches the source part, runs your `INSERT ... SELECT`, freezes the produced destination parts, uploads one uncompressed tarball per part, and marks the row `COMPACT_READY`. When no rewrite work is left, workers opportunistically compact the eligible finished artifact with the most physical parts first before promoting artifacts to `FINISHED`. diff --git a/e2e/expected-chain.tsv b/e2e/expected-chain.tsv new file mode 100644 index 0000000..7868fe4 --- /dev/null +++ b/e2e/expected-chain.tsv @@ -0,0 +1,4 @@ +1 ALPHA 20 2024-01-01 +2 BRAVO 40 2024-01-02 +3 CHARLIE 60 2024-01-03 +4 DELTA 80 2024-01-04 diff --git a/e2e/run.sh b/e2e/run.sh index 4b574a3..f99b681 100755 --- a/e2e/run.sh +++ b/e2e/run.sh @@ -10,6 +10,7 @@ JOB_ID="e2e-job" COPY_JOB_ID="e2e-copy-job" BACKUP_JOB_ID="e2e-backup-job" BACKUP_COPY_JOB_ID="e2e-backup-copy-job" +CHAIN_JOB_ID="e2e-chain-job" JOB_NAME="E2E import" CLICKHOUSE_BUILD="${PARTFORGE_E2E_CLICKHOUSE_BUILD:-26.6}" @@ -626,6 +627,73 @@ if [[ "$output_part_count" != "1" ]]; then exit 1 fi +# Chain a second rewrite from the imported job's finished artifacts. +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm \ + --workdir /work \ + -v "$ROOT:/work:ro" \ + worker \ + upload-chain \ + -source-job-id="$JOB_ID" \ + -destination-schema-file=e2e/sql/chain-destination.sql \ + -insert-select-file=e2e/sql/chain-insert.sql \ + -bucket=partforge \ + -prefix=e2e \ + -job-id="$CHAIN_JOB_ID" \ + -s3-endpoint=http://localstack:4566 \ + -postgres-url="$POSTGRES_URL" + +chain_part_count="$(docker compose exec -T postgres psql -U partforge -d partforge -Atc \ + "SELECT count(*) FROM partforge_state WHERE job_id = '$CHAIN_JOB_ID' AND status = 'READY'")" +parent_artifact_count="$(docker compose exec -T postgres psql -U partforge -d partforge -Atc \ + "SELECT count(*) FROM partforge_state WHERE job_id = '$JOB_ID' AND status = 'IMPORTED' AND COALESCE(data->>'empty_output', 'false') <> 'true'")" +if [[ "$chain_part_count" == "0" || "$chain_part_count" != "$parent_artifact_count" ]]; then + echo "chain READY parts=$chain_part_count, expected $parent_artifact_count parent artifacts" >&2 + exit 1 +fi +if docker compose exec -T localstack awslocal s3 ls \ + "s3://partforge/e2e/jobs/$CHAIN_JOB_ID/source/" --recursive | + grep -vE '/manifest\.json$' | grep -E '[^[:space:]]' >/dev/null; then + echo "upload-chain unexpectedly copied source part data" >&2 + exit 1 +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 \ + worker \ + -role=inserter \ + -insert-chunk-min-rows=1 \ + -merge-max-runtime=1ns \ + -s3-endpoint=http://localstack:4566 \ + -postgres-url="$POSTGRES_URL" \ + -once 2>&1 | tee "$chain_worker_log" + if ! grep -F "downloading chained source artifact" "$chain_worker_log" >/dev/null; then + echo "chain worker did not read the parent job's finished artifact" >&2 + exit 1 + fi +done + +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \ + finalize-compaction -job-id="$CHAIN_JOB_ID" -all -force -postgres-url="$POSTGRES_URL" + +docker compose exec -T clickhouse clickhouse-client --multiquery < e2e/sql/chain-destination.sql +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm --user "$clickhouse_owner" \ + -v "$DATA_DIR:/var/lib/clickhouse" worker import-finished \ + -database=dst -table=events_v3 -job-id="$CHAIN_JOB_ID" \ + -clickhouse-url="$CH_HTTP_DOCKER" -s3-endpoint=http://localstack:4566 \ + -postgres-url="$POSTGRES_URL" +docker compose exec -T clickhouse clickhouse-client --query \ + "SELECT id, name, amount, event_date FROM dst.events_v3 ORDER BY id FORMAT TSV" \ + > "$ROOT/.e2e/actual-chain.tsv" +diff -u e2e/expected-chain.tsv "$ROOT/.e2e/actual-chain.tsv" + +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \ + delete-job \ + -job-id="$CHAIN_JOB_ID" \ + -delete-s3 \ + -s3-endpoint=http://localstack:4566 \ + -postgres-url="$POSTGRES_URL" + CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker \ delete-job \ -job-id="$JOB_ID" \ diff --git a/e2e/sql/chain-destination.sql b/e2e/sql/chain-destination.sql new file mode 100644 index 0000000..8297d52 --- /dev/null +++ b/e2e/sql/chain-destination.sql @@ -0,0 +1,10 @@ +CREATE TABLE dst.events_v3 +( + id UInt64, + name String, + amount UInt32, + event_date Date +) +ENGINE = MergeTree +PARTITION BY toYYYYMM(event_date) +ORDER BY id diff --git a/e2e/sql/chain-insert.sql b/e2e/sql/chain-insert.sql new file mode 100644 index 0000000..534bb08 --- /dev/null +++ b/e2e/sql/chain-insert.sql @@ -0,0 +1,7 @@ +INSERT INTO dst.events_v3 +SELECT + id, + upper(name) AS name, + toUInt32(amount_text) * 2 AS amount, + event_date +FROM dst.events_new diff --git a/internal/manifest/manifest.go b/internal/manifest/manifest.go index 4f22caf..1351015 100644 --- a/internal/manifest/manifest.go +++ b/internal/manifest/manifest.go @@ -52,6 +52,9 @@ type S3Refs struct { SourceKey string `json:"source_key"` FinishedKey string `json:"finished_key"` SourceObjects []SourceObject `json:"source_objects,omitempty"` + // SourceFinishedKey points at another job's finished artifact prefix. The + // worker extracts its part tarballs and attaches every part as source. + SourceFinishedKey string `json:"source_finished_key,omitempty"` } type SourceObject struct { @@ -79,6 +82,14 @@ func (m Manifest) Validate() error { if m.S3.Bucket == "" || m.S3.SourceKey == "" || m.S3.FinishedKey == "" { return Error("manifest is missing S3 references") } + if m.S3.SourceFinishedKey != "" { + if len(m.S3.SourceObjects) > 0 { + return Error("manifest cannot set both source_objects and source_finished_key") + } + if strings.ContainsAny(m.S3.SourceFinishedKey, "\x00\r\n*?[]{}") { + return Error("manifest contains an unsafe source_finished_key") + } + } for _, object := range m.S3.SourceObjects { if object.Bucket == "" || object.Key == "" || object.Path == "" { return Error("manifest source object is missing bucket, key, or path") diff --git a/internal/manifest/manifest_test.go b/internal/manifest/manifest_test.go index c49f4fa..4c3e89e 100644 --- a/internal/manifest/manifest_test.go +++ b/internal/manifest/manifest_test.go @@ -66,3 +66,36 @@ func TestManifestRejectsUnsafeSourceObjectPath(t *testing.T) { t.Fatal("expected unsafe source object path error") } } + +func TestManifestValidatesSourceFinishedKey(t *testing.T) { + m := Manifest{ + Version: Version, + JobID: "job-2", + PartID: "part-2", + Source: TableRef{Database: "dst", Table: "events"}, + Dest: TableRef{Database: "dst2", Table: "events"}, + Part: SourcePart{Disk: "chain", Name: "part-1", RelativePath: "jobs/job-1/finished/part-1"}, + SQL: SQLBundle{SourceSchema: "source", DestinationSchema: "destination", InsertSelect: "insert"}, + S3: S3Refs{ + Bucket: "partforge", + SourceKey: "jobs/job-2/source/part-2", + FinishedKey: "jobs/job-2/finished/part-2", + SourceFinishedKey: "jobs/job-1/finished/part-1", + }, + } + if err := m.Validate(); err != nil { + t.Fatalf("valid chained manifest: %v", err) + } + + unsafe := m + unsafe.S3.SourceFinishedKey = "jobs/job-1/finished/*" + if err := unsafe.Validate(); err == nil { + t.Fatal("expected unsafe source_finished_key error") + } + + both := m + both.S3.SourceObjects = []SourceObject{{Bucket: "b", Key: "k", Path: "data.bin"}} + if err := both.Validate(); err == nil { + t.Fatal("expected error when source_objects and source_finished_key are both set") + } +} diff --git a/internal/rewrite/processor.go b/internal/rewrite/processor.go index 40a0840..58989b6 100644 --- a/internal/rewrite/processor.go +++ b/internal/rewrite/processor.go @@ -185,6 +185,11 @@ type ProcessResult struct { DestinationPartitions []PartPartitionStats } +type sourcePartDir struct { + Name string + Path string +} + type frozenPartGlob struct { Disk string Glob string @@ -396,7 +401,29 @@ func (p Processor) ProcessPart(ctx context.Context, item WorkItem) (result Proce return ProcessResult{}, fmt.Errorf("download zero-copy source objects: %w", err) } } - sourceStats, err := fileutil.StatDir(sourceRoot) + sourceParts := []sourcePartDir{{Name: m.Part.Name, Path: sourceRoot}} + sourceStatsRoot := sourceRoot + if m.S3.SourceFinishedKey != "" { + finishedRoot := filepath.Join(root, "source-finished") + extractRoot := filepath.Join(root, "source-parts") + slog.Info("downloading chained source artifact", "stage", "download_source", "job_id", item.JobID, "part_id", item.PartID, "bucket", m.S3.Bucket, "source_finished_key", m.S3.SourceFinishedKey) + if err := p.S3Copy.DownloadPrefix(ctx, m.S3.Bucket, m.S3.SourceFinishedKey, finishedRoot); err != nil { + return ProcessResult{}, fmt.Errorf("download chained source artifact %s: %w", m.S3.SourceFinishedKey, err) + } + sourceStatsRoot = finishedRoot + partNames, err := artifact.ExtractFinishedTarballsFromDirContext(ctx, finishedRoot, extractRoot) + if err != nil { + return ProcessResult{}, fmt.Errorf("extract chained source artifact %s: %w", m.S3.SourceFinishedKey, err) + } + if len(partNames) == 0 { + return ProcessResult{}, fmt.Errorf("chained source artifact %s contains no part tarballs", m.S3.SourceFinishedKey) + } + sourceParts = sourceParts[:0] + for _, partName := range partNames { + sourceParts = append(sourceParts, sourcePartDir{Name: partName, Path: filepath.Join(extractRoot, partName)}) + } + } + sourceStats, err := fileutil.StatDir(sourceStatsRoot) if err != nil { return ProcessResult{}, fmt.Errorf("stat downloaded source artifact %s: %w", item.SourceKey, err) } @@ -431,7 +458,7 @@ func (p Processor) ProcessPart(ctx context.Context, item WorkItem) (result Proce } slog.Info("rewriting source part", "stage", "rewrite_part", "job_id", m.JobID, "part_id", m.PartID) rewriteStartedAt := time.Now() - rewriteResult, err := p.rewritePart(ctx, m, sourceRoot, stageTracker) + rewriteResult, err := p.rewritePart(ctx, m, sourceParts, stageTracker) if err != nil { recorder.ForgeFailed(m) return ProcessResult{}, err @@ -513,7 +540,7 @@ func downloadManifestSourceObjects(ctx context.Context, copier s3copy.Copier, ob return os.RemoveAll(segmentRoot) } -func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourcePartRoot string, stageTracker *rewriteStageTracker) (result rewriteResult, err error) { +func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourceParts []sourcePartDir, stageTracker *rewriteStageTracker) (result rewriteResult, err error) { if m.Source.Database == m.Dest.Database && m.Source.Table == m.Dest.Table { return rewriteResult{}, fmt.Errorf("source and destination table names must differ inside the worker") } @@ -562,15 +589,19 @@ func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourceP if err := os.MkdirAll(sourceDetached, 0o755); err != nil { return rewriteResult{}, err } - if err := fileutil.MoveDir(sourcePartRoot, filepath.Join(sourceDetached, m.Part.Name)); err != nil { - return rewriteResult{}, fmt.Errorf("move source part to detached: %w", err) + for _, sourcePart := range sourceParts { + if err := fileutil.MoveDir(sourcePart.Path, filepath.Join(sourceDetached, sourcePart.Name)); err != nil { + return rewriteResult{}, fmt.Errorf("move source part %s to detached: %w", sourcePart.Name, err) + } } if err := p.reportStageProgress(ctx, m, stageTracker, stageAttachSourcePart); err != nil { return rewriteResult{}, err } - slog.Info("attaching source part", "stage", "attach_source_part", "job_id", m.JobID, "part_id", m.PartID, "part", m.Part.Name, "source_table", chhttp.TableSQL(m.Source.Database, m.Source.Table)) - if err := p.ClickHouse.Exec(ctx, "ALTER TABLE "+chhttp.TableSQL(m.Source.Database, m.Source.Table)+" ATTACH PART "+chhttp.StringLiteral(m.Part.Name)); err != nil { - return rewriteResult{}, fmt.Errorf("attach source part %s: %w", m.Part.Name, err) + for _, sourcePart := range sourceParts { + slog.Info("attaching source part", "stage", "attach_source_part", "job_id", m.JobID, "part_id", m.PartID, "part", sourcePart.Name, "source_table", chhttp.TableSQL(m.Source.Database, m.Source.Table)) + if err := p.ClickHouse.Exec(ctx, "ALTER TABLE "+chhttp.TableSQL(m.Source.Database, m.Source.Table)+" ATTACH PART "+chhttp.StringLiteral(sourcePart.Name)); err != nil { + return rewriteResult{}, fmt.Errorf("attach source part %s: %w", sourcePart.Name, err) + } } sourceStats, err := p.activePartStats(ctx, m.Source.Database, m.Source.Table) if err != nil { @@ -585,6 +616,12 @@ func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourceP if err := p.reportStageProgress(ctx, m, stageTracker, stageInsertSelect); err != nil { return rewriteResult{}, err } + if sourceStats.Count > 1 && p.InsertChunkMinRows > 0 { + // Chunk ranges use _part_offset, which only partitions a single part. + // Chained sources can attach several parts from one finished artifact. + slog.Info("disabling insert chunks for multi-part source", "stage", "insert_select", "job_id", m.JobID, "part_id", m.PartID, "source_parts", sourceStats.Count) + p.InsertChunkMinRows = 0 + } slog.Info("running insert-select", "stage", "insert_select", "job_id", m.JobID, "part_id", m.PartID) insertStartedAt := time.Now() if err := p.runInsertSelectWithRetries(ctx, m, sourceStats.Rows); err != nil {