#!/usr/bin/env python

# Maps one datasource to another, using a map spreadsheet if needed

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

# Multi-safe (supports an input appearing multiple times).

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

# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml

# Duplicate-column safe (supports multiple columns of the same name, which will

# be combined)

# Case- and punctuation-insensitive.

import copy

import csv

import itertools

import os.path

import warnings

import sys

import xml.dom.minidom as minidom

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

import csvs

import db_xml

import exc

import iters

import maps

import opts

import parallelproc

import Parser

import profiling

import sql

import sql_gen

import sql_io

import streams

import strings

import term

import util

import xpath

import xml_dom

import xml_func

import xml_parse

metadata_prefix = ':'

collision_suffix = '/_alt/'

def get_with_prefix(map_, prefixes, key):

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

    @return tuple(found_key, found_value)

    '''

    values = []

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

        try: value = map_[key_]

        except KeyError, e: continue # keep going

        values.append((key_, value))

    if values != []: return values

    else: raise e # re-raise last KeyError

def is_metadata(str_): return str_.startswith(metadata_prefix)

def metadata_value(name):

    removed_ref = [False]

    name = strings.remove_prefix(metadata_prefix, name, removed_ref)

    if removed_ref[0]: return name

    else: return None

def cleanup(val):

    if val == None: return val

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

def main_():

    env_names = []

    def usage_err():

        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', False, env_names) and not commit

        # never redo in commit mode (run `make schemas/reinstall` instead)

    # Ranges

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

    # Make start interally 0-based.

    # It's 1-based to the user to match up with the staging table row #s.

    start -= 1

    if test: n_default = 1

    else: n_default = None

    n = util.cast(int, util.none_if(opts.get_env_var('n', n_default, env_names),

        u''))

    end = n

    if end != None: end += start

    # Debugging

    verbosity = util.cast(float, opts.get_env_var('verbosity', None, env_names))

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

    # DB

    def get_db_config(prefix):

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

    in_db_config = get_db_config('in')

    out_db_config = get_db_config('out')

    in_is_db = 'engine' in in_db_config

    out_is_db = 'engine' in out_db_config

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

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

    if in_schema != None:

        for config in [in_db_config, out_db_config]:

            config['schemas'] += ','+in_schema

    # Optimization

    cache_sql = opts.env_flag('cache_sql', True, env_names)

    by_col = in_db_config == out_db_config and opts.env_flag('by_col', False,

        env_names) # by-column optimization only applies if mapping to same DB

    if test: cpus_default = 0

    else: cpus_default = 0 # or None to use parallel processing by default

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

        env_names), u''))

    # Set default verbosity. Must happen after by_col is set.

    if verbosity == None:

        if test: verbosity = 0.5 # automated tests should not be verbose

        elif by_col: verbosity = 3 # show all queries to assist profiling

        else: verbosity = 1.1 # just show row progress

    ##

    # Logging

    verbose_errors = test and verbosity > 0

    debug = verbosity >= 1.5

    def log(msg, level=1):

        '''Higher level -> more verbose'''

        if level <= verbosity:

            if verbosity <= 2:

                if level == 1.5: msg = '# '+msg # msg is Redmine list item

                elif msg.startswith('DB query:'): # remove extra debug info

                    first_line, nl, msg = msg.partition('\n')

            elif level > 1: msg = '['+str(level)+'] '+msg # include level in msg

            sys.stderr.write(strings.to_raw_str(msg.rstrip('\n')+'\n'))

    if debug: log_debug = lambda msg, level=2: log(msg, level)

    else: log_debug = sql.log_debug_none

    # Parse args

    map_paths = sys.argv[1:]

    if map_paths == []:

        if in_is_db or not out_is_db: usage_err()

        else: map_paths = [None]

    def connect_db(db_config):

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

        return sql.connect(db_config, caching=cache_sql, autocommit=commit,

            debug_temp=verbosity > 3 and commit, log_debug=log_debug)

    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')

    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')

    # Parallel processing

    pool = parallelproc.MultiProducerPool(cpus)

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

    # Set up DB access

    row_ins_ct_ref = [0]

    if out_is_db:

        out_db = connect_db(out_db_config)

        def is_rel_func(name):

            return (name in db_xml.put_special_funcs

                or sql.function_exists(out_db, sql_gen.Function(name)))

    doc = xml_dom.create_doc()

    root = doc.documentElement

    out_is_xml_ref = [False]

    in_label_ref = [None]

    col_defaults = {}

    def update_in_label():

        if in_label_ref[0] != None:

            os.environ['source'] = in_label_ref[0]

    def prep_root():

        root.clear()

        update_in_label()

    prep_root()

    def process_input(root, row_ready, map_path):

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

        # Load map header

        in_is_xpaths = True

        out_is_xpaths = True

        out_label = None

        if map_path != None:

            metadata = []

            mappings = []

            stream = open(map_path, 'rb')

            reader = csv.reader(stream)

            in_label, out_label = reader.next()[:2]

            def split_col_name(name):

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

                return label, sep != '', root, []

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

            in_is_xpaths = in_root != None

            in_label_ref[0] = in_label

            update_in_label()

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

            out_is_xpaths = out_root != None

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

                # outer elements are types

            for row in reader:

                in_, out = row[:2]

                if out != '': mappings.append([in_, out_root+out])

            stream.close()

            root.ownerDocument.documentElement.tagName = out_label

        in_is_xml = in_is_xpaths and not in_is_db

        out_is_xml_ref[0] = out_is_xpaths and not out_is_db

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

            '''Processes input rows

            @param process_row(in_row, i)

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

                only if the pre-start rows have already been skipped.

            '''

            rows = iter(rows)

            if end != None: row_nums = xrange(rows_start, end)

            else: row_nums = itertools.count(rows_start)

            i = -1

            for i in row_nums:

                try: row = rows.next()

                except StopIteration:

                    i -= 1 # last row # didn't count

                    break # no more rows

                if i < start: continue # not at start row yet

                # Row # is interally 0-based, but 1-based to the user

                log('Processing input row #'+str(i+1), level=1.1)

                process_row(row, i)

                row_ready(i, row)

            row_ct = i-start+1

            return row_ct

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

            '''Maps input rows

            @param get_value(in_, row):str

            '''

            # Prevent collisions if multiple inputs mapping to same output

            outputs_idxs = dict()

            for i, mapping in enumerate(mappings):

                in_, out = mapping

                default = util.NamedTuple(count=1, first=i)

                idxs = outputs_idxs.setdefault(out, default)

                if idxs is not default: # key existed, so there was a collision

                    if idxs.count == 1: # first key does not yet have suffix

                        mappings[idxs.first][1] += collision_suffix+'0'

                    mappings[i][1] += collision_suffix+str(idxs.count)

                    idxs.count += 1

            id_node = None

            if out_is_db:

                mappings_orig = mappings[:] # save a copy

                mappings[:] = [] # empty existing elements

                for in_, out in mappings_orig:

                    in_str = strings.ustr(in_)

                    is_metadata_ = is_metadata(in_str)

                    if is_metadata_: value = metadata_value(in_str)

                    else: value = '$'+in_str # mark as name

                    # All put_obj()s should return the same id_node

                    nodes, id_node = xpath.put_obj(root, out, '-1', has_types,

                        value) # value is placeholder that documents name

                    if not is_metadata_: mappings.append([in_, nodes])

                if id_node == None:

                    warnings.warn(UserWarning('Map warning: No mappings'))

                xml_func.simplify(root)

                sys.stdout.write(strings.to_raw_str('Put template:\n'

                    +strings.ustr(root)))

                sys.stdout.flush()

            def process_row(row, i):

                row_id = str(i)

                if id_node != None: xml_dom.set_value(id_node, row_id)

                for in_, out in mappings:

                    log_debug('Getting '+strings.ustr(in_))

                    value = cleanup(get_value(in_, row))

                    log_debug('Putting '+strings.urepr(value)+' to '

                        +strings.ustr(out))

                    if out_is_db: # out is list of XML nodes

                        for node in out: xml_dom.set_value(node, value)

                    elif value != None: # out is XPath

                        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)

            col_names_map = dict(zip(col_names, col_names))

            prefixes_simp = map(maps.simplify, prefixes)

            # Resolve prefixes

            mappings_orig = mappings[:] # save a copy

            mappings[:] = [] # empty existing elements

            for in_, out in mappings_orig:

                if is_metadata(in_): mappings.append([in_, out])

                else:

                    try:

                        cols = get_with_prefix(col_names_map, prefixes_simp,

                            in_)

                    except KeyError: pass

                    else:

                        cols = [(orig, col_idxs[orig]) for simp, orig in cols]

                        mappings[len(mappings):] = [[db_xml.ColRef(*col), out]

                            for col in cols] # can't use += because that uses =

            def get_value(in_, row): return row.list[in_.idx]

            def wrap_row(row):

                return util.ListDict(util.list_as_length(row, col_names_ct),

                    col_names, col_idxs) # handle CSV rows of different lengths

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

        if in_is_db:

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

            if by_col: in_db = out_db

            else: in_db = connect_db(in_db_config)

            # Get table and schema name

            schema = in_schema # modified, so can't have same name as outer var

            table = in_table # modified, so can't have same name as outer var

            if table == None:

                assert in_is_xpaths

                schema, sep, table = in_root.partition('.')

                if sep == '': # only the table name was specified

                    table = schema

                    schema = None

            table = sql_gen.Table(table, schema)

            # Fetch rows

            if by_col: limit = 0 # only fetch column names

            else: limit = n

            try:

                cur = sql.select(in_db, table, limit=limit, start=start,

                    recover=True, cacheable=False)

            except sql.DoesNotExistException:

                table = None

                col_names = []

                rows = []

            else:

                col_names = list(sql.col_names(cur))

                rows = sql.rows(cur)

            if by_col:

                map_table(col_names, []) # just create the template

                if table != None and start == 0 and n == None: # full re-import

                    log('Clearing errors table')

                    errors_table_ = sql_io.errors_table(in_db, table)

                    if errors_table_ != None:

                        sql.drop_table(in_db, errors_table_)

                # Strip XML functions not in the DB

                xml_func.process(root, is_rel_func=is_rel_func)

                if debug: log_debug('Putting stripped:\n'+strings.ustr(root))

                    # only calc if debug

                # Import rows

                in_row_ct_ref = [0]

                db_xml.put_table(in_db, root.firstChild, table, in_row_ct_ref,

                    row_ins_ct_ref, n, start, on_error, col_defaults)

                row_ct = in_row_ct_ref[0]

            else:

                # Use normal by-row method

                row_ct = map_table(col_names, rows, rows_start=start)

                    # rows_start: pre-start rows have been skipped

                in_db.db.close()

        elif in_is_xml:

            stdin = streams.LineCountStream(sys.stdin)

            def on_error(e):

                exc.add_msg(e, term.emph('input line #:')+' '

                    +str(stdin.line_num))

                ex_tracker.track(e)

            def get_rows(doc2rows):

                return iters.flatten(itertools.imap(doc2rows,

                    xml_parse.docs_iter(stdin, on_error)))

            if map_path == None:

                def doc2rows(in_xml_root):

                    iter_ = xml_dom.NodeElemIter(in_xml_root)

                    util.skip(iter_, xml_dom.is_text) # skip metadata

                    return iter_

                row_ct = process_rows(lambda row, i: root.appendChild(row),

                    get_rows(doc2rows))

            else:

                def doc2rows(in_xml_root):

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

                    if rows == []: warnings.warn(UserWarning('Map warning: '

                        'Root "'+in_root+'" not found in input'))

                    return rows

                def get_value(in_, row):

                    in_ = './{'+(','.join(strings.with_prefixes(

                        ['']+prefixes, in_)))+'}' # also with no prefix

                    nodes = xpath.get(row, in_, allow_rooted=False)

                    if nodes != []: return xml_dom.value(nodes[0])

                    else: return None

                row_ct = map_rows(get_value, get_rows(doc2rows))

        else: # input is CSV

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

            row_ct = map_table(col_names, reader)

        return row_ct

    def process_inputs(root, row_ready):

        row_ct = 0

        for map_path in map_paths:

            row_ct += process_input(root, row_ready, map_path)

        return row_ct

    pool.share_vars(locals())

    if out_is_db:

        try:

            if redo: sql.empty_db(out_db)

            pool.share_vars(locals())

            def row_ready(row_num, input_row):

                row_str_ = [None]

                def row_str():

                    if row_str_[0] == None:

                        # Row # is interally 0-based, but 1-based to the user

                        row_str_[0] = (term.emph('row #:')+' '+str(row_num+1)

                            +'\n'+term.emph('input row:')+'\n'

                            +strings.ustr(input_row))

                        if verbose_errors: row_str_[0] += ('\n'

                            +term.emph('output row:')+'\n'+strings.ustr(root))

                    return row_str_[0]

                if debug: log_debug(row_str()) # only calc if debug

                def on_error(e):

                    exc.add_msg(e, row_str())

                    ex_tracker.track(e, row_num, detail=verbose_errors)

                pool.share_vars(locals())

                row_root = root.cloneNode(True) # deep copy so don't modify root

                xml_func.process(row_root, on_error, is_rel_func, out_db)

                if debug: log_debug('Putting processed:\n'

                    +strings.ustr(row_root)) # only calc if debug

                if not xml_dom.is_empty(row_root):

                    assert xml_dom.has_one_child(row_root)

                    try:

                        sql.with_savepoint(out_db,

                            lambda: db_xml.put(out_db, row_root.firstChild,

                                row_ins_ct_ref, on_error, col_defaults))

                    except sql.DatabaseErrors, e: on_error(e)

            row_ct = process_inputs(root, row_ready)

            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+

                ' new rows into database\n')

            sys.stdout.flush()

            # Consume asynchronous tasks

            pool.main_loop()

        finally: out_db.close()

    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]:

            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)

        else: # output is CSV

            raise NotImplementedError('CSV output not supported yet')

    # Consume any asynchronous tasks not already consumed above

    pool.main_loop()

    profiler.stop(row_ct)

    if not by_col: ex_tracker.add_iters(row_ct) # only if errors are done by row

    log('Processed '+str(row_ct)+' input rows')

    log(profiler.msg())

    log(ex_tracker.msg())

    ex_tracker.exit()

def main():

    try: main_()

    except Parser.SyntaxError, e: raise SystemExit(strings.ustr(e))

if __name__ == '__main__':

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

    if profile_to != None:

        import cProfile

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

        cProfile.run(main.func_code, profile_to)

    else: main()