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()
|