# Database import/export

import exc

import dicts

import sql

import sql_gen

import strings

import util

def put(db, table, row, pkey_=None, row_ct_ref=None):

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    '''

    row = sql_gen.ColDict(db, table, row)

    if pkey_ == None: pkey_ = sql.pkey(db, table, recover=True)

    try:

        cur = sql.insert(db, table, row, pkey_, recover=True)

        if row_ct_ref != None and cur.rowcount >= 0:

            row_ct_ref[0] += cur.rowcount

        return sql.value(cur)

    except sql.DuplicateKeyException, e:

        row = sql_gen.ColDict(db, table,

            util.dict_subset_right_join(row, e.cols))

        return sql.value(sql.select(db, table, [pkey_], row, recover=True))

def get(db, table, row, pkey, row_ct_ref=None, create=False):

    '''Recovers from errors'''

    try:

        return sql.value(sql.select(db, table, [pkey], row, limit=1,

            recover=True))

    except StopIteration:

        if not create: raise

        return put(db, table, row, pkey, row_ct_ref) # insert new row

def is_func_result(col):

    return col.table.name.find('(') >= 0 and col.name == 'result'

def into_table_name(out_table, in_tables0, mapping, is_func):

    def in_col_str(in_col):

        in_col = sql_gen.remove_col_rename(in_col)

        if isinstance(in_col, sql_gen.Col):

            table = in_col.table

            if table == in_tables0:

                in_col = sql_gen.to_name_only_col(in_col)

            elif is_func_result(in_col): in_col = table # omit col name

        return str(in_col)

    str_ = str(out_table)

    if is_func:

        str_ += '('

        try: value_in_col = mapping['value']

        except KeyError:

            str_ += ', '.join((str(k)+'='+in_col_str(v)

                for k, v in mapping.iteritems()))

        else: str_ += in_col_str(value_in_col)

        str_ += ')'

    else:

        out_col = 'rank'

        try: in_col = mapping[out_col]

        except KeyError: str_ += '_pkeys'

        else: # has a rank column, so hierarchical

            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'

    return str_

def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,

    default=None, is_func=False, on_error=exc.raise_):

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    IMPORTANT: Must be run at the *beginning* of a transaction.

    @param in_tables The main input table to select from, followed by a list of

        tables to join with it using the main input table's pkey

    @param mapping dict(out_table_col=in_table_col, ...)

        * out_table_col: str (*not* sql_gen.Col)

        * in_table_col: sql_gen.Col|literal-value

    @param into The table to contain the output and input pkeys.

        Defaults to `out_table.name+'_pkeys'`.

    @param default The *output* column to use as the pkey for missing rows.

        If this output column does not exist in the mapping, uses None.

    @param is_func Whether out_table is the name of a SQL function, not a table

    @return sql_gen.Col Where the output pkeys are made available

    '''

    out_table = sql_gen.as_Table(out_table)

    def log_debug(msg): db.log_debug(msg, level=1.5)

    def col_ustr(str_):

        return strings.repr_no_u(sql_gen.remove_col_rename(str_))

    log_debug('********** New iteration **********')

    log_debug('Inserting these input columns into '+strings.as_tt(

        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))

    is_function = sql.function_exists(db, out_table)

    if is_function: out_pkey = 'result'

    else: out_pkey = sql.pkey(db, out_table, recover=True)

    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)

    if mapping == {}: # need at least one column for INSERT SELECT

        mapping = {out_pkey: None} # ColDict will replace with default value

    # Create input joins from list of input tables

    in_tables_ = in_tables[:] # don't modify input!

    in_tables0 = in_tables_.pop(0) # first table is separate

    errors_table_ = sql.errors_table(db, in_tables0)

    in_pkey = sql.pkey(db, in_tables0, recover=True)

    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)

    input_joins = [in_tables0]+[sql_gen.Join(v,

        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]

    if into == None:

        into = into_table_name(out_table, in_tables0, mapping, is_func)

    into = sql_gen.as_Table(into)

    # Set column sources

    in_cols = filter(sql_gen.is_table_col, mapping.values())

    for col in in_cols:

        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)

    log_debug('Joining together input tables into temp table')

    # Place in new table for speed and so don't modify input if values edited

    in_table = sql_gen.Table('in')

    mapping = dicts.join(mapping, sql.flatten(db, in_table, input_joins,

        in_cols, preserve=[in_pkey_col], start=0))

    input_joins = [in_table]

    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)

    mapping = sql_gen.ColDict(db, out_table, mapping)

        # after applying dicts.join() because that returns a plain dict

    # Resolve default value column

    if default != None:

        try: default = mapping[default]

        except KeyError:

            db.log_debug('Default value column '

                +strings.as_tt(strings.repr_no_u(default))

                +' does not exist in mapping, falling back to None', level=2.1)

            default = None

    pkeys_names = [in_pkey, out_pkey]

    pkeys_cols = [in_pkey_col, out_pkey_col]

    pkeys_table_exists_ref = [False]

    def insert_into_pkeys(joins, cols, distinct=False):

        kw_args = {}

        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))

        query = sql.mk_select(db, joins, cols, order_by=None, start=0,

            **kw_args)

        if pkeys_table_exists_ref[0]:

            sql.insert_select(db, into, pkeys_names, query)

        else:

            sql.run_query_into(db, query, into=into)

            pkeys_table_exists_ref[0] = True

    limit_ref = [None]

    conds = set()

    distinct_on = sql_gen.ColDict(db, out_table)

    def mk_main_select(joins, cols):

        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]

        return sql.mk_select(db, joins, cols, conds, distinct_on_cols,

            limit=limit_ref[0], start=0)

    exc_strs = set()

    def log_exc(e):

        e_str = exc.str_(e, first_line_only=True)

        log_debug('Caught exception: '+e_str)

        assert e_str not in exc_strs # avoid infinite loops

        exc_strs.add(e_str)

    def remove_all_rows():

        log_debug('Ignoring all rows')

        limit_ref[0] = 0 # just create an empty pkeys table

    def ignore(in_col, value, e):

        sql.track_data_error(db, errors_table_, in_col.srcs, value,

            e.cause.pgcode, e.cause.pgerror)

        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '

            +strings.as_tt(repr(value)))

    def remove_rows(in_col, value, e):

        ignore(in_col, value, e)

        cond = (in_col, sql_gen.CompareCond(value, '!='))

        assert cond not in conds # avoid infinite loops

        conds.add(cond)

    def invalid2null(in_col, value, e):

        ignore(in_col, value, e)

        sql.update(db, in_table, [(in_col, None)],

            sql_gen.ColValueCond(in_col, value))

    def insert_pkeys_table(which):

        return sql_gen.Table(sql_gen.concat(in_table.name,

            '_insert_'+which+'_pkeys'))

    insert_out_pkeys = insert_pkeys_table('out')

    insert_in_pkeys = insert_pkeys_table('in')

    # Do inserts and selects

    join_cols = sql_gen.ColDict(db, out_table)

    while True:

        if limit_ref[0] == 0: # special case

            log_debug('Creating an empty pkeys table')

            cur = sql.run_query_into(db, sql.mk_select(db, out_table,

                [out_pkey], limit=limit_ref[0]), into=insert_out_pkeys)

            break # don't do main case

        has_joins = join_cols != {}

        log_debug('Trying to insert new rows')

        # Prepare to insert new rows

        insert_joins = input_joins[:] # don't modify original!

        insert_args = dict(recover=True, cacheable=False)

        if has_joins:

            insert_args.update(dict(ignore=True))

        else:

            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))

        main_select = mk_main_select(insert_joins, mapping.values())

        def main_insert():

            if is_function:

                log_debug('Calling function on input rows')

                args = dict(((k.name, v) for k, v in mapping.iteritems()))

                func_call = sql_gen.NamedCol(out_pkey,

                    sql_gen.FunctionCall(out_table, **args))

                insert_into_pkeys(input_joins, [in_pkey_col, func_call])

                return None

            else:

                return sql.insert_select(db, out_table, mapping.keys(),

                    main_select, **insert_args)

        try:

            cur = sql.with_savepoint(db, main_insert)

            break # insert successful

        except sql.MissingCastException, e:

            log_exc(e)

            out_col = e.col

            type_ = e.type

            log_debug('Casting '+strings.as_tt(out_col)+' input to '

                +strings.as_tt(type_))

            mapping[out_col] = sql.cast_temp_col(db, type_, mapping[out_col],

                errors_table_)

        except sql.DuplicateKeyException, e:

            log_exc(e)

            old_join_cols = join_cols.copy()

            distinct_on.update(util.dict_subset(mapping, e.cols))

            join_cols.update(util.dict_subset_right_join(mapping, e.cols))

            log_debug('Ignoring existing rows, comparing on these columns:\n'

                +strings.as_inline_table(join_cols, ustr=col_ustr))

            assert join_cols != old_join_cols # avoid infinite loops

        except sql.NullValueException, e:

            log_exc(e)

            out_col, = e.cols

            try: in_col = mapping[out_col]

            except KeyError:

                log_debug('Missing mapping for NOT NULL column '+out_col)

                remove_all_rows()

            else: remove_rows(in_col, None, e)

        except sql.FunctionValueException, e:

            log_exc(e)

            func_name = e.name

            value = e.value

            for out_col, in_col in mapping.iteritems():

                in_col = sql_gen.unwrap_func_call(in_col, func_name)

                invalid2null(in_col, value, e)

        except sql.DatabaseErrors, e:

            log_exc(e)

            log_debug('No handler for exception')

            on_error(e)

            remove_all_rows()

        # after exception handled, rerun loop with additional constraints

    if cur != None and row_ct_ref != None and cur.rowcount >= 0:

        row_ct_ref[0] += cur.rowcount

    if is_function: pass # pkeys table already created

    elif has_joins:

        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]

        log_debug('Getting output table pkeys of existing/inserted rows')

        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)

    else:

        sql.add_row_num(db, insert_out_pkeys) # for joining with input pkeys

        log_debug('Getting input table pkeys of inserted rows')

        sql.run_query_into(db, mk_main_select(input_joins, [in_pkey]),

            into=insert_in_pkeys)

        sql.add_row_num(db, insert_in_pkeys) # for joining with output pkeys

        assert sql.table_row_count(db, insert_out_pkeys) == sql.table_row_count(

            db, insert_in_pkeys)

        log_debug('Combining output and input pkeys in inserted order')

        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,

            {sql.row_num_col: sql_gen.join_same_not_null})]

        insert_into_pkeys(pkey_joins, pkeys_names)

        sql.empty_temp(db, [insert_out_pkeys, insert_in_pkeys])

    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)

    sql.add_pkey(db, into)

    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))

    missing_rows_joins = input_joins+[sql_gen.Join(into,

        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]

        # must use join_same_not_null or query will take forever

    insert_into_pkeys(missing_rows_joins,

        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])

    assert sql.table_row_count(db, into) == sql.table_row_count(db, in_table)

    sql.empty_temp(db, in_table)

    srcs = []

    if is_func: srcs = sql_gen.cols_srcs(in_cols)

    return sql_gen.Col(out_pkey, into, srcs)