Project

General

Profile

1 53 aaronmk
#!/usr/bin/env python
2
# Maps one datasource to another, using a map spreadsheet if needed
3 986 aaronmk
# Exit status is the # of errors in the import, up to the maximum exit status
4 53 aaronmk
# 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 1014 aaronmk
import csv
8 1714 aaronmk
import itertools
9 53 aaronmk
import os.path
10
import sys
11 299 aaronmk
import xml.dom.minidom as minidom
12 53 aaronmk
13 266 aaronmk
sys.path.append(os.path.dirname(__file__)+"/../lib")
14 53 aaronmk
15 1389 aaronmk
import csvs
16 2040 aaronmk
import db_xml
17 344 aaronmk
import exc
18 1714 aaronmk
import iters
19 1705 aaronmk
import maps
20 64 aaronmk
import opts
21 2035 aaronmk
import parallelproc
22 281 aaronmk
import Parser
23 982 aaronmk
import profiling
24 131 aaronmk
import sql
25 2288 aaronmk
import sql_gen
26 1758 aaronmk
import streams
27 715 aaronmk
import strings
28 828 aaronmk
import term
29 310 aaronmk
import util
30 1014 aaronmk
import xpath
31 133 aaronmk
import xml_dom
32 86 aaronmk
import xml_func
33 1713 aaronmk
import xml_parse
34 53 aaronmk
35 1404 aaronmk
def get_with_prefix(map_, prefixes, key):
36 2040 aaronmk
    '''Gets all entries for the given key with any of the given prefixes
37
    @return tuple(found_key, found_value)
38
    '''
39 1484 aaronmk
    values = []
40 1681 aaronmk
    for key_ in strings.with_prefixes(['']+prefixes, key): # also with no prefix
41
        try: value = map_[key_]
42 1484 aaronmk
        except KeyError, e: continue # keep going
43 2040 aaronmk
        values.append((key_, value))
44 1484 aaronmk
45
    if values != []: return values
46
    else: raise e # re-raise last KeyError
47 1404 aaronmk
48 1018 aaronmk
def metadata_value(name): return None # this feature has been removed
49 84 aaronmk
50 1360 aaronmk
def cleanup(val):
51
    if val == None: return val
52 1374 aaronmk
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
53 1360 aaronmk
54 847 aaronmk
def main_():
55 131 aaronmk
    env_names = []
56
    def usage_err():
57 944 aaronmk
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '
58 1945 aaronmk
            +sys.argv[0]+' [map_path...] [<input] [>output]\n'
59 1946 aaronmk
            'Note: Row #s start with 1')
60 838 aaronmk
61 1570 aaronmk
    ## Get config from env vars
62 838 aaronmk
63 1570 aaronmk
    # Modes
64 946 aaronmk
    test = opts.env_flag('test', False, env_names)
65 947 aaronmk
    commit = opts.env_flag('commit', False, env_names) and not test
66 944 aaronmk
        # never commit in test mode
67 947 aaronmk
    redo = opts.env_flag('redo', test, env_names) and not commit
68
        # never redo in commit mode (manually run `make empty_db` instead)
69 1570 aaronmk
70
    # Ranges
71 1946 aaronmk
    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 1945 aaronmk
    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 1570 aaronmk
    if end != None: end += start
81
82
    # Debugging
83 946 aaronmk
    debug = opts.env_flag('debug', False, env_names)
84 859 aaronmk
    sql.run_raw_query.debug = debug
85 1574 aaronmk
    verbose = debug or opts.env_flag('verbose', not test, env_names)
86 2118 aaronmk
    verbose_errors = opts.env_flag('verbose_errors', test or debug, env_names)
87 944 aaronmk
    opts.get_env_var('profile_to', None, env_names) # add to env_names
88 131 aaronmk
89 1570 aaronmk
    # DB
90
    def get_db_config(prefix):
91 1926 aaronmk
        return opts.get_env_vars(sql.db_config_names, prefix, env_names)
92 1570 aaronmk
    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 1982 aaronmk
    in_schema = opts.get_env_var('in_schema', None, env_names)
97
    in_table = opts.get_env_var('in_table', None, env_names)
98 1570 aaronmk
99 1989 aaronmk
    # Optimization
100 2050 aaronmk
    cache_sql = opts.env_flag('cache_sql', True, env_names)
101 1989 aaronmk
    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 2044 aaronmk
    else: cpus_default = 0 # or None to use parallel processing by default
105 1989 aaronmk
    cpus = util.cast(int, util.none_if(opts.get_env_var('cpus', cpus_default,
106
        env_names), u''))
107
108 1570 aaronmk
    ##
109
110 838 aaronmk
    # Logging
111 662 aaronmk
    def log(msg, on=verbose):
112 1901 aaronmk
        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 662 aaronmk
116 53 aaronmk
    # Parse args
117 510 aaronmk
    map_paths = sys.argv[1:]
118 512 aaronmk
    if map_paths == []:
119
        if in_is_db or not out_is_db: usage_err()
120
        else: map_paths = [None]
121 53 aaronmk
122 646 aaronmk
    def connect_db(db_config):
123 1901 aaronmk
        log('Connecting to '+sql.db_config_str(db_config))
124 2192 aaronmk
        return sql.connect(db_config, caching=cache_sql,
125
            autocommit=debug and commit, log_debug=log_debug)
126 646 aaronmk
127 1945 aaronmk
    if end != None: end_str = str(end-1) # end is one past the last #
128 1573 aaronmk
    else: end_str = 'end'
129 1901 aaronmk
    log('Processing input rows '+str(start)+'-'+end_str)
130 1573 aaronmk
131 1014 aaronmk
    ex_tracker = exc.ExPercentTracker(iter_text='row')
132
    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
133
134 1856 aaronmk
    # Parallel processing
135 2035 aaronmk
    pool = parallelproc.MultiProducerPool(cpus)
136 1901 aaronmk
    log('Using '+str(pool.process_ct)+' parallel CPUs')
137 1846 aaronmk
138 1014 aaronmk
    doc = xml_dom.create_doc()
139
    root = doc.documentElement
140 751 aaronmk
    out_is_xml_ref = [False]
141 1014 aaronmk
    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 751 aaronmk
150 1991 aaronmk
    # Define before the out_is_db section because it's used by by_col
151
    row_ins_ct_ref = [0]
152
153 838 aaronmk
    def process_input(root, row_ready, map_path):
154 512 aaronmk
        '''Inputs datasource to XML tree, mapping if needed'''
155
        # Load map header
156
        in_is_xpaths = True
157 751 aaronmk
        out_is_xpaths = True
158 512 aaronmk
        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 1402 aaronmk
166 512 aaronmk
            def split_col_name(name):
167 1402 aaronmk
                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 1198 aaronmk
                    # extract datasrc from "datasrc[data_format]"
173 1402 aaronmk
174 1705 aaronmk
            in_label, in_root, prefixes = maps.col_info(in_label)
175
            in_is_xpaths = in_root != None
176 1014 aaronmk
            in_label_ref[0] = in_label
177
            update_in_label()
178 1705 aaronmk
            out_label, out_root = maps.col_info(out_label)[:2]
179
            out_is_xpaths = out_root != None
180 1841 aaronmk
            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0
181 1705 aaronmk
                # outer elements are types
182 1402 aaronmk
183 512 aaronmk
            for row in reader:
184
                in_, out = row[:2]
185 2029 aaronmk
                if out != '': mappings.append([in_, out_root+out])
186 1402 aaronmk
187 512 aaronmk
            stream.close()
188
189
            root.ownerDocument.documentElement.tagName = out_label
190
        in_is_xml = in_is_xpaths and not in_is_db
191 751 aaronmk
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
192 56 aaronmk
193 1897 aaronmk
        def process_rows(process_row, rows, rows_start=0):
194
            '''Processes input rows
195 838 aaronmk
            @param process_row(in_row, i)
196 1945 aaronmk
            @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 297 aaronmk
            '''
199 1897 aaronmk
            rows = iter(rows)
200
201
            if end != None: row_nums = xrange(rows_start, end)
202
            else: row_nums = itertools.count(rows_start)
203 2038 aaronmk
            i = -1
204 1897 aaronmk
            for i in row_nums:
205 1718 aaronmk
                try: row = rows.next()
206 2038 aaronmk
                except StopIteration:
207
                    i -= 1 # last row # didn't count
208
                    break # no more rows
209 1718 aaronmk
                if i < start: continue # not at start row yet
210
211 838 aaronmk
                process_row(row, i)
212
                row_ready(i, row)
213 2033 aaronmk
            row_ct = i-start+1
214 982 aaronmk
            return row_ct
215 838 aaronmk
216 1898 aaronmk
        def map_rows(get_value, rows, **kw_args):
217 838 aaronmk
            '''Maps input rows
218
            @param get_value(in_, row):str
219
            '''
220 2030 aaronmk
            # 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 2026 aaronmk
            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 2032 aaronmk
                    mappings[i] = [in_, nodes]
240
                assert id_node != None
241 2026 aaronmk
242 2119 aaronmk
                if debug: # only calc if debug
243 2026 aaronmk
                    log_debug('Put template:\n'+str(root))
244
245 838 aaronmk
            def process_row(row, i):
246 316 aaronmk
                row_id = str(i)
247 2032 aaronmk
                if id_node != None: xml_dom.set_value(id_node, row_id)
248 316 aaronmk
                for in_, out in mappings:
249 2027 aaronmk
                    log_debug('Getting '+str(in_))
250 316 aaronmk
                    value = metadata_value(in_)
251 2027 aaronmk
                    if value == None: value = cleanup(get_value(in_, row))
252
                    log_debug('Putting '+repr(value)+' to '+str(out))
253 2032 aaronmk
                    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 1360 aaronmk
                        xpath.put_obj(root, out, row_id, has_types, value)
257 1898 aaronmk
            return process_rows(process_row, rows, **kw_args)
258 297 aaronmk
259 1898 aaronmk
        def map_table(col_names, rows, **kw_args):
260 1416 aaronmk
            col_names_ct = len(col_names)
261 1403 aaronmk
            col_idxs = util.list_flip(col_names)
262
263 2029 aaronmk
            # Resolve prefixes
264 2025 aaronmk
            mappings_orig = mappings[:] # save a copy
265
            mappings[:] = [] # empty existing elements
266
            for in_, out in mappings_orig:
267 1403 aaronmk
                if metadata_value(in_) == None:
268 2040 aaronmk
                    try: cols = get_with_prefix(col_idxs, prefixes, in_)
269 2025 aaronmk
                    except KeyError: pass
270 2040 aaronmk
                    else: mappings[len(mappings):] = [[db_xml.ColRef(*col), out]
271
                        for col in cols] # can't use += because that uses =
272 1403 aaronmk
273 2040 aaronmk
            def get_value(in_, row): return row.list[in_.idx]
274 1416 aaronmk
            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 1403 aaronmk
278 1898 aaronmk
            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)
279 1403 aaronmk
280 1758 aaronmk
        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 1700 aaronmk
        if in_is_db:
286 1982 aaronmk
            in_db = connect_db(in_db_config)
287 126 aaronmk
288 1982 aaronmk
            # 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 2313 aaronmk
            table = sql_gen.Table(table, schema)
298 1982 aaronmk
299 1991 aaronmk
            # Fetch rows
300
            if by_col: limit = 0 # only fetch column names
301
            else: limit = n
302 2313 aaronmk
            cur = sql.select(in_db, table, limit=limit, start=start,
303
                cacheable=False)
304 1991 aaronmk
            col_names = list(sql.col_names(cur))
305
306 1989 aaronmk
            if by_col:
307 2042 aaronmk
                map_table(col_names, []) # just create the template
308 1993 aaronmk
                xml_func.strip(root)
309 2103 aaronmk
                if debug: log_debug('Putting stripped:\n'+str(root))
310 2119 aaronmk
                    # only calc if debug
311 2313 aaronmk
                db_xml.put_table(in_db, root.firstChild, table, n, start,
312
                    commit, row_ins_ct_ref)
313 2174 aaronmk
                row_ct = 0 # unknown for now
314 1984 aaronmk
            else:
315
                # Use normal by-row method
316 1991 aaronmk
                row_ct = map_table(col_names, sql.rows(cur), rows_start=start)
317
                    # rows_start: pre-start rows have been skipped
318 1136 aaronmk
319 1849 aaronmk
            in_db.db.close()
320 161 aaronmk
        elif in_is_xml:
321 1715 aaronmk
            def get_rows(doc2rows):
322 1758 aaronmk
                return iters.flatten(itertools.imap(doc2rows,
323
                    xml_parse.docs_iter(stdin, on_error)))
324 1715 aaronmk
325 1714 aaronmk
            if map_path == None:
326 1715 aaronmk
                def doc2rows(in_xml_root):
327 1701 aaronmk
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
328
                    util.skip(iter_, xml_dom.is_text) # skip metadata
329 1715 aaronmk
                    return iter_
330
331
                row_ct = process_rows(lambda row, i: root.appendChild(row),
332
                    get_rows(doc2rows))
333 1714 aaronmk
            else:
334 1715 aaronmk
                def doc2rows(in_xml_root):
335 1701 aaronmk
                    rows = xpath.get(in_xml_root, in_root, limit=end)
336 1714 aaronmk
                    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 1715 aaronmk
                row_ct = map_rows(get_value, get_rows(doc2rows))
348 56 aaronmk
        else: # input is CSV
349 133 aaronmk
            map_ = dict(mappings)
350 1389 aaronmk
            reader, col_names = csvs.reader_and_header(sys.stdin)
351 1403 aaronmk
            row_ct = map_table(col_names, reader)
352 982 aaronmk
353
        return row_ct
354 53 aaronmk
355 838 aaronmk
    def process_inputs(root, row_ready):
356 982 aaronmk
        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 512 aaronmk
361 1886 aaronmk
    pool.share_vars(locals())
362 130 aaronmk
    if out_is_db:
363 646 aaronmk
        out_db = connect_db(out_db_config)
364 53 aaronmk
        try:
365 947 aaronmk
            if redo: sql.empty_db(out_db)
366 1886 aaronmk
            pool.share_vars(locals())
367 449 aaronmk
368 838 aaronmk
            def row_ready(row_num, input_row):
369 2119 aaronmk
                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 452 aaronmk
                def on_error(e):
382 2119 aaronmk
                    exc.add_msg(e, row_str())
383 2046 aaronmk
                    ex_tracker.track(e, row_num, detail=verbose_errors)
384 1886 aaronmk
                pool.share_vars(locals())
385 452 aaronmk
386 2032 aaronmk
                row_root = root.cloneNode(True) # deep copy so don't modify root
387 2106 aaronmk
                xml_func.process(row_root, on_error, out_db)
388 2032 aaronmk
                if not xml_dom.is_empty(row_root):
389
                    assert xml_dom.has_one_child(row_root)
390 442 aaronmk
                    try:
391 982 aaronmk
                        sql.with_savepoint(out_db,
392 2032 aaronmk
                            lambda: db_xml.put(out_db, row_root.firstChild,
393 1850 aaronmk
                                row_ins_ct_ref, on_error))
394 1849 aaronmk
                        if commit: out_db.db.commit()
395 449 aaronmk
                    except sql.DatabaseErrors, e: on_error(e)
396
397 982 aaronmk
            row_ct = process_inputs(root, row_ready)
398
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
399 460 aaronmk
                ' new rows into database\n')
400 1868 aaronmk
401
            # Consume asynchronous tasks
402
            pool.main_loop()
403 53 aaronmk
        finally:
404 2166 aaronmk
            if out_db.connected():
405
                out_db.db.rollback()
406
                out_db.db.close()
407 751 aaronmk
    else:
408 759 aaronmk
        def on_error(e): ex_tracker.track(e)
409 838 aaronmk
        def row_ready(row_num, input_row): pass
410 982 aaronmk
        row_ct = process_inputs(root, row_ready)
411 759 aaronmk
        xml_func.process(root, on_error)
412 751 aaronmk
        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 985 aaronmk
417 1868 aaronmk
    # Consume any asynchronous tasks not already consumed above
418 1862 aaronmk
    pool.main_loop()
419
420 982 aaronmk
    profiler.stop(row_ct)
421
    ex_tracker.add_iters(row_ct)
422 990 aaronmk
    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 985 aaronmk
    ex_tracker.exit()
427 53 aaronmk
428 847 aaronmk
def main():
429
    try: main_()
430 1719 aaronmk
    except Parser.SyntaxError, e: raise SystemExit(str(e))
431 847 aaronmk
432 846 aaronmk
if __name__ == '__main__':
433 847 aaronmk
    profile_to = opts.get_env_var('profile_to', None)
434
    if profile_to != None:
435
        import cProfile
436 1584 aaronmk
        sys.stderr.write('Profiling to '+profile_to+'\n')
437 847 aaronmk
        cProfile.run(main.func_code, profile_to)
438
    else: main()