| 1 |
1942
|
aaronmk
|
#!/usr/bin/env python
|
| 2 |
|
|
# Loads a command's CSV output stream into a PostgreSQL table.
|
| 3 |
4446
|
aaronmk
|
# When no command is specified, just cleans up the specified table.
|
| 4 |
1942
|
aaronmk
|
# The command may be run more than once.
|
| 5 |
|
|
|
| 6 |
|
|
import os.path
|
| 7 |
|
|
import subprocess
|
| 8 |
|
|
import sys
|
| 9 |
|
|
|
| 10 |
|
|
sys.path.append(os.path.dirname(__file__)+"/../lib")
|
| 11 |
|
|
|
| 12 |
|
|
import exc
|
| 13 |
|
|
import opts
|
| 14 |
|
|
import sql
|
| 15 |
3080
|
aaronmk
|
import sql_io
|
| 16 |
2680
|
aaronmk
|
import sql_gen
|
| 17 |
1942
|
aaronmk
|
import streams
|
| 18 |
1963
|
aaronmk
|
import strings
|
| 19 |
1965
|
aaronmk
|
import util
|
| 20 |
1942
|
aaronmk
|
|
| 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 |
4446
|
aaronmk
|
+' [input_cmd args...]')
|
| 27 |
1942
|
aaronmk
|
|
| 28 |
|
|
# Parse args
|
| 29 |
|
|
input_cmd = sys.argv[1:]
|
| 30 |
|
|
|
| 31 |
|
|
# Get config from env vars
|
| 32 |
|
|
table = opts.get_env_var('table', None, env_names)
|
| 33 |
|
|
schema = opts.get_env_var('schema', 'public', env_names)
|
| 34 |
4257
|
aaronmk
|
has_row_num = opts.env_flag('has_row_num', True, env_names)
|
| 35 |
1942
|
aaronmk
|
db_config = opts.get_env_vars(sql.db_config_names, None, env_names)
|
| 36 |
3271
|
aaronmk
|
verbosity = util.cast(float, opts.get_env_var('verbosity', 3, env_names))
|
| 37 |
1942
|
aaronmk
|
|
| 38 |
4446
|
aaronmk
|
if not (table != None and 'engine' in db_config): usage_err()
|
| 39 |
2890
|
aaronmk
|
|
| 40 |
1942
|
aaronmk
|
# Connect to DB
|
| 41 |
2680
|
aaronmk
|
def log(msg, level=1):
|
| 42 |
|
|
'''Higher level -> more verbose'''
|
| 43 |
3610
|
aaronmk
|
if level <= verbosity:
|
| 44 |
|
|
sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))
|
| 45 |
2680
|
aaronmk
|
db = sql.connect(db_config, log_debug=log)
|
| 46 |
1942
|
aaronmk
|
|
| 47 |
3138
|
aaronmk
|
table = sql_gen.Table(table, schema)
|
| 48 |
|
|
|
| 49 |
4996
|
aaronmk
|
use_copy_from = True
|
| 50 |
1963
|
aaronmk
|
|
| 51 |
|
|
# Loads data into the table using the currently-selected approach.
|
| 52 |
2680
|
aaronmk
|
def load():
|
| 53 |
1942
|
aaronmk
|
# Open input stream
|
| 54 |
|
|
proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
|
| 55 |
|
|
in_ = proc.stdout
|
| 56 |
4996
|
aaronmk
|
line_in = streams.ProgressInputStream(in_, sys.stderr, n=1000)
|
| 57 |
1942
|
aaronmk
|
|
| 58 |
4996
|
aaronmk
|
# Import data
|
| 59 |
|
|
try: sql_io.import_csv(db, table, line_in, use_copy_from, has_row_num)
|
| 60 |
|
|
finally:
|
| 61 |
|
|
line_in.close() # also closes proc.stdout
|
| 62 |
|
|
proc.wait()
|
| 63 |
1942
|
aaronmk
|
|
| 64 |
4446
|
aaronmk
|
if input_cmd != []:
|
| 65 |
|
|
try: load()
|
| 66 |
|
|
except sql.DatabaseErrors, e:
|
| 67 |
4996
|
aaronmk
|
if use_copy_from: # first try
|
| 68 |
4446
|
aaronmk
|
exc.print_ex(e, plain=True)
|
| 69 |
4996
|
aaronmk
|
use_copy_from = False
|
| 70 |
4446
|
aaronmk
|
load() # try again with different approach
|
| 71 |
|
|
else: raise
|
| 72 |
4996
|
aaronmk
|
else: sql_io.cleanup_table(db, table)
|
| 73 |
1942
|
aaronmk
|
|
| 74 |
|
|
main()
|
