Move identity_lookup to session

This performance critical method is on Query needlessly, just to appease
the horizontal sharding API.   Have the performance impact of invoking
Query only incur if horizontal sharding is actually used.

Change-Id: I03db2befe2f5614380258927a62ed389a6ba0fae

4 files changed, 98 insertions, 90 deletions

diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py
index f56381685..129e5b488 100644
--- a/lib/sqlalchemy/ext/horizontal_shard.py
+++ b/lib/sqlalchemy/ext/horizontal_shard.py
@@ -86,40 +86,6 @@ class ShardedQuery(Query):
 
             return ShardedResult(results, rowcount)
 
-    def _identity_lookup(
-        self,
-        mapper,
-        primary_key_identity,
-        identity_token=None,
-        lazy_loaded_from=None,
-        **kw
-    ):
-        """override the default Query._identity_lookup method so that we
-        search for a given non-token primary key identity across all
-        possible identity tokens (e.g. shard ids).
-
-        """
-
-        if identity_token is not None:
-            return super(ShardedQuery, self)._identity_lookup(
-                mapper,
-                primary_key_identity,
-                identity_token=identity_token,
-                **kw
-            )
-        else:
-            q = self.session.query(mapper)
-            if lazy_loaded_from:
-                q = q._set_lazyload_from(lazy_loaded_from)
-            for shard_id in self.id_chooser(q, primary_key_identity):
-                obj = super(ShardedQuery, self)._identity_lookup(
-                    mapper, primary_key_identity, identity_token=shard_id, **kw
-                )
-                if obj is not None:
-                    return obj
-
-            return None
-
     def _get_impl(self, primary_key_identity, db_load_fn, identity_token=None):
         """Override the default Query._get_impl() method so that we emit
         a query to the DB for each possible identity token, if we don't
@@ -218,6 +184,47 @@ class ShardedSession(Session):
             for k in shards:
                 self.bind_shard(k, shards[k])
 
+    def _identity_lookup(
+        self,
+        mapper,
+        primary_key_identity,
+        identity_token=None,
+        lazy_loaded_from=None,
+        **kw
+    ):
+        """override the default :meth:`.Session._identity_lookup` method so that we
+        search for a given non-token primary key identity across all
+        possible identity tokens (e.g. shard ids).
+
+        .. versionchanged:: 1.4  Moved :meth:`.Session._identity_lookup` from
+           the :class:`.Query` object to the :class:`.Session`.
+
+        """
+
+        if identity_token is not None:
+            return super(ShardedSession, self)._identity_lookup(
+                mapper,
+                primary_key_identity,
+                identity_token=identity_token,
+                **kw
+            )
+        else:
+            q = self.query(mapper)
+            if lazy_loaded_from:
+                q = q._set_lazyload_from(lazy_loaded_from)
+            for shard_id in self.id_chooser(q, primary_key_identity):
+                obj = super(ShardedSession, self)._identity_lookup(
+                    mapper,
+                    primary_key_identity,
+                    identity_token=shard_id,
+                    lazy_loaded_from=lazy_loaded_from,
+                    **kw
+                )
+                if obj is not None:
+                    return obj
+
+            return None
+
     def _choose_shard_and_assign(self, mapper, instance, **kw):
         if instance is not None:
             state = inspect(instance)
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index a30cfa7e4..37bd77f63 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -975,60 +975,6 @@ class Query(Generative):
         """
         return self._get_impl(ident, loading.load_on_pk_identity)
 
-    def _identity_lookup(
-        self,
-        mapper,
-        primary_key_identity,
-        identity_token=None,
-        passive=attributes.PASSIVE_OFF,
-        lazy_loaded_from=None,
-    ):
-        """Locate an object in the identity map.
-
-        Given a primary key identity, constructs an identity key and then
-        looks in the session's identity map.  If present, the object may
-        be run through unexpiration rules (e.g. load unloaded attributes,
-        check if was deleted).
-
-        For performance reasons, while the :class:`.Query` must be
-        instantiated, it may be instantiated with no entities, and the
-        mapper is passed::
-
-            obj = session.query()._identity_lookup(inspect(SomeClass), (1, ))
-
-        :param mapper: mapper in use
-        :param primary_key_identity: the primary key we are searching for, as
-         a tuple.
-        :param identity_token: identity token that should be used to create
-         the identity key.  Used as is, however overriding subclasses can
-         repurpose this in order to interpret the value in a special way,
-         such as if None then look among multiple target tokens.
-        :param passive: passive load flag passed to
-         :func:`.loading.get_from_identity`, which impacts the behavior if
-         the object is found; the object may be validated and/or unexpired
-         if the flag allows for SQL to be emitted.
-        :param lazy_loaded_from: an :class:`.InstanceState` that is
-         specifically asking for this identity as a related identity.  Used
-         for sharding schemes where there is a correspondence between an object
-         and a related object being lazy-loaded (or otherwise
-         relationship-loaded).
-
-         .. versionadded:: 1.2.9
-
-        :return: None if the object is not found in the identity map, *or*
-         if the object was unexpired and found to have been deleted.
-         if passive flags disallow SQL and the object is expired, returns
-         PASSIVE_NO_RESULT.   In all other cases the instance is returned.
-
-        .. versionadded:: 1.2.7
-
-        """
-
-        key = mapper.identity_key_from_primary_key(
-            primary_key_identity, identity_token=identity_token
-        )
-        return loading.get_from_identity(self.session, key, passive)
-
     def _get_impl(self, primary_key_identity, db_load_fn, identity_token=None):
         # convert composite types to individual args
         if hasattr(primary_key_identity, "__composite_values__"):
@@ -1071,7 +1017,7 @@ class Query(Generative):
             and self._for_update_arg is None
         ):
 
-            instance = self._identity_lookup(
+            instance = self.session._identity_lookup(
                 mapper, primary_key_identity, identity_token=identity_token
             )
 
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 5ff390e19..c345c4281 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -1542,6 +1542,60 @@ class Session(_SessionClassMethods):
 
         return self._query_cls(entities, self, **kwargs)
 
+    def _identity_lookup(
+        self,
+        mapper,
+        primary_key_identity,
+        identity_token=None,
+        passive=attributes.PASSIVE_OFF,
+        lazy_loaded_from=None,
+    ):
+        """Locate an object in the identity map.
+
+        Given a primary key identity, constructs an identity key and then
+        looks in the session's identity map.  If present, the object may
+        be run through unexpiration rules (e.g. load unloaded attributes,
+        check if was deleted).
+
+        e.g.::
+
+            obj = session._identity_lookup(inspect(SomeClass), (1, ))
+
+        :param mapper: mapper in use
+        :param primary_key_identity: the primary key we are searching for, as
+         a tuple.
+        :param identity_token: identity token that should be used to create
+         the identity key.  Used as is, however overriding subclasses can
+         repurpose this in order to interpret the value in a special way,
+         such as if None then look among multiple target tokens.
+        :param passive: passive load flag passed to
+         :func:`.loading.get_from_identity`, which impacts the behavior if
+         the object is found; the object may be validated and/or unexpired
+         if the flag allows for SQL to be emitted.
+        :param lazy_loaded_from: an :class:`.InstanceState` that is
+         specifically asking for this identity as a related identity.  Used
+         for sharding schemes where there is a correspondence between an object
+         and a related object being lazy-loaded (or otherwise
+         relationship-loaded).
+
+        :return: None if the object is not found in the identity map, *or*
+         if the object was unexpired and found to have been deleted.
+         if passive flags disallow SQL and the object is expired, returns
+         PASSIVE_NO_RESULT.   In all other cases the instance is returned.
+
+        .. versionchanged:: 1.4.0 - the :meth:`.Session._identity_lookup`
+           method was moved from :class:`.Query` to
+           :class:`.Session`, to avoid having to instantiate the
+           :class:`.Query` object.
+
+
+        """
+
+        key = mapper.identity_key_from_primary_key(
+            primary_key_identity, identity_token=identity_token
+        )
+        return loading.get_from_identity(self, key, passive)
+
     @property
     @util.contextmanager
     def no_autoflush(self):
diff --git a/lib/sqlalchemy/orm/strategies.py b/lib/sqlalchemy/orm/strategies.py
index 49b5b4f64..cdc5c10c9 100644
--- a/lib/sqlalchemy/orm/strategies.py
+++ b/lib/sqlalchemy/orm/strategies.py
@@ -709,7 +709,8 @@ class LazyLoader(AbstractRelationshipLoader, util.MemoizedSlots):
             # Query class in use, as it may have special rules for how it
             # does this, including how it decides what the correct
             # identity_token would be for this identity.
-            instance = session.query()._identity_lookup(
+
+            instance = session._identity_lookup(
                 self.entity,
                 primary_key_identity,
                 passive=passive,