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
# Multi-safe (supports an input appearing multiple times).
5
# For outputting an XML file to a PostgreSQL database, use the general format of
6
# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml
7
# Duplicate-column safe (supports multiple columns of the same name, which will
8
# be combined)
9
# Case- and punctuation-insensitive.
10

    
11
import copy
12
import csv
13
import itertools
14
import os.path
15
import warnings
16
import sys
17
import xml.dom.minidom as minidom
18

    
19
sys.path.append(os.path.dirname(__file__)+"/../lib")
20

    
21
import csvs
22
import db_xml
23
import exc
24
import iters
25
import maps
26
import opts
27
import parallelproc
28
import Parser
29
import profiling
30
import sql
31
import sql_gen
32
import sql_io
33
import streams
34
import strings
35
import term
36
import util
37
import xpath
38
import xml_dom
39
import xml_func
40
import xml_parse
41

    
42
metadata_prefix = ':'
43
collision_suffix = '/_alt/'
44

    
45
def get_with_prefix(map_, prefixes, key):
46
    '''Gets all entries for the given key with any of the given prefixes
47
    @return tuple(found_key, found_value)
48
    '''
49
    values = []
50
    for key_ in strings.with_prefixes(['']+prefixes, key): # also with no prefix
51
        try: value = map_[key_]
52
        except KeyError, e: continue # keep going
53
        values.append((key_, value))
54
    
55
    if values != []: return values
56
    else: raise e # re-raise last KeyError
57

    
58
def is_metadata(str_): return str_.startswith(metadata_prefix)
59

    
60
def metadata_value(name):
61
    removed_ref = [False]
62
    name = strings.remove_prefix(metadata_prefix, name, removed_ref)
63
    if removed_ref[0]: return name
64
    else: return None
65

    
66
def cleanup(val):
67
    if val == None: return val
68
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
69

    
70
def main_():
71
    env_names = []
72
    def usage_err():
73
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '
74
            +sys.argv[0]+' [map_path...] [<input] [>output]\n'
75
            'Note: Row #s start with 1')
76
    
77
    ## Get config from env vars
78
    
79
    # Modes
80
    test = opts.env_flag('test', False, env_names)
81
    commit = opts.env_flag('commit', False, env_names) and not test
82
        # never commit in test mode
83
    redo = opts.env_flag('redo', test, env_names) and not commit
84
        # never redo in commit mode (run `make schemas/reinstall` instead)
85
    
86
    # Ranges
87
    start = util.cast(int, opts.get_env_var('start', 1, env_names)) # 1-based
88
    # Make start interally 0-based.
89
    # It's 1-based to the user to match up with the staging table row #s.
90
    start -= 1
91
    if test: n_default = 1
92
    else: n_default = None
93
    n = util.cast(int, util.none_if(opts.get_env_var('n', n_default, env_names),
94
        u''))
95
    end = n
96
    if end != None: end += start
97
    
98
    # Debugging
99
    verbosity = util.cast(float, opts.get_env_var('verbosity', None, env_names))
100
    opts.get_env_var('profile_to', None, env_names) # add to env_names
101
    
102
    # DB
103
    def get_db_config(prefix):
104
        return opts.get_env_vars(sql.db_config_names, prefix, env_names)
105
    in_db_config = get_db_config('in')
106
    out_db_config = get_db_config('out')
107
    in_is_db = 'engine' in in_db_config
108
    out_is_db = 'engine' in out_db_config
109
    in_schema = opts.get_env_var('in_schema', None, env_names)
110
    in_table = opts.get_env_var('in_table', None, env_names)
111
    if in_schema != None:
112
        for config in [in_db_config, out_db_config]:
113
            config['schemas'] += ','+in_schema
114
    
115
    # Optimization
116
    cache_sql = opts.env_flag('cache_sql', True, env_names)
117
    by_col = in_db_config == out_db_config and opts.env_flag('by_col', False,
118
        env_names) # by-column optimization only applies if mapping to same DB
119
    if test: cpus_default = 0
120
    else: cpus_default = 0 # or None to use parallel processing by default
121
    cpus = util.cast(int, util.none_if(opts.get_env_var('cpus', cpus_default,
122
        env_names), u''))
123
    
124
    # Set default verbosity. Must happen after by_col is set.
125
    if verbosity == None:
126
        if test: verbosity = 0.5 # automated tests should not be verbose
127
        elif by_col: verbosity = 3 # show all queries to assist profiling
128
        else: verbosity = 1.1 # just show row progress
129
    
130
    ##
131
    
132
    # Logging
133
    verbose_errors = test and verbosity > 0
134
    debug = verbosity >= 1.5
135
    def log(msg, level=1):
136
        '''Higher level -> more verbose'''
137
        if level <= verbosity:
138
            if verbosity <= 2:
139
                if level == 1.5: msg = '# '+msg # msg is Redmine list item
140
                elif msg.startswith('DB query:'): # remove extra debug info
141
                    first_line, nl, msg = msg.partition('\n')
142
            elif level > 1: msg = '['+str(level)+'] '+msg # include level in msg
143
            
144
            sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))
145
    if debug: log_debug = lambda msg, level=2: log(msg, level)
146
    else: log_debug = sql.log_debug_none
147
    
148
    # Parse args
149
    map_paths = sys.argv[1:]
150
    if map_paths == []:
151
        if in_is_db or not out_is_db: usage_err()
152
        else: map_paths = [None]
153
    
154
    def connect_db(db_config):
155
        log('Connecting to '+sql.db_config_str(db_config))
156
        return sql.connect(db_config, caching=cache_sql, autocommit=commit,
157
            debug_temp=verbosity > 3 and commit, log_debug=log_debug)
158
    
159
    if end != None: end_str = str(end-1) # end is one past the last #
160
    else: end_str = 'end'
161
    log('Processing input rows '+str(start)+'-'+end_str)
162
    
163
    ex_tracker = exc.ExPercentTracker(iter_text='row')
164
    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
165
    
166
    # Parallel processing
167
    pool = parallelproc.MultiProducerPool(cpus)
168
    log('Using '+str(pool.process_ct)+' parallel CPUs')
169
    
170
    # Set up DB access
171
    row_ins_ct_ref = [0]
172
    if out_is_db:
173
        out_db = connect_db(out_db_config)
174
        def is_rel_func(name):
175
            return (name in db_xml.put_special_funcs
176
                or sql.function_exists(out_db, sql_gen.Function(name)))
177
    
178
    doc = xml_dom.create_doc()
179
    root = doc.documentElement
180
    out_is_xml_ref = [False]
181
    
182
    in_label_ref = [None]
183
    col_defaults = {}
184
    def update_in_label():
185
        if in_label_ref[0] != None and out_is_db:
186
            # TODO: Move this to the mappings as some kind of metadata
187
            col_defaults['source_id'] = db_xml.put(out_db, xpath.path2xml(
188
                'source/shortname="'+in_label_ref[0]+'"'), row_ins_ct_ref)
189
    
190
    def prep_root():
191
        root.clear()
192
        update_in_label()
193
    prep_root()
194
    
195
    def process_input(root, row_ready, map_path):
196
        '''Inputs datasource to XML tree, mapping if needed'''
197
        # Load map header
198
        in_is_xpaths = True
199
        out_is_xpaths = True
200
        out_label = None
201
        if map_path != None:
202
            metadata = []
203
            mappings = []
204
            stream = open(map_path, 'rb')
205
            reader = csv.reader(stream)
206
            in_label, out_label = reader.next()[:2]
207
            
208
            def split_col_name(name):
209
                label, sep, root = name.partition(':')
210
                return label, sep != '', root, []
211
            
212
            in_label, in_root, prefixes = maps.col_info(in_label)
213
            in_is_xpaths = in_root != None
214
            in_label_ref[0] = in_label
215
            update_in_label()
216
            out_label, out_root = maps.col_info(out_label)[:2]
217
            out_is_xpaths = out_root != None
218
            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0
219
                # outer elements are types
220
            
221
            for row in reader:
222
                in_, out = row[:2]
223
                if out != '': mappings.append([in_, out_root+out])
224
            
225
            stream.close()
226
            
227
            root.ownerDocument.documentElement.tagName = out_label
228
        in_is_xml = in_is_xpaths and not in_is_db
229
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
230
        
231
        def process_rows(process_row, rows, rows_start=0):
232
            '''Processes input rows
233
            @param process_row(in_row, i)
234
            @rows_start The (0-based) row # of the first row in rows. Set this
235
                only if the pre-start rows have already been skipped.
236
            '''
237
            rows = iter(rows)
238
            
239
            if end != None: row_nums = xrange(rows_start, end)
240
            else: row_nums = itertools.count(rows_start)
241
            i = -1
242
            for i in row_nums:
243
                try: row = rows.next()
244
                except StopIteration:
245
                    i -= 1 # last row # didn't count
246
                    break # no more rows
247
                if i < start: continue # not at start row yet
248
                
249
                # Row # is interally 0-based, but 1-based to the user
250
                log('Processing input row #'+str(i+1), level=1.1)
251
                process_row(row, i)
252
                row_ready(i, row)
253
            row_ct = i-start+1
254
            return row_ct
255
        
256
        def map_rows(get_value, rows, **kw_args):
257
            '''Maps input rows
258
            @param get_value(in_, row):str
259
            '''
260
            # Prevent collisions if multiple inputs mapping to same output
261
            outputs_idxs = dict()
262
            for i, mapping in enumerate(mappings):
263
                in_, out = mapping
264
                default = util.NamedTuple(count=1, first=i)
265
                idxs = outputs_idxs.setdefault(out, default)
266
                if idxs is not default: # key existed, so there was a collision
267
                    if idxs.count == 1: # first key does not yet have suffix
268
                        mappings[idxs.first][1] += collision_suffix+'0'
269
                    mappings[i][1] += collision_suffix+str(idxs.count)
270
                    idxs.count += 1
271
            
272
            id_node = None
273
            if out_is_db:
274
                mappings_orig = mappings[:] # save a copy
275
                mappings[:] = [] # empty existing elements
276
                for in_, out in mappings_orig:
277
                    in_str = strings.ustr(in_)
278
                    is_metadata_ = is_metadata(in_str)
279
                    if is_metadata_: value = metadata_value(in_str)
280
                    else: value = '$'+in_str # mark as name
281
                    
282
                    # All put_obj()s should return the same id_node
283
                    nodes, id_node = xpath.put_obj(root, out, '-1', has_types,
284
                        value) # value is placeholder that documents name
285
                    if not is_metadata_: mappings.append([in_, nodes])
286
                if id_node == None:
287
                    warnings.warn(UserWarning('Map warning: No mappings'))
288
                xml_func.simplify(root)
289
                sys.stdout.write(strings.to_raw_str('Put template:\n'
290
                    +strings.ustr(root)))
291
                sys.stdout.flush()
292
            
293
            def process_row(row, i):
294
                row_id = str(i)
295
                if id_node != None: xml_dom.set_value(id_node, row_id)
296
                for in_, out in mappings:
297
                    log_debug('Getting '+strings.ustr(in_))
298
                    value = cleanup(get_value(in_, row))
299
                    log_debug('Putting '+strings.urepr(value)+' to '
300
                        +strings.ustr(out))
301
                    if out_is_db: # out is list of XML nodes
302
                        for node in out: xml_dom.set_value(node, value)
303
                    elif value != None: # out is XPath
304
                        xpath.put_obj(root, out, row_id, has_types, value)
305
            return process_rows(process_row, rows, **kw_args)
306
        
307
        def map_table(col_names, rows, **kw_args):
308
            col_names_ct = len(col_names)
309
            col_idxs = util.list_flip(col_names)
310
            col_names_simp = map(maps.simplify, col_names)
311
            col_names_map = dict(zip(col_names_simp, col_names))
312
            prefixes_simp = map(maps.simplify, prefixes)
313
            
314
            # Resolve prefixes
315
            mappings_orig = mappings[:] # save a copy
316
            mappings[:] = [] # empty existing elements
317
            for in_, out in mappings_orig:
318
                if is_metadata(in_): mappings.append((in_, out))
319
                else:
320
                    try:
321
                        cols = get_with_prefix(col_names_map, prefixes_simp,
322
                            maps.simplify(in_))
323
                    except KeyError: pass
324
                    else:
325
                        cols = [(orig, col_idxs[orig]) for simp, orig in cols]
326
                        mappings[len(mappings):] = [[db_xml.ColRef(*col), out]
327
                            for col in cols] # can't use += because that uses =
328
            
329
            def get_value(in_, row): return row.list[in_.idx]
330
            def wrap_row(row):
331
                return util.ListDict(util.list_as_length(row, col_names_ct),
332
                    col_names, col_idxs) # handle CSV rows of different lengths
333
            
334
            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)
335
        
336
        if in_is_db:
337
            def on_error(e): ex_tracker.track(e)
338
            
339
            if by_col: in_db = out_db
340
            else: in_db = connect_db(in_db_config)
341
            
342
            # Get table and schema name
343
            schema = in_schema # modified, so can't have same name as outer var
344
            table = in_table # modified, so can't have same name as outer var
345
            if table == None:
346
                assert in_is_xpaths
347
                schema, sep, table = in_root.partition('.')
348
                if sep == '': # only the table name was specified
349
                    table = schema
350
                    schema = None
351
            table = sql_gen.Table(table, schema)
352
            
353
            # Fetch rows
354
            if by_col: limit = 0 # only fetch column names
355
            else: limit = n
356
            cur = sql.select(in_db, table, limit=limit, start=start,
357
                cacheable=False)
358
            col_names = list(sql.col_names(cur))
359
            
360
            if by_col:
361
                map_table(col_names, []) # just create the template
362
                
363
                if start == 0 and n == None: # doing full re-import
364
                    log('Clearing errors table')
365
                    errors_table_ = sql_io.errors_table(in_db, table)
366
                    if errors_table_ != None:
367
                        sql.drop_table(in_db, errors_table_)
368
                
369
                # Strip XML functions not in the DB
370
                xml_func.process(root, is_rel_func=is_rel_func)
371
                if debug: log_debug('Putting stripped:\n'+strings.ustr(root))
372
                    # only calc if debug
373
                
374
                # Import rows
375
                in_row_ct_ref = [0]
376
                db_xml.put_table(in_db, root.firstChild, table, in_row_ct_ref,
377
                    row_ins_ct_ref, n, start, on_error, col_defaults)
378
                row_ct = in_row_ct_ref[0]
379
            else:
380
                # Use normal by-row method
381
                row_ct = map_table(col_names, sql.rows(cur), rows_start=start)
382
                    # rows_start: pre-start rows have been skipped
383
                
384
                in_db.db.close()
385
        elif in_is_xml:
386
            stdin = streams.LineCountStream(sys.stdin)
387
            def on_error(e):
388
                exc.add_msg(e, term.emph('input line #:')+' '
389
                    +str(stdin.line_num))
390
                ex_tracker.track(e)
391
            
392
            def get_rows(doc2rows):
393
                return iters.flatten(itertools.imap(doc2rows,
394
                    xml_parse.docs_iter(stdin, on_error)))
395
            
396
            if map_path == None:
397
                def doc2rows(in_xml_root):
398
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
399
                    util.skip(iter_, xml_dom.is_text) # skip metadata
400
                    return iter_
401
                
402
                row_ct = process_rows(lambda row, i: root.appendChild(row),
403
                    get_rows(doc2rows))
404
            else:
405
                def doc2rows(in_xml_root):
406
                    rows = xpath.get(in_xml_root, in_root, limit=end)
407
                    if rows == []: warnings.warn(UserWarning('Map warning: '
408
                        'Root "'+in_root+'" not found in input'))
409
                    return rows
410
                
411
                def get_value(in_, row):
412
                    in_ = './{'+(','.join(strings.with_prefixes(
413
                        ['']+prefixes, in_)))+'}' # also with no prefix
414
                    nodes = xpath.get(row, in_, allow_rooted=False)
415
                    if nodes != []: return xml_dom.value(nodes[0])
416
                    else: return None
417
                
418
                row_ct = map_rows(get_value, get_rows(doc2rows))
419
        else: # input is CSV
420
            reader, col_names = csvs.reader_and_header(sys.stdin)
421
            row_ct = map_table(col_names, reader)
422
        
423
        return row_ct
424
    
425
    def process_inputs(root, row_ready):
426
        row_ct = 0
427
        for map_path in map_paths:
428
            row_ct += process_input(root, row_ready, map_path)
429
        return row_ct
430
    
431
    pool.share_vars(locals())
432
    if out_is_db:
433
        try:
434
            if redo: sql.empty_db(out_db)
435
            pool.share_vars(locals())
436
            
437
            def row_ready(row_num, input_row):
438
                row_str_ = [None]
439
                def row_str():
440
                    if row_str_[0] == None:
441
                        # Row # is interally 0-based, but 1-based to the user
442
                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)
443
                            +'\n'+term.emph('input row:')+'\n'
444
                            +strings.ustr(input_row))
445
                        if verbose_errors: row_str_[0] += ('\n'
446
                            +term.emph('output row:')+'\n'+strings.ustr(root))
447
                    return row_str_[0]
448
                
449
                if debug: log_debug(row_str()) # only calc if debug
450
                
451
                def on_error(e):
452
                    exc.add_msg(e, row_str())
453
                    ex_tracker.track(e, row_num, detail=verbose_errors)
454
                pool.share_vars(locals())
455
                
456
                row_root = root.cloneNode(True) # deep copy so don't modify root
457
                xml_func.process(row_root, on_error, is_rel_func, out_db)
458
                if debug: log_debug('Putting processed:\n'
459
                    +strings.ustr(row_root)) # only calc if debug
460
                if not xml_dom.is_empty(row_root):
461
                    assert xml_dom.has_one_child(row_root)
462
                    try:
463
                        sql.with_savepoint(out_db,
464
                            lambda: db_xml.put(out_db, row_root.firstChild,
465
                                row_ins_ct_ref, on_error, col_defaults))
466
                    except sql.DatabaseErrors, e: on_error(e)
467
            
468
            row_ct = process_inputs(root, row_ready)
469
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
470
                ' new rows into database\n')
471
            sys.stdout.flush()
472
            
473
            # Consume asynchronous tasks
474
            pool.main_loop()
475
        finally: out_db.close()
476
    else:
477
        def on_error(e): ex_tracker.track(e)
478
        def row_ready(row_num, input_row): pass
479
        row_ct = process_inputs(root, row_ready)
480
        xml_func.process(root, on_error)
481
        if out_is_xml_ref[0]:
482
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
483
        else: # output is CSV
484
            raise NotImplementedError('CSV output not supported yet')
485
    
486
    # Consume any asynchronous tasks not already consumed above
487
    pool.main_loop()
488
    
489
    profiler.stop(row_ct)
490
    if not by_col: ex_tracker.add_iters(row_ct) # only if errors are done by row
491
    log('Processed '+str(row_ct)+' input rows')
492
    log(profiler.msg())
493
    log(ex_tracker.msg())
494
    ex_tracker.exit()
495

    
496
def main():
497
    try: main_()
498
    except Parser.SyntaxError, e: raise SystemExit(strings.ustr(e))
499

    
500
if __name__ == '__main__':
501
    profile_to = opts.get_env_var('profile_to', None)
502
    if profile_to != None:
503
        import cProfile
504
        sys.stderr.write('Profiling to '+profile_to+'\n')
505
        cProfile.run(main.func_code, profile_to)
506
    else: main()
(34-34/66)