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 |
3271
|
aaronmk
|
verbosity = util.cast(float, opts.get_env_var('verbosity', 3, 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 |
3138
|
aaronmk
|
table = sql_gen.Table(table, schema)
|
50 |
|
|
|
51 |
3134
|
aaronmk
|
log('Creating errors table')
|
52 |
|
|
errors_table = sql_io.errors_table(db, table, if_exists=False)
|
53 |
|
|
sql.drop_table(db, errors_table)
|
54 |
|
|
typed_cols = [
|
55 |
|
|
sql_gen.TypedCol('column', 'text', nullable=False),
|
56 |
|
|
sql_gen.TypedCol('value', 'text'),
|
57 |
|
|
sql_gen.TypedCol('error_code', 'character varying(5)',
|
58 |
|
|
nullable=False),
|
59 |
|
|
sql_gen.TypedCol('error', 'text', nullable=False),
|
60 |
|
|
]
|
61 |
|
|
sql.create_table(db, errors_table, typed_cols, has_pkey=False)
|
62 |
|
|
index_cols = ['column', 'value', 'error_code', 'error']
|
63 |
|
|
sql.add_index(db, index_cols, errors_table, unique=True)
|
64 |
2890
|
aaronmk
|
|
65 |
1963
|
aaronmk
|
use_copy_from = [True]
|
66 |
|
|
|
67 |
|
|
# Loads data into the table using the currently-selected approach.
|
68 |
2680
|
aaronmk
|
def load():
|
69 |
1942
|
aaronmk
|
# Open input stream
|
70 |
|
|
proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
|
71 |
|
|
in_ = proc.stdout
|
72 |
|
|
|
73 |
|
|
# Get format info
|
74 |
|
|
info = csvs.stream_info(in_, parse_header=True)
|
75 |
|
|
dialect = info.dialect
|
76 |
1963
|
aaronmk
|
if csvs.is_tsv(dialect): use_copy_from[0] = False
|
77 |
2680
|
aaronmk
|
col_names = info.header
|
78 |
|
|
for i, col in enumerate(col_names): # replace empty column names
|
79 |
|
|
if col == '': col_names[i] = 'column_'+str(i)
|
80 |
1942
|
aaronmk
|
|
81 |
1963
|
aaronmk
|
# Select schema and escape names
|
82 |
2062
|
aaronmk
|
def esc_name(name): return sql.esc_name(db, name)
|
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 |
3149
|
aaronmk
|
log('Creating table')
|
88 |
|
|
sql.create_table(db, table, typed_cols, col_indexes=False)
|
89 |
|
|
|
90 |
|
|
# Remove rows from any failed COPY FROM
|
91 |
|
|
sql.truncate(db, table)
|
92 |
|
|
|
93 |
2680
|
aaronmk
|
def load_():
|
94 |
|
|
# Create COPY FROM statement
|
95 |
|
|
if use_copy_from[0]:
|
96 |
|
|
cur = db.db.cursor()
|
97 |
3139
|
aaronmk
|
copy_from = ('COPY '+table.to_str(db)+' ('
|
98 |
2680
|
aaronmk
|
+(', '.join(map(esc_name, col_names)))
|
99 |
|
|
+') FROM STDIN DELIMITER %(delimiter)s NULL %(null)s')
|
100 |
|
|
assert not csvs.is_tsv(dialect)
|
101 |
|
|
copy_from += ' CSV'
|
102 |
|
|
if dialect.quoting != csv.QUOTE_NONE:
|
103 |
|
|
copy_from += ' QUOTE %(quotechar)s'
|
104 |
|
|
if dialect.doublequote: copy_from += ' ESCAPE %(quotechar)s'
|
105 |
|
|
copy_from += ';\n'
|
106 |
|
|
copy_from = cur.mogrify(copy_from, dict(delimiter=
|
107 |
|
|
dialect.delimiter, null='', quotechar=dialect.quotechar))
|
108 |
|
|
|
109 |
|
|
# Load the data
|
110 |
2892
|
aaronmk
|
line_in = streams.ProgressInputStream(in_, sys.stderr, n=1000)
|
111 |
2680
|
aaronmk
|
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 |
3270
|
aaronmk
|
sql.insert(db, table, row, cacheable=False, log_level=5)
|
124 |
2680
|
aaronmk
|
finally:
|
125 |
|
|
line_in.close() # also closes proc.stdout
|
126 |
|
|
proc.wait()
|
127 |
|
|
sql.with_savepoint(db, load_)
|
128 |
1963
|
aaronmk
|
|
129 |
2680
|
aaronmk
|
log('Cleaning up table')
|
130 |
3081
|
aaronmk
|
sql_io.cleanup_table(db, table, col_names)
|
131 |
1942
|
aaronmk
|
|
132 |
3147
|
aaronmk
|
log('Vacuuming and reanalyzing table')
|
133 |
|
|
sql.vacuum(db, table)
|
134 |
1942
|
aaronmk
|
|
135 |
3134
|
aaronmk
|
if not errors_table_only:
|
136 |
2890
|
aaronmk
|
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 |
1942
|
aaronmk
|
|
144 |
|
|
main()
|