#!/usr/bin/env python

import argparse
import os
import psycopg2

parser = argparse.ArgumentParser(description='Make a new now() that returns an offset value (without clobbering pg_catalog.now())')

# Format: as used by psycopg2 - for example,
# 'host=MY_HOST dbname=MY_DB user=MY_USER'
parser.add_argument('--connect-string',
                    required=True,
                    )

# Format: 'YYYY-MM-DD HH:MM:SS.xxx'
parser.add_argument('--log-start',
                    required=True,
                    )

# What to name the new DEFAULT value-generating function we're installing.
# Format: no parentheses, e.g., 'public.NOW'.
parser.add_argument('--function-name',
                    default='public.NOW',
                    )

# Format: whitespace-separated string
parser.add_argument('--schemas',
                    required=True,
                    )

parser.add_argument('--verbose',
                    action='store_true',
                    )

parser.add_argument('--dry-run',
                    action='store_true',
                    )

args = parser.parse_args()
connection_string = args.connect_string
log_start = args.log_start
function_name = args.function_name
schemas = args.schemas.split()

# Let the --verbose option override the corresponding environment variable.
verbose = os.getenv('NOWTHEN_VERBOSE', False)
verbose = args.verbose

# But be safe, and do a dry run if either the environment variable
# or the command line requests a dry run, even if the two conflict.
dry_run = args.dry_run or os.getenv('NOWTHEN_DRY_RUN', False)

conn = psycopg2.connect(connection_string)
conn.set_session(autocommit=True)

def get_script_start():
    start_cur = conn.cursor()
    start_cur.execute('SELECT now()')
    now = start_cur.fetchone()[0]
    return now


def install_function(script_start):
    function_defining_cur = conn.cursor()
    function_sql = '''
        CREATE OR REPLACE FUNCTION {function_name}() RETURNS timestamptz LANGUAGE 'sql' AS $$
            SELECT pg_catalog.NOW() - ('{script_start}'::timestamptz - '{log_start}'::timestamptz);
        $$;
    '''.format(function_name=function_name,
               script_start=script_start,
               log_start=log_start)
    if not dry_run:
        function_defining_cur.execute(function_sql)


def _get_columns_matching(search_clause):
    """
    Do not call this function directly. Call find_and_alter_dl instead.

    Get information_schema.columns rows for columns that match
    'WHERE <search_clause>'. For details, see find_and_alter_dl.

    """
    schemas_csv = ', '.join(map(lambda x: "'" + x + "'", schemas))
    defaults_cur = conn.cursor()
    defaults_query = '''
          SELECT table_schema, table_name, column_name, column_default
          FROM information_schema.columns
         WHERE table_schema IN ({schemas})
           AND {clause}
    '''.format(schemas=schemas_csv,
               clause=search_clause)
    if verbose:
        print defaults_query
    defaults_cur.execute(defaults_query)
    results = defaults_cur.fetchall()
    return results


def _set_column_defaults(columns, new_default):
    """
    Do not call this function directly. Call find_and_alter_dl instead.

    Set the default for each given column to the string <default>.

    <columns> is a list of lists, where each sub-list looks like this:
    [schema, table, column, old_default].

    <default> is a string giving the new default value.

    """

    alter_cur = conn.cursor()
    for col in columns:
        (schema, table, column, default) = col
        alter_query = '''
          ALTER TABLE {schema}.{table} ALTER COLUMN {column} SET DEFAULT {to}
        '''.format(
            schema=schema,
            table=table,
            column=column,
            to=new_default)
        if verbose:
            print alter_query
        if not dry_run:
            alter_cur.execute(alter_query)


def find_and_alter_dl(match_clause, new_default):
    """
    Find columns with default defaults matching 'WHERE <match_clause>',
    and change them to use <new_default>. The actual query that contains
    <match_clause> is run against information_schema.columns, so it
    should reference the columns of that table, such as column_default,
    e.g., find_and_alter_dl("column_default = 'old_default'", 'new_default')

    """
    columns = _get_columns_matching(match_clause)
    _set_column_defaults(columns, new_default)


def alter_default_now_columns():
    find_and_alter_dl("column_default IN ( 'now()', $$('now'::text)::date$$ )",
                      '{}()'.format(function_name))

def alter_default_epochy_columns():
    find_and_alter_dl("column_default = $$date_part('epoch'::text, now())$$",
                      "date_part('epoch'::text, {}())".format(function_name))

def alter_default_offset_columns():
    find_and_alter_dl("column_default = $$(now() + '01:00:00'::interval)$$",
                      "({}() + '01:00:00'::interval)".format(function_name))


def main():
    now = get_script_start()
    install_function(script_start=now)

    alter_default_now_columns()
    alter_default_epochy_columns()
    alter_default_offset_columns()


main()
