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 exc
17
import iters
18
import maps
19
import opts
20
import parallel
21
import Parser
22
import profiling
23
import sql
24
import streams
25
import strings
26
import term
27
import util
28
import xpath
29
import xml_dom
30
import xml_func
31
import xml_parse
32

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

    
44
def metadata_value(name): return None # this feature has been removed
45

    
46
def cleanup(val):
47
    if val == None: return val
48
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
49

    
50
def main_():
51
    env_names = []
52
    def usage_err():
53
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '
54
            +sys.argv[0]+' [map_path...] [<input] [>output]')
55
    
56
    ## Get config from env vars
57
    
58
    # Modes
59
    test = opts.env_flag('test', False, env_names)
60
    commit = opts.env_flag('commit', False, env_names) and not test
61
        # never commit in test mode
62
    redo = opts.env_flag('redo', test, env_names) and not commit
63
        # never redo in commit mode (manually run `make empty_db` instead)
64
    
65
    # Ranges
66
    start = util.cast(int, opts.get_env_var('start', '0', env_names))
67
    if test: end_default = 1
68
    else: end_default = None
69
    end = util.cast(int, util.none_if(
70
        opts.get_env_var('n', end_default, env_names), u''))
71
    if end != None: end += start
72
    
73
    # Optimization
74
    if test: cpus_default = 0
75
    else: cpus_default = None
76
    cpus = util.cast(int, util.none_if(opts.get_env_var('cpus', cpus_default,
77
        env_names), u''))
78
    
79
    # Debugging
80
    debug = opts.env_flag('debug', False, env_names)
81
    sql.run_raw_query.debug = debug
82
    verbose = debug or opts.env_flag('verbose', not test, env_names)
83
    opts.get_env_var('profile_to', None, env_names) # add to env_names
84
    
85
    # DB
86
    db_config_names = ['engine', 'host', 'user', 'password', 'database']
87
    def get_db_config(prefix):
88
        return opts.get_env_vars(db_config_names, prefix, env_names)
89
    in_db_config = get_db_config('in')
90
    out_db_config = get_db_config('out')
91
    in_is_db = 'engine' in in_db_config
92
    out_is_db = 'engine' in out_db_config
93
    
94
    ##
95
    
96
    # Logging
97
    def log(msg, on=verbose):
98
        if on: sys.stderr.write(msg)
99
    def log_start(action, on=verbose): log(action+'...\n', on)
100
    
101
    # Parse args
102
    map_paths = sys.argv[1:]
103
    if map_paths == []:
104
        if in_is_db or not out_is_db: usage_err()
105
        else: map_paths = [None]
106
    
107
    def connect_db(db_config):
108
        log_start('Connecting to '+sql.db_config_str(db_config))
109
        return sql.connect(db_config)
110
    
111
    if end != None: end_str = str(end-1)
112
    else: end_str = 'end'
113
    log_start('Processing input rows '+str(start)+'-'+end_str)
114
    
115
    ex_tracker = exc.ExPercentTracker(iter_text='row')
116
    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
117
    
118
    # Parallel processing
119
    pool = parallel.MultiProducerPool(cpus)
120
    log_start('Using '+str(pool.process_ct)+' parallel CPUs')
121
    
122
    doc = xml_dom.create_doc()
123
    root = doc.documentElement
124
    out_is_xml_ref = [False]
125
    in_label_ref = [None]
126
    def update_in_label():
127
        if in_label_ref[0] != None:
128
            xpath.get(root, '/_ignore/inLabel="'+in_label_ref[0]+'"', True)
129
    def prep_root():
130
        root.clear()
131
        update_in_label()
132
    prep_root()
133
    
134
    def process_input(root, row_ready, map_path):
135
        '''Inputs datasource to XML tree, mapping if needed'''
136
        # Load map header
137
        in_is_xpaths = True
138
        out_is_xpaths = True
139
        out_label = None
140
        if map_path != None:
141
            metadata = []
142
            mappings = []
143
            stream = open(map_path, 'rb')
144
            reader = csv.reader(stream)
145
            in_label, out_label = reader.next()[:2]
146
            
147
            def split_col_name(name):
148
                label, sep, root = name.partition(':')
149
                label, sep2, prefixes_str = label.partition('[')
150
                prefixes_str = strings.remove_suffix(']', prefixes_str)
151
                prefixes = strings.split(',', prefixes_str)
152
                return label, sep != '', root, prefixes
153
                    # extract datasrc from "datasrc[data_format]"
154
            
155
            in_label, in_root, prefixes = maps.col_info(in_label)
156
            in_is_xpaths = in_root != None
157
            in_label_ref[0] = in_label
158
            update_in_label()
159
            out_label, out_root = maps.col_info(out_label)[:2]
160
            out_is_xpaths = out_root != None
161
            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0
162
                # outer elements are types
163
            
164
            for row in reader:
165
                in_, out = row[:2]
166
                if out != '':
167
                    if out_is_xpaths: out = xpath.parse(out_root+out)
168
                    mappings.append((in_, out))
169
            
170
            stream.close()
171
            
172
            root.ownerDocument.documentElement.tagName = out_label
173
        in_is_xml = in_is_xpaths and not in_is_db
174
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
175
        
176
        def process_rows(process_row, rows):
177
            '''Processes input rows
178
            @param process_row(in_row, i)
179
            '''
180
            i = 0
181
            while end == None or i < end:
182
                try: row = rows.next()
183
                except StopIteration: break # no more rows
184
                if i < start: continue # not at start row yet
185
                
186
                process_row(row, i)
187
                row_ready(i, row)
188
                i += 1
189
            row_ct = i-start
190
            return row_ct
191
        
192
        def map_rows(get_value, rows):
193
            '''Maps input rows
194
            @param get_value(in_, row):str
195
            '''
196
            def process_row(row, i):
197
                row_id = str(i)
198
                for in_, out in mappings:
199
                    value = metadata_value(in_)
200
                    if value == None:
201
                        log_start('Getting '+str(in_), debug)
202
                        value = cleanup(get_value(in_, row))
203
                    if value != None:
204
                        log_start('Putting '+str(out), debug)
205
                        xpath.put_obj(root, out, row_id, has_types, value)
206
            return process_rows(process_row, rows)
207
        
208
        def map_table(col_names, rows):
209
            col_names_ct = len(col_names)
210
            col_idxs = util.list_flip(col_names)
211
            
212
            i = 0
213
            while i < len(mappings): # mappings len changes in loop
214
                in_, out = mappings[i]
215
                if metadata_value(in_) == None:
216
                    try: mappings[i] = (
217
                        get_with_prefix(col_idxs, prefixes, in_), out)
218
                    except KeyError:
219
                        del mappings[i]
220
                        continue # keep i the same
221
                i += 1
222
            
223
            def get_value(in_, row):
224
                return util.coalesce(*util.list_subset(row.list, in_))
225
            def wrap_row(row):
226
                return util.ListDict(util.list_as_length(row, col_names_ct),
227
                    col_names, col_idxs) # handle CSV rows of different lengths
228
            
229
            return map_rows(get_value, util.WrapIter(wrap_row, rows))
230
        
231
        stdin = streams.LineCountStream(sys.stdin)
232
        def on_error(e):
233
            exc.add_msg(e, term.emph('input line #:')+' '+str(stdin.line_num))
234
            ex_tracker.track(e)
235
        
236
        if in_is_db:
237
            assert in_is_xpaths
238
            
239
            in_db = connect_db(in_db_config)
240
            cur = sql.select(in_db, table=in_root, fields=None, conds=None,
241
                limit=end, start=0)
242
            row_ct = map_table(list(sql.col_names(cur)), sql.rows(cur))
243
            
244
            in_db.db.close()
245
        elif in_is_xml:
246
            def get_rows(doc2rows):
247
                return iters.flatten(itertools.imap(doc2rows,
248
                    xml_parse.docs_iter(stdin, on_error)))
249
            
250
            if map_path == None:
251
                def doc2rows(in_xml_root):
252
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
253
                    util.skip(iter_, xml_dom.is_text) # skip metadata
254
                    return iter_
255
                
256
                row_ct = process_rows(lambda row, i: root.appendChild(row),
257
                    get_rows(doc2rows))
258
            else:
259
                def doc2rows(in_xml_root):
260
                    rows = xpath.get(in_xml_root, in_root, limit=end)
261
                    if rows == []: raise SystemExit('Map error: Root "'
262
                        +in_root+'" not found in input')
263
                    return rows
264
                
265
                def get_value(in_, row):
266
                    in_ = './{'+(','.join(strings.with_prefixes(
267
                        ['']+prefixes, in_)))+'}' # also with no prefix
268
                    nodes = xpath.get(row, in_, allow_rooted=False)
269
                    if nodes != []: return xml_dom.value(nodes[0])
270
                    else: return None
271
                
272
                row_ct = map_rows(get_value, get_rows(doc2rows))
273
        else: # input is CSV
274
            map_ = dict(mappings)
275
            reader, col_names = csvs.reader_and_header(sys.stdin)
276
            row_ct = map_table(col_names, reader)
277
        
278
        return row_ct
279
    
280
    def process_inputs(root, row_ready):
281
        row_ct = 0
282
        for map_path in map_paths:
283
            row_ct += process_input(root, row_ready, map_path)
284
        return row_ct
285
    
286
    if out_is_db:
287
        import db_xml
288
        
289
        out_db = connect_db(out_db_config)
290
        try:
291
            if redo: sql.empty_db(out_db)
292
            row_ins_ct_ref = [0]
293
            
294
            def row_ready(row_num, input_row):
295
                def on_error(e):
296
                    exc.add_msg(e, term.emph('row #:')+' '+str(row_num))
297
                    exc.add_msg(e, term.emph('input row:')+'\n'+str(input_row))
298
                    exc.add_msg(e, term.emph('output row:')+'\n'+str(root))
299
                    ex_tracker.track(e, row_num)
300
                
301
                xml_func.process(root, on_error)
302
                if not xml_dom.is_empty(root):
303
                    assert xml_dom.has_one_child(root)
304
                    try:
305
                        sql.with_savepoint(out_db,
306
                            lambda: db_xml.put(out_db, root.firstChild,
307
                                row_ins_ct_ref, on_error))
308
                        if commit: out_db.db.commit()
309
                    except sql.DatabaseErrors, e: on_error(e)
310
                prep_root()
311
            
312
            row_ct = process_inputs(root, row_ready)
313
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
314
                ' new rows into database\n')
315
            
316
            # Consume asynchronous tasks
317
            pool.main_loop()
318
        finally:
319
            out_db.db.rollback()
320
            out_db.db.close()
321
    else:
322
        def on_error(e): ex_tracker.track(e)
323
        def row_ready(row_num, input_row): pass
324
        row_ct = process_inputs(root, row_ready)
325
        xml_func.process(root, on_error)
326
        if out_is_xml_ref[0]:
327
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
328
        else: # output is CSV
329
            raise NotImplementedError('CSV output not supported yet')
330
    
331
    # Consume any asynchronous tasks not already consumed above
332
    pool.main_loop()
333
    
334
    profiler.stop(row_ct)
335
    ex_tracker.add_iters(row_ct)
336
    if verbose:
337
        sys.stderr.write('Processed '+str(row_ct)+' input rows\n')
338
        sys.stderr.write(profiler.msg()+'\n')
339
        sys.stderr.write(ex_tracker.msg()+'\n')
340
    ex_tracker.exit()
341

    
342
def main():
343
    try: main_()
344
    except Parser.SyntaxError, e: raise SystemExit(str(e))
345

    
346
if __name__ == '__main__':
347
    profile_to = opts.get_env_var('profile_to', None)
348
    if profile_to != None:
349
        import cProfile
350
        sys.stderr.write('Profiling to '+profile_to+'\n')
351
        cProfile.run(main.func_code, profile_to)
352
    else: main()
(23-23/43)