Merge pull request #99 from dimagi/union-cte

Fix UNION query with common table expression

Author: millerdev

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 = []

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()
+
+        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),
+        ])