ClickHouse/utils/zero_copy/zero_copy_schema_converter.py

#!/usr/bin/env python3
import argparse
import socket
import uuid
from kazoo.client import KazooClient


def parse_args():
    """
    Parse command-line arguments.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--hosts', default=socket.getfqdn() + ':2181', help='ZooKeeper hosts (host:port,host:port,...)')
    parser.add_argument('-s', '--secure', default=False, action='store_true', help='Use secure connection')
    parser.add_argument('--cert', default='', help='Client TLS certificate file')
    parser.add_argument('--key', default='', help='Client TLS key file')
    parser.add_argument('--ca', default='', help='Client TLS ca file')
    parser.add_argument('-u', '--user', default='', help='ZooKeeper ACL user')
    parser.add_argument('-p', '--password', default='', help='ZooKeeper ACL password')
    parser.add_argument('-r', '--root', default='/clickhouse', help='ZooKeeper root path for ClickHouse')
    parser.add_argument('-z', '--zcroot', default='clickhouse/zero_copy', help='ZooKeeper node for new zero-copy data')
    parser.add_argument('--dryrun', default=False, action='store_true', help='Do not perform any actions')
    parser.add_argument('--cleanup', default=False, action='store_true', help='Clean old nodes')
    parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Verbose mode')

    return parser.parse_args()


# Several folders to heuristic that zookeepr node is folder node
# May be false positive when someone creates set of tables with same paths
table_nodes = ['alter_partition_version', 'block_numbers', 'blocks', 'columns', 'leader_election']
zc_nodes = ['zero_copy_s3', 'zero_copy_hdfs']


def convert_node(client, args, path, zc_node):
    base_path = f'{path}/{zc_node}/shared'
    parts = client.get_children(base_path)
    table_id_path = f'{path}/table_shared_id'
    table_id = ''
    if client.exists(table_id_path):
        table_id = client.get(table_id_path)[0].decode('UTF-8')
    else:
        table_id = str(uuid.uuid4())
        if args.verbose:
            print(f'Make table_id "{table_id_path}" = "{table_id}"')
        if not args.dryrun:
            client.create(table_id_path, bytes(table_id, 'UTF-8'))
    for part in parts:
        part_path = f'{base_path}/{part}'
        uniq_ids = client.get_children(part_path)
        for uniq_id in uniq_ids:
            uniq_path = f'{part_path}/{uniq_id}'
            replicas = client.get_children(uniq_path)
            for replica in replicas:
                replica_path = f'{uniq_path}/{replica}'
                new_path = f'{args.root}/{args.zcroot}/{zc_node}/{table_id}/{part}/{uniq_id}/{replica}'
                if not client.exists(new_path):
                    if args.verbose:
                        print(f'Make node "{new_path}"')
                    if not args.dryrun:
                        client.ensure_path(f'{args.root}/{args.zcroot}/{zc_node}/{table_id}/{part}/{uniq_id}')
                        client.create(new_path, value=b'lock')
                if args.cleanup:
                    if args.verbose:
                        print(f'Remove node "{replica_path}"')
                    if not args.dryrun:
                        client.delete(replica_path)
            if args.cleanup and not args.dryrun:
                client.delete(uniq_path)
        if args.cleanup and not args.dryrun:
            client.delete(part_path)
    if args.cleanup and not args.dryrun:
        client.delete(base_path)
        client.delete(f'{path}/{zc_node}')


def convert_table(client, args, path, nodes):
    print(f'Convert table nodes by path "{path}"')
    for zc_node in zc_nodes:
        if zc_node in nodes:
            convert_node(client, args, path, zc_node)


def is_like_a_table(nodes):
    for tn in table_nodes:
        if tn not in nodes:
            return False
    return True


def scan_recursive(client, args, path):
    nodes = client.get_children(path)
    if is_like_a_table(nodes):
        convert_table(client, args, path, nodes)
    else:
        for node in nodes:
            scan_recursive(client, args, f'{path}/{node}')


def scan(client, args):
    nodes = client.get_children(args.root)
    for node in nodes:
        if node != args.zcroot:
            scan_recursive(client, args, f'{args.root}/{node}')


def get_client(args):
    client = KazooClient(connection_retry=3,
                         command_retry=3,
                         timeout=1,
                         hosts=args.hosts,
                         use_ssl=args.secure,
                         certfile=args.cert,
                         keyfile=args.key,
                         ca=args.ca
                         )
    client.start()
    if (args.user and args.password):
        client.add_auth('digest', f'{args.user}:{args.password}')
    return client


def main():
    args = parse_args()
    client = get_client(args)
    scan(client, args)


if __name__ == '__main__':
    main()