refactor: Use qs._result_cache instead of len() for safety reasons

strawberry-graphql · Dec 18, 2024 · a98263b · a98263b
1 parent 51a6291
commit a98263b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/strawberry_django/fields/field.py b/strawberry_django/fields/field.py
@@ -31,7 +31,7 @@
     ReverseOneToOneDescriptor,
 )
 from django.db.models.manager import BaseManager
-from django.db.models.query import MAX_GET_RESULTS
+from django.db.models.query import MAX_GET_RESULTS  # type: ignore
 from django.db.models.query_utils import DeferredAttribute
 from strawberry import UNSET, relay
 from strawberry.annotation import StrawberryAnnotation
@@ -292,10 +292,10 @@ def qs_hook(qs: models.QuerySet):
                 # Don't use qs.get() if the queryset is optimized by prefetching.
                 # Calling get in that case would disregard the prefetched results, because get implicitly
                 # adds a limit to the query
-                if is_optimized_by_prefetching(qs):
+                if (result_cache := qs._result_cache) is not None:  # type: ignore
                     # mimic behavior of get()
                     # the queryset is already prefetched, no issue with just using len()
-                    qs_len = len(qs)
+                    qs_len = len(result_cache)
                     if qs_len == 0:
                         raise qs.model.DoesNotExist(
                             f"{qs.model._meta.object_name} matching query does not exist."
@@ -305,7 +305,8 @@ def qs_hook(qs: models.QuerySet):
                             f"get() returned more than one {qs.model._meta.object_name} -- it returned "
                             f"{qs_len if qs_len < MAX_GET_RESULTS else f'more than {qs_len - 1}'}!"
                         )
-                    return qs[0]
+                    return result_cache[0]
+
                 return qs.get()
 
         return qs_hook

diff --git a/strawberry_django/pagination.py b/strawberry_django/pagination.py
@@ -6,7 +6,7 @@
 from django.db import DEFAULT_DB_ALIAS
 from django.db.models import Count, QuerySet, Window
 from django.db.models.functions import RowNumber
-from django.db.models.query import MAX_GET_RESULTS
+from django.db.models.query import MAX_GET_RESULTS  # type: ignore
 from strawberry.types import Info
 from strawberry.types.arguments import StrawberryArgument
 from strawberry.types.unset import UNSET, UnsetType

diff --git a/tests/projects/schema.py b/tests/projects/schema.py
@@ -159,6 +159,8 @@ class MilestoneType(relay.Node, Named):
         order=IssueOrder,
         pagination=True,
     )
+    first_issue: Optional["IssueType"] = strawberry_django.field(field_name="issues")
+    first_issue_required: "IssueType" = strawberry_django.field(field_name="issues")
     issues_paginated: OffsetPaginated["IssueType"] = strawberry_django.offset_paginated(
         field_name="issues",
         order=IssueOrder,
@@ -169,8 +171,6 @@ class MilestoneType(relay.Node, Named):
             filters=IssueFilter,
         )
     )
-    first_issue: Optional["IssueType"] = strawberry_django.field(field_name="issues")
-    first_issue_required: "IssueType" = strawberry_django.field(field_name="issues")
 
     @strawberry_django.field(
         prefetch_related=[