zero-copy migration script

[MikhailBurdukov]

Im pretty sure that the logic will be never merged but it looks like useful script to migrate cluster from one zk path to another one.

another great feature that it is resumable, so if smth goes wrong you can resume the execution.

Summary by Sourcery

Add a zero-copy migration script to the chadmin CLI that supports detaching tables, cleaning ZooKeeper nodes, restarting ClickHouse, and restoring replicas in parallel, with progress persisted to a status file for resumability.

New Features:

• Add a zero-copy migration CLI group with a resumable “migrate” command for migrating replicated tables between ZooKeeper paths
• Persist migration progress in a YAML status file to allow interruption and resumption of the migration workflow

Enhancements:

• Extend execute_tasks_in_parallel to accept an optional callback for tracking task completion and updating migration status

[Alex-Burmak]

@sourcery-ai review

[sourcery-ai[bot]]

Reviewer's Guide

This PR integrates a resumable zero-copy migration workflow into the chadmin CLI by adding a new command group that orchestrates table detachment, ZK cleanup, server restart, and parallel replica restoration using a status file and enhanced process pool callbacks.

Sequence Diagram for Parallel Replica Restoration and Status Update

sequenceDiagram
    participant CLI as "migrate command\n(via restore_replica_step)"
    participant ETP as "execute_tasks_in_parallel()"
    participant WT1 as "WorkerTask 1\n(restore_replica for Table A)"
    participant WT2 as "WorkerTask 2\n(restore_replica for Table B)"
    participant CB as "callback_update_status_file()"
    participant SF as "StatusFile.yaml"

    CLI ->> ETP: Submit tasks [TableA_task, TableB_task, ...], callback_update_status_file
    activate ETP
    ETP ->> WT1: Run target_func(restore_replica)
    activate WT1
    ETP ->> WT2: Run target_func(restore_replica)
    activate WT2

    WT1 -->> ETP: Result for Table A (success/failure)
    deactivate WT1
    ETP ->> CB: callback(TableA_ID)
    activate CB
    CB ->> SF: Update status of Table A to RESTORED
    deactivate CB

    WT2 -->> ETP: Result for Table B (success/failure)
    deactivate WT2
    ETP ->> CB: callback(TableB_ID)
    activate CB
    CB ->> SF: Update status of Table B to RESTORED
    deactivate CB

    ETP -->> CLI: All tasks completed, results returned
    deactivate ETP

Class Diagram for Zero-Copy Migration Components

classDiagram
    class STATUS {
        <<enumeration>>
        INIT
        DETACHED
        ZK_CLEANED
        RESTORED
    }
    class TableMeta {
        +db: str
        +table: str
        +status: STATUS
        +zk_path: str
    }
    TableMeta *-- STATUS : uses

    class ZeroCopyMigrationCLI {
        <<CLI Group>>
        +migrate(ctx, status_file_path, dry_run, do_restore)
        .. internal functions ..
        #generate_status_file(ctx, status_file_path)
        #load_statuses(ctx, status_file_path)
        #update_status_file(status_file_path, tables_stats)
        #detach_tables(ctx, tables_stat, status_file_path, dry_run)
        #remove_zk_nodes(ctx, tables_stat, status_file_path, dry_run)
        #restart_clickhouse_server(ctx, tables_stat, dry_run)
        #restore_replica_step(ctx, tables_stat, status_file_path, dry_run)
    }
    ZeroCopyMigrationCLI ..> TableMeta : manages
    ZeroCopyMigrationCLI ..> STATUS : manages state via TableMeta

    class ProcessPool {
        <<module>>
        +execute_tasks_in_parallel(tasks: List~WorkerTask~, max_workers: int, keep_going: bool, callback: function) : Dict~str, Any~
    }
    class WorkerTask {
        +identifier: str
        +target_func: function
        +args: dict
    }
    ProcessPool o-- WorkerTask : executes
    ZeroCopyMigrationCLI ..> ProcessPool : uses for restore_replica_step

Flow Diagram for the Zero-Copy Migration Process

graph TD
    Start["User executes 'migrate' command"] --> CheckStatusFile{"Status file exists?"};
    CheckStatusFile -- No --> GenerateStatusFile["generate_status_file()\n(All tables: INIT)"];
    GenerateStatusFile --> LoadStatusFile;
    CheckStatusFile -- Yes --> LoadStatusFile["load_statuses()"];

    LoadStatusFile --> CheckRestoreFlag{"--do-restore specified?"};

    CheckRestoreFlag -- No --> DetachPhase["1. Detach Tables (INIT -> DETACHED)"];
    DetachPhase --> UpdateStatusFile1["Update Status File"];
    UpdateStatusFile1 --> ZKCleanPhase["2. Clean ZooKeeper Nodes (DETACHED -> ZK_CLEANED)"];
    ZKCleanPhase --> UpdateStatusFile2["Update Status File"];
    UpdateStatusFile2 --> RestartCH["3. Restart ClickHouse Server (if needed)"];
    RestartCH --> EndPreRestore["Pre-Restore Steps Complete"];

    CheckRestoreFlag -- Yes --> RestorePhase["4. Restore Replicas (ZK_CLEANED -> RESTORED)\n(Uses execute_tasks_in_parallel with callback)"];
    RestorePhase --> UpdateStatusFile3["Update Status File (via callback)"];
    UpdateStatusFile3 --> EndRestore["Restore Steps Complete"];

File-Level Changes

Change	Details	Files
Enhanced parallel task execution with callback support	Added optional callback parameter to execute_tasks_in_parallel signature Invoke callback(idf) after each task completes successfully	ch_tools/common/process_pool.py
Registered zero-copy migration in the main CLI	Imported zero_copy_migration_group into the CLI entrypoint Appended zero_copy_migration_group to the commands list	ch_tools/chadmin/chadmin_cli.py
Added a new zero-copy migration CLI group with resumable workflow	Defined zero_copy_migration_group and migrate subcommand using click Implemented status file generation, loading, and updates via YAML Built multi-step migration: detach tables, delete ZK nodes, restart server, and restore replicas Leveraged execute_tasks_in_parallel for parallel replica restoration with progress callbacks	ch_tools/chadmin/cli/zero_copy_migration_group.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an issue from a review comment by replying to it. You can also reply to a review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull request title to generate a title at any time. You can also comment @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in the pull request body to generate a PR summary at any time exactly where you want it. You can also comment @sourcery-ai summary on the pull request to (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the pull request to resolve all Sourcery comments. Useful if you've already addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull request to dismiss all existing Sourcery reviews. Especially useful if you want to start fresh with a new review - don't forget to comment @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.