Fixed #36210, Refs #36181 -- Allowed Subquery usage in further lookups against composite pks.

Follow-up to 8561100425876bde3be4b2a22324655f74ff9609.

Co-authored-by: Simon Charette <charette.s@gmail.com>

2 files changed, 64 insertions, 5 deletions

diff --git a/tests/composite_pk/test_filter.py b/tests/composite_pk/test_filter.py
index 902bd832ff..d7ecfbec11 100644
--- a/tests/composite_pk/test_filter.py
+++ b/tests/composite_pk/test_filter.py
@@ -1,6 +1,6 @@
 from unittest.mock import patch
 
-from django.db import connection
+from django.db import NotSupportedError, connection
 from django.db.models import (
     Case,
     F,
@@ -14,7 +14,7 @@ from django.db.models import (
 )
 from django.db.models.functions import Cast
 from django.db.models.lookups import Exact
-from django.test import TestCase, skipUnlessDBFeature
+from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
 
 from .models import Comment, Tenant, User
 
@@ -492,6 +492,39 @@ class CompositePKFilterTests(TestCase):
                 queryset = Comment.objects.filter(**{f"id{lookup}": subquery})
                 self.assertEqual(queryset.count(), expected_count)
 
+    def test_outer_ref_pk_filter_on_pk_exact(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        qs = Comment.objects.filter(pk=subquery)
+        self.assertEqual(qs.count(), 2)
+
+    @skipUnlessDBFeature("supports_tuple_comparison_against_subquery")
+    def test_outer_ref_pk_filter_on_pk_comparison(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        tests = [
+            ("gt", 0),
+            ("gte", 2),
+            ("lt", 0),
+            ("lte", 2),
+        ]
+        for lookup, expected_count in tests:
+            with self.subTest(f"pk__{lookup}"):
+                qs = Comment.objects.filter(**{f"pk__{lookup}": subquery})
+                self.assertEqual(qs.count(), expected_count)
+
+    @skipIfDBFeature("supports_tuple_comparison_against_subquery")
+    def test_outer_ref_pk_filter_on_pk_comparison_unsupported(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        tests = ["gt", "gte", "lt", "lte"]
+        for lookup in tests:
+            with self.subTest(f"pk__{lookup}"):
+                qs = Comment.objects.filter(**{f"pk__{lookup}": subquery})
+                with self.assertRaisesMessage(
+                    NotSupportedError,
+                    f'"{lookup}" cannot be used to target composite fields '
+                    "through subqueries on this backend",
+                ):
+                    qs.count()
+
     def test_unsupported_rhs(self):
         pk = Exact(F("tenant_id"), 1)
         msg = (
@@ -561,7 +594,11 @@ class CompositePKFilterTests(TestCase):
 @skipUnlessDBFeature("supports_tuple_lookups")
 class CompositePKFilterTupleLookupFallbackTests(CompositePKFilterTests):
     def setUp(self):
-        feature_patch = patch.object(
+        feature_patch_1 = patch.object(
             connection.features, "supports_tuple_lookups", False
         )
-        self.enterContext(feature_patch)
+        feature_patch_2 = patch.object(
+            connection.features, "supports_tuple_comparison_against_subquery", False
+        )
+        self.enterContext(feature_patch_1)
+        self.enterContext(feature_patch_2)
diff --git a/tests/expressions/tests.py b/tests/expressions/tests.py
index 228d73c568..27d88be621 100644
--- a/tests/expressions/tests.py
+++ b/tests/expressions/tests.py
@@ -988,11 +988,24 @@ class BasicExpressionsTests(TestCase):
                 )
                 .order_by("-salary_raise")
                 .values("salary_raise")[:1],
-                output_field=IntegerField(),
             ),
         ).get(pk=self.gmbh.pk)
         self.assertEqual(gmbh_salary.max_ceo_salary_raise, 2332)
 
+    def test_annotation_with_outerref_and_output_field(self):
+        gmbh_salary = Company.objects.annotate(
+            max_ceo_salary_raise=Subquery(
+                Company.objects.annotate(
+                    salary_raise=OuterRef("num_employees") + F("num_employees"),
+                )
+                .order_by("-salary_raise")
+                .values("salary_raise")[:1],
+                output_field=DecimalField(),
+            ),
+        ).get(pk=self.gmbh.pk)
+        self.assertEqual(gmbh_salary.max_ceo_salary_raise, 2332.0)
+        self.assertIsInstance(gmbh_salary.max_ceo_salary_raise, Decimal)
+
     def test_annotation_with_nested_outerref(self):
         self.gmbh.point_of_contact = Employee.objects.get(lastname="Meyer")
         self.gmbh.save()
@@ -2542,6 +2555,15 @@ class ExistsTests(TestCase):
         self.assertSequenceEqual(qs, [manager])
         self.assertIs(qs.get().exists, False)
 
+    def test_annotate_by_empty_custom_exists(self):
+        class CustomExists(Exists):
+            template = Subquery.template
+
+        manager = Manager.objects.create()
+        qs = Manager.objects.annotate(exists=CustomExists(Manager.objects.none()))
+        self.assertSequenceEqual(qs, [manager])
+        self.assertIs(qs.get().exists, False)
+
 
 class FieldTransformTests(TestCase):
     @classmethod