Project

General

Profile

1
#!/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
import xml.dom.minidom as minidom
9

    
10
sys.path.append(os.path.dirname(__file__)+"/../lib")
11

    
12
import exc
13
import opts
14
import Parser
15
import sql
16
import util
17
import xml_dom
18
import xml_func
19

    
20
def metadata_value(name):
21
    if type(name) == str and name.startswith(':'): return name[1:]
22
    else: return None
23

    
24
def main():
25
    ex_tracker = exc.ExTracker()
26
    
27
    env_names = []
28
    def usage_err():
29
        raise SystemExit('Usage: ' + opts.env_usage(env_names, True)
30
            +' [commit=1] [test=1] [verbose=1] [debug=1] '+sys.argv[0]
31
            +' [map_path...] [<input] [>output]')
32
    limit = opts.get_env_var('n', None, env_names)
33
    if limit != None: limit = int(limit)
34
    test = opts.env_flag('test')
35
    commit = not test and opts.env_flag('commit') # never commit in test mode
36
    debug = opts.env_flag('debug')
37
    verbose = debug or opts.env_flag('verbose')
38
    
39
    def log(msg, on=verbose):
40
        if on: sys.stderr.write(msg)
41
    def log_start(action, on=verbose): log(action+'...')
42
    def log_done(on=verbose): log('Done\n')
43
    
44
    # Get db config from env vars
45
    db_config_names = ['engine', 'host', 'user', 'password', 'database']
46
    def get_db_config(prefix):
47
        return opts.get_env_vars(db_config_names, prefix, env_names)
48
    in_db_config = get_db_config('in')
49
    out_db_config = get_db_config('out')
50
    in_is_db = 'engine' in in_db_config
51
    out_is_db = 'engine' in out_db_config
52
    
53
    # Parse args
54
    map_paths = sys.argv[1:]
55
    if map_paths == []:
56
        if in_is_db or not out_is_db: usage_err()
57
        else: map_paths = [None]
58
    
59
    def connect_db(db_config):
60
        log_start('Connecting to '+sql.db_config_str(db_config))
61
        db = sql.connect(db_config)
62
        log_done()
63
        return db
64
    
65
    def process_input(root, process_row, map_path):
66
        '''Inputs datasource to XML tree, mapping if needed'''
67
        # Load map header
68
        in_is_xpaths = True
69
        out_label = None
70
        if map_path != None:
71
            import copy
72
            import csv
73
            
74
            import xpath
75
            
76
            metadata = []
77
            mappings = []
78
            stream = open(map_path, 'rb')
79
            reader = csv.reader(stream)
80
            in_label, out_label = reader.next()[:2]
81
            def split_col_name(name):
82
                name, sep, root = name.partition(':')
83
                return name, sep != '', root
84
            in_label, in_is_xpaths, in_root = split_col_name(in_label)
85
            out_label, out_is_xpaths, out_root = split_col_name(out_label)
86
            assert out_is_xpaths # CSV output not supported yet
87
            has_types = out_root.startswith('/*s/') # outer elements are types
88
            for row in reader:
89
                in_, out = row[:2]
90
                if out != '':
91
                    if out_is_xpaths: out = xpath.parse(out_root+out)
92
                    mappings.append((in_, out))
93
            stream.close()
94
            
95
            root.ownerDocument.documentElement.tagName = out_label
96
        in_is_xml = in_is_xpaths and not in_is_db
97
        
98
        if in_is_xml:
99
            doc0 = minidom.parse(sys.stdin)
100
            if out_label == None: out_label = doc0.documentElement.tagName
101
        
102
        def process_rows(get_value, rows):
103
            '''Processes input values
104
            @param get_value f(in_, row):str
105
            '''
106
            for i, row in enumerate(rows):
107
                if not (limit == None or i < limit): break
108
                row_id = str(i)
109
                for in_, out in mappings:
110
                    value = metadata_value(in_)
111
                    if value == None:
112
                        log_start('Getting '+str(in_), debug)
113
                        value = get_value(in_, row)
114
                        log_done(debug)
115
                    if value != None:
116
                        xpath.put_obj(root, out, row_id, has_types, value)
117
                process_row(row)
118
            sys.stderr.write('Processed '+str(i+1)+' input rows\n')
119
        
120
        if map_path == None:
121
            iter_ = xml_dom.NodeElemIter(doc0.documentElement)
122
            util.skip(iter_, xml_dom.is_text) # skip metadata
123
            for child in iter_:
124
                root.appendChild(child)
125
                process_row(child)
126
        elif in_is_db:
127
            assert in_is_xpaths
128
            
129
            import db_xml
130
            
131
            in_root_xml = xpath.path2xml(in_root)
132
            for i, mapping in enumerate(mappings):
133
                in_, out = mapping
134
                if metadata_value(in_) == None:
135
                    mappings[i] = (xpath.path2xml(in_root+'/'+in_), out)
136
            
137
            in_db = connect_db(in_db_config)
138
            in_pkeys = {}
139
            def get_value(in_, row):
140
                pkey, = row
141
                in_ = in_.cloneNode(True) # don't modify orig value!
142
                xml_dom.set_id(xpath.get(in_, in_root), pkey)
143
                value = sql.value_or_none(db_xml.get(in_db, in_, in_pkeys))
144
                if value != None: return str(value)
145
                else: return None
146
            process_rows(get_value, sql.rows(db_xml.get(in_db, in_root_xml,
147
                in_pkeys, limit)))
148
            in_db.close()
149
        elif in_is_xml:
150
            def get_value(in_, row):
151
                node = xpath.get(row, in_)
152
                if node != None: return xml_dom.value(node)
153
                else: return None
154
            row0 = xpath.get(doc0.documentElement, in_root)
155
            process_rows(get_value, xml_dom.NodeElemIter(row0.parentNode))
156
        else: # input is CSV
157
            map_ = dict(mappings)
158
            reader = csv.reader(sys.stdin)
159
            cols = reader.next()
160
            col_idxs = dict([(value, idx) for idx, value in enumerate(cols)])
161
            for i, mapping in enumerate(mappings):
162
                in_, out = mapping
163
                if metadata_value(in_) == None:
164
                    try: mappings[i] = (col_idxs[in_], out)
165
                    except KeyError: pass
166
            
167
            def get_value(in_, row):
168
                value = row[in_]
169
                if value != '': return value
170
                else: return None
171
            process_rows(get_value, reader)
172
    
173
    def process_inputs(root, process_row):
174
        for map_path in map_paths: process_input(root, process_row, map_path)
175
    
176
    # Output XML tree
177
    doc = xml_dom.create_doc()
178
    root = doc.documentElement
179
    if out_is_db:
180
        import db_xml
181
        
182
        out_db = connect_db(out_db_config)
183
        out_pkeys = {}
184
        try:
185
            if test: sql.empty_db(out_db)
186
            row_ct_ref = [0]
187
            
188
            def process_row(input_row):
189
                def on_error(e):
190
                    exc.add_msg(e, 'output row:\n'+str(root))
191
                    exc.add_msg(e, 'input row:\n'+str(input_row))
192
                    ex_tracker.track(e)
193
                
194
                xml_func.process(root, on_error)
195
                if not xml_dom.is_empty(root):
196
                    assert xml_dom.has_one_child(root)
197
                    try:
198
                        sql.with_savepoint(out_db, lambda: db_xml.put(out_db,
199
                            root.firstChild, out_pkeys, row_ct_ref, on_error))
200
                        if commit: out_db.commit()
201
                    except sql.DatabaseErrors, e: on_error(e)
202
                root.clear()
203
            
204
            process_inputs(root, process_row)
205
            sys.stderr.write('Inserted '+str(row_ct_ref[0])+
206
                ' new rows into database\n')
207
        finally:
208
            out_db.rollback()
209
            out_db.close()
210
    else: # output is XML
211
        def process_row(input_row): pass
212
        process_inputs(root, process_row)
213
        xml_func.process(root)
214
        doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
215

    
216
try: main()
217
except Parser.SyntaxException, e: raise SystemExit(str(e))
(9-9/17)