# orm/query.py
# Copyright (C) 2005-2018 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

"""The Query class and support.

Defines the :class:`.Query` class, the central
construct used by the ORM to construct database queries.

The :class:`.Query` class should not be confused with the
:class:`.Select` class, which defines database
SELECT operations at the SQL (non-ORM) level.  ``Query`` differs from
``Select`` in that it returns ORM-mapped objects and interacts with an
ORM session, whereas the ``Select`` construct interacts directly with the
database to return iterable result sets.

"""

from itertools import chain

from . import (
    attributes, interfaces, object_mapper, persistence,
    exc as orm_exc, loading
)
from .base import _entity_descriptor, _is_aliased_class, \
    _is_mapped_class, _orm_columns, _generative, InspectionAttr
from .path_registry import PathRegistry
from .util import (
    AliasedClass, ORMAdapter, join as orm_join, with_parent, aliased,
    _entity_corresponds_to
)
from .. import sql, util, log, exc as sa_exc, inspect, inspection
from ..sql.expression import _interpret_as_from
from ..sql import (
    util as sql_util,
    expression, visitors
)
from ..sql.base import ColumnCollection
from . import properties

__all__ = ['Query', 'QueryContext', 'aliased']


_path_registry = PathRegistry.root


@inspection._self_inspects
@log.class_logger
class Query(object):
    """ORM-level SQL construction object.

    :class:`.Query` is the source of all SELECT statements generated by the
    ORM, both those formulated by end-user query operations as well as by
    high level internal operations such as related collection loading.  It
    features a generative interface whereby successive calls return a new
    :class:`.Query` object, a copy of the former with additional
    criteria and options associated with it.

    :class:`.Query` objects are normally initially generated using the
    :meth:`~.Session.query` method of :class:`.Session`, and in
    less common cases by instantiating the :class:`.Query` directly and
    associating with a :class:`.Session` using the :meth:`.Query.with_session`
    method.

    For a full walkthrough of :class:`.Query` usage, see the
    :ref:`ormtutorial_toplevel`.

    """

    _only_return_tuples = False
    _enable_eagerloads = True
    _enable_assertions = True
    _with_labels = False
    _criterion = None
    _yield_per = None
    _order_by = False
    _group_by = False
    _having = None
    _distinct = False
    _prefixes = None
    _suffixes = None
    _offset = None
    _limit = None
    _for_update_arg = None
    _statement = None
    _correlate = frozenset()
    _populate_existing = False
    _invoke_all_eagers = True
    _version_check = False
    _autoflush = True
    _only_load_props = None
    _refresh_state = None
    _refresh_identity_token = None
    _from_obj = ()
    _join_entities = ()
    _select_from_entity = None
    _mapper_adapter_map = {}
    _filter_aliases = None
    _from_obj_alias = None
    _joinpath = _joinpoint = util.immutabledict()
    _execution_options = util.immutabledict()
    _params = util.immutabledict()
    _attributes = util.immutabledict()
    _with_options = ()
    _with_hints = ()
    _enable_single_crit = True
    _orm_only_adapt = True
    _orm_only_from_obj_alias = True
    _current_path = _path_registry
    _has_mapper_entities = False

    lazy_loaded_from = None
    """An :class:`.InstanceState` that is using this :class:`.Query` for a
    lazy load operation.

    This can be used for extensions like the horizontal sharding extension
    as well as event handlers and custom mapper options to determine
    when a query is being used to lazy load a relationship on an object.

    .. versionadded:: 1.2.9

    """

    def __init__(self, entities, session=None):
        """Construct a :class:`.Query` directly.

        E.g.::

            q = Query([User, Address], session=some_session)

        The above is equivalent to::

            q = some_session.query(User, Address)

        :param entities: a sequence of entities and/or SQL expressions.

        :param session: a :class:`.Session` with which the :class:`.Query`
         will be associated.   Optional; a :class:`.Query` can be associated
         with a :class:`.Session` generatively via the
         :meth:`.Query.with_session` method as well.

        .. seealso::

            :meth:`.Session.query`

            :meth:`.Query.with_session`

        """
        self.session = session
        self._polymorphic_adapters = {}
        self._set_entities(entities)

    def _set_entities(self, entities, entity_wrapper=None):
        if entity_wrapper is None:
            entity_wrapper = _QueryEntity
        self._entities = []
        self._primary_entity = None
        self._has_mapper_entities = False

        # 1. don't run util.to_list() or _set_entity_selectables
        #    if no entities were passed - major performance bottleneck
        #    from lazy loader implementation when it seeks to use Query
        #    class for an identity lookup, causes test_orm.py to fail
        #    with thousands of extra function calls, see issue #4228
        #    for why this use had to be added
        # 2. can't use classmethod on Query because session.query_cls
        #    is an arbitrary callable in some user recipes, not
        #    necessarily a class, so we don't have the class available.
        #    see issue #4256
        # 3. can't do "if entities is not None" because we usually get here
        #    from session.query() which takes in *entities.
        # 4. can't do "if entities" because users make use of undocumented
        #    to_list() behavior here and they pass clause expressions that
        #    can't be evaluated as boolean.  See issue #4269.
        # 5. the empty tuple is a singleton in cPython, take advantage of this
        #    so that we can skip for the empty "*entities" case without using
        #    any Python overloadable operators.
        #
        if entities is not ():
            for ent in util.to_list(entities):
                entity_wrapper(self, ent)

            self._set_entity_selectables(self._entities)

    def _set_entity_selectables(self, entities):
        self._mapper_adapter_map = d = self._mapper_adapter_map.copy()

        for ent in entities:
            for entity in ent.entities:
                if entity not in d:
                    ext_info = inspect(entity)
                    if not ext_info.is_aliased_class and \
                            ext_info.mapper.with_polymorphic:
                        if ext_info.mapper.mapped_table not in \
                                self._polymorphic_adapters:
                            self._mapper_loads_polymorphically_with(
                                ext_info.mapper,
                                sql_util.ColumnAdapter(
                                    ext_info.selectable,
                                    ext_info.mapper._equivalent_columns
                                )
                            )
                        aliased_adapter = None
                    elif ext_info.is_aliased_class:
                        aliased_adapter = ext_info._adapter
                    else:
                        aliased_adapter = None

                    d[entity] = (
                        ext_info,
                        aliased_adapter
                    )
                ent.setup_entity(*d[entity])

    def _mapper_loads_polymorphically_with(self, mapper, adapter):
        for m2 in mapper._with_polymorphic_mappers or [mapper]:
            self._polymorphic_adapters[m2] = adapter
            for m in m2.iterate_to_root():
                self._polymorphic_adapters[m.local_table] = adapter

    def _set_select_from(self, obj, set_base_alias):
        fa = []
        select_from_alias = None

        for from_obj in obj:
            info = inspect(from_obj)
            if hasattr(info, 'mapper') and \
                    (info.is_mapper or info.is_aliased_class):
                self._select_from_entity = info
                if set_base_alias and not info.is_aliased_class:
                    raise sa_exc.ArgumentError(
                        "A selectable (FromClause) instance is "
                        "expected when the base alias is being set.")
                fa.append(info.selectable)
            elif not info.is_selectable:
                raise sa_exc.ArgumentError(
                    "argument is not a mapped class, mapper, "
                    "aliased(), or FromClause instance.")
            else:
                if isinstance(from_obj, expression.SelectBase):
                    from_obj = from_obj.alias()
                if set_base_alias:
                    select_from_alias = from_obj
                fa.append(from_obj)

        self._from_obj = tuple(fa)

        if set_base_alias and \
                len(self._from_obj) == 1 and \
                isinstance(select_from_alias, expression.Alias):
            equivs = self.__all_equivs()
            self._from_obj_alias = sql_util.ColumnAdapter(
                self._from_obj[0], equivs)
        elif set_base_alias and \
                len(self._from_obj) == 1 and \
                hasattr(info, "mapper") and \
                info.is_aliased_class:
            self._from_obj_alias = info._adapter

    def _reset_polymorphic_adapter(self, mapper):
        for m2 in mapper._with_polymorphic_mappers:
            self._polymorphic_adapters.pop(m2, None)
            for m in m2.iterate_to_root():
                self._polymorphic_adapters.pop(m.local_table, None)

    def _adapt_polymorphic_element(self, element):
        if "parententity" in element._annotations:
            search = element._annotations['parententity']
            alias = self._polymorphic_adapters.get(search, None)
            if alias:
                return alias.adapt_clause(element)

        if isinstance(element, expression.FromClause):
            search = element
        elif hasattr(element, 'table'):
            search = element.table
        else:
            return None

        alias = self._polymorphic_adapters.get(search, None)
        if alias:
            return alias.adapt_clause(element)

    def _adapt_col_list(self, cols):
        return [
            self._adapt_clause(
                expression._literal_as_label_reference(o),
                True, True)
            for o in cols
        ]

    @_generative()
    def _set_lazyload_from(self, state):
        self.lazy_loaded_from = state

    @_generative()
    def _adapt_all_clauses(self):
        self._orm_only_adapt = False

    def _adapt_clause(self, clause, as_filter, orm_only):
        """Adapt incoming clauses to transformations which
        have been applied within this query."""

        adapters = []
        # do we adapt all expression elements or only those
        # tagged as 'ORM' constructs ?
        if not self._orm_only_adapt:
            orm_only = False

        if as_filter and self._filter_aliases:
            for fa in self._filter_aliases._visitor_iterator:
                adapters.append(
                    (
                        orm_only, fa.replace
                    )
                )

        if self._from_obj_alias:
            # for the "from obj" alias, apply extra rule to the
            # 'ORM only' check, if this query were generated from a
            # subquery of itself, i.e. _from_selectable(), apply adaption
            # to all SQL constructs.
            adapters.append(
                (
                    orm_only if self._orm_only_from_obj_alias else False,
                    self._from_obj_alias.replace
                )
            )

        if self._polymorphic_adapters:
            adapters.append(
                (
                    orm_only, self._adapt_polymorphic_element
                )
            )

        if not adapters:
            return clause

        def replace(elem):
            for _orm_only, adapter in adapters:
                # if 'orm only', look for ORM annotations