# Database import/export

import dicts

import sql

import sql_gen

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 not col_names: col_names = [None] # need at least one entry

    for col_name in col_names:

                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

        value_param.name)

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

of the function, outside any block.) */

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

row_num_col_def.name = 'row_num'

    assert sql.table_col_names(db, table) == header

        n=1000))

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

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

    except Exception, e: sql.parse_exception(db, e, recover=True)

    # Get format info

        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]

    header.insert(0, row_num_col_def.name)

    reader = csvs.RowNumFilter(reader)

    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.

    cleanup_table(db, table)

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

    Warning: This function's insert/if not exists get algorithm does not support

    database triggers that populate fields covered by a unique constraint. Such

    fields must be populated by the mappings instead. (Other fields are not

    affected by this restriction on triggers.)

        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.

    '''

    if out_table == '_map':

        import psycopg2.extras

        psycopg2.extras.register_hstore(db.db)

        # Parse args

        try: value = mapping.pop('value')

        except KeyError: return None # value required

            for k, v in mapping.iteritems()]) # unwrap literal value

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

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

            sql.run_query_into(db, query, into=into, **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()

        log_debug('Ignoring all rows')

        log_debug('No handler for exception')

        on_error(e)

        remove_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({},

                    sql.table_col_names(db, out_table)))

            in_cols = []

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

            if cond == sql_gen.false_expr:

                assert failed

                remove_all_rows()

            elif cond == sql_gen.true_expr: assert passed

                while True:

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

                    try:

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

                        break

                    except sql.DoesNotExistException, e:

                        if e.type != 'column': raise

                        last_cond = cond

                        cond = sql_gen.map_expr(db, cond, {e.name: None})

                        if cond == last_cond: raise # not fixable

            if failed or cur != None and cur.rowcount > 0:

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

                    +'condition: '+orig_cond+'\ntranslated condition: '+cond)

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

    while True:

        # Prepare to insert new rows

                log_debug('Calling function')

                func_call, args = mk_func_call()

            if has_joins:

                insert_args.update(dict(ignore=True))

            else:

                if not is_literals:

                    insert_args.update(dict(into=insert_out_pkeys))