Fixed #29338 -- Allowed using combined queryset in Subquery.

Thanks Eugene Kovalev for the initial patch, Simon Charette for the
review, and Chetan Khanna for help.

1 files changed, 35 insertions, 2 deletions

diff --git a/tests/queries/test_qs_combinators.py b/tests/queries/test_qs_combinators.py
index 89d08ef6db..37ce37c2c1 100644
--- a/tests/queries/test_qs_combinators.py
+++ b/tests/queries/test_qs_combinators.py
@@ -1,11 +1,11 @@
 import operator
 
 from django.db import DatabaseError, NotSupportedError, connection
-from django.db.models import Exists, F, IntegerField, OuterRef, Value
+from django.db.models import Exists, F, IntegerField, OuterRef, Subquery, Value
 from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
 from django.test.utils import CaptureQueriesContext
 
-from .models import Celebrity, Number, ReservedName
+from .models import Author, Celebrity, ExtraInfo, Number, ReservedName
 
 
 @skipUnlessDBFeature('supports_select_union')
@@ -252,6 +252,39 @@ class QuerySetSetOperationTests(TestCase):
             [reserved_name.pk],
         )
 
+    def test_union_in_subquery(self):
+        ReservedName.objects.bulk_create([
+            ReservedName(name='rn1', order=8),
+            ReservedName(name='rn2', order=1),
+            ReservedName(name='rn3', order=5),
+        ])
+        qs1 = Number.objects.filter(num__gt=7, num=OuterRef('order'))
+        qs2 = Number.objects.filter(num__lt=2, num=OuterRef('order'))
+        self.assertCountEqual(
+            ReservedName.objects.annotate(
+                number=Subquery(qs1.union(qs2).values('num')),
+            ).filter(number__isnull=False).values_list('order', flat=True),
+            [8, 1],
+        )
+
+    def test_union_in_subquery_related_outerref(self):
+        e1 = ExtraInfo.objects.create(value=7, info='e3')
+        e2 = ExtraInfo.objects.create(value=5, info='e2')
+        e3 = ExtraInfo.objects.create(value=1, info='e1')
+        Author.objects.bulk_create([
+            Author(name='a1', num=1, extra=e1),
+            Author(name='a2', num=3, extra=e2),
+            Author(name='a3', num=2, extra=e3),
+        ])
+        qs1 = ExtraInfo.objects.order_by().filter(value=OuterRef('num'))
+        qs2 = ExtraInfo.objects.order_by().filter(value__lt=OuterRef('extra__value'))
+        qs = Author.objects.annotate(
+            info=Subquery(qs1.union(qs2).values('info')[:1]),
+        ).filter(info__isnull=False).values_list('name', flat=True)
+        self.assertCountEqual(qs, ['a1', 'a2'])
+        # Combined queries don't mutate.
+        self.assertCountEqual(qs, ['a1', 'a2'])
+
     def test_count_union(self):
         qs1 = Number.objects.filter(num__lte=1).values('num')
         qs2 = Number.objects.filter(num__gte=2, num__lte=3).values('num')