ClickHouse/tests/integration/test_backup_restore_keeper_map/test.py

from time import sleep
import pytest
from helpers.cluster import ClickHouseCluster


cluster = ClickHouseCluster(__file__)

main_configs = [
    "configs/remote_servers.xml",
    "configs/backups_disk.xml",
    "configs/keeper_map_path_prefix.xml",
]

user_configs = [
    "configs/zookeeper_retries.xml",
]

node1 = cluster.add_instance(
    "node1",
    main_configs=main_configs,
    user_configs=user_configs,
    external_dirs=["/backups/"],
    macros={"replica": "node1", "shard": "shard1"},
    with_zookeeper=True,
    stay_alive=True,
)

node2 = cluster.add_instance(
    "node2",
    main_configs=main_configs,
    user_configs=user_configs,
    external_dirs=["/backups/"],
    macros={"replica": "node2", "shard": "shard1"},
    with_zookeeper=True,
    stay_alive=True,
)


node3 = cluster.add_instance(
    "node3",
    main_configs=main_configs,
    user_configs=user_configs,
    external_dirs=["/backups/"],
    macros={"replica": "node3", "shard": "shard2"},
    with_zookeeper=True,
    stay_alive=True,
)


@pytest.fixture(scope="module", autouse=True)
def start_cluster():
    try:
        cluster.start()
        yield cluster
    finally:
        cluster.shutdown()

backup_id_counter = 0

def new_backup_name(base_name):
    global backup_id_counter
    backup_id_counter += 1
    return f"Disk('backups', '{base_name}{backup_id_counter}')"

def test_on_cluster():
    node1.query_with_retry("CREATE DATABASE keeper_backup ON CLUSTER cluster")
    node1.query_with_retry("CREATE TABLE keeper_backup.keeper1 ON CLUSTER cluster (key UInt64, value String) Engine=KeeperMap('/test_on_cluster1') PRIMARY KEY key")
    node1.query_with_retry("CREATE TABLE keeper_backup.keeper2 ON CLUSTER cluster (key UInt64, value String) Engine=KeeperMap('/test_on_cluster1') PRIMARY KEY key")
    node1.query_with_retry("CREATE TABLE keeper_backup.keeper3 ON CLUSTER cluster (key UInt64, value String) Engine=KeeperMap('/test_on_cluster2') PRIMARY KEY key")
    node1.query_with_retry("INSERT INTO keeper_backup.keeper2 SELECT number, 'test' || toString(number) FROM system.numbers LIMIT 5")
    node1.query_with_retry("INSERT INTO keeper_backup.keeper3 SELECT number, 'test' || toString(number) FROM system.numbers LIMIT 5")

    expected_result = ''.join(f'{i}\ttest{i}\n' for i in range(5))

    def verify_data():
        for node in [node1, node2, node3]:
            for i in range(1, 4):
                result = node.query_with_retry(f'SELECT key, value FROM keeper_backup.keeper{i} ORDER BY key FORMAT TSV')
                assert result == expected_result

    verify_data()

    backup_name = new_backup_name('test_on_cluster')
    node1.query(f"BACKUP DATABASE keeper_backup ON CLUSTER cluster TO {backup_name} SETTINGS async = false;")

    node1.query("DROP DATABASE keeper_backup ON CLUSTER cluster SYNC;")

    def apply_for_all_nodes(f):
        for node in [node1, node2, node3]:
            f(node)

    def change_keeper_map_prefix(node):
        node.replace_config(
            "/etc/clickhouse-server/config.d/keeper_map_path_prefix.xml", """
<clickhouse>
    <keeper_map_path_prefix>/different_path/keeper_map</keeper_map_path_prefix>
</clickhouse>
""")

    apply_for_all_nodes(lambda node: node.stop_clickhouse())
    apply_for_all_nodes(change_keeper_map_prefix)
    apply_for_all_nodes(lambda node: node.start_clickhouse())

    node1.query(f"RESTORE DATABASE keeper_backup ON CLUSTER cluster FROM {backup_name} SETTINGS async = false;")
    
    verify_data()
    
    node1.query("DROP TABLE keeper_backup.keeper3 ON CLUSTER cluster SYNC;")
    node1.query(f"RESTORE TABLE keeper_backup.keeper3 ON CLUSTER cluster FROM {backup_name} SETTINGS async = false;")

    verify_data()