#!/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 subprocess

import sys

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

import csvs

import exc

import opts

import sql

import streams

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:]

    if input_cmd == []: usage_err()

    # 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)

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

    if not (table != None and 'engine' in db_config): usage_err()

    # Connect to DB

    db = sql.connect(db_config)

    def try_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

        # Escape names

        def esc_name(name): return sql.esc_name(db, name, preserve_case=True)

        qual_table = esc_name(schema)+'.'+esc_name(table)

        esc_cols = map(esc_name, info.header)

        # Create CREATE TABLE statement

        pkey = esc_name(table+'_pkey')

        create_table = 'CREATE TABLE '+qual_table+' (\n'

        create_table += '    row_num serial NOT NULL,\n'

        for esc_col in esc_cols: create_table += '    '+esc_col+' text,\n'

        create_table += '    CONSTRAINT '+pkey+' PRIMARY KEY (row_num)\n'

        create_table += ');\n'

        if debug: sys.stderr.write(create_table)

        # Create COPY FROM statement

        cur = db.db.cursor()

        copy_from = ('COPY '+qual_table+' ('+(', '.join(esc_cols))

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

        if 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=r'\N', quotechar=dialect.quotechar))

        if debug: sys.stderr.write(copy_from)

        # Create table

        sql.run_query(db, create_table)

        # COPY FROM the data

        line_in = streams.ProgressInputStream(in_, sys.stderr,

            'Processed %d row(s)', n=10000)

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

        finally:

            line_in.close() # also closes proc.stdout

            proc.wait()

    for encoding in ['UTF8', 'LATIN1']:

        db.db.set_client_encoding(encoding)

        try: sql.with_savepoint(db, try_load)

        except sql.DatabaseErrors, e:

            if str(e).find('invalid byte sequence for encoding') >= 0:

                exc.print_ex(e, plain=True)

                # now, continue to next encoding

            else: raise e

        else:

            db.db.commit() # commit must occur outside of with_savepoint()

            break # don't try further encodings

main()