# orm/persistence.py
# Copyright (C) 2005-2016 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: http://www.opensource.org/licenses/mit-license.php

"""private module containing functions used to emit INSERT, UPDATE
and DELETE statements on behalf of a :class:`.Mapper` and its descending
mappers.

The functions here are called only by the unit of work functions
in unitofwork.py.

"""

import operator
from itertools import groupby, chain
from .. import sql, util, exc as sa_exc
from . import attributes, sync, exc as orm_exc, evaluator
from .base import state_str, _attr_as_key, _entity_descriptor
from ..sql import expression
from ..sql.base import _from_objects
from . import loading


def _bulk_insert(
        mapper, mappings, session_transaction, isstates, return_defaults,
        render_nulls):
    base_mapper = mapper.base_mapper

    cached_connections = _cached_connection_dict(base_mapper)

    if session_transaction.session.connection_callable:
        raise NotImplementedError(
            "connection_callable / per-instance sharding "
            "not supported in bulk_insert()")

    if isstates:
        if return_defaults:
            states = [(state, state.dict) for state in mappings]
            mappings = [dict_ for (state, dict_) in states]
        else:
            mappings = [state.dict for state in mappings]
    else:
        mappings = list(mappings)

    connection = session_transaction.connection(base_mapper)
    for table, super_mapper in base_mapper._sorted_tables.items():
        if not mapper.isa(super_mapper):
            continue

        records = (
            (None, state_dict, params, mapper,
                connection, value_params, has_all_pks, has_all_defaults)
            for
            state, state_dict, params, mp,
            conn, value_params, has_all_pks,
            has_all_defaults in _collect_insert_commands(table, (
                (None, mapping, mapper, connection)
                for mapping in mappings),
                bulk=True, return_defaults=return_defaults,
                render_nulls=render_nulls
            )
        )
        _emit_insert_statements(base_mapper, None,
                                cached_connections,
                                super_mapper, table, records,
                                bookkeeping=return_defaults)

    if return_defaults and isstates:
        identity_cls = mapper._identity_class
        identity_props = [p.key for p in mapper._identity_key_props]
        for state, dict_ in states:
            state.key = (
                identity_cls,
                tuple([dict_[key] for key in identity_props])
            )


def _bulk_update(mapper, mappings, session_transaction,
                 isstates, update_changed_only):
    base_mapper = mapper.base_mapper

    cached_connections = _cached_connection_dict(base_mapper)

    search_keys = mapper._primary_key_propkeys
    if mapper._version_id_prop:
        search_keys = set([mapper._version_id_prop.key]).union(search_keys)

    def _changed_dict(mapper, state):
        return dict(
            (k, v)
            for k, v in state.dict.items() if k in state.committed_state or k
            in search_keys

        )

    if isstates:
        if update_changed_only:
            mappings = [_changed_dict(mapper, state) for state in mappings]
        else:
            mappings = [state.dict for state in mappings]
    else:
        mappings = list(mappings)

    if session_transaction.session.connection_callable:
        raise NotImplementedError(
            "connection_callable / per-instance sharding "
            "not supported in bulk_update()")

    connection = session_transaction.connection(base_mapper)

    for table, super_mapper in base_mapper._sorted_tables.items():
        if not mapper.isa(super_mapper):
            continue

        records = _collect_update_commands(None, table, (
            (None, mapping, mapper, connection,
                (mapping[mapper._version_id_prop.key]
                    if mapper._version_id_prop else None))
            for mapping in mappings
        ), bulk=True)

        _emit_update_statements(base_mapper, None,
                                cached_connections,
                                super_mapper, table, records,
                                bookkeeping=False)


def save_obj(
        base_mapper, states, uowtransaction, single=False):
    """Issue ``INSERT`` and/or ``UPDATE`` statements for a list
    of objects.

    This is called within the context of a UOWTransaction during a
    flush operation, given a list of states to be flushed.  The
    base mapper in an inheritance hierarchy handles the inserts/
    updates for all descendant mappers.

    """

    # if batch=false, call _save_obj separately for each object
    if not single and not base_mapper.batch:
        for state in _sort_states(states):
            save_obj(base_mapper, [state], uowtransaction, single=True)
        return

    states_to_update = []
    states_to_insert = []
    cached_connections = _cached_connection_dict(base_mapper)

    for (state, dict_, mapper, connection,
            has_identity,
            row_switch, update_version_id) in _organize_states_for_save(
            base_mapper, states, uowtransaction
    ):
        if has_identity or row_switch:
            states_to_update.append(
                (state, dict_, mapper, connection, update_version_id)
            )
        else:
            states_to_insert.append(
                (state, dict_, mapper, connection)
            )

    for table, mapper in base_mapper._sorted_tables.items():
        if table not in mapper._pks_by_table:
            continue
        insert = _collect_insert_commands(table, states_to_insert)

        update = _collect_update_commands(
            uowtransaction, table, states_to_update)

        _emit_update_statements(base_mapper, uowtransaction,
                                cached_connections,
                                mapper, table, update)

        _emit_insert_statements(base_mapper, uowtransaction,
                                cached_connections,
                                mapper, table, insert)

    _finalize_insert_update_commands(
        base_mapper, uowtransaction,
        chain(
            (
                (state, state_dict, mapper, connection, False)
                for state, state_dict, mapper, connection in states_to_insert
            ),
            (
                (state, state_dict, mapper, connection, True)
                for state, state_dict, mapper, connection,
                update_version_id in states_to_update
            )
        )
    )


def post_update(base_mapper, states, uowtransaction, post_update_cols):
    """Issue UPDATE statements on behalf of a relationship() which
    specifies post_update.

    """
    cached_connections = _cached_connection_dict(base_mapper)

    states_to_update = list(_organize_states_for_post_update(
        base_mapper,
        states, uowtransaction))

    for table, mapper in base_mapper._sorted_tables.items():
        if table not in mapper._pks_by_table:
            continue

        update = (
            (state, state_dict, sub_mapper, connection)
            for
            state, state_dict, sub_mapper, connection in states_to_update
            if table in sub_mapper._pks_by_table
        )

        update = _collect_post_update_commands(base_mapper, uowtransaction,
                                               table, update,
                                               post_update_cols)

        _emit_post_update_statements(base_mapper, uowtransaction,
                                     cached_connections,
                                     mapper, table, update)


def delete_obj(base_mapper, states, uowtransaction):
    """Issue ``DELETE`` statements for a list of objects.

    This is called within the context of a UOWTransaction during a
    flush operation.

    """

    cached_connections = _cached_connection_dict(base_mapper)

    states_to_delete = list(_organize_states_for_delete(
        base_mapper,
        states,
        uowtransaction))

    table_to_mapper = base_mapper._sorted_tables

    for table in reversed(list(table_to_mapper.keys())):
        mapper = table_to_mapper[table]
        if table not in mapper._pks_by_table:
            continue
        elif mapper.inherits and mapper.passive_deletes:
            continue

        delete = _collect_delete_commands(base_mapper, uowtransaction,
                                          table, states_to_delete)

        _emit_delete_statements(base_mapper, uowtransaction,
                                cached_connections, mapper, table, delete)

    for state, state_dict, mapper, connection, \
            update_version_id in states_to_delete:
        mapper.dispatch.after_delete(mapper, connection, state)


def _organize_states_for_save(base_mapper, states, uowtransaction):
    """Make an initial pass across a set of states for INSERT or
    UPDATE.

    This includes splitting out into distinct lists for
    each, calling before_insert/before_update, obtaining
    key information for each state including its dictionary,
    mapper, the connection to use for the execution per state,
    and the identity flag.

    """

    for state, dict_, mapper, connection in _connections_for_states(
            base_mapper, uowtransaction,
            states):

        has_identity = bool(state.key)

        instance_key = state.key or mapper._identity_key_from_state(state)

        row_switch = update_version_id = None

        # call before_XXX extensions
        if not has_identity:
            mapper.dispatch.before_insert(mapper, connection, state)
        else:
            mapper.dispatch.before_update(mapper, connection, state)

        if mapper._validate_polymorphic_identity:
            mapper._validate_polymorphic_identity(mapper, state, dict_)

        # detect if we have a "pending" instance (i.e. has
        # no instance_key attached to it), and another instance
        # with the same identity key already exists as persistent.
        # convert to an UPDATE if so.
        if not has_identity and \
                instance_key in uowtransaction.session.identity_map:
            instance = \
                uowtransaction.session.identity_map[instance_key]
            existing = attributes.instance_state(instance)

            if not uowtransaction.was_already_deleted(existing):
                if not uowtransaction.is_deleted(existing):
                    raise orm_exc.FlushError(
                        "New instance %s with identity key %s conflicts "
                        "with persistent instance %s" %
                        (state_str(state), instance_key,
                         state_str(existing)))

                base_mapper._log_debug(
                    "detected row switch for identity %s.  "
                    "will update %s, remove %s from "
                    "transaction", instance_key,
                    state_str(state), state_str(existing))

                # remove the "delete" flag from the existing element
                uowtransaction.remove_state_actions(existing)
                row_switch = existing

        if (has_identity or row_switch) and mapper.version_id_col is not None:
            update_version_id = mapper._get_committed_state_attr_by_column(
                row_switch if row_switch else state,
                row_switch.dict if row_switch else dict_,
                mapper.version_id_col)

        yield (state, dict_, mapper, connection,
               has_identity, row_switch, update_version_id)


def _organize_states_for_post_update(base_mapper, states,
                                     uowtransaction):
    """Make an initial pass across a set of states for UPDATE
    corresponding to post_update.

    This includes obtaining key information for each state
    including its dictionary, mapper, the connection to use for
    the execution per state.

    """
    return _connections_for_states(base_mapper, uowtransaction, states)


def _organize_states_for_delete(base_mapper, states, uowtransaction):
    """Make an initial pass across a set of states for DELETE.

    This includes calling out before_delete and obtaining
    key information for each state including its dictionary,
    mapper, the connection to use for the execution per state.

    """
    for state, dict_, mapper, connection in _connections_for_states(
            base_mapper, uowtransaction,
            states):

        mapper.dispatch.before_delete(mapper, connection, state)

        if mapper.version_id_col is not None:
            update_version_id = \
                mapper._get_committed_state_attr_by_column(
                    state, dict_,
                    mapper.version_id_col)
        else:
            update_version_id = None

        yield (
            state, dict_, mapper, connection, update_version_id)


def _collect_insert_commands(
        table, states_to_insert,
        bulk=False, return_defaults=False, render_nulls=False):
    """Identify sets of values to use in INSERT statements for a
    list of states.

    """
    for state, state_dict, mapper, connection in states_to_insert:
        if table not in mapper._pks_by_table:
            continue

        params = {}
        value_params = {}

        propkey_to_col = mapper._propkey_to_col[table]

        eval_none = mapper._insert_cols_evaluating_none[table]

        for propkey in set(propkey_to_col).intersection(state_dict):
            value = state_dict[propkey]
            col = propkey_to_col[propkey]
            if value is None and propkey not in eval_none and not render_nulls:
                continue
            elif not bulk and hasattr(value, '__clause_element__') or \
                    isinstance(value, sql.ClauseElement):
                value_params[col.key] = value.__clause_element__() \
                    if hasattr(value, '__clause_element__') else value
            else:
                params[col.key] = value

        if not bulk:
            # for all the columns that have no default and we don't have
            # a value and where "None" is not a special value, add
            # explicit None to the INSERT.   This is a legacy behavior
            # which might be worth removing, as it should not be necessary
            # and also produces confusion, given that "missing" and None
            # now have distinct meanings
            for colkey in mapper._insert_cols_as_none[table].\
                    difference(params).difference(value_params):
                params[colkey] = None

        if not bulk or return_defaults:
            has_all_pks = mapper._pk_keys_by_table[table].issubset(params)

            if mapper.base_mapper.eager_defaults:
                has_all_defaults = mapper._server_default_cols[table].\
                    issubset(params)
            else:
                has_all_defaults = True
        else:
            has_all_defaults = has_all_pks = True

        if mapper.version_id_generator is not False \
                and mapper.version_id_col is not None and \
                mapper.version_id_col in mapper._cols_by_table[table]:
            params[mapper.version_id_col.key] = \
                mapper.version_id_generator(None)

        yield (
            state, state_dict, params, mapper,
            connection, value_params, has_all_pks,
            has_all_defaults)


def _collect_update_commands(
        uowtransaction, table, states_to_update,
        bulk=False):
    """Identify sets of values to use in UPDATE statements for a
    list of states.

    This function works intricately with the history system
    to determine exactly what values should be updated
    as well as how the row should be matched within an UPDATE
    statement.  Includes some tricky scenarios where the primary
    key of an object might have been changed.

    """

    for state, state_dict, mapper, connection, \
            update_version_id in states_to_update:

        if table not in mapper._pks_by_table:
            continue

        pks = mapper._pks_by_table[table]

        value_params = {}

        propkey_to_col = mapper._propkey_to_col[table]

        if bulk:
            params = dict(
                (propkey_to_col[propkey].key, state_dict[propkey])
                for propkey in
                set(propkey_to_col).intersection(state_dict).difference(
                    mapper._pk_keys_by_table[table])
            )
            has_all_defaults = True
        else:
            params = {}
            for propkey in set(propkey_to_col).intersection(
                    state.committed_state):
                value = state_dict[propkey]
                col = propkey_to_col[propkey]

                if hasattr(value, '__clause_element__') or \
                        isinstance(value, sql.ClauseElement):
                    value_params[col] = value.__clause_element__() \
                        if hasattr(value, '__clause_element__') else value
                # guard against values that generate non-__nonzero__
                # objects for __eq__()
                elif state.manager[propkey].impl.is_equal(
                        value, state.committed_state[propkey]) is not True:
                    params[col.key] = value

            if mapper.base_mapper.eager_defaults:
                has_all_defaults = mapper._server_onupdate_default_cols[table].\
                    issubset(params)
            else:
                has_all_defaults = True

        if update_version_id is not None and \
                mapper.version_id_col in mapper._cols_by_table[table]:

            if not bulk and not (params or value_params):
                # HACK: check for history in other tables, in case the
                # history is only in a different table than the one
                # where the version_id_col is.  This logic was lost
                # from 0.9 -> 1.0.0 and restored in 1.0.6.
                for prop in mapper._columntoproperty.values():
                    history = (
                        state.manager[prop.key].impl.get_history(
                            state, state_dict,
                            attributes.PASSIVE_NO_INITIALIZE))
                    if history.added:
                        break
                else:
                    # no net change, break
                    continue

            col = mapper.version_id_col
            params[col._label] = update_version_id

            if (bulk or col.key not in params) and \
                    mapper.version_id_generator is not False:
                val = mapper.version_id_generator(update_version_id)
                params[col.key] = val

        elif not (params or value_params):
            continue

        has_all_pks = True
        if bulk:
            pk_params = dict(
                (propkey_to_col[propkey]._label, state_dict.get(propkey))
                for propkey in
                set(propkey_to_col).
                intersection(mapper._pk_attr_keys_by_table[table])
            )
        else:
            pk_params = {}
            for col in pks:
                propkey = mapper._columntoproperty[col].key

                history = state.manager[propkey].impl.get_history(
                    state, state_dict, attributes.PASSIVE_OFF)

                if history.added:
                    if not history.deleted or \
                            ("pk_cascaded", state, col) in \
                            uowtransaction.attributes:
                        pk_params[col._label] = history.added[0]
                        params.pop(col.key, None)
                    else:
                        # else, use the old value to locate the row
                        pk_params[col._label] = history.deleted[0]
                        if col in value_params:
                            has_all_pks = False
                else:
                    pk_params[col._label] = history.unchanged[0]
                if pk_params[col._label] is None:
                    raise orm_exc.FlushError(
                        "Can't update table %s using NULL for primary "
                        "key value on column %s" % (table, col))

        if params or value_params:
            params.update(pk_params)
            yield (
                state, state_dict, params, mapper,
                connection, value_params, has_all_defaults, has_all_pks)


def _collect_post_update_commands(base_mapper, uowtransaction, table,
                                  states_to_update, post_update_cols):
    """Identify sets of values to use in UPDATE statements for a
    list of states within a post_update operation.

    """

    for state, state_dict, mapper, connection in states_to_update:

        # assert table in mapper._pks_by_table

        pks = mapper._pks_by_table[table]
        params = {}
        hasdata = False

        for col in mapper._cols_by_table[table]:
            if col in pks:
                params[col._label] = \
                    mapper._get_state_attr_by_column(
                        state,
                        state_dict, col, passive=attributes.PASSIVE_OFF)

            elif col in post_update_cols:
                prop = mapper._columntoproperty[col]
                history = state.manager[prop.key].impl.get_history(
                    state, state_dict,
                    attributes.PASSIVE_NO_INITIALIZE)
                if history.added:
                    value = history.added[0]
                    params[col.key] = value
                    hasdata = True
        if hasdata:
            yield params, connection


def _collect_delete_commands(base_mapper, uowtransaction, table,
                             states_to_delete):
    """Identify values to use in DELETE statements for a list of
    states to be deleted."""

    for state, state_dict, mapper, connection, \
            update_version_id in states_to_delete:

        if table not in mapper._pks_by_table:
            continue

        params = {}
        for col in mapper._pks_by_table[table]:
            params[col.key] = \
                value = \
                mapper._get_committed_state_attr_by_column(
                    state, state_dict, col)
            if value is None:
                raise orm_exc.FlushError(
                    "Can't delete from table %s "
                    "using NULL for primary "
                    "key value on column %s" % (table, col))

        if update_version_id is not None and \
                mapper.version_id_col in mapper._cols_by_table[table]:
            params[mapper.version_id_col.key] = update_version_id
        yield params, connection


def _emit_update_statements(base_mapper, uowtransaction,
                            cached_connections, mapper, table, update,
                            bookkeeping=True):
    """Emit UPDATE statements corresponding to value lists collected
    by _collect_update_commands()."""

    needs_version_id = mapper.version_id_col is not None and \
        mapper.version_id_col in mapper._cols_by_table[table]

    def update_stmt():
        clause = sql.and_()

        for col in mapper._pks_by_table[table]:
            clause.clauses.append(col == sql.bindparam(col._label,
                                                       type_=col.type))

        if needs_version_id:
            clause.clauses.append(
                mapper.version_id_col == sql.bindparam(
                    mapper.version_id_col._label,
                    type_=mapper.version_id_col.type))

        stmt = table.update(clause)
        return stmt

    cached_stmt = base_mapper._memo(('update', table), update_stmt)

    for (connection, paramkeys, hasvalue, has_all_defaults, has_all_pks), \
        records in groupby(
            update,
            lambda rec: (
                rec[4],  # connection
                set(rec[2]),  # set of parameter keys
                bool(rec[5]),  # whether or not we have "value" parameters
                rec[6],  # has_all_defaults
                rec[7]  # has all pks
            )
    ):
        rows = 0
        records = list(records)

        statement = cached_stmt

        # TODO: would be super-nice to not have to determine this boolean
        # inside the loop here, in the 99.9999% of the time there's only
        # one connection in use
        assert_singlerow = connection.dialect.supports_sane_rowcount
        assert_multirow = assert_singlerow and \
            connection.dialect.supports_sane_multi_rowcount
        allow_multirow = has_all_defaults and not needs_version_id

        if not has_all_pks:
            statement = statement.return_defaults()
        elif bookkeeping and not has_all_defaults and \
                mapper.base_mapper.eager_defaults:
            statement = statement.return_defaults()
        elif mapper.version_id_col is not None:
            statement = statement.return_defaults(mapper.version_id_col)

        if hasvalue: