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
            universal_newlines=True)
54
        in_ = proc.stdout
55
        
56
        # Import data
57
        try: sql_io.import_csv(db, table, *csvs.reader_and_header(in_))
58
        finally:
59
            in_.close() # also closes proc.stdout
60
            proc.wait()
61
    
62
    if input_cmd != []:
63
        try: load()
64
        except sql.EncodingException, e:
65
            exc.print_ex(e, plain=True)
66
            assert e.name == 'UTF8'
67
            
68
            db.set_encoding('LATIN1')
69
            load() # try again with new encoding
70
    else: sql_io.cleanup_table(db, table)
71

    
72
main()
(11-11/80)