From 45c4f5d076b13f74203f730cb663f0c721a84745 Mon Sep 17 00:00:00 2001
From: Daniel Miller <dmiller@dimagi.com>
Date: Thu, 2 Jan 2025 15:59:18 -0500
Subject: [PATCH 1/2] Fix UNION query with common table expression

---
 django_cte/query.py |  2 +-
 tests/test_cte.py   | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/django_cte/query.py b/django_cte/query.py
index 5201bfd..64294d1 100644
--- a/django_cte/query.py
+++ b/django_cte/query.py
@@ -68,7 +68,7 @@ class CTECompiler(object):
 
     @classmethod
     def generate_sql(cls, connection, query, as_sql):
-        if query.combinator:
+        if not query._with_ctes:
             return as_sql()
 
         ctes = []
diff --git a/tests/test_cte.py b/tests/test_cte.py
index 49f00cd..d3fcaeb 100644
--- a/tests/test_cte.py
+++ b/tests/test_cte.py
@@ -596,3 +596,34 @@ def test_left_outer_join_on_empty_result_set_cte(self):
         )
 
         self.assertEqual(len(orders), 22)
+
+    def test_union_query_with_cte(self):
+        orders = (
+            Order.objects
+            .filter(region__parent="sun")
+            .only("region", "amount")
+        )
+        orders_cte = With(orders, name="orders_cte")
+        orders_cte_queryset = orders_cte.queryset().only("region", "amount")
+
+        earth_orders = orders_cte_queryset.filter(region="earth")
+        mars_orders = orders_cte_queryset.filter(region="mars")
+
+        earth_mars = earth_orders.union(mars_orders, all=True)
+        earth_mars_cte = (
+            earth_mars
+            .with_cte(orders_cte)
+            .order_by("region", "amount")
+            .values_list("region", "amount")
+        )
+        print(earth_mars_cte.query)
+
+        self.assertEqual(list(earth_mars_cte), [
+            ('earth', 30),
+            ('earth', 31),
+            ('earth', 32),
+            ('earth', 33),
+            ('mars', 40),
+            ('mars', 41),
+            ('mars', 42),
+        ])

From 0e7f07460d2d75ef172de2b5bb14ccf38c3fbd3e Mon Sep 17 00:00:00 2001
From: Daniel Miller <dmiller@dimagi.com>
Date: Fri, 3 Jan 2025 12:08:20 -0500
Subject: [PATCH 2/2] Remove redundant .only(...)

---
 tests/test_cte.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_cte.py b/tests/test_cte.py
index d3fcaeb..ed7a112 100644
--- a/tests/test_cte.py
+++ b/tests/test_cte.py
@@ -604,7 +604,7 @@ def test_union_query_with_cte(self):
             .only("region", "amount")
         )
         orders_cte = With(orders, name="orders_cte")
-        orders_cte_queryset = orders_cte.queryset().only("region", "amount")
+        orders_cte_queryset = orders_cte.queryset()
 
         earth_orders = orders_cte_queryset.filter(region="earth")
         mars_orders = orders_cte_queryset.filter(region="mars")