Fixed #36210 -- Allowed Subquery usage in futher lookups against composite pks.

jacobtylerwalls · charettes · jacobtylerwalls · commit 82bdc8888efb · 2025-05-11T22:16:12.000-04:00
Follow-up to 8561100. co-authored-by: Simon Charette <charette.s@gmail.com>
diff --git a/django/db/models/expressions.py b/django/db/models/expressions.py
@@ -1756,6 +1756,7 @@ def __init__(self, queryset, output_field=None, **extra):
         # Allow the usage of both QuerySet and sql.Query objects.
         self.query = getattr(queryset, "query", queryset).clone()
         self.query.subquery = True
+        self.template = extra.pop("template", self.template)
         self.extra = extra
         super().__init__(output_field)
 
@@ -1768,6 +1769,14 @@ def set_source_expressions(self, exprs):
     def _resolve_output_field(self):
         return self.query.output_field
 
+    def resolve_expression(self, *args, **kwargs):
+        resolved = super().resolve_expression(*args, **kwargs)
+        if self.template == Subquery.template:
+            # Subquery is an unnecessary shim for a resolved query as it
+            # complexifies the lookup's right-hand-side introspection.
+            return resolved.query
+        return resolved
+
     def copy(self):
         clone = super().copy()
         clone.query = clone.query.clone()
diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py
@@ -235,6 +235,7 @@ class Query(BaseExpression):
 
     filter_is_sticky = False
     subquery = False
+    contains_subquery = True
 
     # SQL-related attributes.
     # Select and related select clauses are expressions to use in the SELECT
@@ -1226,7 +1227,7 @@ def add_annotation(self, annotation, alias, select=True):
 
     @property
     def _subquery_fields_len(self):
-        if self.has_select_fields:
+        if self.has_select_fields and tuple(self.selected) != ("pk",):
             return len(self.selected)
         return len(self.model._meta.pk_fields)
 
diff --git a/tests/composite_pk/test_filter.py b/tests/composite_pk/test_filter.py
@@ -476,6 +476,20 @@ def test_outer_ref_pk(self):
                 queryset = Comment.objects.filter(**{f"id{lookup}": subquery})
                 self.assertEqual(queryset.count(), expected_count)
 
+    def test_outer_ref_pk_select_pk(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        tests = [
+            ("", 2),
+            ("__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)
+
     def test_unsupported_rhs(self):
         pk = Exact(F("tenant_id"), 1)
         msg = (
