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

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

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

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

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

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

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

    
428
def main():
429
    try: main_()
430
    except Parser.SyntaxError, e: raise SystemExit(str(e))
431

    
432
if __name__ == '__main__':
433
    profile_to = opts.get_env_var('profile_to', None)
434
    if profile_to != None:
435
        import cProfile
436
        sys.stderr.write('Profiling to '+profile_to+'\n')
437
        cProfile.run(main.func_code, profile_to)
438
    else: main()
(25-25/48)