# Database import/export

import dicts

import sql

import sql_gen

import strings

import util

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)

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

    @param errors_table None|sql_gen.Table|str

    '''

    col = sql_gen.as_Col(col)

    if isinstance(col, sql_gen.Literal): return sql_gen.Cast(type_, 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))

    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 += '''\

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 += '''\

        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]),

                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;

'''

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

    # Add column

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

    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

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

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

            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)

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

    except IndexError: in_tables0 = None

    else:

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

        in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)

    # Determine if can use optimization for only literal values

    is_literals = not reduce(operator.or_, map(sql_gen.is_table_col,

    else: errors_table_ = errors_table(db, in_tables0)

    # Create input joins from list of input tables

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

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

        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 so don't modify input and for speed

        in_table = sql_gen.Table('in')

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

            in_cols, preserve=[in_pkey_col]))

        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

    elif is_function: full_in_table = in_table

        full_in_table = sql_gen.suffixed_table(in_table, '_full')

        full_in_table_cols = [in_pkey_col]

        if default != None:

            full_in_table_cols.append(default)

            default = sql_gen.with_table(default, full_in_table)

        sql.run_query_into(db, sql.mk_select(db, in_table, full_in_table_cols,

            order_by=None), into=full_in_table, add_pkey_=True)

        pkeys_names = [in_pkey, out_pkey]

        pkeys_cols = [in_pkey_col, out_pkey_col]

    pkeys_table_exists_ref = [False]

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

    limit_ref = [None]

        return sql.mk_select(db, joins, cols, limit=limit_ref[0], order_by=None)

    else:

        insert_in_table = in_table

        insert_in_tables = [insert_in_table]

    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

        out_table_cols = sql_gen.ColDict(db, out_table)

        out_table_cols.update(util.dict_subset_right_join({},

            sql.table_cols(db, out_table)))

        in_cols = []

        cond = sql.map_expr(db, cond, mapping, in_cols)

        cond = sql.map_expr(db, cond, out_table_cols)

        track_data_error(db, errors_table_, sql_gen.cols_srcs(in_cols), None,

            e.cause.pgcode,

            strings.ensure_newl(e.cause.pgerror)+'condition: '+cond)

        not_cond = sql_gen.NotCond(sql_gen.CustomCode(cond))

        log_debug('Ignoring rows where '+strings.as_tt(not_cond.to_str(db)))

        sql.delete(db, insert_in_table, not_cond)

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

            +strings.as_tt(repr(value)))

            # Try just mapping the value to NULL

            sql.update(db, insert_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

    while True:

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

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

        has_joins = join_cols != {}

        log_debug('Trying to insert new rows')

        # Prepare to insert new rows

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

            if has_joins:

                insert_args.update(dict(ignore=True))

            else:

                if not is_literals:

                    insert_args.update(dict(into=insert_out_pkeys))

                c, insert_in_table) for c in mapping.values()])

            cur = None

                else:

                    insert_into_pkeys(input_joins, [in_pkey_col, func_call],

                        recover=True)

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

            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)

            log_exc(e)

            # supported

            assert not join_cols

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

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

                return sql.value(sql.select(db, out_table, [out_pkey_col],

                    mapping, order_by=None))

            insert_in_table = sql.distinct_table(db, insert_in_table,

            log_exc(e)

            out_col, = e.cols

            try: in_col = mapping[out_col]

            except KeyError:

                log_debug(msg)

                if default == None: on_error(SyntaxError(msg)) # required col

            log_exc(e)

            ignore_cond(e.cond, e)

            log_exc(e)

            for in_col in mapping.values(): ignore(in_col, e.value, 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

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

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

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

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

        # since the SELECT query is identical to the one used in INSERT SELECT,

        # its rows will be retrieved in the same order.

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

            +strings.as_tt(repr(default)))

            # must use join_same_not_null or query will take forever

        insert_into_pkeys(missing_rows_joins,

            sql_gen.NamedCol(out_pkey, default)])

    else:

        assert (sql.table_row_count(db, into)

            == sql.table_row_count(db, full_in_table))

        sql.empty_temp(db, insert_in_tables+[full_in_table])

        srcs = []

        if is_func: srcs = sql_gen.cols_srcs(in_cols)

        return sql_gen.Col(out_pkey, into, srcs)