From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 20 Jul 2020 20:51:24 +0000 (-0400)
Subject: Ensure _distinct_on removed within from_selectable
X-Git-Tag: rel_1_4_0b1~217
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=e385fc6215aea605e4b09babbb2c21702000b84f;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Ensure _distinct_on removed within from_selectable

Fixes: #5469
Change-Id: I85498d3380c2efb595782d2ba92cfd692c3efc48
---

diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 7bf69f99f3..d991f62293 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -1326,6 +1326,7 @@ class Query(
             "_last_joined_entity",
             "_legacy_setup_joins",
             "_distinct",
+            "_distinct_on",
             "_having_criteria",
             "_prefixes",
             "_suffixes",
diff --git a/test/orm/test_query.py b/test/orm/test_query.py
index 93d1e88a5a..b6d014936b 100644
--- a/test/orm/test_query.py
+++ b/test/orm/test_query.py
@@ -4016,6 +4016,66 @@ class DistinctTest(QueryTest, AssertsCompiledSQL):
             ") AS anon_1",
         )
 
+    def test_columns_augmented_sql_union_one(self):
+        User, Address = self.classes.User, self.classes.Address
+
+        sess = create_session()
+
+        q = (
+            sess.query(
+                User.id,
+                User.name.label("foo"),
+                Address.id,
+                Address.email_address,
+            )
+            .distinct()
+            .order_by(User.id, User.name, Address.email_address)
+        )
+        q2 = sess.query(
+            User.id, User.name.label("foo"), Address.id, Address.email_address,
+        )
+
+        self.assert_compile(
+            q.union(q2),
+            "SELECT anon_1.users_id AS anon_1_users_id, "
+            "anon_1.foo AS anon_1_foo, anon_1.addresses_id AS "
+            "anon_1_addresses_id, anon_1.addresses_email_address AS "
+            "anon_1_addresses_email_address FROM "
+            "((SELECT DISTINCT users.id AS users_id, users.name AS foo, "
+            "addresses.id AS addresses_id, addresses.email_address "
+            "AS addresses_email_address FROM users, addresses "
+            "ORDER BY users.id, users.name, addresses.email_address) "
+            "UNION SELECT users.id AS users_id, users.name AS foo, "
+            "addresses.id AS addresses_id, addresses.email_address AS "
+            "addresses_email_address FROM users, addresses) AS anon_1",
+        )
+
+    def test_columns_augmented_sql_union_two(self):
+        User, Address = self.classes.User, self.classes.Address
+
+        sess = create_session()
+
+        q = (
+            sess.query(User.id, User.name.label("foo"), Address.id,)
+            .distinct(Address.email_address)
+            .order_by(User.id, User.name)
+        )
+        q2 = sess.query(User.id, User.name.label("foo"), Address.id)
+
+        self.assert_compile(
+            q.union(q2),
+            "SELECT anon_1.users_id AS anon_1_users_id, "
+            "anon_1.foo AS anon_1_foo, anon_1.addresses_id AS "
+            "anon_1_addresses_id FROM "
+            "((SELECT DISTINCT ON (addresses.email_address) users.id "
+            "AS users_id, users.name AS foo, "
+            "addresses.id AS addresses_id FROM users, addresses "
+            "ORDER BY users.id, users.name) "
+            "UNION SELECT users.id AS users_id, users.name AS foo, "
+            "addresses.id AS addresses_id FROM users, addresses) AS anon_1",
+            dialect="postgresql",
+        )
+
     def test_columns_augmented_sql_two(self):
         User, Address = self.classes.User, self.classes.Address
 
@@ -4072,6 +4132,40 @@ class DistinctTest(QueryTest, AssertsCompiledSQL):
             dialect="postgresql",
         )
 
+    def test_columns_augmented_distinct_on(self):
+        User, Address = self.classes.User, self.classes.Address
+
+        sess = create_session()
+
+        q = (
+            sess.query(
+                User.id,
+                User.name.label("foo"),
+                Address.id,
+                Address.email_address,
+            )
+            .distinct(Address.email_address)
+            .order_by(User.id, User.name, Address.email_address)
+            .from_self(User.id, User.name.label("foo"), Address.id)
+        )
+
+        # Address.email_address is added because of DISTINCT,
+        # however User.id, User.name are not b.c. they're already there,
+        # even though User.name is labeled
+        self.assert_compile(
+            q,
+            "SELECT anon_1.users_id AS anon_1_users_id, anon_1.foo AS foo, "
+            "anon_1.addresses_id AS anon_1_addresses_id "
+            "FROM ("
+            "SELECT DISTINCT ON (addresses.email_address) "
+            "users.id AS users_id, users.name AS foo, "
+            "addresses.id AS addresses_id, addresses.email_address AS "
+            "addresses_email_address FROM users, addresses ORDER BY "
+            "users.id, users.name, addresses.email_address"
+            ") AS anon_1",
+            dialect="postgresql",
+        )
+
     def test_columns_augmented_sql_three_using_label_reference(self):
         User, Address = self.classes.User, self.classes.Address