# Database import/export

import dicts

import sql

import sql_gen

import strings

import util

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

class SyntaxError(Exception): pass

        sql.table_cols(db, table))

    except sql.DoesNotExistException: pass

    else:

        try: cols.remove(pkey_col)

        except ValueError: pass

    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]

        try:

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

            break # successful

        except sql.NullValueException, e:

            sql.drop_not_null(db, col)

    db.log_debug('Vacuuming and reanalyzing table', level=1.5)

    sql.vacuum(db, table)

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 an exception by saving it or converting it to a warning.'''

        @param return_ See sql_gen.ExcToWarning

        @param srcs The column names for the errors table

        @param errors_table None|sql_gen.Table

        '''

        sql_gen.ExcToWarning.__init__(self, return_)

        self.errors_table = errors_table

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

            return sql_gen.ExcToWarning.to_str(self, db)

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

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

            errors_table_cols,

DECLARE

    error_code text := SQLSTATE;

    error text := SQLERRM;

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

    For params, see ExcToErrorsTable().

    '''

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

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

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

    modifiers = 'STRICT'

    if not save_errors: modifiers = 'IMMUTABLE '+modifiers

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

        value_param.name)

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'

    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)

        data_exception_handler(return_, srcs, errors_table, value)

        , sql_gen.plpythonu_error_handler

        )

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

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

    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)

    dialect = stream_info.dialect

    if csvs.is_tsv(dialect): use_copy_from = False

    if use_copy_from:

        log('Using COPY FROM')

        # Create COPY FROM statement

        copy_from = ('COPY '+table.to_str(db)+' FROM STDIN DELIMITER '

            +db.esc_value(dialect.delimiter)+' NULL '+db.esc_value(''))

        assert not csvs.is_tsv(dialect)

        copy_from += ' CSV'

        if dialect.quoting != csv.QUOTE_NONE:

            quote_str = db.esc_value(dialect.quotechar)

            copy_from += ' QUOTE '+quote_str

            if dialect.doublequote: copy_from += ' ESCAPE '+quote_str

        copy_from += ';\n'

        log(copy_from, level=2)

        db.db.cursor().copy_expert(copy_from, stream)

    else:

        log('Using INSERT')

        cols_ct = len(stream_info.header)

        inserted_row_ct = 0

            util.list_set_length(row, cols_ct) # truncate extra cols

                log_level=5) != None:

                inserted_row_ct += 1

        log('Inserted '+str(inserted_row_ct)+' of '+str(row_ct)+' rows')

    # Get format info

    info = csvs.stream_info(stream, parse_header=True)

    col_names = map(strings.to_unicode, info.header)

    for i, col in enumerate(col_names): # replace empty column names

        if col == '': col_names[i] = 'column_'+str(i)

    # Select schema and escape names

    def esc_name(name): return db.esc_name(name)

    typed_cols = [sql_gen.TypedCol(v, 'text') for v in col_names]

    log('Creating table')

    # cached, and will not be re-run if the import is retried.

    # This MUST be run so that the rows will be stored in inserted order, and

    # the row_num added after import will match up with the CSV's row order.

    if use_copy_from: sql.with_savepoint(db, load)

    else: load()

    cleanup_table(db, table)

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    '''

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

        str_ += '('

        try: value_in_col = mapping['value']

        except KeyError:

        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

    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.

    '''

    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 not mapping and not is_function: # functions with no args OK

        warnings.warn(UserWarning('Inserting empty table row(s)'))

    except IndexError: in_tables0 = None

    else:

    # 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

            is_function))

        else: into_out_pkey = strings.ustr(out_pkey_col)

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

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

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

    elif is_function: full_in_table = in_table

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

                **kw_args)

    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)

            log_debug('Exception already seen, handler broken')

            on_error(e)

            remove_all_rows()

        else: exc_strs.add(e_str)

        log_debug('Ignoring all rows')

            if passed: pass

            elif failed: remove_all_rows()

            else: raise NotImplementedError()

            out_table_cols = sql_gen.ColDict(db, out_table)

            out_table_cols.update(util.dict_subset_right_join({},

            in_cols = []

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

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

            if cond == sql_gen.false_expr:

                assert failed

                remove_all_rows()

            elif cond == sql_gen.true_expr: assert passed

            else: cur = sql.delete(db, insert_in_table, not_cond)

            if failed or cur.rowcount > 0: # only if any rows failed cond

                    sql_gen.cross_join_srcs(in_cols), None, e.cause.pgcode,

            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:

                    +strings.as_tt(strings.urepr(value))+' with NULL in column '

                    +strings.as_tt(in_col.to_str(db)))

                    sql_gen.ColValueCond(in_col, value))

            else:

                    sql_gen.ColValueCond(in_col, value))

                if value == None: not_null_cols.add(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

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

        log_debug('Defining wrapper function')

        func_call, args = mk_func_call()

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

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

                exc_handler=exc_handler)

        wrapper_table = sql_gen.FunctionCall(wrapper)

    while True:

        # Prepare to insert new rows

            log_debug('Calling function on input rows')

            if is_literals: func_call, args = mk_func_call()

            if has_joins:

                insert_args.update(dict(ignore=True))

            else:

                if not is_literals: