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

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

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

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

    
41
collision_suffix = '/_alt/'
42

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

    
56
def metadata_value(name): return None # this feature has been removed
57

    
58
def cleanup(val):
59
    if val == None: return val
60
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
61

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

    
474
def main():
475
    try: main_()
476
    except Parser.SyntaxError, e: raise SystemExit(str(e))
477

    
478
if __name__ == '__main__':
479
    profile_to = opts.get_env_var('profile_to', None)
480
    if profile_to != None:
481
        import cProfile
482
        sys.stderr.write('Profiling to '+profile_to+'\n')
483
        cProfile.run(main.func_code, profile_to)
484
    else: main()
(28-28/54)