fix: issues caused by #620 change of polymorphic_primary_key_name

Author: pgammans

src/polymorphic/query.py

@@ -8,7 +8,7 @@
 from collections import defaultdict
 
 from django.contrib.contenttypes.models import ContentType
-from django.core.exceptions import FieldDoesNotExist
+from django.core.exceptions import FieldDoesNotExist, FieldError
 from django.db import connections, models
 from django.db.models import FilteredRelation, Manager
 from django.db.models.constants import LOOKUP_SEP
@@ -240,8 +240,19 @@ def pivot_onto_cached_subclass(self, from_obj, obj, model_target_cls):
                     self.remote_setter(obj, from_obj)
             return None, None
 
-        pk_name = self.model_cls.polymorphic_primary_key_name
-        return model_target_cls, (getattr(original, pk_name), self.field.name)
+        local_pk_name = original.__class__.polymorphic_primary_key_name
+        target_pk_name = original.__class__.polymorphic_primary_key_name
+        original_pk = getattr(original, local_pk_name)
+
+        # NOTE: We could use a recursive function on model_target_cls._meta.parents
+        # PolymorphicModel.much _get_inheritance_relation_fields_and_models.like add_all_sub_models
+        for field in model_target_cls._meta.fields:
+            if field.is_relation is True:
+                for rel_field in field.foreign_related_fields:
+                    if rel_field.name is local_pk_name and rel_field.model is original._meta.model:
+                        target_pk_name = field.attname
+
+        return model_target_cls, (original_pk, self.field.name, target_pk_name)
 
 
 def get_related_populators(klass_info, select, db):
@@ -409,8 +420,8 @@ def fetch_polymorphic(self, post_actions, base_result_objects):
         for action, populate_fn in post_actions:
             target_class, pk_info = action()
             if target_class:
-                pk, name = pk_info
-                idlist_per_model[target_class].append((pk, name))
+                pk, name, pk_name = pk_info
+                idlist_per_model[target_class].append(pk_info)
                 update_fn_per_model[target_class].append((populate_fn, pk))
 
         # For each model in "idlist_per_model" request its objects (the real model)
@@ -419,16 +430,26 @@ def fetch_polymorphic(self, post_actions, base_result_objects):
         # Then we copy the extra() select fields from the base objects to the real objects.
         # TODO: defer(), only(): support for these would be around here
         for real_concrete_class, data in idlist_per_model.items():
-            idlist, names = zip(*data)
+            idlist, names, pk_attr_names = zip(*data)
             updates = update_fn_per_model[real_concrete_class]
-            pk_name = real_concrete_class.polymorphic_primary_key_name
+
+            if len(set(pk_attr_names)) != 1:
+                raise FieldError(
+                    "PolymorphicModel: cannot convert model type as non "
+                    f"upk_namesnique related key names {pk_attr_names}"
+                )
+
+            pk_attr_name = pk_attr_names[0]
+            # FIXME: this seams to get extra field already fetch in base
+            # initial query, we may need to add defer?
+
             real_objects = real_concrete_class._base_objects.db_manager(self.queryset.db).filter(
-                **{("%s__in" % pk_name): idlist}
+                **{("%s__in" % pk_attr_name): idlist},
             )
 
             real_objects = self.apply_select_related(real_objects, set(names))
             real_objects_dict = {
-                getattr(real_object, pk_name): real_object for real_object in real_objects
+                getattr(real_object, pk_attr_name): real_object for real_object in real_objects
             }
 
             for populate_fn, o_pk in updates:

src/polymorphic/tests/migrations/0001_initial.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.2 on 2025-12-23 09:40
+# Generated by Django 4.2 on 2025-12-23 08:21
 
 from django.conf import settings
 from django.db import migrations, models
@@ -1433,6 +1433,30 @@ class Migration(migrations.Migration):
             },
             bases=('tests.mrobase2', 'tests.mrobase3'),
         ),
+        migrations.CreateModel(
+            name='NonAutoPKChild',
+            fields=[
+                ('altchildmodel_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, to='tests.altchildmodel')),
+                ('uuid_primary_key', models.UUIDField(default=uuid.uuid1, primary_key=True, serialize=False)),
+            ],
+            options={
+                'abstract': False,
+                'base_manager_name': 'objects',
+            },
+            bases=('tests.altchildmodel',),
+        ),
+        migrations.CreateModel(
+            name='NonUUIDArtProject',
+            fields=[
+                ('uuidresearchproject_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, to='tests.uuidresearchproject')),
+                ('idkey', models.AutoField(primary_key=True, serialize=False)),
+            ],
+            options={
+                'abstract': False,
+                'base_manager_name': 'objects',
+            },
+            bases=('tests.uuidresearchproject',),
+        ),
         migrations.CreateModel(
             name='PlainC',
             fields=[

src/polymorphic/tests/models.py

@@ -346,6 +346,10 @@ class UUIDResearchProject(UUIDProject):
     supervisor = models.CharField(max_length=30)
 
 
+class NonUUIDArtProject(UUIDResearchProject):
+    idkey = models.AutoField(primary_key=True)
+
+
 class UUIDArtProjectA(UUIDArtProject): ...
 
 
@@ -833,6 +837,10 @@ class AltChildAsBaseModel(AltChildModel):
     more_name = models.CharField(max_length=10)
 
 
+class NonAutoPKChild(AltChildModel):
+    uuid_primary_key = models.UUIDField(primary_key=True, default=uuid.uuid1)
+
+
 class PlainModel(models.Model):
     relation = models.ForeignKey(ParentModel, on_delete=models.CASCADE)
     objects = models.Manager.from_queryset(PolymorphicRelatedQuerySet)()

src/polymorphic/tests/test_orm.py

@@ -74,12 +74,14 @@
     MultiTableDerived,
     MyManager,
     MyManagerQuerySet,
+    NonAutoPKChild,
     NonPolymorphicParent,
     NonProxyChild,
     NonSymRelationA,
     NonSymRelationB,
     NonSymRelationBase,
     NonSymRelationBC,
+    NonUUIDArtProject,
     One2OneRelatingModel,
     One2OneRelatingModelDerived,
     ParentModel,
@@ -112,7 +114,6 @@
     SubclassSelectorProxyBaseModel,
     SubclassSelectorProxyConcreteModel,
     ParentLinkAndRelatedName,
-    TestParentLinkAndRelatedName,
     UUIDArtProject,
     UUIDArtProjectA,
     UUIDArtProjectB,
@@ -1997,38 +1998,144 @@ def test_infinite_recursion_with_only(self):
     def test_normal_django_to_poly_related_give_poly_type(self):
         obj1 = ParentModel.objects.create(name="m1")
         obj2 = ChildModel.objects.create(name="m2", other_name="m2")
-        obj3 = ChildModel.objects.create(name="m1")
+        obj3 = ChildModel.objects.create(name="m3")
+        obj4 = ChildModel.objects.create(name="m3")
+        obj5 = AltChildModel.objects.create(name="m4")
 
         PlainModel.objects.create(relation=obj1)
         PlainModel.objects.create(relation=obj2)
         PlainModel.objects.create(relation=obj3)
+        PlainModel.objects.create(relation=obj4)
+        PlainModel.objects.create(relation=obj5)
 
-        ContentType.objects.get_for_model(AltChildModel)
-
-        with self.assertNumQueries(6):
+        with self.assertNumQueries(10):
             # Queries will be
             #    * 1 for All PlainModels object (1)
-            #    * 1 for each relations ParentModel (4)
-            #    * 1 for each relations ChilModel is needed (3)
+            #    * 1 for each relations ParentModel (5)
+            #    * 1 for each relations ChildModel is needed (3)
+            #    * 1 for each relations AltChildModel is needed (1)
             multi_q = [
                 # these obj.relation values will have their proper sub type
                 obj.relation
                 for obj in PlainModel.objects.all()
             ]
             multi_q_types = [type(obj) for obj in multi_q]
 
-        with self.assertNumQueries(2):
+        with self.assertNumQueries(3):
             grouped_q = [
                 # these obj.relation values will all be ParentModel's
                 # unless we fix select related but should be their proper
                 # sub type by using PolymorphicRelatedQuerySetMixin
+                # 1 query for each relation type
                 obj.relation
                 for obj in PlainModel.objects.select_related("relation")
             ]
             grouped_q_types = [type(obj) for obj in grouped_q]
 
         self.assertListEqual(multi_q_types, grouped_q_types)
-        self.assertListEqual(grouped_q, [obj1, obj2, obj3])
+        self.assertListEqual(grouped_q, [obj1, obj2, obj3, obj4, obj5])
+
+    def test_normal_django_to_multi_level_poly_related_give_poly_type(self):
+        obj1 = ParentModel.objects.create(name="m1")
+        obj2 = ChildModel.objects.create(name="m2", other_name="c1")
+        obj3 = AltChildModel.objects.create(name="m3")
+        obj4 = AltChildAsBaseModel.objects.create(name="m4", more_name="acab1")
+
+        PlainModel.objects.create(relation=obj1)
+        PlainModel.objects.create(relation=obj2)
+        PlainModel.objects.create(relation=obj3)
+        PlainModel.objects.create(relation=obj4)
+
+        with self.assertNumQueries(4):
+            grouped_q = [
+                # these obj.relation values will all be ParentModel's
+                # unless we fix select related but should be their proper
+                # sub type by using PolymorphicRelatedQuerySetMixin
+                obj.relation
+                for obj in PlainModel.objects.select_related("relation")
+            ]
+        self.assertListEqual(grouped_q, [obj1, obj2, obj3, obj4])
+
+    def test_related_fetch_of_different_type_pks(self):
+        "pk on child is not same field type as pk on parent and thus prt field"
+        obj1 = ChildModel.objects.create(name="m1", other_name="c1")
+        obj2 = ParentModel.objects.create(name="m2")
+        obj3 = NonAutoPKChild.objects.create(name="m3", other_name="napk1")
+        obj4 = AltChildModel.objects.create(name="m4", other_name="acm1")
+        obj5 = ChildModel.objects.create(name="m5", other_name="c3")
+        obj6 = AltChildAsBaseModel.objects.create(name="m6", more_name="acab1")
+        obj7 = NonAutoPKChild.objects.create(name="m7", other_name="napk2")
+
+        PlainModel.objects.create(relation=obj1)
+        PlainModel.objects.create(relation=obj2)
+        PlainModel.objects.create(relation=obj3)
+        PlainModel.objects.create(relation=obj4)
+        PlainModel.objects.create(relation=obj5)
+        PlainModel.objects.create(relation=obj6)
+        PlainModel.objects.create(relation=obj7)
+
+        def object_info(obj):
+            return {
+                "pk": obj.pk,
+                "parentmodel_ptr": getattr(obj, "parentmodel_ptr_id", None),
+                "altchildmodel_ptr": getattr(obj, "altchildmodel_ptr_id", None),
+            }
+
+        with self.assertNumQueries(5):
+            grouped_q = [
+                # these obj.relation values will all be ParentModel's
+                # unless we fix select related but should be their proper
+                # sub type by using PolymorphicRelatedQuerySetMixin
+                obj.relation
+                for obj in PlainModel.objects.select_related("relation").order_by("pk")
+            ]
+            grouped_info = [object_info(obj) for obj in grouped_q]
+        self.assertListEqual(
+            grouped_info, [object_info(obj) for obj in [obj1, obj2, obj3, obj4, obj5, obj6, obj7]]
+        )
+        self.assertListEqual(grouped_q, [obj1, obj2, obj3, obj4, obj5, obj6, obj7])
+
+    def test_related_fetch_of_non_sequential_pks(self):
+        obj1 = ChildModel.objects.create(name="m1", other_name="c1")
+        obj2 = ParentModel.objects.create(name="m2")
+
+        # FIXME use PK from table to get in sequential PKS
+        # from django.db import connection
+        # with connection.cursor() as cursor:
+        #    cursor.execute('INSERT INTO "tests_childmodel" ("other_name", "parentmodel_ptr_id") VALUES (%s, %s)', ['fake', 1])
+
+        obj3 = ChildModel.objects.create(name="m3", other_name="c2")
+        obj4 = AltChildModel.objects.create(name="m4", other_name="acm1")
+        obj5 = ChildModel.objects.create(name="m5", other_name="c3")
+        obj6 = AltChildAsBaseModel.objects.create(name="m6", more_name="acab1")
+
+        PlainModel.objects.create(relation=obj1)
+        PlainModel.objects.create(relation=obj2)
+        PlainModel.objects.create(relation=obj3)
+        PlainModel.objects.create(relation=obj4)
+        PlainModel.objects.create(relation=obj5)
+        PlainModel.objects.create(relation=obj6)
+
+        def object_info(obj):
+            return {
+                "pk": obj.pk,
+                "parentmodel_ptr": getattr(obj, "parentmodel_ptr_id", None),
+                "altchildmodel_ptr": getattr(obj, "altchildmodel_ptr_id", None),
+            }
+
+        with self.assertNumQueries(4):
+            grouped_q = [
+                # these obj.relation values will all be ParentModel's
+                # unless we fix select related but should be their proper
+                # sub type by using PolymorphicRelatedQuerySetMixin
+                obj.relation
+                for obj in PlainModel.objects.select_related("relation")
+            ]
+            grouped_info = [object_info(obj) for obj in grouped_q]
+        self.assertListEqual(
+            grouped_info, [object_info(obj) for obj in [obj1, obj2, obj3, obj4, obj5, obj6]]
+        )
+        self.assertListEqual(grouped_q, [obj1, obj2, obj3, obj4, obj5, obj6])
 
     def test_normal_django_to_poly_related_give_poly_type_using_select_related_true(self):
         obj1 = ParentModel.objects.create(name="m1")
@@ -2335,7 +2442,12 @@ def test_select_related_fecth_all_poly_classes_indirect_related(self):
 
         # Prefetch content_types
         ContentType.objects.get_for_models(
-            PlainModel, PlainA, ModelExtraExternal, AltChildAsBaseModel, AltChildWithM2MModel
+            AltChildAsBaseModel,
+            AltChildWithM2MModel,
+            ModelExtraExternal,
+            NonAutoPKChild,
+            PlainA,
+            PlainModel,
         )
 
         with self.assertNumQueries(1):