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', False, 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:
186
            os.environ['source'] = in_label_ref[0]
187
    
188
    def prep_root():
189
        root.clear()
190
        update_in_label()
191
    prep_root()
192
    
193
    def process_input(root, row_ready, map_path):
194
        '''Inputs datasource to XML tree, mapping if needed'''
195
        # Load map header
196
        in_is_xpaths = True
197
        out_is_xpaths = True
198
        out_label = None
199
        if map_path != None:
200
            metadata = []
201
            mappings = []
202
            stream = open(map_path, 'rb')
203
            reader = csv.reader(stream)
204
            in_label, out_label = reader.next()[:2]
205
            
206
            def split_col_name(name):
207
                label, sep, root = name.partition(':')
208
                return label, sep != '', root, []
209
            
210
            in_label, in_root, prefixes = maps.col_info(in_label)
211
            in_is_xpaths = in_root != None
212
            in_label_ref[0] = in_label
213
            update_in_label()
214
            out_label, out_root = maps.col_info(out_label)[:2]
215
            out_is_xpaths = out_root != None
216
            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0
217
                # outer elements are types
218
            
219
            for row in reader:
220
                in_, out = row[:2]
221
                if out != '': mappings.append([in_, out_root+out])
222
            
223
            stream.close()
224
            
225
            root.ownerDocument.documentElement.tagName = out_label
226
        in_is_xml = in_is_xpaths and not in_is_db
227
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
228
        
229
        def process_rows(process_row, rows, rows_start=0):
230
            '''Processes input rows
231
            @param process_row(in_row, i)
232
            @rows_start The (0-based) row # of the first row in rows. Set this
233
                only if the pre-start rows have already been skipped.
234
            '''
235
            rows = iter(rows)
236
            
237
            if end != None: row_nums = xrange(rows_start, end)
238
            else: row_nums = itertools.count(rows_start)
239
            i = -1
240
            for i in row_nums:
241
                try: row = rows.next()
242
                except StopIteration:
243
                    i -= 1 # last row # didn't count
244
                    break # no more rows
245
                if i < start: continue # not at start row yet
246
                
247
                # Row # is interally 0-based, but 1-based to the user
248
                log('Processing input row #'+str(i+1), level=1.1)
249
                process_row(row, i)
250
                row_ready(i, row)
251
            row_ct = i-start+1
252
            return row_ct
253
        
254
        def map_rows(get_value, rows, **kw_args):
255
            '''Maps input rows
256
            @param get_value(in_, row):str
257
            '''
258
            # Prevent collisions if multiple inputs mapping to same output
259
            outputs_idxs = dict()
260
            for i, mapping in enumerate(mappings):
261
                in_, out = mapping
262
                default = util.NamedTuple(count=1, first=i)
263
                idxs = outputs_idxs.setdefault(out, default)
264
                if idxs is not default: # key existed, so there was a collision
265
                    if idxs.count == 1: # first key does not yet have suffix
266
                        mappings[idxs.first][1] += collision_suffix+'0'
267
                    mappings[i][1] += collision_suffix+str(idxs.count)
268
                    idxs.count += 1
269
            
270
            id_node = None
271
            if out_is_db:
272
                mappings_orig = mappings[:] # save a copy
273
                mappings[:] = [] # empty existing elements
274
                for in_, out in mappings_orig:
275
                    in_str = strings.ustr(in_)
276
                    is_metadata_ = is_metadata(in_str)
277
                    if is_metadata_: value = metadata_value(in_str)
278
                    else: value = '$'+in_str # mark as name
279
                    
280
                    # All put_obj()s should return the same id_node
281
                    nodes, id_node = xpath.put_obj(root, out, '-1', has_types,
282
                        value) # value is placeholder that documents name
283
                    if not is_metadata_: mappings.append([in_, nodes])
284
                if id_node == None:
285
                    warnings.warn(UserWarning('Map warning: No mappings'))
286
                xml_func.simplify(root)
287
                sys.stdout.write(strings.to_raw_str('Put template:\n'
288
                    +strings.ustr(root)))
289
                sys.stdout.flush()
290
            
291
            def process_row(row, i):
292
                row_id = str(i)
293
                if id_node != None: xml_dom.set_value(id_node, row_id)
294
                for in_, out in mappings:
295
                    log_debug('Getting '+strings.ustr(in_))
296
                    value = cleanup(get_value(in_, row))
297
                    log_debug('Putting '+strings.urepr(value)+' to '
298
                        +strings.ustr(out))
299
                    if out_is_db: # out is list of XML nodes
300
                        for node in out: xml_dom.set_value(node, value)
301
                    elif value != None: # out is XPath
302
                        xpath.put_obj(root, out, row_id, has_types, value)
303
            return process_rows(process_row, rows, **kw_args)
304
        
305
        def map_table(col_names, rows, **kw_args):
306
            col_names_ct = len(col_names)
307
            col_idxs = util.list_flip(col_names)
308
            col_names_map = dict(zip(col_names, col_names))
309
            prefixes_simp = map(maps.simplify, prefixes)
310
            
311
            # Resolve prefixes
312
            mappings_orig = mappings[:] # save a copy
313
            mappings[:] = [] # empty existing elements
314
            for in_, out in mappings_orig:
315
                if is_metadata(in_): mappings.append([in_, out])
316
                else:
317
                    try:
318
                        cols = get_with_prefix(col_names_map, prefixes_simp,
319
                            in_)
320
                    except KeyError: pass
321
                    else:
322
                        cols = [(orig, col_idxs[orig]) for simp, orig in cols]
323
                        mappings[len(mappings):] = [[db_xml.ColRef(*col), out]
324
                            for col in cols] # can't use += because that uses =
325
            
326
            def get_value(in_, row): return row.list[in_.idx]
327
            def wrap_row(row):
328
                return util.ListDict(util.list_as_length(row, col_names_ct),
329
                    col_names, col_idxs) # handle CSV rows of different lengths
330
            
331
            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)
332
        
333
        if in_is_db:
334
            def on_error(e): ex_tracker.track(e)
335
            
336
            if by_col: in_db = out_db
337
            else: in_db = connect_db(in_db_config)
338
            
339
            # Get table and schema name
340
            schema = in_schema # modified, so can't have same name as outer var
341
            table = in_table # modified, so can't have same name as outer var
342
            if table == None:
343
                assert in_is_xpaths
344
                schema, sep, table = in_root.partition('.')
345
                if sep == '': # only the table name was specified
346
                    table = schema
347
                    schema = None
348
            table = sql_gen.Table(table, schema)
349
            
350
            # Fetch rows
351
            if by_col: limit = 0 # only fetch column names
352
            else: limit = n
353
            try:
354
                cur = sql.select(in_db, table, limit=limit, start=start,
355
                    recover=True, cacheable=False)
356
            except sql.DoesNotExistException:
357
                table = None
358
                col_names = []
359
                rows = []
360
            else:
361
                col_names = list(sql.col_names(cur))
362
                rows = sql.rows(cur)
363
            
364
            if by_col:
365
                map_table(col_names, []) # just create the template
366
                
367
                if table != None and start == 0 and n == None: # full re-import
368
                    log('Clearing errors table')
369
                    errors_table_ = sql_io.errors_table(in_db, table)
370
                    if errors_table_ != None:
371
                        sql.drop_table(in_db, errors_table_)
372
                
373
                # Strip XML functions not in the DB
374
                xml_func.process(root, is_rel_func=is_rel_func)
375
                if debug: log_debug('Putting stripped:\n'+strings.ustr(root))
376
                    # only calc if debug
377
                
378
                # Import rows
379
                in_row_ct_ref = [0]
380
                db_xml.put_table(in_db, root.firstChild, table, in_row_ct_ref,
381
                    row_ins_ct_ref, n, start, on_error, col_defaults)
382
                row_ct = in_row_ct_ref[0]
383
            else:
384
                # Use normal by-row method
385
                row_ct = map_table(col_names, rows, rows_start=start)
386
                    # rows_start: pre-start rows have been skipped
387
                
388
                in_db.db.close()
389
        elif in_is_xml:
390
            stdin = streams.LineCountStream(sys.stdin)
391
            def on_error(e):
392
                exc.add_msg(e, term.emph('input line #:')+' '
393
                    +str(stdin.line_num))
394
                ex_tracker.track(e)
395
            
396
            def get_rows(doc2rows):
397
                return iters.flatten(itertools.imap(doc2rows,
398
                    xml_parse.docs_iter(stdin, on_error)))
399
            
400
            if map_path == None:
401
                def doc2rows(in_xml_root):
402
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
403
                    util.skip(iter_, xml_dom.is_text) # skip metadata
404
                    return iter_
405
                
406
                row_ct = process_rows(lambda row, i: root.appendChild(row),
407
                    get_rows(doc2rows))
408
            else:
409
                def doc2rows(in_xml_root):
410
                    rows = xpath.get(in_xml_root, in_root, limit=end)
411
                    if rows == []: warnings.warn(UserWarning('Map warning: '
412
                        'Root "'+in_root+'" not found in input'))
413
                    return rows
414
                
415
                def get_value(in_, row):
416
                    in_ = './{'+(','.join(strings.with_prefixes(
417
                        ['']+prefixes, in_)))+'}' # also with no prefix
418
                    nodes = xpath.get(row, in_, allow_rooted=False)
419
                    if nodes != []: return xml_dom.value(nodes[0])
420
                    else: return None
421
                
422
                row_ct = map_rows(get_value, get_rows(doc2rows))
423
        else: # input is CSV
424
            reader, col_names = csvs.reader_and_header(sys.stdin)
425
            row_ct = map_table(col_names, reader)
426
        
427
        return row_ct
428
    
429
    def process_inputs(root, row_ready):
430
        row_ct = 0
431
        for map_path in map_paths:
432
            row_ct += process_input(root, row_ready, map_path)
433
        return row_ct
434
    
435
    pool.share_vars(locals())
436
    if out_is_db:
437
        try:
438
            if redo: sql.empty_db(out_db)
439
            pool.share_vars(locals())
440
            
441
            def row_ready(row_num, input_row):
442
                row_str_ = [None]
443
                def row_str():
444
                    if row_str_[0] == None:
445
                        # Row # is interally 0-based, but 1-based to the user
446
                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)
447
                            +'\n'+term.emph('input row:')+'\n'
448
                            +strings.ustr(input_row))
449
                        if verbose_errors: row_str_[0] += ('\n'
450
                            +term.emph('output row:')+'\n'+strings.ustr(root))
451
                    return row_str_[0]
452
                
453
                if debug: log_debug(row_str()) # only calc if debug
454
                
455
                def on_error(e):
456
                    exc.add_msg(e, row_str())
457
                    ex_tracker.track(e, row_num, detail=verbose_errors)
458
                pool.share_vars(locals())
459
                
460
                row_root = root.cloneNode(True) # deep copy so don't modify root
461
                xml_func.process(row_root, on_error, is_rel_func, out_db)
462
                if debug: log_debug('Putting processed:\n'
463
                    +strings.ustr(row_root)) # only calc if debug
464
                if not xml_dom.is_empty(row_root):
465
                    assert xml_dom.has_one_child(row_root)
466
                    try:
467
                        sql.with_savepoint(out_db,
468
                            lambda: db_xml.put(out_db, row_root.firstChild,
469
                                row_ins_ct_ref, on_error, col_defaults))
470
                    except sql.DatabaseErrors, e: on_error(e)
471
            
472
            row_ct = process_inputs(root, row_ready)
473
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
474
                ' new rows into database\n')
475
            sys.stdout.flush()
476
            
477
            # Consume asynchronous tasks
478
            pool.main_loop()
479
        finally: out_db.close()
480
    else:
481
        def on_error(e): ex_tracker.track(e)
482
        def row_ready(row_num, input_row): pass
483
        row_ct = process_inputs(root, row_ready)
484
        xml_func.process(root, on_error)
485
        if out_is_xml_ref[0]:
486
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
487
        else: # output is CSV
488
            raise NotImplementedError('CSV output not supported yet')
489
    
490
    # Consume any asynchronous tasks not already consumed above
491
    pool.main_loop()
492
    
493
    profiler.stop(row_ct)
494
    if not by_col: ex_tracker.add_iters(row_ct) # only if errors are done by row
495
    log('Processed '+str(row_ct)+' input rows')
496
    log(profiler.msg())
497
    log(ex_tracker.msg())
498
    ex_tracker.exit()
499

    
500
def main():
501
    try: main_()
502
    except Parser.SyntaxError, e: raise SystemExit(strings.ustr(e))
503

    
504
if __name__ == '__main__':
505
    profile_to = opts.get_env_var('profile_to', None)
506
    if profile_to != None:
507
        import cProfile
508
        sys.stderr.write('Profiling to '+profile_to+'\n')
509
        cProfile.run(main.func_code, profile_to)
510
    else: main()
(38-38/72)