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_path = None
47
    try: _prog_name, map_path = sys.argv
48
    except ValueError:
49
        if in_is_db or not out_is_db: usage_err()
50
    
51
    # Load map header
52
    in_is_xpaths = True
53
    out_label = None
54
    if map_path != None:
55
        import copy
56
        import csv
57
        
58
        import xpath
59
        
60
        metadata = []
61
        mappings = []
62
        stream = open(map_path, 'rb')
63
        reader = csv.reader(stream)
64
        in_label, out_label = reader.next()[:2]
65
        def split_col_name(name):
66
            name, sep, root = name.partition(':')
67
            return name, sep != '', root
68
        in_label, in_is_xpaths, in_root = split_col_name(in_label)
69
        out_label, out_is_xpaths, out_root = split_col_name(out_label)
70
        assert out_is_xpaths # CSV output not supported yet
71
        has_types = out_root.startswith('/*s/') # outer elements are types
72
        for row in reader:
73
            in_, out = row[:2]
74
            if out != '':
75
                if out_is_xpaths: out = xpath.parse(out_root+out)
76
                mappings.append((in_, out))
77
        stream.close()
78
    in_is_xml = in_is_xpaths and not in_is_db
79
    
80
    if in_is_xml:
81
        doc0 = minidom.parse(sys.stdin)
82
        if out_label == None: out_label = doc0.documentElement.tagName
83
    
84
    def process_input(root, process_row):
85
        '''Inputs datasource to XML tree, mapping if needed'''
86
        def process_rows(get_value, rows):
87
            '''Processes input values
88
            @param get_value f(in_, row):str
89
            '''
90
            for i, row in enumerate(rows):
91
                if not (limit == None or i < limit): break
92
                row_id = str(i)
93
                for in_, out in mappings:
94
                    value = metadata_value(in_)
95
                    if value == None: value = get_value(in_, row)
96
                    if value != None:
97
                        xpath.put_obj(root, out, row_id, has_types, value)
98
                process_row()
99
        
100
        if map_path == None:
101
            iter_ = xml_dom.NodeElemIter(doc0.documentElement)
102
            util.skip(iter_, xml_dom.is_text) # skip metadata
103
            for child in iter_:
104
                root.appendChild(child)
105
                process_row()
106
        elif in_is_db:
107
            assert in_is_xpaths
108
            
109
            import db_xml
110
            
111
            in_root_xml = xpath.path2xml(in_root)
112
            for i, mapping in enumerate(mappings):
113
                in_, out = mapping
114
                if metadata_value(in_) == None:
115
                    mappings[i] = (xpath.path2xml(in_root+'/'+in_), out)
116
            
117
            in_db = sql.connect(in_db_config)
118
            in_pkeys = {}
119
            def get_value(in_, row):
120
                pkey, = row
121
                in_ = in_.cloneNode(True) # don't modify orig value!
122
                xml_dom.set_id(xpath.get(in_, in_root), pkey)
123
                value = sql.value_or_none(db_xml.get(in_db, in_, in_pkeys))
124
                if value != None: return str(value)
125
                else: return None
126
            process_rows(get_value, sql.rows(db_xml.get(in_db, in_root_xml,
127
                in_pkeys, limit)))
128
            in_db.close()
129
        elif in_is_xml:
130
            def get_value(in_, row):
131
                node = xpath.get(row, in_)
132
                if node != None: return xml_dom.value(node)
133
                else: return None
134
            row0 = xpath.get(doc0.documentElement, in_root)
135
            process_rows(get_value, xml_dom.NodeElemIter(row0.parentNode))
136
        else: # input is CSV
137
            map_ = dict(mappings)
138
            reader = csv.reader(sys.stdin)
139
            cols = reader.next()
140
            col_idxs = dict([(value, idx) for idx, value in enumerate(cols)])
141
            for i, mapping in enumerate(mappings):
142
                in_, out = mapping
143
                if metadata_value(in_) == None:
144
                    try: mappings[i] = (col_idxs[in_], out)
145
                    except KeyError: pass
146
            
147
            def get_value(in_, row):
148
                value = row[in_]
149
                if value != '': return value
150
                else: return None
151
            process_rows(get_value, reader)
152
    
153
    # Output XML tree
154
    doc = xml_dom.create_doc(out_label)
155
    root = doc.documentElement
156
    if out_is_db:
157
        from psycopg2.extensions import ISOLATION_LEVEL_SERIALIZABLE
158
        import db_xml
159
        
160
        out_db = sql.connect(out_db_config)
161
        out_db.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
162
        out_pkeys = {}
163
        try:
164
            if test: sql.empty_db(out_db)
165
            row_ct_ref = [0]
166
            def process_row():
167
                try: xml_func.process(root)
168
                except xml_func.SyntaxException, e: ex_tracker.track(e, False)
169
                else:
170
                    assert xml_dom.has_one_child(root)
171
                    child = root.firstChild
172
                    try:
173
                        sql.with_savepoint(out_db, lambda: db_xml.put(out_db,
174
                            child, False, row_ct_ref, out_pkeys))
175
                        if commit: out_db.commit()
176
                    except sql.DatabaseErrors, e: ex_tracker.track(e)
177
                root.clear()
178
            process_input(root, process_row)
179
            print 'Inserted '+str(row_ct_ref[0])+' rows'
180
        finally:
181
            out_db.rollback()
182
            out_db.close()
183
    else: # output is XML
184
        def process_row(): pass
185
        process_input(root, process_row)
186
        xml_func.process(root)
187
        doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
188

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