# Database import/export

import exc

import dicts

import sql

import sql_gen

import strings

import util

##### Data cleanup

def cleanup_table(db, table, cols):

    table = sql_gen.as_Table(table)

    cols = map(sql_gen.as_Col, cols)

    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '

        +db.esc_value(r'\N')+')')

    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))

        for v in cols]

    sql.update(db, table, changes, in_place=True)

##### Error tracking

def track_data_error(db, errors_table, cols, value, error_code, error):

    '''

    @param errors_table If None, does nothing.

    '''

    if errors_table == None or cols == (): return

    for col in cols:

        try:

            sql.insert(db, errors_table, dict(column=col.name, value=value,

                error_code=error_code, error=error), recover=True,

                cacheable=True, log_level=4)

        except sql.DuplicateKeyException: pass

def cast(db, type_, col, errors_table=None):

    '''Casts an (unrenamed) column or value.

    If errors_table set and col has srcs, saves errors in errors_table (using

    col's srcs attr as the source columns) and converts errors to warnings.

    @param col str|sql_gen.Col|sql_gen.Literal

    @param errors_table None|sql_gen.Table|str

    '''

    col = sql_gen.as_Col(col)

    # Don't convert exceptions to warnings for user-supplied constants

    if isinstance(col, sql_gen.Literal): return sql_gen.Cast(type_, col)

    assert not isinstance(col, sql_gen.NamedCol)

    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)

        and col.srcs != ())

    function_name = type_

    if save_errors:

        errors_table = sql_gen.as_Table(errors_table)

        srcs = map(sql_gen.to_name_only_col, col.srcs)

        function_name = str(sql_gen.FunctionCall(function_name, *srcs))

    function = db.TempFunction(function_name)

    while True:

        # Create function definition

        query = '''\

CREATE FUNCTION '''+function.to_str(db)+'''(value text)

RETURNS '''+type_+'''

LANGUAGE plpgsql

'''

        if not save_errors: query += 'IMMUTABLE '

        query += '''\

STRICT

AS $$

BEGIN

    /* The explicit cast to the return type is needed to make the cast happen

    inside the try block. (Implicit casts to the return type happen at the end

    of the function, outside any block.) */

    RETURN value::'''+type_+''';

EXCEPTION

    WHEN data_exception THEN

'''

        if save_errors:

            errors_table_cols = map(sql_gen.Col,

                ['column', 'value', 'error_code', 'error'])

            query += '''\

        -- Save error in errors table.

        DECLARE

            error_code text := SQLSTATE;

            error text := SQLERRM;

            "column" text;

        BEGIN

            -- Insert the value and error for *each* source column.

            FOR "column" IN

'''+sql.mk_select(db, sql_gen.NamedValues('c', None, [[c.name] for c in srcs]),

    order_by=None, start=0)+'''

            LOOP

                BEGIN

'''+sql.mk_insert_select(db, errors_table, errors_table_cols,

    sql_gen.Values(errors_table_cols).to_str(db))+''';

                EXCEPTION

                    WHEN unique_violation THEN NULL; -- continue to next row

                END;

            END LOOP;

        END;

'''

        query += '''\

        RAISE WARNING '%', SQLERRM;

        RETURN NULL;

END;

$$;

'''

        # Create function

        try:

            sql.run_query(db, query, recover=True, cacheable=True,

                log_ignore_excs=(sql.DuplicateException,))

            break # successful

        except sql.DuplicateException:

            function.name = sql.next_version(function.name)

            # try again with next version of name

    return sql_gen.FunctionCall(function, col)

def cast_temp_col(db, type_, col, errors_table=None):

    '''Like cast(), but creates a new column with the cast values if the input

    is a column.

    @return The new column or cast value

    '''

    def cast_(col): return cast(db, type_, col, errors_table)

    try: col = sql_gen.underlying_col(col)

    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)

    table = col.table

    new_col = sql_gen.suffixed_col(col, '::'+strings.first_word(type_))

    expr = cast_(col)

    # Add column

    new_typed_col = sql_gen.TypedCol(new_col.name, type_)

    sql.add_col(db, table, new_typed_col, comment=repr(col)+'::'+type_)

    new_col.name = new_typed_col.name # propagate any renaming

    sql.update(db, table, [(new_col, expr)], in_place=True, recover=True)

    return new_col

def errors_table(db, table, if_exists=True):

    '''

    @param if_exists If set, returns None if the errors table doesn't exist

    @return None|sql_gen.Table

    '''

    table = sql_gen.as_Table(table)

    if table.srcs != (): table = table.srcs[0]

    errors_table = sql_gen.suffixed_table(table, '.errors')

    if if_exists and not sql.table_exists(db, errors_table): return None

    return errors_table

##### Import

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, log_level=3.5)

        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,

            log_level=3.5))

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_ = 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]

    insert_in_table = in_table

    insert_in_tables = [insert_in_table]

    conds = set()

    distinct_on = sql_gen.ColDict(db, out_table)

    def mk_main_select(joins, cols):

        conds_ = [(sql_gen.with_table(k, insert_in_table), v) for k, v in conds]

        return sql.mk_select(db, joins, cols, conds_, 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):

        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_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_in_table],

            [sql_gen.with_table(c, insert_in_table) for c in 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_))

            in_col = mapping[out_col]

            while True:

                try:

                    mapping[out_col] = cast_temp_col(db, type_, in_col,

                        errors_table_)

                    break # cast successful

                except sql.InvalidValueException, e:

                    log_exc(e)

                    invalid2null(in_col, e.value, e)

        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

            # Uniquify input table to avoid internal duplicate keys

            insert_in_table = sql.distinct_table(db, insert_in_table,

                filter(sql_gen.is_table_col, distinct_on.values()))

            insert_in_tables.append(insert_in_table)

        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.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, insert_in_tables)

    srcs = []

    if is_func: srcs = sql_gen.cols_srcs(in_cols)

    return sql_gen.Col(out_pkey, into, srcs)