# 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

    '''Handles a data_exception by saving the error or converting it to a

    warning, and returning NULL.

    @param srcs The column names for the errors table

    @param errors_table None|sql_gen.Table

    '''

    save_errors = errors_table != None and srcs

    handler = ''

    if save_errors:

        errors_table_cols = map(sql_gen.Col,

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

        handler += '''\

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

    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;

'''

    else:

        handler += '''\

RAISE WARNING '%', SQLERRM;

'''

    handler += '''\

RETURN NULL;

'''

    return sql_gen.ExcHandler('data_exception', handler)

    '''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 != ())

    if save_errors:

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

        query += data_exception_handler(db, srcs, errors_table).to_str(db, '''\

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

    of the function, outside any block.) */

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

'''

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