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
298 1991 aaronmk
            # Fetch rows
299
            if by_col: limit = 0 # only fetch column names
300
            else: limit = n
301 2288 aaronmk
            cur = sql.select(in_db, sql_gen.Table(table, schema), limit=limit,
302
                start=start, cacheable=False)
303 1991 aaronmk
            col_names = list(sql.col_names(cur))
304
305 1989 aaronmk
            if by_col:
306 2042 aaronmk
                map_table(col_names, []) # just create the template
307 1993 aaronmk
                xml_func.strip(root)
308 2103 aaronmk
                if debug: log_debug('Putting stripped:\n'+str(root))
309 2119 aaronmk
                    # only calc if debug
310 2195 aaronmk
                db_xml.put_table(in_db, root.firstChild, table, schema, n,
311
                    start, commit, row_ins_ct_ref)
312 2174 aaronmk
                row_ct = 0 # unknown for now
313 1984 aaronmk
            else:
314
                # Use normal by-row method
315 1991 aaronmk
                row_ct = map_table(col_names, sql.rows(cur), rows_start=start)
316
                    # rows_start: pre-start rows have been skipped
317 1136 aaronmk
318 1849 aaronmk
            in_db.db.close()
319 161 aaronmk
        elif in_is_xml:
320 1715 aaronmk
            def get_rows(doc2rows):
321 1758 aaronmk
                return iters.flatten(itertools.imap(doc2rows,
322
                    xml_parse.docs_iter(stdin, on_error)))
323 1715 aaronmk
324 1714 aaronmk
            if map_path == None:
325 1715 aaronmk
                def doc2rows(in_xml_root):
326 1701 aaronmk
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
327
                    util.skip(iter_, xml_dom.is_text) # skip metadata
328 1715 aaronmk
                    return iter_
329
330
                row_ct = process_rows(lambda row, i: root.appendChild(row),
331
                    get_rows(doc2rows))
332 1714 aaronmk
            else:
333 1715 aaronmk
                def doc2rows(in_xml_root):
334 1701 aaronmk
                    rows = xpath.get(in_xml_root, in_root, limit=end)
335 1714 aaronmk
                    if rows == []: raise SystemExit('Map error: Root "'
336
                        +in_root+'" not found in input')
337
                    return rows
338
339
                def get_value(in_, row):
340
                    in_ = './{'+(','.join(strings.with_prefixes(
341
                        ['']+prefixes, in_)))+'}' # also with no prefix
342
                    nodes = xpath.get(row, in_, allow_rooted=False)
343
                    if nodes != []: return xml_dom.value(nodes[0])
344
                    else: return None
345
346 1715 aaronmk
                row_ct = map_rows(get_value, get_rows(doc2rows))
347 56 aaronmk
        else: # input is CSV
348 133 aaronmk
            map_ = dict(mappings)
349 1389 aaronmk
            reader, col_names = csvs.reader_and_header(sys.stdin)
350 1403 aaronmk
            row_ct = map_table(col_names, reader)
351 982 aaronmk
352
        return row_ct
353 53 aaronmk
354 838 aaronmk
    def process_inputs(root, row_ready):
355 982 aaronmk
        row_ct = 0
356
        for map_path in map_paths:
357
            row_ct += process_input(root, row_ready, map_path)
358
        return row_ct
359 512 aaronmk
360 1886 aaronmk
    pool.share_vars(locals())
361 130 aaronmk
    if out_is_db:
362 646 aaronmk
        out_db = connect_db(out_db_config)
363 53 aaronmk
        try:
364 947 aaronmk
            if redo: sql.empty_db(out_db)
365 1886 aaronmk
            pool.share_vars(locals())
366 449 aaronmk
367 838 aaronmk
            def row_ready(row_num, input_row):
368 2119 aaronmk
                row_str_ = [None]
369
                def row_str():
370
                    if row_str_[0] == None:
371
                        # Row # is interally 0-based, but 1-based to the user
372
                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)
373
                            +'\n'+term.emph('input row:')+'\n'+str(input_row))
374
                        if verbose_errors: row_str_[0] += ('\n'
375
                            +term.emph('output row:')+'\n'+str(root))
376
                    return row_str_[0]
377
378
                if debug: log_debug(row_str()) # only calc if debug
379
380 452 aaronmk
                def on_error(e):
381 2119 aaronmk
                    exc.add_msg(e, row_str())
382 2046 aaronmk
                    ex_tracker.track(e, row_num, detail=verbose_errors)
383 1886 aaronmk
                pool.share_vars(locals())
384 452 aaronmk
385 2032 aaronmk
                row_root = root.cloneNode(True) # deep copy so don't modify root
386 2106 aaronmk
                xml_func.process(row_root, on_error, out_db)
387 2032 aaronmk
                if not xml_dom.is_empty(row_root):
388
                    assert xml_dom.has_one_child(row_root)
389 442 aaronmk
                    try:
390 982 aaronmk
                        sql.with_savepoint(out_db,
391 2032 aaronmk
                            lambda: db_xml.put(out_db, row_root.firstChild,
392 1850 aaronmk
                                row_ins_ct_ref, on_error))
393 1849 aaronmk
                        if commit: out_db.db.commit()
394 449 aaronmk
                    except sql.DatabaseErrors, e: on_error(e)
395
396 982 aaronmk
            row_ct = process_inputs(root, row_ready)
397
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
398 460 aaronmk
                ' new rows into database\n')
399 1868 aaronmk
400
            # Consume asynchronous tasks
401
            pool.main_loop()
402 53 aaronmk
        finally:
403 2166 aaronmk
            if out_db.connected():
404
                out_db.db.rollback()
405
                out_db.db.close()
406 751 aaronmk
    else:
407 759 aaronmk
        def on_error(e): ex_tracker.track(e)
408 838 aaronmk
        def row_ready(row_num, input_row): pass
409 982 aaronmk
        row_ct = process_inputs(root, row_ready)
410 759 aaronmk
        xml_func.process(root, on_error)
411 751 aaronmk
        if out_is_xml_ref[0]:
412
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
413
        else: # output is CSV
414
            raise NotImplementedError('CSV output not supported yet')
415 985 aaronmk
416 1868 aaronmk
    # Consume any asynchronous tasks not already consumed above
417 1862 aaronmk
    pool.main_loop()
418
419 982 aaronmk
    profiler.stop(row_ct)
420
    ex_tracker.add_iters(row_ct)
421 990 aaronmk
    if verbose:
422
        sys.stderr.write('Processed '+str(row_ct)+' input rows\n')
423
        sys.stderr.write(profiler.msg()+'\n')
424
        sys.stderr.write(ex_tracker.msg()+'\n')
425 985 aaronmk
    ex_tracker.exit()
426 53 aaronmk
427 847 aaronmk
def main():
428
    try: main_()
429 1719 aaronmk
    except Parser.SyntaxError, e: raise SystemExit(str(e))
430 847 aaronmk
431 846 aaronmk
if __name__ == '__main__':
432 847 aaronmk
    profile_to = opts.get_env_var('profile_to', None)
433
    if profile_to != None:
434
        import cProfile
435 1584 aaronmk
        sys.stderr.write('Profiling to '+profile_to+'\n')
436 847 aaronmk
        cProfile.run(main.func_code, profile_to)
437
    else: main()