# Exit status is the # of errors in the import, up to the maximum exit status

# For outputting an XML file to a PostgreSQL database, use the general format of

import csv

import itertools

import os.path

import xml.dom.minidom as minidom

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

import csvs

import exc

import iters

import maps

import opts

import parallel

import Parser

import profiling

import sql

import streams

import strings

import term

import util

import xpath

import xml_dom

import xml_func

import xml_parse

def get_with_prefix(map_, prefixes, key):

    '''Gets all entries for the given key with any of the given prefixes'''

    for key_ in strings.with_prefixes(['']+prefixes, key): # also with no prefix

        except KeyError, e: continue # keep going

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

def cleanup(val):

    return util.none_if(strings.cleanup(strings.ustr(val)), u'', u'\\N')

def main_():

    env_names = []

        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '

            +sys.argv[0]+' [map_path...] [<input] [>output]\n'

            'Note: Row #s start with 1')

    ## Get config from env vars

    # Modes

    test = opts.env_flag('test', False, env_names)

    commit = opts.env_flag('commit', False, env_names) and not test

        # never commit in test mode

    redo = opts.env_flag('redo', test, env_names) and not commit

    start = util.cast(int, opts.get_env_var('start', 1, env_names)) # 1-based

    if test: n_default = 1

    if end != None: end += start

    # Optimization

    by_col = opts.env_flag('by_col', False, env_names)

    if test: cpus_default = 0

    cpus = util.cast(int, util.none_if(opts.get_env_var('cpus', cpus_default,

    # Debugging

    debug = opts.env_flag('debug', False, env_names)

    sql.run_raw_query.debug = debug

    verbose = debug or opts.env_flag('verbose', not test, env_names)

    opts.get_env_var('profile_to', None, env_names) # add to env_names

    # DB

        return opts.get_env_vars(sql.db_config_names, prefix, env_names)

    in_db_config = get_db_config('in')

    in_schema = opts.get_env_var('in_schema', None, env_names)

    # Logging

    def log(msg, on=verbose):

        if on: sys.stderr.write(msg+'\n')

    # Parse args

    map_paths = sys.argv[1:]

    if map_paths == []:

    def connect_db(db_config):

        log('Connecting to '+sql.db_config_str(db_config))

    if end != None: end_str = str(end-1) # end is one past the last #

    else: end_str = 'end'

    log('Processing input rows '+str(start)+'-'+end_str)

    ex_tracker = exc.ExPercentTracker(iter_text='row')

    # Parallel processing

    pool = parallel.MultiProducerPool(cpus)

    log('Using '+str(pool.process_ct)+' parallel CPUs')

    doc = xml_dom.create_doc()

    out_is_xml_ref = [False]

    in_label_ref = [None]

    def process_input(root, row_ready, map_path):

        '''Inputs datasource to XML tree, mapping if needed'''

        out_is_xpaths = True

        out_label = None

            def split_col_name(name):

                label, sep, root = name.partition(':')

                    # extract datasrc from "datasrc[data_format]"

            in_label, in_root, prefixes = maps.col_info(in_label)

            in_label_ref[0] = in_label

            out_label, out_root = maps.col_info(out_label)[:2]

            if out_is_xpaths: has_types = out_root.find('/*s/') >= 0

                # outer elements are types

            for row in reader:

            stream.close()

        out_is_xml_ref[0] = out_is_xpaths and not out_is_db

        def process_rows(process_row, rows, rows_start=0):

            @param process_row(in_row, i)

            @rows_start The (0-based) row # of the first row in rows. Set this

            '''

            rows = iter(rows)

                try: row = rows.next()

                process_row(row, i)

            row_ct = i-start

            return row_ct

        def map_rows(get_value, rows, **kw_args):

            '''Maps input rows

                row_id = str(i)

                    if value == None:

                        log_debug('Getting '+str(in_))

                        value = cleanup(get_value(in_, row))

                    if value != None:

                        log_debug('Putting '+str(out))

                        xpath.put_obj(root, out, row_id, has_types, value)

            return process_rows(process_row, rows, **kw_args)

        def map_table(col_names, rows, **kw_args):

            col_names_ct = len(col_names)

            col_idxs = util.list_flip(col_names)

                    try: mappings[i] = (

                i += 1

                return util.coalesce(*util.list_subset(row.list, in_))

            def wrap_row(row):

            return map_rows(get_value, util.WrapIter(wrap_row, rows), **kw_args)

        stdin = streams.LineCountStream(sys.stdin)

        if in_is_db:

            in_db = connect_db(in_db_config)

            # Get table and schema name

            if by_col and in_db_config == out_db_config:

                raise NotImplementedError(

            in_db.db.close()

        elif in_is_xml:

            def get_rows(doc2rows):

                return iters.flatten(itertools.imap(doc2rows,

            if map_path == None:

                def doc2rows(in_xml_root):

                    iter_ = xml_dom.NodeElemIter(in_xml_root)

                    return iter_

            else:

                def doc2rows(in_xml_root):

                    rows = xpath.get(in_xml_root, in_root, limit=end)

                    if rows == []: raise SystemExit('Map error: Root "'

                row_ct = map_rows(get_value, get_rows(doc2rows))

        else: # input is CSV

            map_ = dict(mappings)

            reader, col_names = csvs.reader_and_header(sys.stdin)

            row_ct = map_table(col_names, reader)

    def process_inputs(root, row_ready):

        row_ct = 0

    pool.share_vars(locals())

    if out_is_db:

        import db_xml

        out_db = connect_db(out_db_config)

        try:

            if redo: sql.empty_db(out_db)

            row_ins_ct_ref = [0]

            pool.share_vars(locals())

            def row_ready(row_num, input_row):

                def on_error(e):

                    exc.add_msg(e, term.emph('row #:')+' '+str(row_num+1))

                    exc.add_msg(e, term.emph('input row:')+'\n'+str(input_row))

                    ex_tracker.track(e, row_num)

                pool.share_vars(locals())

                xml_func.process(root, on_error)

                if not xml_dom.is_empty(root):

                        sql.with_savepoint(out_db,

                                row_ins_ct_ref, on_error))

                        if commit: out_db.db.commit()

                    except sql.DatabaseErrors, e: on_error(e)

                prep_root()

            row_ct = process_inputs(root, row_ready)

                ' new rows into database\n')

        finally:

            out_db.db.rollback()

    else:

        def on_error(e): ex_tracker.track(e)

        def row_ready(row_num, input_row): pass

        row_ct = process_inputs(root, row_ready)

        xml_func.process(root, on_error)

        if out_is_xml_ref[0]:

    # Consume any asynchronous tasks not already consumed above

    pool.main_loop()

    profiler.stop(row_ct)

    if verbose:

    ex_tracker.exit()

def main():

    except Parser.SyntaxError, e: raise SystemExit(str(e))

if __name__ == '__main__':

    profile_to = opts.get_env_var('profile_to', None)

        sys.stderr.write('Profiling to '+profile_to+'\n')

        cProfile.run(main.func_code, profile_to)