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

    
21
def main():
22
    # Usage
23
    env_names = []
24
    def usage_err():
25
        raise SystemExit('Usage: '+opts.env_usage(env_names)+' '+sys.argv[0]
26
            +' input_cmd [args...]')
27
    
28
    # Parse args
29
    input_cmd = sys.argv[1:]
30
    if input_cmd == []: usage_err()
31
    
32
    # Get config from env vars
33
    table = opts.get_env_var('table', None, env_names)
34
    schema = opts.get_env_var('schema', 'public', env_names)
35
    db_config = opts.get_env_vars(sql.db_config_names, None, env_names)
36
    debug = opts.env_flag('debug', False, env_names)
37
    if not (table != None and 'engine' in db_config): usage_err()
38
    
39
    # Connect to DB
40
    db = sql.connect(db_config)
41
    
42
    use_copy_from = [True]
43
    
44
    # Loads data into the table using the currently-selected approach.
45
    def load_():
46
        # Open input stream
47
        proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
48
        in_ = proc.stdout
49
        
50
        # Get format info
51
        info = csvs.stream_info(in_, parse_header=True)
52
        dialect = info.dialect
53
        if csvs.is_tsv(dialect): use_copy_from[0] = False
54
        cols = info.header
55
        for i, col in enumerate(cols): # replace empty column names
56
            if col == '': cols[i] = 'column_'+str(i)
57
        
58
        # Select schema and escape names
59
        def esc_name(name): return sql.esc_name(db, name)
60
        sql.run_query(db, 'SET search_path TO '+esc_name(schema))
61
        esc_table = esc_name(table)
62
        esc_cols = map(esc_name, cols)
63
        
64
        # Create CREATE TABLE statement
65
        pkey = esc_name(table+'_pkey')
66
        create_table = 'CREATE TABLE '+esc_table+' (\n'
67
        create_table += '    row_num serial NOT NULL PRIMARY KEY\n'
68
        for esc_col in esc_cols: create_table += '    , '+esc_col+' text\n'
69
        create_table += ');\n'
70
        if debug: sys.stderr.write(create_table)
71
        
72
        # Create table
73
        sql.run_query(db, create_table)
74
        
75
        # Create COPY FROM statement
76
        if use_copy_from[0]:
77
            cur = db.db.cursor()
78
            copy_from = ('COPY '+esc_table+' ('+(', '.join(esc_cols))
79
                +') FROM STDIN DELIMITER %(delimiter)s NULL %(null)s')
80
            assert not csvs.is_tsv(dialect)
81
            copy_from += ' CSV'
82
            if dialect.quoting != csv.QUOTE_NONE:
83
                copy_from += ' QUOTE %(quotechar)s'
84
                if dialect.doublequote: copy_from += ' ESCAPE %(quotechar)s'
85
            copy_from += ';\n'
86
            copy_from = cur.mogrify(copy_from, dict(delimiter=dialect.delimiter,
87
                null='', quotechar=dialect.quotechar))
88
            if debug: sys.stderr.write(copy_from)
89
        
90
        # Load the data
91
        line_in = streams.ProgressInputStream(in_, sys.stderr,
92
            'Processed %d row(s)', n=1000)
93
        try:
94
            if use_copy_from[0]:
95
                sys.stderr.write('Using COPY FROM\n')
96
                db.db.cursor().copy_expert(copy_from, line_in)
97
            else:
98
                sys.stderr.write('Using INSERT\n')
99
                cols_ct = len(cols)+1 # +1 for row_num
100
                for row in csvs.make_reader(line_in, dialect):
101
                    row = map(strings.to_unicode, row)
102
                    row.insert(0, sql.default) # leave space for autogen row_num
103
                    util.list_set_length(row, cols_ct) # truncate extra cols
104
                    sql.insert(db, table, row)
105
        finally:
106
            line_in.close() # also closes proc.stdout
107
            proc.wait()
108
        
109
        # Clean up the data
110
        sys.stderr.write('Cleaning up table\n')
111
        sql.cleanup_table(db, table, cols)
112
    load = lambda: sql.with_savepoint(db, load_)
113
    
114
    try: load()
115
    except sql.DatabaseErrors, e:
116
        if use_copy_from[0]: # first try
117
            exc.print_ex(e, plain=True)
118
            use_copy_from[0] = False
119
            load() # try again with different approach
120
        else: raise e
121
    db.db.commit()
122

    
123
main()
(7-7/50)