select_related() doesn't work when targeting composite primary keys

[Jacob Walls]

The ​suggested workaround until #35956 implements foreign keys to models with composite primary keys is to use ForeignObject, but it causes failures in select_related(). (Of note: prefetch_related() works fine.)

tests/composite_pk/tests.py

@@ -184 +184 @@
         with self.assertNumQueries(1):
             self.assertEqual(user.email, self.user.email)
 
+    def test_select_related(self):
+        with self.assertNumQueries(1):
+            for comment in Comment.objects.select_related():
+                comment.user
+
     def test_model_forms(self):
         fields = ["tenant", "id", "user_id", "text", "integer"]
         self.assertEqual(list(CommentForm.base_fields), fields)

======================================================================
ERROR: test_select_related (composite_pk.tests.CompositePKTests.test_select_related)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/.../django/tests/composite_pk/tests.py", line 189, in test_select_related
    for comment in Comment.objects.select_related():
                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/.../django/django/db/models/query.py", line 403, in __iter__
    self._fetch_all()
    ~~~~~~~~~~~~~~~^^
  File "/Users/.../django/django/db/models/query.py", line 1966, in _fetch_all
    self._result_cache = list(self._iterable_class(self))
                         ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/query.py", line 107, in __iter__
    related_populators = get_related_populators(klass_info, select, db, fetch_mode)
  File "/Users/.../django/django/db/models/query.py", line 2742, in get_related_populators
    rel_cls = RelatedPopulator(rel_klass_info, select, db, fetch_mode)
  File "/Users/.../django/django/db/models/query.py", line 2710, in __init__
    self.pk_idx = self.init_list.index(self.model_cls._meta.pk.attname)
                  ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: 'pk' is not in list

[Simon Charette]

Thanks for the report, it's worth pointing out that the issue is not specific to calling select_related without a subset of relationship the following test fails as well and is a more common usage of select_related

tests/composite_pk/tests.py

@@ -184 +184 @@
         with self.assertNumQueries(1):
             self.assertEqual(user.email, self.user.email)
 
+    def test_select_related(self):
+        with self.assertNumQueries(1):
+            for comment in Comment.objects.select_related("user"):
+                comment.user
+
     def test_model_forms(self):
         fields = ["tenant", "id", "user_id", "text", "integer"]
         self.assertEqual(list(CommentForm.base_fields), fields)

[Simon Charette]

Looks like something like the following should do

django/db/models/query.py

@@ -2680 +2680 @@
             )
 
         self.model_cls = klass_info["model"]
-        self.pk_idx = self.init_list.index(self.model_cls._meta.pk.attname)
+        pk_fields = self.model_cls._meta.pk_fields
+        pk_idx = self.init_list.index(pk_fields[0].attname)
+        if (pk_fields_len := len(pk_fields)) > 1:
+            self.pk_idx = slice(pk_idx, pk_fields_len)
+        else:
+            self.pk_idx = pk_idx
         self.related_populators = get_related_populators(klass_info, select, self.db)
         self.local_setter = klass_info["local_setter"]
         self.remote_setter = klass_info["remote_setter"]

It might be worth adding a test with a model where pk fields are not defined next to each other or just alter the existing models to cover for that

tests/composite_pk/models/tenant.py

@@ -17 +17 @@
 class BaseModel(models.Model):
     pk = models.CompositePrimaryKey("tenant_id", "id")
     tenant = models.ForeignKey(Tenant, on_delete=models.CASCADE)
+    between = models.IntegerField(null=True)
     id = models.SmallIntegerField(unique=True)
 
     class Meta:
@@ -34 +35 @@
         on_delete=models.CASCADE,
         related_name="comments",
     )
-    id = models.SmallIntegerField(unique=True, db_column="comment_id")
     user_id = models.SmallIntegerField()
+    id = models.SmallIntegerField(unique=True, db_column="comment_id")
     user = models.ForeignObject(
         User,
         on_delete=models.CASCADE,

[Sarah Boyce <42296566+sarahboyce@…>]

In 8be0c0d:

Fixed #36373 -- Fixed select_related() crash on foreign object for a composite pk.

Thanks Jacob Walls for the report and Sarah for the in-depth review.

[Sarah Boyce <42296566+sarahboyce@…>]

In e23dd728:

[5.2.x] Fixed #36373 -- Fixed select_related() crash on foreign object for a composite pk.

Thanks Jacob Walls for the report and Sarah for the in-depth review.

Backport of 8be0c0d6901669661fca578f474cd51cd284d35a from main.