set _render_for_subquery for legacy set ops

Adjusted :meth:`_orm.Query.union` and similar set operations to be
correctly compatible with the new capabilities just added in
:ticket:`6661`, with SQLAlchemy 1.4.19, such that the SELECT statements
rendered as elements of the UNION or other set operation will include
directly mapped columns that are mapped as deferred; this both fixes a
regression involving unions with multiple levels of nesting that would
produce a column mismatch, and also allows the :func:`_orm.undefer` option
to be used at the top level of such a :class:`_orm.Query` without having to
apply the option to each of the elements within the UNION.

Fixes: #6678
Change-Id: Iba97ce7fd8a965499853256fd2eb7f61512db60f

diff --git a/doc/build/changelog/unreleased_14/6678.rst b/doc/build/changelog/unreleased_14/6678.rst
new file mode 100644 (file)
index 0000000..db461ee
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/6678.rst
@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, regression, orm
+    :tickets: 6678
+
+    Adjusted :meth:`_orm.Query.union` and similar set operations to be
+    correctly compatible with the new capabilities just added in
+    :ticket:`6661`, with SQLAlchemy 1.4.19, such that the SELECT statements
+    rendered as elements of the UNION or other set operation will include
+    directly mapped columns that are mapped as deferred; this both fixes a
+    regression involving unions with multiple levels of nesting that would
+    produce a column mismatch, and also allows the :func:`_orm.undefer` option
+    to be used at the top level of such a :class:`_orm.Query` without having to
+    apply the option to each of the elements within the UNION.

diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 7ba31fa7a0e01f75f152064f4158c29a3712853d..d8f4b4ea7c8db2fdee69ce29d16925fc55c00739 100644 (file)
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -611,7 +611,9 @@ class Query(
 
     def __clause_element__(self):
         return (
-            self.enable_eagerloads(False)
+            self._with_compile_options(
+                _enable_eagerloads=False, _render_for_subquery=True
+            )
             .set_label_style(LABEL_STYLE_TABLENAME_PLUS_COL)
             .statement
         )
@@ -675,6 +677,10 @@ class Query(
         """
         self._compile_options += {"_enable_eagerloads": value}
 
+    @_generative
+    def _with_compile_options(self, **opt):
+        self._compile_options += opt
+
     @util.deprecated_20(
         ":meth:`_orm.Query.with_labels` and :meth:`_orm.Query.apply_labels`",
         alternative="Use set_label_style(LABEL_STYLE_TABLENAME_PLUS_COL) "

diff --git a/test/orm/test_core_compilation.py b/test/orm/test_core_compilation.py
index 12cfef339d61296d5c754d20a78db9dddb3833ac..5f25b56e882b5147335a1fe88b02acb2baa93eda 100644 (file)
--- a/test/orm/test_core_compilation.py
+++ b/test/orm/test_core_compilation.py
@@ -8,6 +8,7 @@ from sqlalchemy import or_
 from sqlalchemy import select
 from sqlalchemy import testing
 from sqlalchemy import text
+from sqlalchemy import union
 from sqlalchemy import util
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm import column_property
@@ -611,6 +612,57 @@ class LoadersInSubqueriesTest(QueryTest, AssertsCompiledSQL):
             "FROM users JOIN anon_1 ON users.id = anon_1.id",
         )
 
+    def test_nested_union_deferred(self, deferred_fixture):
+        """test #6678"""
+        User = deferred_fixture
+
+        s1 = select(User).where(User.id == 5)
+        s2 = select(User).where(User.id == 6)
+
+        s3 = select(User).where(User.id == 7)
+
+        stmt = union(s1.union(s2), s3)
+
+        u_alias = aliased(User, stmt.subquery())
+
+        self.assert_compile(
+            select(u_alias),
+            "SELECT anon_1.id FROM ((SELECT users.name, users.id FROM users "
+            "WHERE users.id = :id_1 UNION SELECT users.name, users.id "
+            "FROM users WHERE users.id = :id_2) "
+            "UNION SELECT users.name AS name, users.id AS id "
+            "FROM users WHERE users.id = :id_3) AS anon_1",
+        )
+
+    def test_nested_union_undefer_option(self, deferred_fixture):
+        """test #6678
+
+        in this case we want to see that the unions include the deferred
+        columns so that if we undefer on the outside we can get the
+        column.
+
+        """
+        User = deferred_fixture
+
+        s1 = select(User).where(User.id == 5)
+        s2 = select(User).where(User.id == 6)
+
+        s3 = select(User).where(User.id == 7)
+
+        stmt = union(s1.union(s2), s3)
+
+        u_alias = aliased(User, stmt.subquery())
+
+        self.assert_compile(
+            select(u_alias).options(undefer(u_alias.name)),
+            "SELECT anon_1.name, anon_1.id FROM "
+            "((SELECT users.name, users.id FROM users "
+            "WHERE users.id = :id_1 UNION SELECT users.name, users.id "
+            "FROM users WHERE users.id = :id_2) "
+            "UNION SELECT users.name AS name, users.id AS id "
+            "FROM users WHERE users.id = :id_3) AS anon_1",
+        )
+
 
 class ExtraColsTest(QueryTest, AssertsCompiledSQL):
     __dialect__ = "default"

diff --git a/test/orm/test_query.py b/test/orm/test_query.py
index 77535a9bbb29d81588a67ec5ebf9fcabbb973ace..ed705576617dd7313c00812466a45d06a6ae123b 100644 (file)
--- a/test/orm/test_query.py
+++ b/test/orm/test_query.py
@@ -47,6 +47,7 @@ from sqlalchemy.orm import Bundle
 from sqlalchemy.orm import column_property
 from sqlalchemy.orm import contains_eager
 from sqlalchemy.orm import defer
+from sqlalchemy.orm import deferred
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm import lazyload
 from sqlalchemy.orm import mapper
@@ -56,6 +57,7 @@ from sqlalchemy.orm import selectinload
 from sqlalchemy.orm import Session
 from sqlalchemy.orm import subqueryload
 from sqlalchemy.orm import synonym
+from sqlalchemy.orm import undefer
 from sqlalchemy.orm.context import QueryContext
 from sqlalchemy.orm.util import join
 from sqlalchemy.orm.util import with_parent
@@ -3950,6 +3952,199 @@ class SetOpsTest(QueryTest, AssertsCompiledSQL):
         self.assert_sql_count(testing.db, go, 1)
 
 
+class SetOpsWDeferredTest(QueryTest, AssertsCompiledSQL):
+    __dialect__ = "default"
+
+    run_setup_mappers = None
+
+    @testing.fixture
+    def deferred_fixture(self):
+        User = self.classes.User
+        users = self.tables.users
+
+        mapper(
+            User,
+            users,
+            properties={
+                "name": deferred(users.c.name),
+                "name_upper": column_property(
+                    func.upper(users.c.name), deferred=True
+                ),
+            },
+        )
+
+        return User
+
+    def test_flat_twolevel_union_deferred(self, deferred_fixture):
+        """test #6678
+
+        note that due to #6661, the SELECTs inside the union include the
+        deferred "name" column.  this so we can switch to undeferred on
+        the outside.  this didn't work in 1.3.
+
+        """
+        User = deferred_fixture
+
+        s = fixture_session()
+
+        s1 = s.query(User).filter(User.id == 7)
+        s2 = s.query(User).filter(User.id == 8)
+
+        stmt = s1.union(s2).order_by(User.id)
+        self.assert_compile(
+            stmt,
+            "SELECT anon_1.users_id AS anon_1_users_id FROM "
+            "(SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_1 "
+            "UNION "
+            "SELECT users.name AS users_name, users.id AS users_id FROM users "
+            "WHERE users.id = :id_2) AS anon_1 ORDER BY anon_1.users_id",
+        )
+
+        recs = stmt.all()
+        eq_(recs, [User(id=7), User(id=8)])
+        for rec in recs:
+            assert "name" not in rec.__dict__
+
+        eq_(stmt.count(), 2)
+
+    def test_flat_twolevel_union_undeferred(self, deferred_fixture):
+        """test #6678
+
+        in this case we want to see that the unions include the deferred
+        columns so that if we undefer on the outside we can get the
+        column.   #6661 allows this.
+
+        """
+        User = deferred_fixture
+
+        s = fixture_session()
+
+        s1 = s.query(User).filter(User.id == 7)
+        s2 = s.query(User).filter(User.id == 8)
+
+        stmt = s1.union(s2).options(undefer(User.name)).order_by(User.id)
+        self.assert_compile(
+            stmt,
+            "SELECT anon_1.users_name AS anon_1_users_name, "
+            "anon_1.users_id AS anon_1_users_id FROM "
+            "(SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_1 "
+            "UNION "
+            "SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_2) AS anon_1 "
+            "ORDER BY anon_1.users_id",
+        )
+
+        recs = stmt.all()
+        for rec in recs:
+            assert "name" in rec.__dict__
+        eq_(
+            recs,
+            [
+                User(id=7, name="jack"),
+                User(id=8, name="ed"),
+            ],
+        )
+
+        eq_(stmt.count(), 2)
+
+    def test_nested_union_deferred(self, deferred_fixture):
+        """test #6678
+
+        note that due to #6661, the SELECTs inside the union include the
+        deferred "name" column.  this so we can switch to undeferred on
+        the outside.  this didn't work in 1.3.
+
+        """
+        User = deferred_fixture
+
+        s = fixture_session()
+
+        s1 = s.query(User).filter(User.id == 7)
+        s2 = s.query(User).filter(User.id == 8)
+
+        s3 = s.query(User).filter(User.id == 9)
+
+        stmt = s1.union(s2).union(s3).order_by(User.id)
+        self.assert_compile(
+            stmt,
+            "SELECT anon_1.anon_2_users_id AS anon_1_anon_2_users_id "
+            "FROM ("
+            "SELECT anon_2.users_name AS anon_2_users_name, "
+            "anon_2.users_id AS anon_2_users_id FROM "
+            "(SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_1 UNION "
+            "SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_2) AS anon_2 "
+            "UNION "
+            "SELECT users.name AS users_name, users.id AS users_id FROM users "
+            "WHERE users.id = :id_3) AS anon_1 "
+            "ORDER BY anon_1.anon_2_users_id",
+        )
+
+        recs = stmt.all()
+        eq_(recs, [User(id=7), User(id=8), User(id=9)])
+        for rec in recs:
+            assert "name" not in rec.__dict__
+
+        eq_(stmt.count(), 3)
+
+    def test_nested_union_undeferred(self, deferred_fixture):
+        """test #6678
+
+        in this case we want to see that the unions include the deferred
+        columns so that if we undefer on the outside we can get the
+        column.   #6661 allows this.
+
+        """
+        User = deferred_fixture
+
+        s = fixture_session()
+
+        s1 = s.query(User).filter(User.id == 7)
+        s2 = s.query(User).filter(User.id == 8)
+
+        s3 = s.query(User).filter(User.id == 9)
+
+        stmt = (
+            s1.union(s2)
+            .union(s3)
+            .options(undefer(User.name))
+            .order_by(User.id)
+        )
+        self.assert_compile(
+            stmt,
+            "SELECT anon_1.anon_2_users_name AS anon_1_anon_2_users_name, "
+            "anon_1.anon_2_users_id AS anon_1_anon_2_users_id "
+            "FROM ("
+            "SELECT anon_2.users_name AS anon_2_users_name, "
+            "anon_2.users_id AS anon_2_users_id FROM "
+            "(SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_1 UNION "
+            "SELECT users.name AS users_name, users.id AS users_id "
+            "FROM users WHERE users.id = :id_2) AS anon_2 "
+            "UNION "
+            "SELECT users.name AS users_name, users.id AS users_id FROM users "
+            "WHERE users.id = :id_3) AS anon_1 "
+            "ORDER BY anon_1.anon_2_users_id",
+        )
+
+        recs = stmt.all()
+        for rec in recs:
+            assert "name" in rec.__dict__
+        eq_(
+            recs,
+            [
+                User(id=7, name="jack"),
+                User(id=8, name="ed"),
+                User(id=9, name="fred"),
+            ],
+        )
+
+        eq_(stmt.count(), 3)
+
+
 class AggregateTest(QueryTest):
     def test_sum(self):
         Order = self.classes.Order