Merge "Change query._identity_lookup into a normal instance method"

6 files changed, 123 insertions, 40 deletions

diff --git a/doc/build/changelog/unreleased_12/4256.rst b/doc/build/changelog/unreleased_12/4256.rst
new file mode 100644
index 000000000..9937b9529
--- /dev/null
+++ b/doc/build/changelog/unreleased_12/4256.rst
@@ -0,0 +1,10 @@
+.. change::
+	:tags: bug, orm
+	:tickets: 4256
+
+	Fixed regression in 1.2.7 caused by :ticket:`4228`, which itself was fixing
+	a 1.2-level regression, where the ``query_cls`` callable passed to a
+	:class:`.Session` was assumed to be a subclass of :class:`.Query`  with
+	class method availability, as opposed to an arbitrary callable.    In
+	particular, the dogpile caching example illustrates ``query_cls`` as a
+	function and not a :class:`.Query` subclass.
diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py
index 6516950ed..c7770d195 100644
--- a/lib/sqlalchemy/ext/horizontal_shard.py
+++ b/lib/sqlalchemy/ext/horizontal_shard.py
@@ -64,10 +64,8 @@ class ShardedQuery(Query):
             # were done, this is where it would happen
             return iter(partial)
 
-    @classmethod
     def _identity_lookup(
-            cls, session, mapper, primary_key_identity, identity_token=None,
-            **kw):
+            self, mapper, primary_key_identity, identity_token=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).
@@ -75,18 +73,15 @@ class ShardedQuery(Query):
         """
 
         if identity_token is not None:
-            return super(ShardedQuery, cls)._identity_lookup(
-                session, mapper, primary_key_identity,
-                identity_token=identity_token,
-                **kw
+            return super(ShardedQuery, self)._identity_lookup(
+                mapper, primary_key_identity,
+                identity_token=identity_token, **kw
             )
         else:
-            q = cls([mapper], session)
-            for shard_id in q.id_chooser(q, primary_key_identity):
-                obj = super(ShardedQuery, cls)._identity_lookup(
-                    session, mapper, primary_key_identity,
-                    identity_token=shard_id,
-                    **kw
+            q = self.session.query(mapper)
+            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
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index a5f3d01f6..56e42a702 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -147,10 +147,11 @@ class Query(object):
         self._entities = []
         self._primary_entity = None
         self._has_mapper_entities = False
-        for ent in util.to_list(entities):
-            entity_wrapper(self, ent)
+        if entities:
+            for ent in util.to_list(entities):
+                entity_wrapper(self, ent)
 
-        self._set_entity_selectables(self._entities)
+            self._set_entity_selectables(self._entities)
 
     def _set_entity_selectables(self, entities):
         self._mapper_adapter_map = d = self._mapper_adapter_map.copy()
@@ -885,10 +886,9 @@ class Query(object):
         return self._get_impl(
             ident, loading.load_on_pk_identity)
 
-    @classmethod
-    def _identity_lookup(
-            cls, session, mapper, primary_key_identity, identity_token=None,
-            passive=attributes.PASSIVE_OFF):
+    def _identity_lookup(self, mapper, primary_key_identity,
+                         identity_token=None,
+                         passive=attributes.PASSIVE_OFF):
         """Locate an object in the identity map.
 
         Given a primary key identity, constructs an identity key and then
@@ -896,8 +896,13 @@ class Query(object):
         be run through unexpiration rules (e.g. load unloaded attributes,
         check if was deleted).
 
-        :param session: Session in use
-        :param mapper: target mapper
+        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
@@ -916,10 +921,11 @@ class Query(object):
         .. versionadded:: 1.2.7
 
         """
+
         key = mapper.identity_key_from_primary_key(
             primary_key_identity, identity_token=identity_token)
         return loading.get_from_identity(
-            session, key, passive)
+            self.session, key, passive)
 
     def _get_impl(
             self, primary_key_identity, db_load_fn, identity_token=None):
@@ -942,7 +948,7 @@ class Query(object):
                 self._for_update_arg is None:
 
             instance = self._identity_lookup(
-                self.session, mapper, primary_key_identity,
+                mapper, primary_key_identity,
                 identity_token=identity_token)
 
             if instance is not None:
diff --git a/lib/sqlalchemy/orm/strategies.py b/lib/sqlalchemy/orm/strategies.py
index 00c83cea4..93288c3d6 100644
--- a/lib/sqlalchemy/orm/strategies.py
+++ b/lib/sqlalchemy/orm/strategies.py
@@ -615,10 +615,9 @@ class LazyLoader(AbstractRelationshipLoader, util.MemoizedSlots):
             # look for this identity in the identity map.  Delegate to the
             # 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_cls._identity_lookup(
-                session, self.mapper, primary_key_identity,
-                passive=passive
+            # identity_token would be for this identity.
+            instance = session.query()._identity_lookup(
+                self.mapper, primary_key_identity, passive=passive
             )
 
             if instance is not None:
diff --git a/test/orm/test_query.py b/test/orm/test_query.py
index ed2b17edc..d4a122a23 100644
--- a/test/orm/test_query.py
+++ b/test/orm/test_query.py
@@ -4523,3 +4523,76 @@ class SessionBindTest(QueryTest):
 
         with self._assert_bind_args(session):
             session.query(func.max(User.score)).scalar()
+
+
+class QueryClsTest(QueryTest):
+    def _fn_fixture(self):
+        def query(*arg, **kw):
+            return Query(*arg, **kw)
+        return query
+
+    def _subclass_fixture(self):
+        class MyQuery(Query):
+            pass
+
+        return MyQuery
+
+    def _callable_fixture(self):
+        class MyQueryFactory(object):
+            def __call__(self, *arg, **kw):
+                return Query(*arg, **kw)
+
+        return MyQueryFactory()
+
+    def _test_get(self, fixture):
+        User = self.classes.User
+
+        s = Session(query_cls=fixture())
+
+        assert s.query(User).get(19) is None
+        u = s.query(User).get(7)
+        u2 = s.query(User).get(7)
+        assert u is u2
+
+    def _test_o2m_lazyload(self, fixture):
+        User, Address = self.classes('User', 'Address')
+
+        s = Session(query_cls=fixture())
+
+        u1 = s.query(User).filter(User.id == 7).first()
+        eq_(u1.addresses, [Address(id=1)])
+
+    def _test_m2o_lazyload(self, fixture):
+        User, Address = self.classes('User', 'Address')
+
+        s = Session(query_cls=fixture())
+
+        a1 = s.query(Address).filter(Address.id == 1).first()
+        eq_(a1.user, User(id=7))
+
+    def test_callable_get(self):
+        self._test_get(self._callable_fixture)
+
+    def test_subclass_get(self):
+        self._test_get(self._subclass_fixture)
+
+    def test_fn_get(self):
+        self._test_get(self._fn_fixture)
+
+    def test_callable_o2m_lazyload(self):
+        self._test_o2m_lazyload(self._callable_fixture)
+
+    def test_subclass_o2m_lazyload(self):
+        self._test_o2m_lazyload(self._subclass_fixture)
+
+    def test_fn_o2m_lazyload(self):
+        self._test_o2m_lazyload(self._fn_fixture)
+
+    def test_callable_m2o_lazyload(self):
+        self._test_m2o_lazyload(self._callable_fixture)
+
+    def test_subclass_m2o_lazyload(self):
+        self._test_m2o_lazyload(self._subclass_fixture)
+
+    def test_fn_m2o_lazyload(self):
+        self._test_m2o_lazyload(self._fn_fixture)
diff --git a/test/profiles.txt b/test/profiles.txt
index 7b5386269..394d9ea38 100644
--- a/test/profiles.txt
+++ b/test/profiles.txt
@@ -198,18 +198,18 @@ test.aaa_profiling.test_orm.JoinedEagerLoadTest.test_fetch_results 3.6_sqlite_py
 
 # TEST: test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity
 
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_mysql_mysqldb_dbapiunicode_cextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_mysql_mysqldb_dbapiunicode_nocextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_postgresql_psycopg2_dbapiunicode_cextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_postgresql_psycopg2_dbapiunicode_nocextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_sqlite_pysqlite_dbapiunicode_cextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_sqlite_pysqlite_dbapiunicode_nocextensions 18987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_mysql_mysqldb_dbapiunicode_cextensions 19987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_mysql_mysqldb_dbapiunicode_nocextensions 19987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_postgresql_psycopg2_dbapiunicode_cextensions 19987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_postgresql_psycopg2_dbapiunicode_nocextensions 19987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_sqlite_pysqlite_dbapiunicode_cextensions 19987
-test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_sqlite_pysqlite_dbapiunicode_nocextensions 19987
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_mysql_mysqldb_dbapiunicode_cextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_mysql_mysqldb_dbapiunicode_nocextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_postgresql_psycopg2_dbapiunicode_cextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_postgresql_psycopg2_dbapiunicode_nocextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_sqlite_pysqlite_dbapiunicode_cextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 2.7_sqlite_pysqlite_dbapiunicode_nocextensions 21984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_mysql_mysqldb_dbapiunicode_cextensions 22984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_mysql_mysqldb_dbapiunicode_nocextensions 22984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_postgresql_psycopg2_dbapiunicode_cextensions 22984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_postgresql_psycopg2_dbapiunicode_nocextensions 22984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_sqlite_pysqlite_dbapiunicode_cextensions 22984
+test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_identity 3.6_sqlite_pysqlite_dbapiunicode_nocextensions 22984
 
 # TEST: test.aaa_profiling.test_orm.LoadManyToOneFromIdentityTest.test_many_to_one_load_no_identity