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

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