Project

General

Profile

1 1942 aaronmk
#!/usr/bin/env python
2
# Loads a command's CSV output stream into a PostgreSQL table.
3
# The command may be run more than once.
4
5
import csv
6
import os.path
7 1963 aaronmk
import re
8 1942 aaronmk
import subprocess
9
import sys
10
11
sys.path.append(os.path.dirname(__file__)+"/../lib")
12
13
import csvs
14
import exc
15
import opts
16
import sql
17 3080 aaronmk
import sql_io
18 2680 aaronmk
import sql_gen
19 1942 aaronmk
import streams
20 1963 aaronmk
import strings
21 1965 aaronmk
import util
22 1942 aaronmk
23
def main():
24
    # Usage
25
    env_names = []
26
    def usage_err():
27
        raise SystemExit('Usage: '+opts.env_usage(env_names)+' '+sys.argv[0]
28
            +' input_cmd [args...]')
29
30
    # Parse args
31
    input_cmd = sys.argv[1:]
32
33
    # Get config from env vars
34
    table = opts.get_env_var('table', None, env_names)
35
    schema = opts.get_env_var('schema', 'public', env_names)
36
    db_config = opts.get_env_vars(sql.db_config_names, None, env_names)
37 2890 aaronmk
    errors_table_only = opts.env_flag('errors_table_only', False, env_names)
38 2680 aaronmk
    verbosity = util.cast(float, opts.get_env_var('verbosity', 1, env_names))
39 1942 aaronmk
40 2890 aaronmk
    if not (input_cmd != [] and table != None and 'engine' in db_config):
41
        usage_err()
42
43 1942 aaronmk
    # Connect to DB
44 2680 aaronmk
    def log(msg, level=1):
45
        '''Higher level -> more verbose'''
46
        if level <= verbosity: sys.stderr.write(msg.rstrip('\n')+'\n')
47
    db = sql.connect(db_config, log_debug=log)
48 1942 aaronmk
49 2890 aaronmk
    def mk_errors_table():
50
        log('Creating errors table')
51 3080 aaronmk
        errors_table = sql_io.errors_table(db, table, if_exists=False)
52 2890 aaronmk
        sql.drop_table(db, errors_table)
53
        typed_cols = [
54
            sql_gen.TypedCol('column', 'text', nullable=False),
55
            sql_gen.TypedCol('value', 'text'),
56
            sql_gen.TypedCol('error_code', 'character varying(5)',
57
                nullable=False),
58
            sql_gen.TypedCol('error', 'text', nullable=False),
59
            ]
60
        sql.create_table(db, errors_table, typed_cols, has_pkey=False)
61
        index_cols = ['column', 'value', 'error_code', 'error']
62
        sql.add_index(db, index_cols, errors_table, unique=True)
63
64 1963 aaronmk
    use_copy_from = [True]
65
66
    # Loads data into the table using the currently-selected approach.
67 2680 aaronmk
    def load():
68 1942 aaronmk
        # Open input stream
69
        proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
70
        in_ = proc.stdout
71
72
        # Get format info
73
        info = csvs.stream_info(in_, parse_header=True)
74
        dialect = info.dialect
75 1963 aaronmk
        if csvs.is_tsv(dialect): use_copy_from[0] = False
76 2680 aaronmk
        col_names = info.header
77
        for i, col in enumerate(col_names): # replace empty column names
78
            if col == '': col_names[i] = 'column_'+str(i)
79 1942 aaronmk
80 1963 aaronmk
        # Select schema and escape names
81 2062 aaronmk
        def esc_name(name): return sql.esc_name(db, name)
82 1963 aaronmk
        sql.run_query(db, 'SET search_path TO '+esc_name(schema))
83 1942 aaronmk
84 2876 aaronmk
        typed_cols = [sql_gen.TypedCol('row_num', 'serial', nullable=False)]+[
85 2759 aaronmk
            sql_gen.TypedCol(v, 'text') for v in col_names]
86
87 2680 aaronmk
        def load_():
88
            log('Creating table')
89 2760 aaronmk
            sql.create_table(db, table, typed_cols, col_indexes=False)
90 2680 aaronmk
91
            # Create COPY FROM statement
92
            if use_copy_from[0]:
93
                cur = db.db.cursor()
94
                copy_from = ('COPY '+esc_name(table)+' ('
95
                    +(', '.join(map(esc_name, col_names)))
96
                    +') FROM STDIN DELIMITER %(delimiter)s NULL %(null)s')
97
                assert not csvs.is_tsv(dialect)
98
                copy_from += ' CSV'
99
                if dialect.quoting != csv.QUOTE_NONE:
100
                    copy_from += ' QUOTE %(quotechar)s'
101
                    if dialect.doublequote: copy_from += ' ESCAPE %(quotechar)s'
102
                copy_from += ';\n'
103
                copy_from = cur.mogrify(copy_from, dict(delimiter=
104
                    dialect.delimiter, null='', quotechar=dialect.quotechar))
105
106
            # Load the data
107 2892 aaronmk
            line_in = streams.ProgressInputStream(in_, sys.stderr, n=1000)
108 2680 aaronmk
            try:
109
                if use_copy_from[0]:
110
                    log('Using COPY FROM')
111
                    log(copy_from, level=2)
112
                    db.db.cursor().copy_expert(copy_from, line_in)
113
                else:
114
                    log('Using INSERT')
115
                    cols_ct = len(col_names)+1 # +1 for row_num
116
                    for row in csvs.make_reader(line_in, dialect):
117
                        row = map(strings.to_unicode, row)
118
                        row.insert(0, sql.default) # row_num is autogen
119
                        util.list_set_length(row, cols_ct) # truncate extra cols
120 2926 aaronmk
                        sql.insert(db, table, row, log_level=5)
121 2680 aaronmk
            finally:
122
                line_in.close() # also closes proc.stdout
123
                proc.wait()
124
        sql.with_savepoint(db, load_)
125 1963 aaronmk
126 2680 aaronmk
        log('Cleaning up table')
127 3081 aaronmk
        sql_io.cleanup_table(db, table, col_names)
128 1942 aaronmk
129 2680 aaronmk
        log('Adding indexes')
130 2759 aaronmk
        for col in typed_cols[1:]: # exclude pkey
131
            log('Adding index on '+col.name)
132 3054 aaronmk
            sql.add_index(db, col.name, table, ensure_not_null_=False)
133 2682 aaronmk
134 3069 aaronmk
        log('Reanalyzing table')
135
        sql.analyze(db, table)
136
137 2890 aaronmk
        mk_errors_table()
138 1942 aaronmk
139 2890 aaronmk
    if errors_table_only: mk_errors_table()
140
    else:
141
        try: load()
142
        except sql.DatabaseErrors, e:
143
            if use_copy_from[0]: # first try
144
                exc.print_ex(e, plain=True)
145
                use_copy_from[0] = False
146
                load() # try again with different approach
147
            else: raise
148 1942 aaronmk
149
main()