Add support for polymorphic foreign key optimization

Author: diesieben07

strawberry_django/optimizer.py

@@ -702,6 +702,22 @@ def _get_hints_from_model_property(
     return store
 
 
+def _must_use_prefetch_related(
+    field: StrawberryField,
+    model_field: models.ForeignKey | OneToOneRel,
+) -> bool:
+    f_type = _get_django_type(field)
+    if f_type and hasattr(f_type, "get_queryset"):
+        # If the field has a get_queryset method, change strategy to Prefetch
+        # so it will be respected
+        return True
+    if is_polymorphic_model(model_field.related_model):
+        # If the model is using django-polymorphic, change strategy to Prefetch,
+        # so its custom queryset will be used, returning polymorphic models
+        return True
+    return False
+
+
 def _get_hints_from_django_foreign_key(
     field: StrawberryField,
     field_definition: GraphQLObjectType,
@@ -717,13 +733,9 @@ def _get_hints_from_django_foreign_key(
     cache: dict[type[models.Model], list[tuple[int, OptimizerStore]]],
     level: int = 0,
 ) -> OptimizerStore:
-    f_type = _get_django_type(field)
-    if f_type and hasattr(f_type, "get_queryset"):
-        # If the field has a get_queryset method, change strategy to Prefetch
-        # so it will be respected
+    if _must_use_prefetch_related(field, model_field):
         store = _get_hints_from_django_relation(
             field,
-            field_definition=field_definition,
             field_selection=field_selection,
             model_field=model_field,
             model_fieldname=model_fieldname,
@@ -772,7 +784,6 @@ def _get_hints_from_django_foreign_key(
 
 def _get_hints_from_django_relation(
     field: StrawberryField,
-    field_definition: GraphQLObjectType,
     field_selection: FieldNode,
     model_field: (
         models.ManyToManyField
@@ -961,7 +972,6 @@ def _get_hints_from_django_field(
     elif isinstance(model_field, relation_fields):
         store = _get_hints_from_django_relation(
             field,
-            field_definition=field_definition,
             field_selection=field_selection,
             model_field=model_field,
             model_fieldname=model_fieldname,

tests/polymorphism/models.py

@@ -4,6 +4,10 @@
 
 class Company(models.Model):
     name = models.CharField(max_length=100)
+    main_project = models.ForeignKey('Project', on_delete=models.CASCADE, null=True)
+
+    class Meta:
+        ordering = ('name',)
 
 
 class Project(PolymorphicModel):

tests/polymorphism/schema.py

@@ -26,6 +26,7 @@ class ResearchProjectType(ProjectType):
 class CompanyType:
     name: strawberry.auto
     projects: list[ProjectType]
+    main_project: ProjectType | None
 
 
 @strawberry.type

tests/polymorphism/test_optimizer.py

@@ -164,6 +164,61 @@ def test_polymorphic_offset_paginated_query():
     }
 
 
+@pytest.mark.django_db(transaction=True)
+def test_polymorphic_relation():
+    ap = ArtProject.objects.create(topic="Art", artist="Artist")
+    art_company = Company.objects.create(name="ArtCompany", main_project=ap)
+
+    rp = ResearchProject.objects.create(
+        topic="Research", supervisor="Supervisor"
+    )
+    research_company = Company.objects.create(name="ResearchCompany", main_project=rp)
+
+    query = """\
+    query {
+      companies {
+        name
+        mainProject {
+            __typename
+            topic
+            ... on ArtProjectType {
+              artist
+            }
+            ... on ResearchProjectType {
+              supervisor
+            }
+          }
+      }
+    }
+    """
+
+    # Company, ContentType, base table, two subtables = 5 queries
+    with assert_num_queries(1):
+        result = schema.execute_sync(query)
+    assert not result.errors
+    assert result.data == {
+        "companies": [
+            {
+                "name": art_company.name,
+                "mainProject": {
+                    "__typename": "ArtProjectType",
+                    "topic": ap.topic,
+                    "artist": ap.artist,
+                },
+            },
+            {
+                "name": research_company.name,
+                "mainProject": {
+                    "__typename": "ResearchProjectType",
+                    "topic": rp.topic,
+                    "supervisor": rp.supervisor,
+                },
+            }
+        ]
+    }
+
+
+
 @pytest.mark.django_db(transaction=True)
 def test_polymorphic_nested_list():
     company = Company.objects.create(name="Company")