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_gen
18
import streams
19
import strings
20
import util
21

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

    
145
main()
(7-7/50)