Project

General

Profile

1
#!/usr/bin/env python
2
# Maps one datasource to another, using a map spreadsheet if needed
3
# Exit status is the # of errors in the import, up to the maximum exit status
4
# For outputting an XML file to a PostgreSQL database, use the general format of
5
# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml
6

    
7
import csv
8
import os.path
9
import sys
10
import xml.dom.minidom as minidom
11

    
12
sys.path.append(os.path.dirname(__file__)+"/../lib")
13

    
14
import csvs
15
import exc
16
import maps
17
import opts
18
import Parser
19
import profiling
20
import sql
21
import strings
22
import term
23
import util
24
import xpath
25
import xml_dom
26
import xml_func
27
import xml_parse
28

    
29
def get_with_prefix(map_, prefixes, key):
30
    '''Gets all entries for the given key with any of the given prefixes'''
31
    values = []
32
    for key_ in strings.with_prefixes(['']+prefixes, key): # also with no prefix
33
        try: value = map_[key_]
34
        except KeyError, e: continue # keep going
35
        values.append(value)
36
    
37
    if values != []: return values
38
    else: raise e # re-raise last KeyError
39

    
40
def metadata_value(name): return None # this feature has been removed
41

    
42
def cleanup(val):
43
    if val == None: return val
44
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
45

    
46
def main_():
47
    env_names = []
48
    def usage_err():
49
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '
50
            +sys.argv[0]+' [map_path...] [<input] [>output]')
51
    
52
    ## Get config from env vars
53
    
54
    # Modes
55
    test = opts.env_flag('test', False, env_names)
56
    commit = opts.env_flag('commit', False, env_names) and not test
57
        # never commit in test mode
58
    redo = opts.env_flag('redo', test, env_names) and not commit
59
        # never redo in commit mode (manually run `make empty_db` instead)
60
    
61
    # Ranges
62
    start = util.cast(int, opts.get_env_var('start', '0', env_names))
63
    if test: end_default = 1
64
    else: end_default = None
65
    end = util.cast(int, util.none_if(
66
        opts.get_env_var('n', end_default, env_names), u''))
67
    if end != None: end += start
68
    
69
    # Debugging
70
    debug = opts.env_flag('debug', False, env_names)
71
    sql.run_raw_query.debug = debug
72
    verbose = debug or opts.env_flag('verbose', not test, env_names)
73
    opts.get_env_var('profile_to', None, env_names) # add to env_names
74
    
75
    # DB
76
    db_config_names = ['engine', 'host', 'user', 'password', 'database']
77
    def get_db_config(prefix):
78
        return opts.get_env_vars(db_config_names, prefix, env_names)
79
    in_db_config = get_db_config('in')
80
    out_db_config = get_db_config('out')
81
    in_is_db = 'engine' in in_db_config
82
    out_is_db = 'engine' in out_db_config
83
    
84
    ##
85
    
86
    # Logging
87
    def log(msg, on=verbose):
88
        if on: sys.stderr.write(msg)
89
    def log_start(action, on=verbose): log(action+'...\n', on)
90
    
91
    # Parse args
92
    map_paths = sys.argv[1:]
93
    if map_paths == []:
94
        if in_is_db or not out_is_db: usage_err()
95
        else: map_paths = [None]
96
    
97
    def connect_db(db_config):
98
        log_start('Connecting to '+sql.db_config_str(db_config))
99
        return sql.connect(db_config)
100
    
101
    if end != None: end_str = str(end-1)
102
    else: end_str = 'end'
103
    log_start('Processing input rows '+str(start)+'-'+end_str)
104
    
105
    ex_tracker = exc.ExPercentTracker(iter_text='row')
106
    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
107
    
108
    doc = xml_dom.create_doc()
109
    root = doc.documentElement
110
    out_is_xml_ref = [False]
111
    in_label_ref = [None]
112
    def update_in_label():
113
        if in_label_ref[0] != None:
114
            xpath.get(root, '/_ignore/inLabel="'+in_label_ref[0]+'"', True)
115
    def prep_root():
116
        root.clear()
117
        update_in_label()
118
    prep_root()
119
    
120
    def process_input(root, row_ready, map_path):
121
        '''Inputs datasource to XML tree, mapping if needed'''
122
        # Load map header
123
        in_is_xpaths = True
124
        out_is_xpaths = True
125
        out_label = None
126
        if map_path != None:
127
            metadata = []
128
            mappings = []
129
            stream = open(map_path, 'rb')
130
            reader = csv.reader(stream)
131
            in_label, out_label = reader.next()[:2]
132
            
133
            def split_col_name(name):
134
                label, sep, root = name.partition(':')
135
                label, sep2, prefixes_str = label.partition('[')
136
                prefixes_str = strings.remove_suffix(']', prefixes_str)
137
                prefixes = strings.split(',', prefixes_str)
138
                return label, sep != '', root, prefixes
139
                    # extract datasrc from "datasrc[data_format]"
140
            
141
            in_label, in_root, prefixes = maps.col_info(in_label)
142
            in_is_xpaths = in_root != None
143
            in_label_ref[0] = in_label
144
            update_in_label()
145
            out_label, out_root = maps.col_info(out_label)[:2]
146
            out_is_xpaths = out_root != None
147
            if out_is_xpaths: has_types = out_root.startswith('/*s/')
148
                # outer elements are types
149
            
150
            for row in reader:
151
                in_, out = row[:2]
152
                if out != '':
153
                    if out_is_xpaths: out = xpath.parse(out_root+out)
154
                    mappings.append((in_, out))
155
            
156
            stream.close()
157
            
158
            root.ownerDocument.documentElement.tagName = out_label
159
        in_is_xml = in_is_xpaths and not in_is_db
160
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
161
        
162
        def process_rows(process_row, rows):
163
            '''Processes input rows
164
            @param process_row(in_row, i)
165
            '''
166
            i = -1 # in case for loop does not execute
167
            for i, row in enumerate(rows):
168
                if i < start: continue
169
                if end != None and i >= end: break
170
                process_row(row, i)
171
                row_ready(i, row)
172
            row_ct = i-start+1
173
            return row_ct
174
        
175
        def map_rows(get_value, rows):
176
            '''Maps input rows
177
            @param get_value(in_, row):str
178
            '''
179
            def process_row(row, i):
180
                row_id = str(i)
181
                for in_, out in mappings:
182
                    value = metadata_value(in_)
183
                    if value == None:
184
                        log_start('Getting '+str(in_), debug)
185
                        value = cleanup(get_value(in_, row))
186
                    if value != None:
187
                        log_start('Putting '+str(out), debug)
188
                        xpath.put_obj(root, out, row_id, has_types, value)
189
            return process_rows(process_row, rows)
190
        
191
        def map_table(col_names, rows):
192
            col_names_ct = len(col_names)
193
            col_idxs = util.list_flip(col_names)
194
            
195
            i = 0
196
            while i < len(mappings): # mappings len changes in loop
197
                in_, out = mappings[i]
198
                if metadata_value(in_) == None:
199
                    try: mappings[i] = (
200
                        get_with_prefix(col_idxs, prefixes, in_), out)
201
                    except KeyError:
202
                        del mappings[i]
203
                        continue # keep i the same
204
                i += 1
205
            
206
            def get_value(in_, row):
207
                return util.coalesce(*util.list_subset(row.list, in_))
208
            def wrap_row(row):
209
                return util.ListDict(util.list_as_length(row, col_names_ct),
210
                    col_names, col_idxs) # handle CSV rows of different lengths
211
            
212
            return map_rows(get_value, util.WrapIter(wrap_row, rows))
213
        
214
        if in_is_db:
215
            assert in_is_xpaths
216
            
217
            in_db = connect_db(in_db_config)
218
            in_pkeys = {}
219
            cur = sql.select(in_db, table=in_root, fields=None, conds=None,
220
                limit=end, start=0)
221
            row_ct = map_table(list(sql.col_names(cur)), sql.rows(cur))
222
            
223
            in_db.close()
224
        elif in_is_xml:
225
            for in_xml_root in xml_parse.docs_iter(sys.stdin):
226
                if map_path == None:
227
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
228
                    util.skip(iter_, xml_dom.is_text) # skip metadata
229
                    row_ct = process_rows(lambda row, i: root.appendChild(row),
230
                        iter_)
231
                else:
232
                    rows = xpath.get(in_xml_root, in_root, limit=end)
233
                    if rows == []: raise SystemExit('Map error: Root "'+in_root
234
                        +'" not found in input')
235
                    
236
                    def get_value(in_, row):
237
                        in_ = './{'+(','.join(strings.with_prefixes(
238
                            ['']+prefixes, in_)))+'}' # also with no prefix
239
                        nodes = xpath.get(row, in_, allow_rooted=False)
240
                        if nodes != []: return xml_dom.value(nodes[0])
241
                        else: return None
242
                    
243
                    row_ct = map_rows(get_value, rows)
244
        else: # input is CSV
245
            map_ = dict(mappings)
246
            reader, col_names = csvs.reader_and_header(sys.stdin)
247
            row_ct = map_table(col_names, reader)
248
        
249
        return row_ct
250
    
251
    def process_inputs(root, row_ready):
252
        row_ct = 0
253
        for map_path in map_paths:
254
            row_ct += process_input(root, row_ready, map_path)
255
        return row_ct
256
    
257
    if out_is_db:
258
        import db_xml
259
        
260
        out_db = connect_db(out_db_config)
261
        out_pkeys = {}
262
        try:
263
            if redo: sql.empty_db(out_db)
264
            row_ins_ct_ref = [0]
265
            
266
            def row_ready(row_num, input_row):
267
                def on_error(e):
268
                    exc.add_msg(e, term.emph('row #:')+' '+str(row_num))
269
                    exc.add_msg(e, term.emph('input row:')+'\n'+str(input_row))
270
                    exc.add_msg(e, term.emph('output row:')+'\n'+str(root))
271
                    ex_tracker.track(e, row_num)
272
                
273
                xml_func.process(root, on_error)
274
                if not xml_dom.is_empty(root):
275
                    assert xml_dom.has_one_child(root)
276
                    try:
277
                        sql.with_savepoint(out_db,
278
                            lambda: db_xml.put(out_db, root.firstChild,
279
                                out_pkeys, row_ins_ct_ref, on_error))
280
                        if commit: out_db.commit()
281
                    except sql.DatabaseErrors, e: on_error(e)
282
                prep_root()
283
            
284
            row_ct = process_inputs(root, row_ready)
285
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
286
                ' new rows into database\n')
287
        finally:
288
            out_db.rollback()
289
            out_db.close()
290
    else:
291
        def on_error(e): ex_tracker.track(e)
292
        def row_ready(row_num, input_row): pass
293
        row_ct = process_inputs(root, row_ready)
294
        xml_func.process(root, on_error)
295
        if out_is_xml_ref[0]:
296
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
297
        else: # output is CSV
298
            raise NotImplementedError('CSV output not supported yet')
299
    
300
    profiler.stop(row_ct)
301
    ex_tracker.add_iters(row_ct)
302
    if verbose:
303
        sys.stderr.write('Processed '+str(row_ct)+' input rows\n')
304
        sys.stderr.write(profiler.msg()+'\n')
305
        sys.stderr.write(ex_tracker.msg()+'\n')
306
    ex_tracker.exit()
307

    
308
def main():
309
    try: main_()
310
    except Parser.SyntaxException, e: raise SystemExit(str(e))
311

    
312
if __name__ == '__main__':
313
    profile_to = opts.get_env_var('profile_to', None)
314
    if profile_to != None:
315
        import cProfile
316
        sys.stderr.write('Profiling to '+profile_to+'\n')
317
        cProfile.run(main.func_code, profile_to)
318
    else: main()
(21-21/40)