# Database import/export

import copy

import operator

import warnings

import exc

import dicts

import sql

import sql_gen

import strings

import util

##### Exceptions

# Can't use built-in SyntaxError because it stringifies to only the first line

class SyntaxError(Exception): pass

##### Data cleanup

null_strs = ['', r'\N', 'NULL', 'UNKNOWN']

def cleanup_table(db, table, cols=None):

    table = sql_gen.as_Table(table)

    cols = [sql_gen.as_Col(c, table) for c in sql.table_cols(db, table)]

    cols = filter(lambda c: sql_gen.is_text_col(db, c), cols)

    expr = 'trim(both from %s)'

    for null in null_strs: expr = 'nullif('+expr+', '+db.esc_value(null)+')'

    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db))) for v in cols]

    while True:

        try:

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

            break # successful

        except sql.NullValueException, e:

            col, = e.cols

            sql.drop_not_null(db, col)

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

class ExcToErrorsTable(sql_gen.ExcToWarning):

    '''Handles an exception by saving it or converting it to a warning.'''

    def __init__(self, return_, srcs, errors_table, value=None):

        '''

        @param return_ See sql_gen.ExcToWarning

        @param srcs The column names for the errors table

        @param errors_table None|sql_gen.Table

        @param value The value (or an expression for it) that caused the error

        @pre The invalid value must be in a local variable "value" of type text.

        '''

        sql_gen.ExcToWarning.__init__(self, return_)

        value = sql_gen.as_Code(value)

        self.srcs = srcs

        self.errors_table = errors_table

        self.value = value

    def to_str(self, db):

        if not self.srcs or self.errors_table == None:

            return sql_gen.ExcToWarning.to_str(self, db)

        errors_table_cols = map(sql_gen.Col,

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

        col_names_query = sql.mk_select(db, sql_gen.NamedValues('c', None,

            [[c.name] for c in self.srcs]), order_by=None)

        insert_query = sql.mk_insert_select(db, self.errors_table,

            errors_table_cols,

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

        return '''\

-- Save error in errors table.

DECLARE

    error_code text := SQLSTATE;

    error text := SQLERRM;

    value text := '''+self.value.to_str(db)+''';

    "column" text;

BEGIN

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

'''+strings.indent(sql_gen.RowExcIgnore(None, col_names_query, insert_query,

    row_var=errors_table_cols[0]).to_str(db))+'''

END;

'''+self.return_.to_str(db)

def data_exception_handler(*args, **kw_args):

    '''Handles a data_exception by saving it or converting it to a warning.

    For params, see ExcToErrorsTable().

    '''

    return sql_gen.data_exception_handler(ExcToErrorsTable(*args, **kw_args))

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 source columns). Otherwise, 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)

    function_name = strings.first_word(type_)

    srcs = col.srcs

    save_errors = errors_table != None and srcs

    if save_errors: # function will be unique for the given srcs

        function_name = strings.ustr(sql_gen.FunctionCall(function_name,

            *map(sql_gen.to_name_only_col, srcs)))

    function = db.TempFunction(function_name)

    # Create function definition

    modifiers = 'STRICT'

    if not save_errors: modifiers = 'IMMUTABLE '+modifiers

    value_param = sql_gen.FunctionParam('value', 'text')

    handler = data_exception_handler('RETURN NULL;\n', srcs, errors_table,

        value_param.name)

    body = sql_gen.CustomCode(handler.to_str(db, '''\

/* 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_+''';

'''))

    body.lang='plpgsql'

    sql.define_func(db, sql_gen.FunctionDef(function, type_, body,

        [value_param], modifiers))

    return sql_gen.FunctionCall(function, col)

def func_wrapper_exception_handler(db, return_, args, errors_table):

    '''Handles a function call's data_exceptions.

    Supports PL/Python functions.

    @param return_ See data_exception_handler()

    @param args [arg...] Function call's args

    @param errors_table See data_exception_handler()

    '''

    args = filter(sql_gen.has_srcs, args)

    srcs = sql_gen.cross_join_srcs(args)

    value = sql_gen.merge_not_null(db, ',', args)

    return sql_gen.NestedExcHandler(

        data_exception_handler(return_, srcs, errors_table, value)

        , sql_gen.plpythonu_error_handler

        )

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=strings.urepr(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

def mk_errors_table(db, table):

    errors_table_ = errors_table(db, table, if_exists=False)

    if sql.table_exists(db, errors_table_): return

    typed_cols = [

        sql_gen.TypedCol('column', 'text', nullable=False),

        sql_gen.TypedCol('value', 'text'),

        sql_gen.TypedCol('error_code', 'character varying(5)', nullable=False),

        sql_gen.TypedCol('error', 'text', nullable=False),

        ]

    sql.create_table(db, errors_table_, typed_cols, has_pkey=False)

    index_cols = ['column', 'value', 'error_code', 'error']

    sql.add_index(db, index_cols, errors_table_, unique=True)

##### Import

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

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    '''

    return put_table(db, table, [], row, row_ct_ref)

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 strings.ustr(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, default=None,

    col_defaults={}, on_error=exc.reraise):

    '''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 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 col_defaults Default values for required columns.

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

    '''

    import psycopg2.extensions

    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!

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

    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,

        mapping.values()), False)

    is_literals_or_function = is_literals or is_function

    if in_tables0 == None: errors_table_ = None

    else: errors_table_ = errors_table(db, in_tables0)

    # Create input joins from list of input tables

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

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

    if mapping == {} and not is_function: # need >= one column for INSERT SELECT

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

    if not is_literals:

        into = sql_gen.as_Table(into_table_name(out_table, in_tables0, mapping,

            is_function))

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

    # Wrap mapping in a sql_gen.ColDict.

    # sql_gen.ColDict sanitizes both keys and values passed into it.

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

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

    # 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

    # Save default values for all rows since in_table may have rows deleted

    if is_literals: pass

    elif is_function: full_in_table = in_table

    else:

        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)

    if not is_literals:

        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=None, limit=None, **kw_args):

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

        if pkeys_table_exists_ref[0]:

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

        else:

            sql.run_query_into(db, query, into=into, add_pkey_=True, **kw_args)

            pkeys_table_exists_ref[0] = True

    limit_ref = [None]

    def mk_main_select(joins, cols):

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

    if is_literals: insert_in_table = None

    else:

        insert_in_table = in_table

        insert_in_tables = [insert_in_table]

    join_cols = sql_gen.ColDict(db, out_table)

    exc_strs = set()

    def log_exc(e):

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

        log_debug('Caught exception: '+e_str)

        if e_str in exc_strs: # avoid infinite loops

            log_debug('Exception already seen, handler broken')

            on_error(e)

            remove_all_rows()

        else: 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_cond(cond, e):

        if is_literals: remove_all_rows()

        else:

            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)

    not_null_cols = set()

    def ignore(in_col, value, e):

        if sql_gen.is_table_col(in_col):

            in_col = sql_gen.with_table(in_col, insert_in_table)

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

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

            sql.add_index(db, in_col, insert_in_table) # enable fast filtering

            if value != None and in_col not in not_null_cols:

                log_debug('Replacing invalid value '+strings.as_tt(repr(value))

                    +' with NULL in column '+strings.as_tt(in_col.to_str(db)))

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

                    sql_gen.ColValueCond(in_col, value))

            else:

                log_debug('Ignoring rows with '+strings.as_tt(in_col.to_str(db))

                    +' = '+strings.as_tt(repr(value)))

                sql.delete(db, insert_in_table,

                    sql_gen.ColValueCond(in_col, value))

                if value == None: not_null_cols.add(in_col)

        else:

            assert isinstance(in_col, sql_gen.NamedCol)

            in_value = sql_gen.remove_col_rename(in_col)

            assert sql_gen.is_literal(in_value)

            if value == in_value.value:

                if value != None:

                    log_debug('Replacing invalid literal '

                        +strings.as_tt(in_col.to_str(db))+' with NULL')

                    mapping[in_col.name] = None

                else:

                    remove_all_rows()

            # otherwise, all columns were being ignore()d because the specific

            # column couldn't be identified, and this was not the invalid column

    if not is_literals:

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

    def mk_func_call():

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

        func_call = sql_gen.NamedCol(out_pkey,

            sql_gen.FunctionCall(out_table, **args))

        return func_call, args

    if is_function and not is_literals:

        log_debug('Defining wrapper function')

        func_call, args = mk_func_call()

        # Create empty pkeys table so its row type can be used

        insert_into_pkeys(input_joins, [in_pkey_col, func_call], limit=0,

            recover=True)

        result_type = db.col_info(sql_gen.Col(out_pkey, into)).type

        ## Create error handling wrapper function

        wrapper = db.TempFunction(sql_gen.concat(into.name, '_wrap'))

        select_cols = [in_pkey_col]+args.values()

        row_var = copy.copy(sql_gen.row_var)

        row_var.set_srcs([in_table])

        in_pkey_var = sql_gen.Col(in_pkey, row_var)

        args = dict(((k, sql_gen.with_table(v, row_var))

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

        func_call = sql_gen.FunctionCall(out_table, **args)

        def mk_return(result):

            return sql_gen.ReturnQuery(sql.mk_select(db,

                fields=[in_pkey_var, result], explain=False))

        exc_handler = func_wrapper_exception_handler(db,

            mk_return(sql_gen.Cast(result_type, None)), args.values(),

            errors_table_)

        sql.define_func(db, sql_gen.FunctionDef(wrapper, sql_gen.SetOf(into),

            sql_gen.RowExcIgnore(sql_gen.RowType(in_table),

                sql.mk_select(db, input_joins, order_by=None),

                mk_return(func_call), exc_handler=exc_handler)

            ))

        wrapper_table = sql_gen.FunctionCall(wrapper)

    # Do inserts and selects

    while True:

        has_joins = join_cols != {}

        # Handle unrecoverable errors in a special case

        if limit_ref[0] == 0:

            if is_literals or default == None:

                default = sql_gen.remove_col_rename(default)

                log_debug('Returning default: '+strings.as_tt(repr(default)))

                return default

            elif is_function: pass # empty pkeys table already created

            else:

                log_debug('Creating an empty output pkeys table')

                has_joins = False # use the no-joins case

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

                    [out_pkey], order_by=None, limit=0), into=insert_out_pkeys)

            break # don't do main case

        # Prepare to insert new rows

        if is_function:

            log_debug('Calling function on input rows')

            if is_literals: func_call, args = mk_func_call()

        else:

            log_debug('Trying 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))

                if not is_literals:

                    insert_args.update(dict(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()])

        try:

            cur = None

            if is_function:

                if is_literals:

                    cur = sql.select(db, fields=[func_call], recover=True,

                        cacheable=True)

                else: insert_into_pkeys(wrapper_table, recover=True)

            else:

                cur = sql.insert_select(db, out_table, mapping.keys(),

                    main_select, **insert_args)

            break # insert successful

        except sql.MissingCastException, e:

            log_exc(e)

            type_ = e.type

            if e.col == None: out_cols = mapping.keys()

            else: out_cols = [e.col]

            for out_col in out_cols:

                log_debug('Casting '+strings.as_tt(strings.repr_no_u(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)

                        ignore(in_col, e.value, e)

        except sql.DuplicateKeyException, e:

            log_exc(e)

            # Different rows violating different unique constraints not

            # supported

            assert not join_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))

            if is_literals:

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

                    join_cols, order_by=None))

            # Uniquify input table to avoid internal duplicate keys

            insert_in_table = sql.distinct_table(db, insert_in_table,

                join_cols.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, e:

                try: in_col = mapping[out_col] = col_defaults[out_col]

                except KeyError:

                    msg = 'Missing mapping for NOT NULL column '+out_col

                    log_debug(msg)

                    if default == None: warnings.warn(UserWarning(msg))

                        # not an error because sometimes the mappings include

                        # extra tables which aren't used by the dataset

                    remove_all_rows()

            else: ignore(in_col, None, e)

        except sql.CheckException, e:

            log_exc(e)

            ignore_cond(e.cond, e)

        except sql.InvalidValueException, e:

            log_exc(e)

            for in_col in mapping.values(): ignore(in_col, e.value, e)

        except psycopg2.extensions.TransactionRollbackError, e:

            log_exc(e)

            # retry

        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_literals: return sql.value(cur)

    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)

    else:

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

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

        # Note that mk_main_select() does not use ORDER BY. Instead, assume that

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

        # its rows will be retrieved in the same order.

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

    if limit_ref[0] == 0 or not is_function: # is_function doesn't leave holes

        log_debug('Setting pkeys of missing rows to '

            +strings.as_tt(repr(default)))

        missing_rows_joins = [full_in_table, 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,

            [sql_gen.Col(in_pkey, full_in_table),

            sql_gen.NamedCol(out_pkey, default)])

    # otherwise, there is already an entry for every row

    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_function: srcs = sql_gen.cols_srcs(in_cols)

    return sql_gen.Col(out_pkey, into, srcs)