~adam-collard/landscape-charm/collect-logs-clean-juju-2

#!/usr/bin/python
import os
import sys
import ConfigParser
import base64
import subprocess
import tempfile
import shutil
import multiprocessing

import psycopg2

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

STORES = ["main", "package", "session", "knowledge", "account-1", "resource-1"]

DEAD_TABLES = {
    "main": [
        "ec2_region", "ec2_region_id_seq", "lds_trial", "lds_trial_id_seq"],
    "account-1": [
        "change_package_lock_request", "computer_rule_state_id_seq",
        "computer_rule_state_pkey", "computer_rule_state_rule_id_key",
        "remove_package_lock_request", "create_package_lock_request"],
}

INCONSISTENT_SEQUENCES = {
    "account-1": [
        ("computer_rule_state_id_seq", "rule_state_id_seq"),
        ("event_log_id_seq", "event_log_entry_id_seq"),
    ],
    "resource-1": [
        ("metadata_id_seq", "annotation_id_seq"),
    ],
}


def parse_args(args):
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "action", metavar="ACTION", nargs=1, choices=["import", "export"],
        help="Either 'export' to export db data, or 'import' to import it")
    parser.add_argument(
        "-c", "--config", action="store", metavar="PATH",
        default="/etc/landscape/service.conf",
        help="Path to service.conf.")
    parser.add_argument(
        "-s", "--schema-script", action="store", metavar="PATH",
        default="/usr/bin/landscape-schema",
        help="Path to the landscape-schema script.")
    parser.add_argument(
        "-t", "--tarball", action="store", metavar="PATH",
        default="landscape-db.tgz",
        help="Path to the data tarball to export or import.")
    return parser.parse_args(args)


def read_stores_info(path):
    config = ConfigParser.ConfigParser()
    with open(path) as fd:
        config.readfp(fd)
    host = config.get("stores", "host")
    port = "5432"
    if ":" in host:
        host, port = host.split(":")

    if not config.has_option("stores", "password"):
        user_password = None
    else:
        user_password = decode_password(config.get("stores", "password"))

    if not config.has_option("schema", "store_password"):
        admin_password = None
    else:
        admin_password = decode_password(
            config.get("schema", "store_password"))

    postgres = {
        "host": host,
        "port": port,
        "user": config.get("stores", "user"),
        "password": user_password,
    }

    admin = {
        "user": config.get("schema", "store_user"),
        "password": admin_password,
    }
    if admin["user"] == "superuser":
        admin["user"] = None

    stores = dict((store, config.get("stores", store)) for store in STORES)

    info = {
        "postgres": postgres,
        "admin": admin,
        "stores": stores,
        "root-url": config.get("global", "root-url"),
    }

    return info


def decode_password(password):
    if password.startswith("b64:"):
        password = base64.b64decode(password[len("b64:"):])
    return password


def connect(stores_info, store_name, as_admin=False):
    params = stores_info["postgres"].copy()
    if as_admin:
        params.update(stores_info["admin"])
        if params["user"] is None:
            params = {}

    db_name = stores_info["stores"][store_name]

    return psycopg2.connect(dbname=db_name, **params)


def set_root_url(stores_info):
    connection = connect(stores_info, "main")
    cursor = connection.cursor()
    cursor.execute(
        "UPDATE system_configuration SET value='\"{}\"'"
        "    WHERE key='landscape.root_url'".format(stores_info["root-url"]))
    connection.commit()
    connection.close()


def drop_data(stores_info):
    for store_name in stores_info["stores"]:
        connection = connect(stores_info, store_name, as_admin=True)
        cursor = connection.cursor()

        # Drop all tables
        cursor.execute(
            "SELECT tablename FROM pg_tables WHERE schemaname='public'")
        for row in cursor.fetchall():
            table = row[0]
            if store_name.startswith("resource-"):
                # Check if this is a partioned table
                if "_" in table:
                    suffix = table.split("_")[-1]
                    if len(suffix) == 8:
                        try:
                            int(suffix)
                        except ValueError:
                            # Not a partioned table
                            pass
                        else:
                            # This is a partioned table, it will
                            # be dropped by CASCADE on the parent,
                            # so let's skip it.
                            continue
            cursor.execute("DROP TABLE {} CASCADE".format(table))

        # Drop all extensions
        cursor.execute(
            "SELECT name FROM pg_available_extensions "
            "WHERE installed_version IS NOT NULL")
        for row in cursor.fetchall():
            cursor.execute("DROP EXTENSION {} CASCADE".format(row[0]))

        # Drop all functions
        cursor.execute(
            "SELECT 'DROP FUNCTION ' || ns.nspname || '.' || proname "
            "        || '(' || oidvectortypes(proargtypes) || ');'"
            "FROM pg_proc INNER JOIN pg_namespace ns "
            "ON (pg_proc.pronamespace = ns.oid)"
            "WHERE ns.nspname = 'public'  order by proname")
        for row in cursor.fetchall():
            cursor.execute(row[0])

        # Drop text search
        cursor.execute(
            "DROP TEXT SEARCH CONFIGURATION IF EXISTS landscape_english")

        # Drop large objects
        cursor.execute(
            "SELECT lo_unlink(l.oid) FROM pg_largeobject_metadata l")

        connection.commit()
        connection.close()


def fresh_schema(schema_script, stores_info):
    subprocess.check_call(schema_script)


def rename_incosistent_sequences(stores_info):
    for store_name, sequences in INCONSISTENT_SEQUENCES.iteritems():
        connection = connect(stores_info, store_name, as_admin=True)
        cursor = connection.cursor()
        for old_name, new_name in sequences:
            cursor.execute(
                "ALTER TABLE IF EXISTS {} RENAME TO {}".format(
                    old_name, new_name))
        connection.commit()
        connection.close()


def pg_dump(stores_info, store_name, dump_dir):
    dump_file = os.path.join(dump_dir, store_name + ".sql")
    options = [
           "--host={}".format(stores_info["postgres"]["host"]),
           "--port={}".format(stores_info["postgres"]["port"]),
           "--format=custom",
           "--data-only",
           "--no-acl",
           "--exclude-table=patch",
           "--file={}".format(dump_file),
    ]
    for table in DEAD_TABLES.get(store_name, []):
        options.append("--exclude-table={}".format(table))
    user = stores_info["admin"]["user"]
    if user:
        options.append("--user={}".format(user))

    cmd = ["pg_dump"] + options + [stores_info["stores"][store_name]]
    password = stores_info["admin"]["password"]
    env = os.environ.copy()
    if password:
        env["PGPASSWORD"] = password
    print("Running: {}".format(" ".join(cmd)))
    subprocess.check_call(cmd, env=env)


def pg_dump_partitioned_schemas(stores_info, dump_dir):
    store_name = "resource-1"
    dump_file = os.path.join(dump_dir, store_name + "-partioned.sql")

    options = [
           "--host={}".format(stores_info["postgres"]["host"]),
           "--port={}".format(stores_info["postgres"]["port"]),
           "--schema-only",
           "--no-acl",
           "--no-owner",
           "--file={}".format(dump_file),
    ]

    parents = {}
    children = []

    connection = connect(stores_info, store_name)
    cursor = connection.cursor()
    cursor.execute(
        "SELECT table_name FROM information_schema.tables "
        "WHERE table_schema='public' AND table_name ~ E'[a-z_]+_[0-9]{8}'")

    for row in cursor.fetchall():
        table = row[0]
        options.append("--table={}".format(table))
        parent = "_".join(table.split("_")[:-1])
        if parent not in parents:
            triggers = {}
            trigger_before = parent + "_insert_trigger"
            trigger_after = parent + "_insert_trigger_after"
            for trigger in [trigger_before, trigger_after]:
                cursor.execute(
                    "SELECT pg_get_functiondef('{}'::regproc)".format(trigger))
                triggers[trigger] = cursor.fetchone()[0]
            parents[parent] = triggers
        children.append(table)
    connection.close()

    user = stores_info["admin"]["user"]
    if user:
        options.append("--user={}".format(user))
    cmd = ["pg_dump"] + options + [stores_info["stores"][store_name]]
    password = stores_info["admin"]["password"]
    env = os.environ.copy()
    if password:
        env["PGPASSWORD"] = password
    subprocess.check_call(cmd, env=env)

    with open(dump_file, "a") as fd:
        for parent, triggers in parents.iteritems():
            for trigger, body in triggers.iteritems():
                fd.write(body.strip() + ";\n")
            trigger_before = parent + "_insert_trigger"
            trigger_after = parent + "_insert_trigger_after"
            fd.write(
                "CREATE TRIGGER {} BEFORE INSERT ON {} "
                "FOR EACH ROW EXECUTE PROCEDURE {}();\n".format(
                    trigger_before, parent, trigger_before))
            fd.write(
                "CREATE TRIGGER {} AFTER INSERT ON {} "
                "FOR EACH ROW EXECUTE PROCEDURE {}();\n".format(
                    trigger_after, parent, trigger_after))
        for child in children:
            fd.write(
                "GRANT SELECT,INSERT,DELETE ON TABLE {} "
                "TO landscape;\n".format(child))


def psql_restore_partitioned_schema(stores_info, dump_dir):
    store_name = "resource-1"
    dump_file = os.path.join(dump_dir, store_name + "-partioned.sql")

    options = [
           "--host={}".format(stores_info["postgres"]["host"]),
           "--port={}".format(stores_info["postgres"]["port"]),
           "--file={}".format(dump_file),
    ]
    user = stores_info["admin"]["user"]
    if user:
        options.append("--user={}".format(user))
    cmd = ["psql"] + options + [stores_info["stores"][store_name]]
    password = stores_info["admin"]["password"]
    env = os.environ.copy()
    if password:
        env["PGPASSWORD"] = password
    subprocess.check_call(cmd, env=env)


def pg_restore(stores_info, store_name, dump_dir):
    dump_file = os.path.join(dump_dir, store_name + ".sql")
    options = [
           "--jobs={}".format(multiprocessing.cpu_count()),
           "--disable-triggers",
           "--no-owner",
           "--data-only",
           "--dbname={}".format(stores_info["stores"][store_name]),
    ]
    user = stores_info["admin"]["user"]
    if user:
        options.extend([
            "--user={}".format(user),
            "--host={}".format(stores_info["postgres"]["host"]),
            "--port={}".format(stores_info["postgres"]["port"]),
        ])

    cmd = ["pg_restore"] + options + [dump_file]
    password = stores_info["admin"]["password"]
    env = os.environ.copy()
    if password:
        env["PGPASSWORD"] = password
    print("Running: {}".format(" ".join(cmd)))
    subprocess.check_output(cmd, env=env)


def action_export(stores_info, tarball_file):
    try:
        dump_dir = tempfile.mkdtemp()
        rename_incosistent_sequences(stores_info)
        dump_files = []
        for store_name in STORES:
            pg_dump(stores_info, store_name, dump_dir)
            dump_files.append(store_name + ".sql")
        pg_dump_partitioned_schemas(stores_info, dump_dir)
        dump_files.append("resource-1-partioned.sql")
        cmd = ["tar", "cfz", tarball_file, "-C", dump_dir] + dump_files
        subprocess.check_call(" ".join(cmd), shell=True)
    finally:
        shutil.rmtree(dump_dir)


def action_import(stores_info, tarball_file, schema_script):
    try:
        dump_dir = tempfile.mkdtemp()
        cmd = ["tar", "xfz", tarball_file, "-C", dump_dir]
        subprocess.check_call(" ".join(cmd), shell=True)
        drop_data(stores_info)
        fresh_schema(schema_script, stores_info)
        psql_restore_partitioned_schema(stores_info, dump_dir)
        for store_name in STORES:
            pg_restore(stores_info, store_name, dump_dir)
        set_root_url(stores_info)
    finally:
        shutil.rmtree(dump_dir)


if __name__ == "__main__":
    args = parse_args(sys.argv[1:])
    stores_info = read_stores_info(args.config)
    action = args.action[0]
    if action == "export":
        action_export(stores_info, args.tarball)
        print("Database export available at {}".format(args.tarball))
    if action == "import":
        action_import(stores_info, args.tarball, args.schema_script)