prefetch_related() selects too many objects from a related table with a composite primary key and hidden related name

[Jacob Walls]

Noticed a "fixme" ​comment in the code that pointed me here. With a hidden related name, e.g. comments+, we hit a branch that expects primary keys to consist of one column only. When a composite primary key violates that assumption, more objects than necessary will be fetched.

With this test:

tests/composite_pk/models/__init__.py

@@ -1 @@
-from .tenant import Comment, Post, Tenant, TimeStamped, Token, User
+from .tenant import (
+    Comment,
+    CommentWithHiddenUserField,
+    Post,
+    Tenant,
+    TimeStamped,
+    Token,
+    User,
+)
 
 __all__ = [
     "Comment",
+    "CommentWithHiddenUserField",
     "Post",
     "Tenant",
     "TimeStamped",

tests/composite_pk/models/tenant.py

@@ -46 +46 @@
     text = models.TextField(default="", blank=True)
 
 
+class CommentWithHiddenUserField(models.Model):
+    pk = models.CompositePrimaryKey("tenant", "id")
+    tenant = models.ForeignKey(
+        Tenant,
+        on_delete=models.CASCADE,
+        related_name="comments+",
+    )
+    id = models.SmallIntegerField(unique=True, db_column="comment_id")
+    user_id = models.SmallIntegerField()
+    user = models.ForeignObject(
+        User,
+        on_delete=models.CASCADE,
+        from_fields=("tenant_id", "user_id"),
+        to_fields=("tenant_id", "id"),
+        related_name="comments+",
+    )
+    text = models.TextField(default="", blank=True)
+
+
 class Post(models.Model):
     pk = models.CompositePrimaryKey("tenant_id", "id")
     tenant = models.ForeignKey(Tenant, on_delete=models.CASCADE, default=1)

tests/composite_pk/tests.py

@@ -16 +16 @@
 from django.db.models import CompositePrimaryKey
 from django.forms import modelform_factory
 from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
 
-from .models import Comment, Post, Tenant, TimeStamped, User
+from .models import Comment, CommentWithHiddenUserField, Post, Tenant, TimeStamped, User
 
 
 class CommentForm(forms.ModelForm):
@@ -38 +39 @@
             email="user0001@example.com",
         )
         cls.comment = Comment.objects.create(tenant=cls.tenant, id=1, user=cls.user)
+        cls.comment_by_hidden_user = CommentWithHiddenUserField.objects.create(
+            tenant=cls.tenant, id=1, user=cls.user
+        )
 
     @staticmethod
     def get_primary_key_columns(table):
@@ -157 +161 @@
         result = list(Comment.objects.iterator())
         self.assertEqual(result, [self.comment])
 
+    def test_prefetch_related_hidden_related_name(self):
+        with CaptureQueriesContext(connection) as queries:
+            CommentWithHiddenUserField.objects.prefetch_related("user").get(
+                pk=self.comment_by_hidden_user.pk
+            )
+        self.assertIn('"comment_id" = 1', queries[1]["sql"])
+
     def test_query(self):
         users = User.objects.values_list("pk").order_by("pk")
         self.assertNotIn('AS "pk"', str(users.query))

The WHERE is too permissive -- would expect both primary key columns to be checked, i.e. also "comment_id".

======================================================================
FAIL: test_prefetch_related_hidden_related_name (composite_pk.tests.CompositePKTests.test_prefetch_related_hidden_related_name)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/source/django/tests/composite_pk/tests.py", line 169, in test_prefetch_related_hidden_related_name
    self.assertIn('"comment_id" = 1', queries[1]["sql"])
AssertionError: '"comment_id" = 1' not found in 'SELECT "composite_pk_user"."tenant_id", "composite_pk_user"."id", "composite_pk_user"."email" FROM "composite_pk_user" WHERE "composite_pk_user"."tenant_id" IN (1)'

----------------------------------------------------------------------
Ran 1 test in 0.006s

FAILED (failures=1)

[Simon Charette]

Thanks for surfacing Jacob, glad we foresaw this one in cb83448891f2920c926f61826d8eae4051a3d8f2.

Tentatively elevating to release blocker as ​we've opted to start documenting ForeignObject's usage with CompositePrimaryKey and I would expect users to run into this problem as well.

[Simon Charette]

Going to assign since I got a patch ready that doesn't require introducing new models and we'll surely see many of these required adjustments until 5.0 gets fully released. Let me know if you want me to hold off for longer when you report these.