From: Eric Masseran <eric.masseran@gmail.com>
Date: Fri, 6 Aug 2021 08:11:09 +0000 (-0400)
Subject: Dispatch independent ctes on compound select
X-Git-Tag: rel_1_4_23~17^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=770d50092f847ef7e79a895c72a680f7626a1164;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Dispatch independent ctes on compound select

Fix issue in :class:`_sql.CTE` where new :meth:`_sql.HasCTE.add_cte` method
added in version 1.4.21 / :ticket:`6752` failed to function correctly for
"compound select" structures such as :func:`_sql.union`,
:func:`_sql.union_all`, :func:`_sql.except`, etc. Pull request courtesy
Eric Masseran.

Fixes: #6752
Closes: #6849
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/6849
Pull-request-sha: 1c4b4d72b2789cf89ff5043ca964ebdd6c9a6617

Change-Id: I49a16a4fc2af8299502011f3a02d8a2ad93255e3
---

diff --git a/doc/build/changelog/unreleased_14/6752.rst b/doc/build/changelog/unreleased_14/6752.rst
new file mode 100644
index 0000000000..9cc07c2a1d
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/6752.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, sql
+    :tickets: 6752
+
+    Fix issue in :class:`_sql.CTE` where new :meth:`_sql.HasCTE.add_cte` method
+    added in version 1.4.21 / :ticket:`6752` failed to function correctly for
+    "compound select" structures such as :func:`_sql.union`,
+    :func:`_sql.union_all`, :func:`_sql.except`, etc. Pull request courtesy
+    Eric Masseran.
\ No newline at end of file
diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index a81507acb9..1e090ef103 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1787,6 +1787,8 @@ class SQLCompiler(Compiled):
         if toplevel and not self.compile_state:
             self.compile_state = compile_state
 
+        compound_stmt = compile_state.statement
+
         entry = self._default_stack_entry if toplevel else self.stack[-1]
         need_result_map = toplevel or (
             not compound_index
@@ -1807,6 +1809,10 @@ class SQLCompiler(Compiled):
             }
         )
 
+        if compound_stmt._independent_ctes:
+            for cte in compound_stmt._independent_ctes:
+                cte._compiler_dispatch(self, **kwargs)
+
         keyword = self.compound_keywords.get(cs.keyword)
 
         text = (" " + keyword + " ").join(
diff --git a/test/sql/test_cte.py b/test/sql/test_cte.py
index 77905cd896..cc663d53fb 100644
--- a/test/sql/test_cte.py
+++ b/test/sql/test_cte.py
@@ -1533,6 +1533,37 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             checkparams={"param_1": 10, "price_1": 50, "price_2": 45},
         )
 
+    def test_compound_select_uses_independent_cte(self):
+        products = table("products", column("id"), column("price"))
+
+        upd_cte = (
+            products.update().values(price=10).where(products.c.price > 50)
+        ).cte()
+
+        stmt = (
+            products.select()
+            .where(products.c.price < 45)
+            .union(products.select().where(products.c.price > 90))
+            .add_cte(upd_cte)
+        )
+
+        self.assert_compile(
+            stmt,
+            "WITH anon_1 AS (UPDATE products SET price=:param_1 "
+            "WHERE products.price > :price_1) "
+            "SELECT products.id, products.price "
+            "FROM products WHERE products.price < :price_2 "
+            "UNION "
+            "SELECT products.id, products.price "
+            "FROM products WHERE products.price > :price_3",
+            checkparams={
+                "param_1": 10,
+                "price_1": 50,
+                "price_2": 45,
+                "price_3": 90,
+            },
+        )
+
     def test_insert_uses_independent_cte(self):
         products = table("products", column("id"), column("price"))