Project

General

Profile

1 11 aaronmk
#!/usr/bin/env python
2
# Imports an XML file into a PostgreSQL database
3
# Format: see http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml
4
5
import os
6 42 aaronmk
import os.path
7 11 aaronmk
import psycopg2
8
from psycopg2.extensions import ISOLATION_LEVEL_SERIALIZABLE
9
import sys
10
import xml.dom.minidom
11
12 43 aaronmk
sys.path.append(os.path.dirname(__file__)+"/lib")
13 46 aaronmk
import db_xml
14 11 aaronmk
15
def env_flag(name): return name in os.environ and os.environ[name] != ''
16
17
def main():
18
    prog_name = sys.argv.pop(0)
19
    try:
20
        db_config = {}
21
        for name in ['host', 'user', 'password', 'database']:
22
            if os.environ[name] != '': db_config[name] = os.environ[name]
23
    except KeyError: raise Exception('Usage: env host=... user=... password=...'
24
        ' database=... [commit=1] '+prog_name+' <dataset')
25
    commit = env_flag('commit')
26
27
    # Process dataset
28
    db = psycopg2.connect(**db_config)
29
    db.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
30
    try:
31
        doc = xml.dom.minidom.parse(sys.stdin)
32 40 aaronmk
        row_ct_ref = [0]
33 46 aaronmk
        db_xml.xml2db(db, doc.documentElement, row_ct_ref)
34 40 aaronmk
        print 'Inserted '+str(row_ct_ref[0])+' rows'
35 11 aaronmk
        if commit: db.commit()
36
    finally:
37
        db.rollback()
38
        db.close()
39
40
main()