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 copy
8
import csv
9
import itertools
10
import os.path
11
import sys
12
import xml.dom.minidom as minidom
13

    
14
sys.path.append(os.path.dirname(__file__)+"/../lib")
15

    
16
import csvs
17
import db_xml
18
import exc
19
import iters
20
import maps
21
import opts
22
import parallelproc
23
import Parser
24
import profiling
25
import sql
26
import sql_gen
27
import streams
28
import strings
29
import term
30
import util
31
import xpath
32
import xml_dom
33
import xml_func
34
import xml_parse
35

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

    
49
def metadata_value(name): return None # this feature has been removed
50

    
51
def cleanup(val):
52
    if val == None: return val
53
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
54

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

    
458
def main():
459
    try: main_()
460
    except Parser.SyntaxError, e: raise SystemExit(str(e))
461

    
462
if __name__ == '__main__':
463
    profile_to = opts.get_env_var('profile_to', None)
464
    if profile_to != None:
465
        import cProfile
466
        sys.stderr.write('Profiling to '+profile_to+'\n')
467
        cProfile.run(main.func_code, profile_to)
468
    else: main()
(26-26/50)