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 3768 aaronmk
# Multi-safe (supports an input appearing multiple times).
5 53 aaronmk
# For outputting an XML file to a PostgreSQL database, use the general format of
6
# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml
7 4213 aaronmk
# Duplicate-column safe (supports multiple columns of the same name, which will
8
# be combined)
9 4505 aaronmk
# Case- and punctuation-insensitive.
10 53 aaronmk
11 2550 aaronmk
import copy
12 1014 aaronmk
import csv
13 1714 aaronmk
import itertools
14 53 aaronmk
import os.path
15 3577 aaronmk
import warnings
16 53 aaronmk
import sys
17 299 aaronmk
import xml.dom.minidom as minidom
18 53 aaronmk
19 266 aaronmk
sys.path.append(os.path.dirname(__file__)+"/../lib")
20 53 aaronmk
21 1389 aaronmk
import csvs
22 2040 aaronmk
import db_xml
23 344 aaronmk
import exc
24 9455 aaronmk
import ints
25 1714 aaronmk
import iters
26 1705 aaronmk
import maps
27 64 aaronmk
import opts
28 2035 aaronmk
import parallelproc
29 281 aaronmk
import Parser
30 982 aaronmk
import profiling
31 131 aaronmk
import sql
32 2288 aaronmk
import sql_gen
33 3103 aaronmk
import sql_io
34 1758 aaronmk
import streams
35 715 aaronmk
import strings
36 828 aaronmk
import term
37 310 aaronmk
import util
38 1014 aaronmk
import xpath
39 133 aaronmk
import xml_dom
40 86 aaronmk
import xml_func
41 1713 aaronmk
import xml_parse
42 53 aaronmk
43 5911 aaronmk
metadata_prefix = ':'
44 4068 aaronmk
collision_suffix = '/_alt/'
45 4047 aaronmk
46 1404 aaronmk
def get_with_prefix(map_, prefixes, key):
47 2040 aaronmk
    '''Gets all entries for the given key with any of the given prefixes
48
    @return tuple(found_key, found_value)
49
    '''
50 1484 aaronmk
    values = []
51 1681 aaronmk
    for key_ in strings.with_prefixes(['']+prefixes, key): # also with no prefix
52
        try: value = map_[key_]
53 1484 aaronmk
        except KeyError, e: continue # keep going
54 2040 aaronmk
        values.append((key_, value))
55 1484 aaronmk
56
    if values != []: return values
57
    else: raise e # re-raise last KeyError
58 1404 aaronmk
59 5927 aaronmk
def is_metadata(str_): return str_.startswith(metadata_prefix)
60
61 5911 aaronmk
def metadata_value(name):
62
    removed_ref = [False]
63
    name = strings.remove_prefix(metadata_prefix, name, removed_ref)
64
    if removed_ref[0]: return name
65
    else: return None
66 84 aaronmk
67 1360 aaronmk
def cleanup(val):
68
    if val == None: return val
69 1374 aaronmk
    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')
70 1360 aaronmk
71 847 aaronmk
def main_():
72 131 aaronmk
    env_names = []
73
    def usage_err():
74 944 aaronmk
        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '
75 1945 aaronmk
            +sys.argv[0]+' [map_path...] [<input] [>output]\n'
76 1946 aaronmk
            'Note: Row #s start with 1')
77 838 aaronmk
78 1570 aaronmk
    ## Get config from env vars
79 838 aaronmk
80 1570 aaronmk
    # Modes
81 946 aaronmk
    test = opts.env_flag('test', False, env_names)
82 947 aaronmk
    commit = opts.env_flag('commit', False, env_names) and not test
83 944 aaronmk
        # never commit in test mode
84 7122 aaronmk
    redo = opts.env_flag('redo', False, env_names) and not commit
85 2977 aaronmk
        # never redo in commit mode (run `make schemas/reinstall` instead)
86 1570 aaronmk
87
    # Ranges
88 1946 aaronmk
    start = util.cast(int, opts.get_env_var('start', 1, env_names)) # 1-based
89
    # Make start interally 0-based.
90
    # It's 1-based to the user to match up with the staging table row #s.
91
    start -= 1
92 1945 aaronmk
    if test: n_default = 1
93
    else: n_default = None
94
    n = util.cast(int, util.none_if(opts.get_env_var('n', n_default, env_names),
95
        u''))
96
    end = n
97 1570 aaronmk
    if end != None: end += start
98
99
    # Debugging
100 3132 aaronmk
    verbosity = util.cast(float, opts.get_env_var('verbosity', None, env_names))
101 944 aaronmk
    opts.get_env_var('profile_to', None, env_names) # add to env_names
102 131 aaronmk
103 1570 aaronmk
    # DB
104
    def get_db_config(prefix):
105 1926 aaronmk
        return opts.get_env_vars(sql.db_config_names, prefix, env_names)
106 1570 aaronmk
    in_db_config = get_db_config('in')
107
    out_db_config = get_db_config('out')
108
    in_is_db = 'engine' in in_db_config
109
    out_is_db = 'engine' in out_db_config
110 1982 aaronmk
    in_schema = opts.get_env_var('in_schema', None, env_names)
111
    in_table = opts.get_env_var('in_table', None, env_names)
112 3339 aaronmk
    if in_schema != None:
113
        for config in [in_db_config, out_db_config]:
114
            config['schemas'] += ','+in_schema
115 1570 aaronmk
116 1989 aaronmk
    # Optimization
117 2050 aaronmk
    cache_sql = opts.env_flag('cache_sql', True, env_names)
118 1989 aaronmk
    by_col = in_db_config == out_db_config and opts.env_flag('by_col', False,
119
        env_names) # by-column optimization only applies if mapping to same DB
120
    if test: cpus_default = 0
121 2044 aaronmk
    else: cpus_default = 0 # or None to use parallel processing by default
122 1989 aaronmk
    cpus = util.cast(int, util.none_if(opts.get_env_var('cpus', cpus_default,
123
        env_names), u''))
124
125 3132 aaronmk
    # Set default verbosity. Must happen after by_col is set.
126 4026 aaronmk
    if verbosity == None:
127 3132 aaronmk
        if test: verbosity = 0.5 # automated tests should not be verbose
128 9453 aaronmk
        elif by_col: verbosity = 3 # show all queries to assist debugging
129 3132 aaronmk
        else: verbosity = 1.1 # just show row progress
130
131 9455 aaronmk
    # fix verbosity
132
    if by_col and not test: verbosity = ints.set_min(verbosity, 2)
133
        # live column-based import MUST be run with verbosity 2+ (3 preferred)
134
        # to provide debugging information for often-complex errors.
135
        # without this, debugging is effectively impossible.
136
        # automated tests are exempt from this because they output to the screen
137
138 1570 aaronmk
    ##
139
140 838 aaronmk
    # Logging
141 2457 aaronmk
    verbose_errors = test and verbosity > 0
142 3243 aaronmk
    debug = verbosity >= 1.5
143 2439 aaronmk
    def log(msg, level=1):
144
        '''Higher level -> more verbose'''
145 2452 aaronmk
        if level <= verbosity:
146 2490 aaronmk
            if verbosity <= 2:
147
                if level == 1.5: msg = '# '+msg # msg is Redmine list item
148
                elif msg.startswith('DB query:'): # remove extra debug info
149
                    first_line, nl, msg = msg.partition('\n')
150
            elif level > 1: msg = '['+str(level)+'] '+msg # include level in msg
151
152 2883 aaronmk
            sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))
153 2439 aaronmk
    if debug: log_debug = lambda msg, level=2: log(msg, level)
154 1901 aaronmk
    else: log_debug = sql.log_debug_none
155 662 aaronmk
156 53 aaronmk
    # Parse args
157 510 aaronmk
    map_paths = sys.argv[1:]
158 512 aaronmk
    if map_paths == []:
159
        if in_is_db or not out_is_db: usage_err()
160
        else: map_paths = [None]
161 53 aaronmk
162 646 aaronmk
    def connect_db(db_config):
163 1901 aaronmk
        log('Connecting to '+sql.db_config_str(db_config))
164 2920 aaronmk
        return sql.connect(db_config, caching=cache_sql, autocommit=commit,
165 2916 aaronmk
            debug_temp=verbosity > 3 and commit, log_debug=log_debug)
166 646 aaronmk
167 1945 aaronmk
    if end != None: end_str = str(end-1) # end is one past the last #
168 1573 aaronmk
    else: end_str = 'end'
169 1901 aaronmk
    log('Processing input rows '+str(start)+'-'+end_str)
170 1573 aaronmk
171 1014 aaronmk
    ex_tracker = exc.ExPercentTracker(iter_text='row')
172
    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
173
174 1856 aaronmk
    # Parallel processing
175 2035 aaronmk
    pool = parallelproc.MultiProducerPool(cpus)
176 1901 aaronmk
    log('Using '+str(pool.process_ct)+' parallel CPUs')
177 1846 aaronmk
178 3661 aaronmk
    # Set up DB access
179
    row_ins_ct_ref = [0]
180
    if out_is_db:
181
        out_db = connect_db(out_db_config)
182 3689 aaronmk
        def is_rel_func(name):
183
            return (name in db_xml.put_special_funcs
184
                or sql.function_exists(out_db, sql_gen.Function(name)))
185 3661 aaronmk
186 1014 aaronmk
    doc = xml_dom.create_doc()
187
    root = doc.documentElement
188 751 aaronmk
    out_is_xml_ref = [False]
189 3653 aaronmk
190 1014 aaronmk
    in_label_ref = [None]
191 3653 aaronmk
    col_defaults = {}
192 1014 aaronmk
    def update_in_label():
193 6399 aaronmk
        if in_label_ref[0] != None:
194
            os.environ['source'] = in_label_ref[0]
195 3653 aaronmk
196 1014 aaronmk
    def prep_root():
197
        root.clear()
198
        update_in_label()
199
    prep_root()
200 751 aaronmk
201 838 aaronmk
    def process_input(root, row_ready, map_path):
202 512 aaronmk
        '''Inputs datasource to XML tree, mapping if needed'''
203
        # Load map header
204
        in_is_xpaths = True
205 751 aaronmk
        out_is_xpaths = True
206 512 aaronmk
        out_label = None
207
        if map_path != None:
208
            metadata = []
209
            mappings = []
210
            stream = open(map_path, 'rb')
211
            reader = csv.reader(stream)
212
            in_label, out_label = reader.next()[:2]
213 1402 aaronmk
214 512 aaronmk
            def split_col_name(name):
215 1402 aaronmk
                label, sep, root = name.partition(':')
216 4652 aaronmk
                return label, sep != '', root, []
217 1402 aaronmk
218 1705 aaronmk
            in_label, in_root, prefixes = maps.col_info(in_label)
219
            in_is_xpaths = in_root != None
220 1014 aaronmk
            in_label_ref[0] = in_label
221
            update_in_label()
222 1705 aaronmk
            out_label, out_root = maps.col_info(out_label)[:2]
223
            out_is_xpaths = out_root != None
224 1841 aaronmk
            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0
225 1705 aaronmk
                # outer elements are types
226 1402 aaronmk
227 512 aaronmk
            for row in reader:
228
                in_, out = row[:2]
229 2029 aaronmk
                if out != '': mappings.append([in_, out_root+out])
230 1402 aaronmk
231 512 aaronmk
            stream.close()
232
233
            root.ownerDocument.documentElement.tagName = out_label
234
        in_is_xml = in_is_xpaths and not in_is_db
235 751 aaronmk
        out_is_xml_ref[0] = out_is_xpaths and not out_is_db
236 56 aaronmk
237 1897 aaronmk
        def process_rows(process_row, rows, rows_start=0):
238 4026 aaronmk
            '''Processes input rows
239 838 aaronmk
            @param process_row(in_row, i)
240 1945 aaronmk
            @rows_start The (0-based) row # of the first row in rows. Set this
241
                only if the pre-start rows have already been skipped.
242 297 aaronmk
            '''
243 1897 aaronmk
            rows = iter(rows)
244
245
            if end != None: row_nums = xrange(rows_start, end)
246
            else: row_nums = itertools.count(rows_start)
247 2038 aaronmk
            i = -1
248 1897 aaronmk
            for i in row_nums:
249 1718 aaronmk
                try: row = rows.next()
250 2038 aaronmk
                except StopIteration:
251
                    i -= 1 # last row # didn't count
252
                    break # no more rows
253 1718 aaronmk
                if i < start: continue # not at start row yet
254
255 2881 aaronmk
                # Row # is interally 0-based, but 1-based to the user
256
                log('Processing input row #'+str(i+1), level=1.1)
257 838 aaronmk
                process_row(row, i)
258
                row_ready(i, row)
259 2033 aaronmk
            row_ct = i-start+1
260 982 aaronmk
            return row_ct
261 838 aaronmk
262 1898 aaronmk
        def map_rows(get_value, rows, **kw_args):
263 838 aaronmk
            '''Maps input rows
264
            @param get_value(in_, row):str
265
            '''
266 2030 aaronmk
            # Prevent collisions if multiple inputs mapping to same output
267
            outputs_idxs = dict()
268
            for i, mapping in enumerate(mappings):
269
                in_, out = mapping
270
                default = util.NamedTuple(count=1, first=i)
271
                idxs = outputs_idxs.setdefault(out, default)
272
                if idxs is not default: # key existed, so there was a collision
273 4048 aaronmk
                    if idxs.count == 1: # first key does not yet have suffix
274 4047 aaronmk
                        mappings[idxs.first][1] += collision_suffix+'0'
275
                    mappings[i][1] += collision_suffix+str(idxs.count)
276 2030 aaronmk
                    idxs.count += 1
277
278 2026 aaronmk
            id_node = None
279
            if out_is_db:
280 5927 aaronmk
                mappings_orig = mappings[:] # save a copy
281
                mappings[:] = [] # empty existing elements
282
                for in_, out in mappings_orig:
283
                    in_str = strings.ustr(in_)
284
                    is_metadata_ = is_metadata(in_str)
285
                    if is_metadata_: value = metadata_value(in_str)
286
                    else: value = '$'+in_str # mark as name
287
288 2026 aaronmk
                    # All put_obj()s should return the same id_node
289
                    nodes, id_node = xpath.put_obj(root, out, '-1', has_types,
290 5927 aaronmk
                        value) # value is placeholder that documents name
291
                    if not is_metadata_: mappings.append([in_, nodes])
292 3580 aaronmk
                if id_node == None:
293 8075 aaronmk
                    warnings.warn(UserWarning('Map warning: No mappings or no '
294
                        'column name matches. Are you importing the correct '
295
                        'input table?'))
296 4042 aaronmk
                xml_func.simplify(root)
297 4491 aaronmk
                sys.stdout.write(strings.to_raw_str('Put template:\n'
298
                    +strings.ustr(root)))
299 3669 aaronmk
                sys.stdout.flush()
300 2026 aaronmk
301 838 aaronmk
            def process_row(row, i):
302 316 aaronmk
                row_id = str(i)
303 2032 aaronmk
                if id_node != None: xml_dom.set_value(id_node, row_id)
304 316 aaronmk
                for in_, out in mappings:
305 4491 aaronmk
                    log_debug('Getting '+strings.ustr(in_))
306 5927 aaronmk
                    value = cleanup(get_value(in_, row))
307 4492 aaronmk
                    log_debug('Putting '+strings.urepr(value)+' to '
308
                        +strings.ustr(out))
309 2032 aaronmk
                    if out_is_db: # out is list of XML nodes
310
                        for node in out: xml_dom.set_value(node, value)
311
                    elif value != None: # out is XPath
312 1360 aaronmk
                        xpath.put_obj(root, out, row_id, has_types, value)
313 1898 aaronmk
            return process_rows(process_row, rows, **kw_args)
314 297 aaronmk
315 1898 aaronmk
        def map_table(col_names, rows, **kw_args):
316 1416 aaronmk
            col_names_ct = len(col_names)
317 1403 aaronmk
            col_idxs = util.list_flip(col_names)
318 6740 aaronmk
            col_names_map = dict(zip(col_names, col_names))
319 4503 aaronmk
            prefixes_simp = map(maps.simplify, prefixes)
320 1403 aaronmk
321 2029 aaronmk
            # Resolve prefixes
322 2025 aaronmk
            mappings_orig = mappings[:] # save a copy
323
            mappings[:] = [] # empty existing elements
324
            for in_, out in mappings_orig:
325 7148 aaronmk
                if is_metadata(in_): mappings.append([in_, out])
326 5928 aaronmk
                else:
327 4503 aaronmk
                    try:
328 4505 aaronmk
                        cols = get_with_prefix(col_names_map, prefixes_simp,
329 6740 aaronmk
                            in_)
330 2025 aaronmk
                    except KeyError: pass
331 4503 aaronmk
                    else:
332 4505 aaronmk
                        cols = [(orig, col_idxs[orig]) for simp, orig in cols]
333 4503 aaronmk
                        mappings[len(mappings):] = [[db_xml.ColRef(*col), out]
334
                            for col in cols] # can't use += because that uses =
335 1403 aaronmk
336 2040 aaronmk
            def get_value(in_, row): return row.list[in_.idx]
337 1416 aaronmk
            def wrap_row(row):
338
                return util.ListDict(util.list_as_length(row, col_names_ct),
339
                    col_names, col_idxs) # handle CSV rows of different lengths
340 1403 aaronmk
341 1898 aaronmk
            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)
342 1403 aaronmk
343 1700 aaronmk
        if in_is_db:
344 2809 aaronmk
            def on_error(e): ex_tracker.track(e)
345
346 3186 aaronmk
            if by_col: in_db = out_db
347
            else: in_db = connect_db(in_db_config)
348 126 aaronmk
349 1982 aaronmk
            # Get table and schema name
350
            schema = in_schema # modified, so can't have same name as outer var
351
            table = in_table # modified, so can't have same name as outer var
352
            if table == None:
353
                assert in_is_xpaths
354
                schema, sep, table = in_root.partition('.')
355
                if sep == '': # only the table name was specified
356
                    table = schema
357
                    schema = None
358 2313 aaronmk
            table = sql_gen.Table(table, schema)
359 1982 aaronmk
360 1991 aaronmk
            # Fetch rows
361
            if by_col: limit = 0 # only fetch column names
362
            else: limit = n
363 6385 aaronmk
            try:
364
                cur = sql.select(in_db, table, limit=limit, start=start,
365
                    recover=True, cacheable=False)
366
            except sql.DoesNotExistException:
367 6445 aaronmk
                table = None
368 6385 aaronmk
                col_names = []
369
                rows = []
370
            else:
371
                col_names = list(sql.col_names(cur))
372
                rows = sql.rows(cur)
373 1991 aaronmk
374 1989 aaronmk
            if by_col:
375 2042 aaronmk
                map_table(col_names, []) # just create the template
376 2550 aaronmk
377 6454 aaronmk
                if table != None and start == 0 and n == None: # full re-import
378 2734 aaronmk
                    log('Clearing errors table')
379 4474 aaronmk
                    errors_table_ = sql_io.errors_table(in_db, table)
380
                    if errors_table_ != None:
381
                        sql.drop_table(in_db, errors_table_)
382 2734 aaronmk
383 2550 aaronmk
                # Strip XML functions not in the DB
384 3427 aaronmk
                xml_func.process(root, is_rel_func=is_rel_func)
385 4491 aaronmk
                if debug: log_debug('Putting stripped:\n'+strings.ustr(root))
386 2119 aaronmk
                    # only calc if debug
387 2550 aaronmk
388
                # Import rows
389 2422 aaronmk
                in_row_ct_ref = [0]
390 2928 aaronmk
                db_xml.put_table(in_db, root.firstChild, table, in_row_ct_ref,
391 3627 aaronmk
                    row_ins_ct_ref, n, start, on_error, col_defaults)
392 2422 aaronmk
                row_ct = in_row_ct_ref[0]
393 1984 aaronmk
            else:
394
                # Use normal by-row method
395 6385 aaronmk
                row_ct = map_table(col_names, rows, rows_start=start)
396 1991 aaronmk
                    # rows_start: pre-start rows have been skipped
397 3186 aaronmk
398
                in_db.db.close()
399 161 aaronmk
        elif in_is_xml:
400 2809 aaronmk
            stdin = streams.LineCountStream(sys.stdin)
401
            def on_error(e):
402
                exc.add_msg(e, term.emph('input line #:')+' '
403
                    +str(stdin.line_num))
404
                ex_tracker.track(e)
405
406 1715 aaronmk
            def get_rows(doc2rows):
407 1758 aaronmk
                return iters.flatten(itertools.imap(doc2rows,
408
                    xml_parse.docs_iter(stdin, on_error)))
409 1715 aaronmk
410 1714 aaronmk
            if map_path == None:
411 1715 aaronmk
                def doc2rows(in_xml_root):
412 1701 aaronmk
                    iter_ = xml_dom.NodeElemIter(in_xml_root)
413
                    util.skip(iter_, xml_dom.is_text) # skip metadata
414 1715 aaronmk
                    return iter_
415
416
                row_ct = process_rows(lambda row, i: root.appendChild(row),
417
                    get_rows(doc2rows))
418 1714 aaronmk
            else:
419 1715 aaronmk
                def doc2rows(in_xml_root):
420 1701 aaronmk
                    rows = xpath.get(in_xml_root, in_root, limit=end)
421 3581 aaronmk
                    if rows == []: warnings.warn(UserWarning('Map warning: '
422
                        'Root "'+in_root+'" not found in input'))
423 1714 aaronmk
                    return rows
424
425
                def get_value(in_, row):
426
                    in_ = './{'+(','.join(strings.with_prefixes(
427
                        ['']+prefixes, in_)))+'}' # also with no prefix
428
                    nodes = xpath.get(row, in_, allow_rooted=False)
429
                    if nodes != []: return xml_dom.value(nodes[0])
430
                    else: return None
431
432 1715 aaronmk
                row_ct = map_rows(get_value, get_rows(doc2rows))
433 56 aaronmk
        else: # input is CSV
434 1389 aaronmk
            reader, col_names = csvs.reader_and_header(sys.stdin)
435 1403 aaronmk
            row_ct = map_table(col_names, reader)
436 982 aaronmk
437
        return row_ct
438 53 aaronmk
439 838 aaronmk
    def process_inputs(root, row_ready):
440 982 aaronmk
        row_ct = 0
441
        for map_path in map_paths:
442
            row_ct += process_input(root, row_ready, map_path)
443
        return row_ct
444 512 aaronmk
445 1886 aaronmk
    pool.share_vars(locals())
446 130 aaronmk
    if out_is_db:
447 53 aaronmk
        try:
448 947 aaronmk
            if redo: sql.empty_db(out_db)
449 1886 aaronmk
            pool.share_vars(locals())
450 449 aaronmk
451 838 aaronmk
            def row_ready(row_num, input_row):
452 2119 aaronmk
                row_str_ = [None]
453
                def row_str():
454
                    if row_str_[0] == None:
455
                        # Row # is interally 0-based, but 1-based to the user
456
                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)
457 4491 aaronmk
                            +'\n'+term.emph('input row:')+'\n'
458
                            +strings.ustr(input_row))
459 2119 aaronmk
                        if verbose_errors: row_str_[0] += ('\n'
460 4491 aaronmk
                            +term.emph('output row:')+'\n'+strings.ustr(root))
461 2119 aaronmk
                    return row_str_[0]
462
463
                if debug: log_debug(row_str()) # only calc if debug
464
465 452 aaronmk
                def on_error(e):
466 2119 aaronmk
                    exc.add_msg(e, row_str())
467 2046 aaronmk
                    ex_tracker.track(e, row_num, detail=verbose_errors)
468 1886 aaronmk
                pool.share_vars(locals())
469 452 aaronmk
470 2032 aaronmk
                row_root = root.cloneNode(True) # deep copy so don't modify root
471 3689 aaronmk
                xml_func.process(row_root, on_error, is_rel_func, out_db)
472 4491 aaronmk
                if debug: log_debug('Putting processed:\n'
473
                    +strings.ustr(row_root)) # only calc if debug
474 2032 aaronmk
                if not xml_dom.is_empty(row_root):
475
                    assert xml_dom.has_one_child(row_root)
476 442 aaronmk
                    try:
477 982 aaronmk
                        sql.with_savepoint(out_db,
478 2032 aaronmk
                            lambda: db_xml.put(out_db, row_root.firstChild,
479 3653 aaronmk
                                row_ins_ct_ref, on_error, col_defaults))
480 449 aaronmk
                    except sql.DatabaseErrors, e: on_error(e)
481
482 982 aaronmk
            row_ct = process_inputs(root, row_ready)
483
            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+
484 460 aaronmk
                ' new rows into database\n')
485 3669 aaronmk
            sys.stdout.flush()
486 1868 aaronmk
487
            # Consume asynchronous tasks
488
            pool.main_loop()
489 3117 aaronmk
        finally: out_db.close()
490 751 aaronmk
    else:
491 759 aaronmk
        def on_error(e): ex_tracker.track(e)
492 838 aaronmk
        def row_ready(row_num, input_row): pass
493 982 aaronmk
        row_ct = process_inputs(root, row_ready)
494 759 aaronmk
        xml_func.process(root, on_error)
495 751 aaronmk
        if out_is_xml_ref[0]:
496
            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)
497
        else: # output is CSV
498
            raise NotImplementedError('CSV output not supported yet')
499 985 aaronmk
500 1868 aaronmk
    # Consume any asynchronous tasks not already consumed above
501 1862 aaronmk
    pool.main_loop()
502
503 982 aaronmk
    profiler.stop(row_ct)
504 3324 aaronmk
    if not by_col: ex_tracker.add_iters(row_ct) # only if errors are done by row
505 2439 aaronmk
    log('Processed '+str(row_ct)+' input rows')
506
    log(profiler.msg())
507
    log(ex_tracker.msg())
508 985 aaronmk
    ex_tracker.exit()
509 53 aaronmk
510 847 aaronmk
def main():
511
    try: main_()
512 4491 aaronmk
    except Parser.SyntaxError, e: raise SystemExit(strings.ustr(e))
513 847 aaronmk
514 846 aaronmk
if __name__ == '__main__':
515 847 aaronmk
    profile_to = opts.get_env_var('profile_to', None)
516
    if profile_to != None:
517
        import cProfile
518 1584 aaronmk
        sys.stderr.write('Profiling to '+profile_to+'\n')
519 847 aaronmk
        cProfile.run(main.func_code, profile_to)
520
    else: main()