Merge "qualify compile_state updates for non-current entities"

zzzeek · Gerrit Code Review · commit 2900c472d7f8 · 2021-08-23T16:29:42.000Z
diff --git a/doc/build/changelog/unreleased_14/6924.rst b/doc/build/changelog/unreleased_14/6924.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, orm, regression
+    :tickets: 6924
+
+    Fixed issue in recently repaired ``Query.with_entities()`` method where the
+    flag that determines automatic uniquing for legacy ORM ``Query`` objects
+    only would be set to ``True`` inappropriately in cases where the
+    ``with_entities()`` call would be setting the ``Query`` to return
+    column-only rows, which are not uniqued.
diff --git a/lib/sqlalchemy/orm/context.py b/lib/sqlalchemy/orm/context.py
@@ -417,7 +417,10 @@ def create_for_statement(cls, statement_container, compiler, **kw):
         )
 
         _QueryEntity.to_compile_state(
-            self, statement_container._raw_columns, self._entities
+            self,
+            statement_container._raw_columns,
+            self._entities,
+            is_current_entities=True,
         )
 
         self.current_path = statement_container._compile_options._current_path
@@ -590,14 +593,18 @@ def create_for_statement(cls, statement, compiler, **kw):
                     self,
                     memoized_entities._raw_columns,
                     [],
+                    is_current_entities=False,
                 )
                 for memoized_entities in (
                     select_statement._memoized_select_entities
                 )
             }
 
         _QueryEntity.to_compile_state(
-            self, select_statement._raw_columns, self._entities
+            self,
+            select_statement._raw_columns,
+            self._entities,
+            is_current_entities=True,
         )
 
         self.current_path = select_statement._compile_options._current_path
@@ -839,7 +846,9 @@ def _create_entities_collection(cls, query, legacy):
 
         # entities will also set up polymorphic adapters for mappers
         # that have with_polymorphic configured
-        _QueryEntity.to_compile_state(self, query._raw_columns, self._entities)
+        _QueryEntity.to_compile_state(
+            self, query._raw_columns, self._entities, is_current_entities=True
+        )
         return self
 
     @classmethod
@@ -2278,13 +2287,18 @@ class _QueryEntity(object):
     use_id_for_hash = False
 
     @classmethod
-    def to_compile_state(cls, compile_state, entities, entities_collection):
+    def to_compile_state(
+        cls, compile_state, entities, entities_collection, is_current_entities
+    ):
 
         for idx, entity in enumerate(entities):
             if entity._is_lambda_element:
                 if entity._is_sequence:
                     cls.to_compile_state(
-                        compile_state, entity._resolved, entities_collection
+                        compile_state,
+                        entity._resolved,
+                        entities_collection,
+                        is_current_entities,
                     )
                     continue
                 else:
@@ -2294,7 +2308,10 @@ def to_compile_state(cls, compile_state, entities, entities_collection):
                 if entity.is_selectable:
                     if "parententity" in entity._annotations:
                         _MapperEntity(
-                            compile_state, entity, entities_collection
+                            compile_state,
+                            entity,
+                            entities_collection,
+                            is_current_entities,
                         )
                     else:
                         _ColumnEntity._for_columns(
@@ -2343,12 +2360,15 @@ class _MapperEntity(_QueryEntity):
         "_polymorphic_discriminator",
     )
 
-    def __init__(self, compile_state, entity, entities_collection):
+    def __init__(
+        self, compile_state, entity, entities_collection, is_current_entities
+    ):
         entities_collection.append(self)
-        if compile_state._primary_entity is None:
-            compile_state._primary_entity = self
-        compile_state._has_mapper_entities = True
-        compile_state._has_orm_entities = True
+        if is_current_entities:
+            if compile_state._primary_entity is None:
+                compile_state._primary_entity = self
+            compile_state._has_mapper_entities = True
+            compile_state._has_orm_entities = True
 
         entity = entity._annotations["parententity"]
         entity._post_inspect
diff --git a/test/orm/test_froms.py b/test/orm/test_froms.py
@@ -29,6 +29,7 @@
 from sqlalchemy.orm import mapper
 from sqlalchemy.orm import relationship
 from sqlalchemy.orm import Session
+from sqlalchemy.orm.context import ORMSelectCompileState
 from sqlalchemy.orm.util import join
 from sqlalchemy.sql import column
 from sqlalchemy.sql import table
@@ -1699,6 +1700,59 @@ def test_correlated_subquery(self):
             ],
         )
 
+    @testing.combinations((True,), (False,))
+    def test_no_uniquing_cols_legacy(self, with_entities):
+        """test #6924"""
+        User = self.classes.User
+        Address = self.classes.Address
+
+        sess = fixture_session()
+
+        if with_entities:
+            q = (
+                sess.query(User)
+                .join(Address)
+                .filter(Address.user_id == 8)
+                .with_entities(User.id, User.name)
+                .order_by(User.id)
+            )
+        else:
+            q = (
+                sess.query(User.id, User.name)
+                .join(Address)
+                .filter(Address.user_id == 8)
+                .order_by(User.id)
+            )
+
+        is_(q._compile_state()._primary_entity, None)
+
+        eq_(q.all(), [(8, "ed"), (8, "ed"), (8, "ed")])
+
+    @testing.combinations((True,), (False,))
+    def test_no_uniquing_cols(self, with_entities):
+        """test #6924"""
+        User = self.classes.User
+        Address = self.classes.Address
+
+        if with_entities:
+            stmt = (
+                select(User)
+                .join(Address)
+                .filter(Address.user_id == 8)
+                .with_only_columns(User.id, User.name)
+                .order_by(User.id)
+            )
+        else:
+            stmt = (
+                select(User.id, User.name)
+                .join(Address)
+                .filter(Address.user_id == 8)
+                .order_by(User.id)
+            )
+
+        compile_state = ORMSelectCompileState.create_for_statement(stmt, None)
+        is_(compile_state._primary_entity, None)
+
     def test_column_queries_one(self):
         User = self.classes.User
 
