#!/usr/bin/env python3

# install:
# #pip3 install clickhouse-connect

import argparse
import re
import json
import clickhouse_connect

DUMPFILE = "/tmp/ch_export"
createre = re.compile("(?P<begin>CREATE (TABLE|(MATERIALIZED )?VIEW)) (?P<database>\w+)\.(?P<table>\w+) (?P<query>.*)")

database_exclude = {
    "default": True,
    "INFORMATION_SCHEMA": True,
    "information_schema": True,
    "system": True,
}

def get_databases_tables(client):
    databases = {}
    q_db = client.query("show databases;")
    for d in q_db.result_rows:
        dbname = d[0]
        if database_exclude.get(dbname):
            continue
        databases[dbname] = {"schema":"", "tables":{}}
        q_tables = client.query(f"show tables from {dbname};")
        for t in q_tables.result_rows:
            tablename = t[0]
            q_table_schema = client.query(f"SELECT create_table_query, uuid from system.tables where database='{dbname}' and table='{tablename}';")
            for schema in q_table_schema.result_rows:
                create_table_database = schema[0]
                uuid = schema[1]
                resre = createre.match(create_table_database)
                begin = resre.group("begin")
                database = resre.group("database")
                table = resre.group("table")
                query = resre.group("query")
                fullquery = f"{begin} {database}.{table} UUID '{uuid}' {query}"
                databases[dbname]["tables"][tablename] = fullquery
    return databases

def export(dumpfile=DUMPFILE):
    client = clickhouse_connect.get_client(host='localhost', username='default', password='')
    databases = get_databases_tables(client)
    with open(dumpfile, "w") as f:
        json.dump(databases, f, indent=4)
    return

def dump(dumpfile=DUMPFILE):
    with open(dumpfile, "r") as f:
        a = json.load(f)
        for database,values in a.items():
            for table,c in values["tables"].items():
                print(f"# {database}.{table}")
                print(f"{c}\n")

def main():
    parser = argparse.ArgumentParser(
                    prog="clickhouse_ddl_export",
                    description="Exports clickhouse DDL, useful when adding replicas in a cluster")
    parser.add_argument("dumpfile", default=DUMPFILE, nargs="?")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--export", action="store_true")
    group.add_argument("--print-statements", action="store_true")
    args = parser.parse_args()
    if args.export:
        export(args.dumpfile)
    if args.print_statements:
        dump(args.dumpfile)

if __name__ == "__main__":
    main()