#!/usr/bin/env python

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

# When no command is specified, just cleans up the specified table.

# The command may be run more than once.

import os.path

import subprocess

import sys

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

import exc

import opts

import sql

import sql_io

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)

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

    if not (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(strings.to_raw_str(msg.rstrip('\n')+'\n'))

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

    table = sql_gen.Table(table, schema)

    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

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

        # Import data

        try: sql_io.import_csv(db, table, line_in, use_copy_from)

        finally:

            line_in.close() # also closes proc.stdout

            proc.wait()

    if input_cmd != []:

        try: load()

        except sql.EncodingException, e:

            exc.print_ex(e, plain=True)

            assert e.name == 'UTF8'

            db.set_encoding('LATIN1')

            load() # try again with new encoding

        except sql.DatabaseErrors, e:

            if use_copy_from: # first try

                exc.print_ex(e, plain=True)

                use_copy_from = False

                load() # try again with different approach

            else: raise

    else: sql_io.cleanup_table(db, table)

main()