Project

General

Profile

1
#!/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
import re
8
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
import sql_io
18
import sql_gen
19
import streams
20
import strings
21
import util
22

    
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
    errors_table_only = opts.env_flag('errors_table_only', False, env_names)
38
    verbosity = util.cast(float, opts.get_env_var('verbosity', 3, env_names))
39
    
40
    if not (input_cmd != [] and table != None and 'engine' in db_config):
41
        usage_err()
42
    
43
    # Connect to DB
44
    def log(msg, level=1):
45
        '''Higher level -> more verbose'''
46
        if level <= verbosity:
47
            sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))
48
    db = sql.connect(db_config, log_debug=log)
49
    
50
    table = sql_gen.Table(table, schema)
51
    
52
    log('Creating errors table')
53
    errors_table = sql_io.errors_table(db, table, if_exists=False)
54
    if errors_table_only: sql.drop_table(db, errors_table)
55
    typed_cols = [
56
        sql_gen.TypedCol('column', 'text', nullable=False),
57
        sql_gen.TypedCol('value', 'text'),
58
        sql_gen.TypedCol('error_code', 'character varying(5)',
59
            nullable=False),
60
        sql_gen.TypedCol('error', 'text', nullable=False),
61
        ]
62
    sql.create_table(db, errors_table, typed_cols, has_pkey=False)
63
    index_cols = ['column', 'value', 'error_code', 'error']
64
    sql.add_index(db, index_cols, errors_table, unique=True)
65
    
66
    use_copy_from = [True]
67
    
68
    # Loads data into the table using the currently-selected approach.
69
    def load():
70
        # Open input stream
71
        proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
72
        in_ = proc.stdout
73
        
74
        # Get format info
75
        info = csvs.stream_info(in_, parse_header=True)
76
        dialect = info.dialect
77
        if csvs.is_tsv(dialect): use_copy_from[0] = False
78
        col_names = map(strings.to_unicode, info.header)
79
        for i, col in enumerate(col_names): # replace empty column names
80
            if col == '': col_names[i] = 'column_'+str(i)
81
        
82
        # Select schema and escape names
83
        def esc_name(name): return db.esc_name(name)
84
        
85
        typed_cols = [sql_gen.TypedCol('row_num', 'serial', nullable=False)]+[
86
            sql_gen.TypedCol(v, 'text') for v in col_names]
87
        
88
        log('Creating table')
89
        sql.create_table(db, table, typed_cols, col_indexes=False)
90
        
91
        # Remove rows from any failed COPY FROM
92
        sql.truncate(db, table)
93
        
94
        def load_():
95
            # Create COPY FROM statement
96
            if use_copy_from[0]:
97
                copy_from = ('COPY '+table.to_str(db)+' ('
98
                    +(', '.join(map(esc_name, col_names)))
99
                    +') FROM STDIN DELIMITER '+db.esc_value(dialect.delimiter)
100
                    +' NULL '+db.esc_value(''))
101
                assert not csvs.is_tsv(dialect)
102
                copy_from += ' CSV'
103
                if dialect.quoting != csv.QUOTE_NONE:
104
                    quote_str = db.esc_value(dialect.quotechar)
105
                    copy_from += ' QUOTE '+quote_str
106
                    if dialect.doublequote: copy_from += ' ESCAPE '+quote_str
107
                copy_from += ';\n'
108
            
109
            # Load the data
110
            line_in = streams.ProgressInputStream(in_, sys.stderr, n=1000)
111
            try:
112
                if use_copy_from[0]:
113
                    log('Using COPY FROM')
114
                    log(copy_from, level=2)
115
                    db.db.cursor().copy_expert(copy_from, line_in)
116
                else:
117
                    log('Using INSERT')
118
                    cols_ct = len(col_names)+1 # +1 for row_num
119
                    for row in csvs.make_reader(line_in, dialect):
120
                        row = map(strings.to_unicode, row)
121
                        row.insert(0, sql.default) # row_num is autogen
122
                        util.list_set_length(row, cols_ct) # truncate extra cols
123
                        sql.insert(db, table, row, cacheable=False, log_level=5)
124
            finally:
125
                line_in.close() # also closes proc.stdout
126
                proc.wait()
127
        sql.with_savepoint(db, load_)
128
        
129
        log('Cleaning up table')
130
        sql_io.cleanup_table(db, table, col_names)
131
        
132
        log('Vacuuming and reanalyzing table')
133
        sql.vacuum(db, table)
134
    
135
    if not errors_table_only:
136
        try: load()
137
        except sql.DatabaseErrors, e:
138
            if use_copy_from[0]: # first try
139
                exc.print_ex(e, plain=True)
140
                use_copy_from[0] = False
141
                load() # try again with different approach
142
            else: raise
143

    
144
main()
(7-7/54)