Project

General

Profile

1
#!/usr/bin/env python
2
# Loads a command's CSV output stream into a PostgreSQL table.
3
# When no command is specified, just cleans up the specified table.
4
# 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 csvs
13
import exc
14
import opts
15
import sql
16
import sql_io
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
    
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
    verbosity = util.cast(float, opts.get_env_var('verbosity', 3, env_names))
37
    
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:
44
            sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))
45
    db = sql.connect(db_config, log_debug=log)
46
    
47
    table = sql_gen.Table(table, schema)
48
    
49
    # Loads data into the table using the currently-selected approach.
50
    def load():
51
        # Open input stream
52
        proc = subprocess.Popen(input_cmd, stdout=subprocess.PIPE, bufsize=-1)
53
        in_ = proc.stdout
54
        
55
        # Import data
56
        try: sql_io.import_csv(db, table, *csvs.reader_and_header(in_))
57
        finally:
58
            in_.close() # also closes proc.stdout
59
            proc.wait()
60
    
61
    if input_cmd != []:
62
        try: load()
63
        except sql.EncodingException, e:
64
            exc.print_ex(e, plain=True)
65
            assert e.name == 'UTF8'
66
            
67
            db.set_encoding('LATIN1')
68
            load() # try again with new encoding
69
    else: sql_io.cleanup_table(db, table)
70

    
71
main()
(10-10/68)