From cac94dd8aa2fb49cd2e06b5b37cf039257284bb0 Mon Sep 17 00:00:00 2001 From: Clément Escolano Date: Tue, 1 Aug 2023 23:31:40 +0200 Subject: Fixed #33651 -- Added support for prefetching GenericForeignKey. Co-authored-by: revanthgss Co-authored-by: Mariusz Felisiak --- django/contrib/contenttypes/fields.py | 60 ++++++++++++++++++--- django/contrib/contenttypes/prefetch.py | 36 +++++++++++++ django/db/models/fields/related_descriptors.py | 68 +++++++++++++++++++++-- django/db/models/query.py | 74 ++++++++++++++++++++------ 4 files changed, 210 insertions(+), 28 deletions(-) create mode 100644 django/contrib/contenttypes/prefetch.py (limited to 'django') 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 -- cgit v1.3