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 or no '
286
                        'column name matches. Are you importing the correct '
287
                        'input table?'))
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_map = dict(zip(col_names, col_names))
311
            prefixes_simp = map(maps.simplify, prefixes)
312
            
313
            # Resolve prefixes
314
            mappings_orig = mappings[:] # save a copy
315
            mappings[:] = [] # empty existing elements
316
            for in_, out in mappings_orig:
317
                if is_metadata(in_): mappings.append([in_, out])
318
                else:
319
                    try:
320
                        cols = get_with_prefix(col_names_map, prefixes_simp,
321
                            in_)
322
                    except KeyError: pass
323
                    else:
324
                        cols = [(orig, col_idxs[orig]) for simp, orig in cols]
325
                        mappings[len(mappings):] = [[db_xml.ColRef(*col), out]
326
                            for col in cols] # can't use += because that uses =
327
            
328
            def get_value(in_, row): return row.list[in_.idx]
329
            def wrap_row(row):
330
                return util.ListDict(util.list_as_length(row, col_names_ct),
331
                    col_names, col_idxs) # handle CSV rows of different lengths
332
            
333
            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)
334
        
335
        if in_is_db:
336
            def on_error(e): ex_tracker.track(e)
337
            
338
            if by_col: in_db = out_db
339
            else: in_db = connect_db(in_db_config)
340
            
341
            # Get table and schema name
342
            schema = in_schema # modified, so can't have same name as outer var
343
            table = in_table # modified, so can't have same name as outer var
344
            if table == None:
345
                assert in_is_xpaths
346
                schema, sep, table = in_root.partition('.')
347
                if sep == '': # only the table name was specified
348
                    table = schema
349
                    schema = None
350
            table = sql_gen.Table(table, schema)
351
            
352
            # Fetch rows
353
            if by_col: limit = 0 # only fetch column names
354
            else: limit = n
355
            try:
356
                cur = sql.select(in_db, table, limit=limit, start=start,
357
                    recover=True, cacheable=False)
358
            except sql.DoesNotExistException:
359
                table = None
360
                col_names = []
361
                rows = []
362
            else:
363
                col_names = list(sql.col_names(cur))
364
                rows = sql.rows(cur)
365
            
366
            if by_col:
367
                map_table(col_names, []) # just create the template
368
                
369
                if table != None and start == 0 and n == None: # full re-import
370
                    log('Clearing errors table')
371
                    errors_table_ = sql_io.errors_table(in_db, table)
372
                    if errors_table_ != None:
373
                        sql.drop_table(in_db, errors_table_)
374
                
375
                # Strip XML functions not in the DB
376
                xml_func.process(root, is_rel_func=is_rel_func)
377
                if debug: log_debug('Putting stripped:\n'+strings.ustr(root))
378
                    # only calc if debug
379
                
380
                # Import rows
381
                in_row_ct_ref = [0]
382
                db_xml.put_table(in_db, root.firstChild, table, in_row_ct_ref,
383
                    row_ins_ct_ref, n, start, on_error, col_defaults)
384
                row_ct = in_row_ct_ref[0]
385
            else:
386
                # Use normal by-row method
387
                row_ct = map_table(col_names, rows, rows_start=start)
388
                    # rows_start: pre-start rows have been skipped
389
                
390
                in_db.db.close()
391
        elif in_is_xml:
392
            stdin = streams.LineCountStream(sys.stdin)
393
            def on_error(e):
394
                exc.add_msg(e, term.emph('input line #:')+' '
395
                    +str(stdin.line_num))
396
                ex_tracker.track(e)
397
            
398
            def get_rows(doc2rows):
399
                return iters.flatten(itertools.imap(doc2rows,
400
                    xml_parse.docs_iter(stdin, on_error)))
401
            
402
            if map_path == None:
403
                def doc2rows(in_xml_root):
404
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
405
                    util.skip(iter_, xml_dom.is_text) # skip metadata
406
                    return iter_
407
                
408
                row_ct = process_rows(lambda row, i: root.appendChild(row),
409
                    get_rows(doc2rows))
410
            else:
411
                def doc2rows(in_xml_root):
412
                    rows = xpath.get(in_xml_root, in_root, limit=end)
413
                    if rows == []: warnings.warn(UserWarning('Map warning: '
414
                        'Root "'+in_root+'" not found in input'))
415
                    return rows
416
                
417
                def get_value(in_, row):
418
                    in_ = './{'+(','.join(strings.with_prefixes(
419
                        ['']+prefixes, in_)))+'}' # also with no prefix
420
                    nodes = xpath.get(row, in_, allow_rooted=False)
421
                    if nodes != []: return xml_dom.value(nodes[0])
422
                    else: return None
423
                
424
                row_ct = map_rows(get_value, get_rows(doc2rows))
425
        else: # input is CSV
426
            reader, col_names = csvs.reader_and_header(sys.stdin)
427
            row_ct = map_table(col_names, reader)
428
        
429
        return row_ct
430
    
431
    def process_inputs(root, row_ready):
432
        row_ct = 0
433
        for map_path in map_paths:
434
            row_ct += process_input(root, row_ready, map_path)
435
        return row_ct
436
    
437
    pool.share_vars(locals())
438
    if out_is_db:
439
        try:
440
            if redo: sql.empty_db(out_db)
441
            pool.share_vars(locals())
442
            
443
            def row_ready(row_num, input_row):
444
                row_str_ = [None]
445
                def row_str():
446
                    if row_str_[0] == None:
447
                        # Row # is interally 0-based, but 1-based to the user
448
                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)
449
                            +'\n'+term.emph('input row:')+'\n'
450
                            +strings.ustr(input_row))
451
                        if verbose_errors: row_str_[0] += ('\n'
452
                            +term.emph('output row:')+'\n'+strings.ustr(root))
453
                    return row_str_[0]
454
                
455
                if debug: log_debug(row_str()) # only calc if debug
456
                
457
                def on_error(e):
458
                    exc.add_msg(e, row_str())
459
                    ex_tracker.track(e, row_num, detail=verbose_errors)
460
                pool.share_vars(locals())
461
                
462
                row_root = root.cloneNode(True) # deep copy so don't modify root
463
                xml_func.process(row_root, on_error, is_rel_func, out_db)
464
                if debug: log_debug('Putting processed:\n'
465
                    +strings.ustr(row_root)) # only calc if debug
466
                if not xml_dom.is_empty(row_root):
467
                    assert xml_dom.has_one_child(row_root)
468
                    try:
469
                        sql.with_savepoint(out_db,
470
                            lambda: db_xml.put(out_db, row_root.firstChild,
471
                                row_ins_ct_ref, on_error, col_defaults))
472
                    except sql.DatabaseErrors, e: on_error(e)
473
            
474
            row_ct = process_inputs(root, row_ready)
475
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
476
                ' new rows into database\n')
477
            sys.stdout.flush()
478
            
479
            # Consume asynchronous tasks
480
            pool.main_loop()
481
        finally: out_db.close()
482
    else:
483
        def on_error(e): ex_tracker.track(e)
484
        def row_ready(row_num, input_row): pass
485
        row_ct = process_inputs(root, row_ready)
486
        xml_func.process(root, on_error)
487
        if out_is_xml_ref[0]:
488
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
489
        else: # output is CSV
490
            raise NotImplementedError('CSV output not supported yet')
491
    
492
    # Consume any asynchronous tasks not already consumed above
493
    pool.main_loop()
494
    
495
    profiler.stop(row_ct)
496
    if not by_col: ex_tracker.add_iters(row_ct) # only if errors are done by row
497
    log('Processed '+str(row_ct)+' input rows')
498
    log(profiler.msg())
499
    log(ex_tracker.msg())
500
    ex_tracker.exit()
501

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

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