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 (file)
index 0000000..9cc07c2
--- /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 a81507acb94415f7f67f9b3656aece679e0ac875..1e090ef103fea7035d10d9c430bade1c7a975d6b 100644 (file)
--- 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 77905cd896dbfd040c7f44fe7991e2983174c20a..cc663d53fbf300810294c53c4fd2f7b3fe56cb5d 100644 (file)
--- 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"))