Fixed #33651 -- Added support for prefetching GenericForeignKey.

Co-authored-by: revanthgss <revanthgss@almabase.com>
Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>

4 files changed, 209 insertions, 27 deletions

diff --git a/django/contrib/contenttypes/fields.py b/django/contrib/contenttypes/fields.py
index f92ecfa6c0..1b6abb9818 100644
--- a/django/contrib/contenttypes/fields.py
+++ b/django/contrib/contenttypes/fields.py
@@ -1,5 +1,6 @@
 import functools
 import itertools
+import warnings
 from collections import defaultdict
 
 from asgiref.sync import sync_to_async
@@ -19,6 +20,7 @@ from django.db.models.query_utils import PathInfo
 from django.db.models.sql import AND
 from django.db.models.sql.where import WhereNode
 from django.db.models.utils import AltersData
+from django.utils.deprecation import RemovedInDjango60Warning
 from django.utils.functional import cached_property
 
 
@@ -163,20 +165,44 @@ class GenericForeignKey(FieldCacheMixin):
     def get_cache_name(self):
         return self.name
 
-    def get_content_type(self, obj=None, id=None, using=None):
+    def get_content_type(self, obj=None, id=None, using=None, model=None):
         if obj is not None:
             return ContentType.objects.db_manager(obj._state.db).get_for_model(
                 obj, for_concrete_model=self.for_concrete_model
             )
         elif id is not None:
             return ContentType.objects.db_manager(using).get_for_id(id)
+        elif model is not None:
+            return ContentType.objects.db_manager(using).get_for_model(
+                model, for_concrete_model=self.for_concrete_model
+            )
         else:
             # This should never happen. I love comments like this, don't you?
             raise Exception("Impossible arguments to GFK.get_content_type!")
 
     def get_prefetch_queryset(self, instances, queryset=None):
-        if queryset is not None:
-            raise ValueError("Custom queryset can't be used for this lookup.")
+        warnings.warn(
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead.",
+            RemovedInDjango60Warning,
+            stacklevel=2,
+        )
+        if queryset is None:
+            return self.get_prefetch_querysets(instances)
+        return self.get_prefetch_querysets(instances, [queryset])
+
+    def get_prefetch_querysets(self, instances, querysets=None):
+        custom_queryset_dict = {}
+        if querysets is not None:
+            for queryset in querysets:
+                ct_id = self.get_content_type(
+                    model=queryset.query.model, using=queryset.db
+                ).pk
+                if ct_id in custom_queryset_dict:
+                    raise ValueError(
+                        "Only one queryset is allowed for each content type."
+                    )
+                custom_queryset_dict[ct_id] = queryset
 
         # For efficiency, group the instances by content type and then do one
         # query per model
@@ -195,9 +221,13 @@ class GenericForeignKey(FieldCacheMixin):
 
         ret_val = []
         for ct_id, fkeys in fk_dict.items():
-            instance = instance_dict[ct_id]
-            ct = self.get_content_type(id=ct_id, using=instance._state.db)
-            ret_val.extend(ct.get_all_objects_for_this_type(pk__in=fkeys))
+            if ct_id in custom_queryset_dict:
+                # Return values from the custom queryset, if provided.
+                ret_val.extend(custom_queryset_dict[ct_id].filter(pk__in=fkeys))
+            else:
+                instance = instance_dict[ct_id]
+                ct = self.get_content_type(id=ct_id, using=instance._state.db)
+                ret_val.extend(ct.get_all_objects_for_this_type(pk__in=fkeys))
 
         # For doing the join in Python, we have to match both the FK val and the
         # content type, so we use a callable that returns a (fk, class) pair.
@@ -616,9 +646,23 @@ def create_generic_related_manager(superclass, rel):
                 return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
+            warnings.warn(
+                "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+                "instead.",
+                RemovedInDjango60Warning,
+                stacklevel=2,
+            )
             if queryset is None:
-                queryset = super().get_queryset()
+                return self.get_prefetch_querysets(instances)
+            return self.get_prefetch_querysets(instances, [queryset])
 
+        def get_prefetch_querysets(self, instances, querysets=None):
+            if querysets and len(querysets) != 1:
+                raise ValueError(
+                    "querysets argument of get_prefetch_querysets() should have a "
+                    "length of 1."
+                )
+            queryset = querysets[0] if querysets else super().get_queryset()
             queryset._add_hints(instance=instances[0])
             queryset = queryset.using(queryset._db or self._db)
             # Group instances by content types.
diff --git a/django/contrib/contenttypes/prefetch.py b/django/contrib/contenttypes/prefetch.py
new file mode 100644
index 0000000000..b02ed3bae5
--- /dev/null
+++ b/django/contrib/contenttypes/prefetch.py
@@ -0,0 +1,36 @@
+from django.db.models import Prefetch
+from django.db.models.query import ModelIterable, RawQuerySet
+
+
+class GenericPrefetch(Prefetch):
+    def __init__(self, lookup, querysets=None, to_attr=None):
+        for queryset in querysets:
+            if queryset is not None and (
+                isinstance(queryset, RawQuerySet)
+                or (
+                    hasattr(queryset, "_iterable_class")
+                    and not issubclass(queryset._iterable_class, ModelIterable)
+                )
+            ):
+                raise ValueError(
+                    "Prefetch querysets cannot use raw(), values(), and values_list()."
+                )
+        self.querysets = querysets
+        super().__init__(lookup, to_attr=to_attr)
+
+    def __getstate__(self):
+        obj_dict = self.__dict__.copy()
+        obj_dict["querysets"] = []
+        for queryset in self.querysets:
+            if queryset is not None:
+                queryset = queryset._chain()
+                # Prevent the QuerySet from being evaluated
+                queryset._result_cache = []
+                queryset._prefetch_done = True
+                obj_dict["querysets"].append(queryset)
+        return obj_dict
+
+    def get_current_querysets(self, level):
+        if self.get_current_prefetch_to(level) == self.prefetch_to:
+            return self.querysets
+        return None
diff --git a/django/db/models/fields/related_descriptors.py b/django/db/models/fields/related_descriptors.py
index 4d6164143b..46a5823647 100644
--- a/django/db/models/fields/related_descriptors.py
+++ b/django/db/models/fields/related_descriptors.py
@@ -62,6 +62,7 @@ and two directions (forward and reverse) for a total of six combinations.
    If you're looking for ``ForwardManyToManyDescriptor`` or
    ``ReverseManyToManyDescriptor``, use ``ManyToManyDescriptor`` instead.
 """
+import warnings
 
 from asgiref.sync import sync_to_async
 
@@ -79,6 +80,7 @@ from django.db.models.lookups import GreaterThan, LessThanOrEqual
 from django.db.models.query import QuerySet
 from django.db.models.query_utils import DeferredAttribute
 from django.db.models.utils import AltersData, resolve_callables
+from django.utils.deprecation import RemovedInDjango60Warning
 from django.utils.functional import cached_property
 
 
@@ -153,8 +155,23 @@ class ForwardManyToOneDescriptor:
         return self.field.remote_field.model._base_manager.db_manager(hints=hints).all()
 
     def get_prefetch_queryset(self, instances, queryset=None):
+        warnings.warn(
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead.",
+            RemovedInDjango60Warning,
+            stacklevel=2,
+        )
         if queryset is None:
-            queryset = self.get_queryset()
+            return self.get_prefetch_querysets(instances)
+        return self.get_prefetch_querysets(instances, [queryset])
+
+    def get_prefetch_querysets(self, instances, querysets=None):
+        if querysets and len(querysets) != 1:
+            raise ValueError(
+                "querysets argument of get_prefetch_querysets() should have a length "
+                "of 1."
+            )
+        queryset = querysets[0] if querysets else self.get_queryset()
         queryset._add_hints(instance=instances[0])
 
         rel_obj_attr = self.field.get_foreign_related_value
@@ -427,8 +444,23 @@ class ReverseOneToOneDescriptor:
         return self.related.related_model._base_manager.db_manager(hints=hints).all()
 
     def get_prefetch_queryset(self, instances, queryset=None):
+        warnings.warn(
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead.",
+            RemovedInDjango60Warning,
+            stacklevel=2,
+        )
         if queryset is None:
-            queryset = self.get_queryset()
+            return self.get_prefetch_querysets(instances)
+        return self.get_prefetch_querysets(instances, [queryset])
+
+    def get_prefetch_querysets(self, instances, querysets=None):
+        if querysets and len(querysets) != 1:
+            raise ValueError(
+                "querysets argument of get_prefetch_querysets() should have a length "
+                "of 1."
+            )
+        queryset = querysets[0] if querysets else self.get_queryset()
         queryset._add_hints(instance=instances[0])
 
         rel_obj_attr = self.related.field.get_local_related_value
@@ -728,9 +760,23 @@ def create_reverse_many_to_one_manager(superclass, rel):
                 return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
+            warnings.warn(
+                "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+                "instead.",
+                RemovedInDjango60Warning,
+                stacklevel=2,
+            )
             if queryset is None:
-                queryset = super().get_queryset()
+                return self.get_prefetch_querysets(instances)
+            return self.get_prefetch_querysets(instances, [queryset])
 
+        def get_prefetch_querysets(self, instances, querysets=None):
+            if querysets and len(querysets) != 1:
+                raise ValueError(
+                    "querysets argument of get_prefetch_querysets() should have a "
+                    "length of 1."
+                )
+            queryset = querysets[0] if querysets else super().get_queryset()
             queryset._add_hints(instance=instances[0])
             queryset = queryset.using(queryset._db or self._db)
 
@@ -1087,9 +1133,23 @@ def create_forward_many_to_many_manager(superclass, rel, reverse):
                 return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
+            warnings.warn(
+                "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+                "instead.",
+                RemovedInDjango60Warning,
+                stacklevel=2,
+            )
             if queryset is None:
-                queryset = super().get_queryset()
+                return self.get_prefetch_querysets(instances)
+            return self.get_prefetch_querysets(instances, [queryset])
 
+        def get_prefetch_querysets(self, instances, querysets=None):
+            if querysets and len(querysets) != 1:
+                raise ValueError(
+                    "querysets argument of get_prefetch_querysets() should have a "
+                    "length of 1."
+                )
+            queryset = querysets[0] if querysets else super().get_queryset()
             queryset._add_hints(instance=instances[0])
             queryset = queryset.using(queryset._db or self._db)
             queryset = _filter_prefetch_queryset(
diff --git a/django/db/models/query.py b/django/db/models/query.py
index 0746dc5c6b..1125302933 100644
--- a/django/db/models/query.py
+++ b/django/db/models/query.py
@@ -33,6 +33,7 @@ from django.db.models.utils import (
     resolve_callables,
 )
 from django.utils import timezone
+from django.utils.deprecation import RemovedInDjango60Warning
 from django.utils.functional import cached_property, partition
 
 # The maximum number of results to fetch in a get() query.
@@ -2236,8 +2237,21 @@ class Prefetch:
         return to_attr, as_attr
 
     def get_current_queryset(self, level):
-        if self.get_current_prefetch_to(level) == self.prefetch_to:
-            return self.queryset
+        warnings.warn(
+            "Prefetch.get_current_queryset() is deprecated. Use "
+            "get_current_querysets() instead.",
+            RemovedInDjango60Warning,
+            stacklevel=2,
+        )
+        querysets = self.get_current_querysets(level)
+        return querysets[0] if querysets is not None else None
+
+    def get_current_querysets(self, level):
+        if (
+            self.get_current_prefetch_to(level) == self.prefetch_to
+            and self.queryset is not None
+        ):
+            return [self.queryset]
         return None
 
     def __eq__(self, other):
@@ -2425,9 +2439,9 @@ async def aprefetch_related_objects(model_instances, *related_lookups):
 def get_prefetcher(instance, through_attr, to_attr):
     """
     For the attribute 'through_attr' on the given instance, find
-    an object that has a get_prefetch_queryset().
+    an object that has a get_prefetch_querysets().
     Return a 4 tuple containing:
-    (the object with get_prefetch_queryset (or None),
+    (the object with get_prefetch_querysets (or None),
      the descriptor object representing this relationship (or None),
      a boolean that is False if the attribute was not found at all,
      a function that takes an instance and returns a boolean that is True if
@@ -2462,8 +2476,12 @@ def get_prefetcher(instance, through_attr, to_attr):
         attr_found = True
         if rel_obj_descriptor:
             # singly related object, descriptor object has the
-            # get_prefetch_queryset() method.
-            if hasattr(rel_obj_descriptor, "get_prefetch_queryset"):
+            # get_prefetch_querysets() method.
+            if (
+                hasattr(rel_obj_descriptor, "get_prefetch_querysets")
+                # RemovedInDjango60Warning.
+                or hasattr(rel_obj_descriptor, "get_prefetch_queryset")
+            ):
                 prefetcher = rel_obj_descriptor
                 # If to_attr is set, check if the value has already been set,
                 # which is done with has_to_attr_attribute(). Do not use the
@@ -2476,7 +2494,11 @@ def get_prefetcher(instance, through_attr, to_attr):
                 # the attribute on the instance rather than the class to
                 # support many related managers
                 rel_obj = getattr(instance, through_attr)
-                if hasattr(rel_obj, "get_prefetch_queryset"):
+                if (
+                    hasattr(rel_obj, "get_prefetch_querysets")
+                    # RemovedInDjango60Warning.
+                    or hasattr(rel_obj, "get_prefetch_queryset")
+                ):
                     prefetcher = rel_obj
                 if through_attr == to_attr:
 
@@ -2497,7 +2519,7 @@ def prefetch_one_level(instances, prefetcher, lookup, level):
     Return the prefetched objects along with any additional prefetches that
     must be done due to prefetch_related lookups found from default managers.
     """
-    # prefetcher must have a method get_prefetch_queryset() which takes a list
+    # prefetcher must have a method get_prefetch_querysets() which takes a list
     # of instances, and returns a tuple:
 
     # (queryset of instances of self.model that are related to passed in instances,
@@ -2510,14 +2532,34 @@ def prefetch_one_level(instances, prefetcher, lookup, level):
     # The 'values to be matched' must be hashable as they will be used
     # in a dictionary.
 
-    (
-        rel_qs,
-        rel_obj_attr,
-        instance_attr,
-        single,
-        cache_name,
-        is_descriptor,
-    ) = prefetcher.get_prefetch_queryset(instances, lookup.get_current_queryset(level))
+    if hasattr(prefetcher, "get_prefetch_querysets"):
+        (
+            rel_qs,
+            rel_obj_attr,
+            instance_attr,
+            single,
+            cache_name,
+            is_descriptor,
+        ) = prefetcher.get_prefetch_querysets(
+            instances, lookup.get_current_querysets(level)
+        )
+    else:
+        warnings.warn(
+            "The usage of get_prefetch_queryset() in prefetch_related_objects() is "
+            "deprecated. Implement get_prefetch_querysets() instead.",
+            RemovedInDjango60Warning,
+            stacklevel=2,
+        )
+        (
+            rel_qs,
+            rel_obj_attr,
+            instance_attr,
+            single,
+            cache_name,
+            is_descriptor,
+        ) = prefetcher.get_prefetch_queryset(
+            instances, lookup.get_current_querysets(level)
+        )
     # We have to handle the possibility that the QuerySet we just got back
     # contains some prefetch_related lookups. We don't want to trigger the
     # prefetch_related functionality by evaluating the query. Rather, we need