#!/usr/bin/env python

import argparse
import os
import subprocess
import sys
import time

# Required so we don't generate tons of logs during restore
disable_logging_sql = "ALTER USER postgres RESET pgaudit.log;"

# Re-enable audit logging after restoring
enable_logging_sql = "ALTER USER postgres SET pgaudit.log='all';"

set_table_owners_sql = """
DO $$
DECLARE
  row record;
BEGIN
  FOR row IN SELECT schemaname, tablename
    FROM pg_tables WHERE NOT schemaname IN ('pg_catalog', 'information_schema') LOOP
    EXECUTE FORMAT('ALTER TABLE %I.%I OWNER TO {pg_user};', row.schemaname, row.tablename);
  END LOOP;
END;
$$;
"""

set_sequence_owners_sql = """
DO $$
DECLARE
  row record;
BEGIN
  FOR row IN SELECT sequence_schema, sequence_name
    FROM information_schema.sequences WHERE NOT sequence_schema IN ('pg_catalog', 'information_schema') LOOP
    EXECUTE FORMAT('ALTER SEQUENCE %I.%I OWNER TO {pg_user};', row.sequence_schema, row.sequence_name);
  END LOOP;
END;
$$;
"""

set_view_owners_sql = """
DO $$
DECLARE
  row record;
BEGIN
  FOR row IN SELECT table_schema, table_name
    FROM information_schema.views WHERE NOT table_schema IN ('pg_catalog', 'information_schema')
    AND NOT table_name IN ('pg_stat_statements', 'pg_stat_statements_info')
    LOOP
    EXECUTE FORMAT('ALTER VIEW %I.%I OWNER TO {pg_user};', row.table_schema, row.table_name);
  END LOOP;
END;
$$;
"""

set_pg_class_owners_sql = """
DO $$
DECLARE
  row record;
BEGIN
  FOR row IN SELECT oid::regclass::text as s_oid
    FROM pg_class WHERE relkind = 'm'
    LOOP
    EXECUTE FORMAT('ALTER TABLE %I OWNER TO {pg_user};', row.s_oid);
  END LOOP;
END;
$$;
"""

backup_and_restore = "pg_dump -Fc ${heroku_pg_url_name}| pg_restore --no-acl --no-owner -d {pg_url} && exit"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Backup and restore from Heroku to new PostgreSQL instance."
    )
    parser.add_argument(
        "--app", dest="heroku_app", help="Heroku app name.", required=True
    )
    parser.add_argument(
        "--db",
        dest="heroku_db",
        default="DATABASE_URL",
        help="Heroku database connection string name (e.g. DATABASE_URL, HEROKU_POSTGRESQL_CRIMSON_URL).",
    )
    parser.add_argument(
        "--size",
        dest="dyno_size",
        default="standard",
        choices=[
            "standard",
            "standard-2x",
            "performance",
            "performance-l",
        ],
        help="Heroku dyno size for running the backup and restore.",
    )
    parser.add_argument(
        "--user",
        dest="pg_user",
        default="application",
        help="New PostgreSQL user that should own the DB resources (default is 'application').",
    )
    parser.add_argument(
        "--target",
        dest="target_pg_url",
        help="PostgreSQL connection string for the restoration target (defaults to PG_URL in environment).",
    )

    args = parser.parse_args()

    start_time = time.time()

    target_pg_url = args.target_pg_url or os.environ.get("PG_URL")
    if not target_pg_url:
        sys.stderr.write(
            "You must provide a connection string in either the --target argument or PG_URL environment variable."
        )
        sys.exit(1)

    bash_command = subprocess.Popen(
        [
            "echo",
            backup_and_restore.format(
                heroku_pg_url_name=args.heroku_db,
                pg_url=target_pg_url,
            ),
        ],
        stdout=subprocess.PIPE,
    )
    subprocess.run(["psql", "-c", disable_logging_sql, "-d", target_pg_url])

    heroku_start = time.time()
    subprocess.run(
        ["heroku", "run", "bash", "--app", args.heroku_app, "--size", args.dyno_size],
        stdin=bash_command.stdout,
    )
    heroku_end = time.time()

    cleanup_sql = (
        "".join(
            [
                set_owner_sql.format(pg_user=args.pg_user)
                for set_owner_sql in [
                    set_table_owners_sql,
                    set_sequence_owners_sql,
                    set_view_owners_sql,
                    set_pg_class_owners_sql,
                ]
            ]
        )
        + enable_logging_sql
    )
    subprocess.run(["psql", "-c", cleanup_sql, "-d", os.environ["target_pg_url"]])
    end_time = time.time()

    print("Total completion time", end_time - start_time)
    print("Backup and restore time", heroku_end - heroku_start)