Project

General

Profile

1 53 aaronmk
#!/usr/bin/env python
2
# Maps one datasource to another, using a map spreadsheet if needed
3
# For outputting an XML file to a PostgreSQL database, use the general format of
4
# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml
5
6
import os.path
7
import sys
8 299 aaronmk
import xml.dom.minidom as minidom
9 53 aaronmk
10 266 aaronmk
sys.path.append(os.path.dirname(__file__)+"/../lib")
11 53 aaronmk
12 344 aaronmk
import exc
13 64 aaronmk
import opts
14 281 aaronmk
import Parser
15 131 aaronmk
import sql
16 310 aaronmk
import util
17 133 aaronmk
import xml_dom
18 86 aaronmk
import xml_func
19 53 aaronmk
20 84 aaronmk
def metadata_value(name):
21 164 aaronmk
    if type(name) == str and name.startswith(':'): return name[1:]
22 84 aaronmk
    else: return None
23
24 53 aaronmk
def main():
25 131 aaronmk
    env_names = []
26
    def usage_err():
27
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)
28
            +' [commit=1] '+sys.argv[0]+' [map_path] [<input] [>output]')
29 146 aaronmk
    limit = opts.get_env_var('n', None, env_names)
30
    if limit != None: limit = int(limit)
31 135 aaronmk
    commit = opts.env_flag('commit')
32 131 aaronmk
33 53 aaronmk
    # Get db config from env vars
34 131 aaronmk
    db_config_names = ['engine', 'host', 'user', 'password', 'database']
35 53 aaronmk
    def get_db_config(prefix):
36 64 aaronmk
        return opts.get_env_vars(db_config_names, prefix, env_names)
37 67 aaronmk
    in_db_config = get_db_config('in')
38
    out_db_config = get_db_config('out')
39 131 aaronmk
    in_is_db = 'engine' in in_db_config
40
    out_is_db = 'engine' in out_db_config
41 53 aaronmk
42
    # Parse args
43 73 aaronmk
    map_path = None
44 67 aaronmk
    try: _prog_name, map_path = sys.argv
45 53 aaronmk
    except ValueError:
46 338 aaronmk
        if in_is_db or not out_is_db: usage_err()
47 53 aaronmk
48 57 aaronmk
    # Load map header
49 130 aaronmk
    in_is_xpaths = True
50 318 aaronmk
    out_label = None
51 73 aaronmk
    if map_path != None:
52 56 aaronmk
        import copy
53
        import csv
54
55 53 aaronmk
        import xpath
56
57 133 aaronmk
        metadata = []
58 73 aaronmk
        mappings = []
59
        stream = open(map_path, 'rb')
60
        reader = csv.reader(stream)
61 161 aaronmk
        in_label, out_label = reader.next()[:2]
62 61 aaronmk
        def split_col_name(name):
63 72 aaronmk
            name, sep, root = name.partition(':')
64
            return name, sep != '', root
65 161 aaronmk
        in_label, in_is_xpaths, in_root = split_col_name(in_label)
66
        out_label, out_is_xpaths, out_root = split_col_name(out_label)
67 133 aaronmk
        assert out_is_xpaths # CSV output not supported yet
68 161 aaronmk
        has_types = out_root.startswith('/*s/') # outer elements are types
69 73 aaronmk
        for row in reader:
70
            in_, out = row[:2]
71
            if out != '':
72 164 aaronmk
                if out_is_xpaths: out = out_root+out
73
                mappings.append((in_, out))
74 73 aaronmk
        stream.close()
75 130 aaronmk
    in_is_xml = in_is_xpaths and not in_is_db
76 56 aaronmk
77 318 aaronmk
    if in_is_xml:
78
        doc0 = minidom.parse(sys.stdin)
79
        if out_label == None: out_label = doc0.documentElement.tagName
80 294 aaronmk
81 316 aaronmk
    def process_input(root, process_row):
82 309 aaronmk
        '''Inputs datasource to XML tree, mapping if needed'''
83
        def process_rows(get_value, rows):
84 297 aaronmk
            '''Processes input values
85
            @param get_value f(in_, row):str
86
            '''
87 314 aaronmk
            for i, row in enumerate(rows):
88 316 aaronmk
                if not (limit == None or i < limit): break
89
                row_id = str(i)
90
                for in_, out in mappings:
91
                    value = metadata_value(in_)
92
                    if value == None: value = get_value(in_, row)
93
                    if value != None:
94
                        xpath.put_obj(root, out, row_id, has_types, value)
95
                process_row()
96 297 aaronmk
97 310 aaronmk
        if map_path == None:
98
            iter_ = xml_dom.NodeElemIter(doc0.documentElement)
99
            util.skip(iter_, xml_dom.is_text) # skip metadata
100 317 aaronmk
            for child in iter_:
101
                root.appendChild(child)
102
                process_row()
103 309 aaronmk
        elif in_is_db:
104 130 aaronmk
            assert in_is_xpaths
105 126 aaronmk
106 117 aaronmk
            import db_xml
107
108 161 aaronmk
            in_root_xml = xpath.path2xml(in_root)
109 164 aaronmk
            for i, mapping in enumerate(mappings):
110
                in_, out = mapping
111
                if metadata_value(in_) == None:
112 168 aaronmk
                    mappings[i] = (xpath.path2xml(in_root+'/'+in_), out)
113 126 aaronmk
114 131 aaronmk
            in_db = sql.connect(in_db_config)
115 133 aaronmk
            in_pkeys = {}
116 297 aaronmk
            def get_value(in_, row):
117 167 aaronmk
                pkey, = row
118 297 aaronmk
                in_ = in_.cloneNode(True) # don't modify orig value!
119
                xml_dom.set_id(xpath.get(in_, in_root), pkey)
120
                value = sql.value_or_none(db_xml.get(in_db, in_, in_pkeys))
121
                if value != None: return str(value)
122
                else: return None
123 309 aaronmk
            process_rows(get_value, sql.rows(db_xml.get(in_db, in_root_xml,
124
                in_pkeys, limit)))
125 117 aaronmk
            in_db.close()
126 161 aaronmk
        elif in_is_xml:
127 297 aaronmk
            def get_value(in_, row):
128
                node = xpath.get(row, in_)
129
                if node != None: return xml_dom.value(node)
130
                else: return None
131
            row0 = xpath.get(doc0.documentElement, in_root)
132 309 aaronmk
            process_rows(get_value, xml_dom.NodeElemIter(row0.parentNode))
133 56 aaronmk
        else: # input is CSV
134 133 aaronmk
            map_ = dict(mappings)
135 59 aaronmk
            reader = csv.reader(sys.stdin)
136 84 aaronmk
            cols = reader.next()
137 162 aaronmk
            col_idxs = dict([(value, idx) for idx, value in enumerate(cols)])
138 164 aaronmk
            for i, mapping in enumerate(mappings):
139
                in_, out = mapping
140
                if metadata_value(in_) == None:
141
                    try: mappings[i] = (col_idxs[in_], out)
142
                    except KeyError: pass
143 162 aaronmk
144 297 aaronmk
            def get_value(in_, row):
145
                value = row[in_]
146
                if value != '': return value
147
                else: return None
148 309 aaronmk
            process_rows(get_value, reader)
149 53 aaronmk
150
    # Output XML tree
151 316 aaronmk
    doc = xml_dom.create_doc(out_label)
152
    root = doc.documentElement
153 130 aaronmk
    if out_is_db:
154 53 aaronmk
        from psycopg2.extensions import ISOLATION_LEVEL_SERIALIZABLE
155
        import db_xml
156
157 133 aaronmk
        out_db = sql.connect(out_db_config)
158
        out_db.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
159 310 aaronmk
        out_pkeys = {}
160 53 aaronmk
        try:
161
            row_ct_ref = [0]
162 316 aaronmk
            def process_row():
163
                try: xml_func.process(root)
164 344 aaronmk
                except xml_func.SyntaxException, e: exc.print_ex(e, False)
165 316 aaronmk
                else:
166
                    assert xml_dom.has_one_child(root)
167
                    child = root.firstChild
168
                    try:
169
                        db_xml.put(out_db, child, False, row_ct_ref, out_pkeys)
170 327 aaronmk
                        if commit: out_db.commit()
171 344 aaronmk
                    except sql.DatabaseErrors, e: exc.print_ex(e)
172
                    out_db.rollback() # clean up for next row
173 316 aaronmk
                root.clear()
174
            process_input(root, process_row)
175 53 aaronmk
            print 'Inserted '+str(row_ct_ref[0])+' rows'
176
        finally:
177 133 aaronmk
            out_db.rollback()
178
            out_db.close()
179 299 aaronmk
    else: # output is XML
180 316 aaronmk
        def process_row(): pass
181
        process_input(root, process_row)
182
        xml_func.process(root)
183
        doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
184 53 aaronmk
185 133 aaronmk
try: main()
186 294 aaronmk
except Parser.SyntaxException, e: raise SystemExit(str(e))