#!/usr/bin/env python

# Loads a command's CSV output stream into a PostgreSQL table.

# The command may be run more than once.

import csv

import os.path

import re

import subprocess

import sys

sys.path.append(os.path.dirname(__file__)+"/../lib")

import csvs

import exc

import opts

import sql

import sql_gen

import streams

import strings

import util

def main():

    # Usage

    env_names = []

    def usage_err():

        raise SystemExit('Usage: '+opts.env_usage(env_names)+' '+sys.argv[0]

            +' input_cmd [args...]')

    # Parse args

    input_cmd = sys.argv[1:]

    # Get config from env vars

    table = opts.get_env_var('table', None, env_names)

    schema = opts.get_env_var('schema', 'public', env_names)

    db_config = opts.get_env_vars(sql.db_config_names, None, env_names)

    errors_table_only = opts.env_flag('errors_table_only', False, env_names)

    verbosity = util.cast(float, opts.get_env_var('verbosity', 1, env_names))

    if not (input_cmd != [] and table != None and 'engine' in db_config):

        usage_err()

    # Connect to DB

    def log(msg, level=1):

        '''Higher level -> more verbose'''

        if level <= verbosity: sys.stderr.write(msg.rstrip('\n')+'\n')

    db = sql.connect(db_config, log_debug=log)

    def mk_errors_table():

        log('Creating errors table')

        errors_table = sql.errors_table(db, table, if_exists=False)

        sql.drop_table(db, errors_table)

        typed_cols = [

            sql_gen.TypedCol('column', 'text', nullable=False),

            sql_gen.TypedCol('value', 'text'),

            sql_gen.TypedCol('error_code', 'character varying(5)',

                nullable=False),

            sql_gen.TypedCol('error', 'text', nullable=False),

            ]

        sql.create_table(db, errors_table, typed_cols, has_pkey=False)

        index_cols = ['column', 'value', 'error_code', 'error']

        sql.add_index(db, index_cols, errors_table, unique=True)

    use_copy_from = [True]

    # Loads data into the table using the currently-selected approach.

    def load():

        # Open input stream

        proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)

        in_ = proc.stdout

        # Get format info

        info = csvs.stream_info(in_, parse_header=True)

        dialect = info.dialect

        if csvs.is_tsv(dialect): use_copy_from[0] = False

        col_names = info.header

        for i, col in enumerate(col_names): # replace empty column names

            if col == '': col_names[i] = 'column_'+str(i)

        # Select schema and escape names

        def esc_name(name): return sql.esc_name(db, name)

        sql.run_query(db, 'SET search_path TO '+esc_name(schema))

        typed_cols = [sql_gen.TypedCol('row_num', 'serial', nullable=False)]+[

            sql_gen.TypedCol(v, 'text') for v in col_names]

        def load_():

            log('Creating table')

            sql.create_table(db, table, typed_cols, col_indexes=False)

            # Create COPY FROM statement

            if use_copy_from[0]:

                cur = db.db.cursor()

                copy_from = ('COPY '+esc_name(table)+' ('

                    +(', '.join(map(esc_name, col_names)))

                    +') FROM STDIN DELIMITER %(delimiter)s NULL %(null)s')

                assert not csvs.is_tsv(dialect)

                copy_from += ' CSV'

                if dialect.quoting != csv.QUOTE_NONE:

                    copy_from += ' QUOTE %(quotechar)s'

                    if dialect.doublequote: copy_from += ' ESCAPE %(quotechar)s'

                copy_from += ';\n'

                copy_from = cur.mogrify(copy_from, dict(delimiter=

                    dialect.delimiter, null='', quotechar=dialect.quotechar))

            # Load the data

            line_in = streams.ProgressInputStream(in_, sys.stderr, n=1000)

            try:

                if use_copy_from[0]:

                    log('Using COPY FROM')

                    log(copy_from, level=2)

                    db.db.cursor().copy_expert(copy_from, line_in)

                else:

                    log('Using INSERT')

                    cols_ct = len(col_names)+1 # +1 for row_num

                    for row in csvs.make_reader(line_in, dialect):

                        row = map(strings.to_unicode, row)

                        row.insert(0, sql.default) # row_num is autogen

                        util.list_set_length(row, cols_ct) # truncate extra cols

                        sql.insert(db, table, row, log_level=5)

            finally:

                line_in.close() # also closes proc.stdout

                proc.wait()

        sql.with_savepoint(db, load_)

        log('Cleaning up table')

        sql.cleanup_table(db, table, col_names)

        log('Adding indexes')

        for col in typed_cols[1:]: # exclude pkey

            log('Adding index on '+col.name)

            sql.add_index(db, col.name, table, ensure_not_null=False)

        log('Vacuuming table')

        db.db.rollback()

        sql.vacuum(db, table)

        mk_errors_table()

    if errors_table_only: mk_errors_table()

    else:

        try: load()

        except sql.DatabaseErrors, e:

            if use_copy_from[0]: # first try

                exc.print_ex(e, plain=True)

                use_copy_from[0] = False

                load() # try again with different approach

            else: raise

main()