From f51c1f590085556abca44fd2a49618162203b2ec Mon Sep 17 00:00:00 2001 From: Loic Bistuer Date: Thu, 7 Nov 2013 00:25:05 +0700 Subject: [PATCH] Fixed #17001 -- Custom querysets for prefetch_related. This patch introduces the Prefetch object which allows customizing prefetch operations. This enables things like filtering prefetched relations, calling select_related from a prefetched relation, or prefetching the same relation multiple times with different querysets. When a Prefetch instance specifies a to_attr argument, the result is stored in a list rather than a QuerySet. This has the fortunate consequence of being significantly faster. The preformance improvement is due to the fact that we save the costly creation of a QuerySet instance. Thanks @akaariai for the original patch and @bmispelon and @timgraham for the reviews. --- django/contrib/contenttypes/generic.py | 18 +- django/db/models/__init__.py | 2 +- django/db/models/fields/related.py | 51 ++-- django/db/models/query.py | 141 ++++++++--- docs/ref/models/queries.txt | 29 +++ docs/ref/models/querysets.txt | 87 ++++++- docs/releases/1.7.txt | 15 ++ tests/prefetch_related/models.py | 10 + tests/prefetch_related/tests.py | 328 ++++++++++++++++++++++++- 9 files changed, 616 insertions(+), 65 deletions(-) diff --git a/django/contrib/contenttypes/generic.py b/django/contrib/contenttypes/generic.py index 8b7161251b..4880214200 100644 --- a/django/contrib/contenttypes/generic.py +++ b/django/contrib/contenttypes/generic.py @@ -76,7 +76,10 @@ class GenericForeignKey(six.with_metaclass(RenameGenericForeignKeyMethods)): # 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): + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is not None: + raise ValueError("Custom queryset can't be used for this lookup.") + # For efficiency, group the instances by content type and then do one # query per model fk_dict = defaultdict(set) @@ -348,17 +351,22 @@ def create_generic_related_manager(superclass): db = self._db or router.db_for_read(self.model, instance=self.instance) return super(GenericRelatedObjectManager, self).get_queryset().using(db).filter(**self.core_filters) - def get_prefetch_queryset(self, instances): - db = self._db or router.db_for_read(self.model, instance=instances[0]) + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is None: + queryset = super(GenericRelatedObjectManager, self).get_queryset() + + queryset._add_hints(instance=instances[0]) + queryset = queryset.using(queryset._db or self._db) + query = { '%s__pk' % self.content_type_field_name: self.content_type.id, '%s__in' % self.object_id_field_name: set(obj._get_pk_val() for obj in instances) } - qs = super(GenericRelatedObjectManager, self).get_queryset().using(db).filter(**query) + # We (possibly) need to convert object IDs to the type of the # instances' PK in order to match up instances: object_id_converter = instances[0]._meta.pk.to_python - return (qs, + return (queryset.filter(**query), lambda relobj: object_id_converter(getattr(relobj, self.object_id_field_name)), lambda obj: obj._get_pk_val(), False, diff --git a/django/db/models/__init__.py b/django/db/models/__init__.py index 13785ed17a..e2f7348ef0 100644 --- a/django/db/models/__init__.py +++ b/django/db/models/__init__.py @@ -4,7 +4,7 @@ from django.core.exceptions import ObjectDoesNotExist, ImproperlyConfigured # N from django.db.models.loading import ( # NOQA get_apps, get_app_path, get_app_paths, get_app, get_models, get_model, register_models, UnavailableApp) -from django.db.models.query import Q, QuerySet # NOQA +from django.db.models.query import Q, QuerySet, Prefetch # NOQA from django.db.models.expressions import F # NOQA from django.db.models.manager import Manager # NOQA from django.db.models.base import Model # NOQA diff --git a/django/db/models/fields/related.py b/django/db/models/fields/related.py index 6299b20fba..0f267f0d7c 100644 --- a/django/db/models/fields/related.py +++ b/django/db/models/fields/related.py @@ -162,7 +162,10 @@ class SingleRelatedObjectDescriptor(six.with_metaclass(RenameRelatedObjectDescri def get_queryset(self, **hints): return self.related.model._base_manager.db_manager(hints=hints) - def get_prefetch_queryset(self, instances): + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is not None: + raise ValueError("Custom queryset can't be used for this lookup.") + rel_obj_attr = attrgetter(self.related.field.attname) instance_attr = lambda obj: obj._get_pk_val() instances_dict = dict((instance_attr(inst), inst) for inst in instances) @@ -264,7 +267,10 @@ class ReverseSingleRelatedObjectDescriptor(six.with_metaclass(RenameRelatedObjec else: return QuerySet(self.field.rel.to, hints=hints) - def get_prefetch_queryset(self, instances): + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is not None: + raise ValueError("Custom queryset can't be used for this lookup.") + rel_obj_attr = self.field.get_foreign_related_value instance_attr = self.field.get_local_related_value instances_dict = dict((instance_attr(inst), inst) for inst in instances) @@ -397,23 +403,26 @@ def create_foreign_related_manager(superclass, rel_field, rel_model): qs._known_related_objects = {rel_field: {self.instance.pk: self.instance}} return qs - def get_prefetch_queryset(self, instances): + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is None: + queryset = super(RelatedManager, self).get_queryset() + + queryset._add_hints(instance=instances[0]) + queryset = queryset.using(queryset._db or self._db) + rel_obj_attr = rel_field.get_local_related_value instance_attr = rel_field.get_foreign_related_value instances_dict = dict((instance_attr(inst), inst) for inst in instances) query = {'%s__in' % rel_field.name: instances} - qs = super(RelatedManager, self).get_queryset() - qs._add_hints(instance=instances[0]) - if self._db: - qs = qs.using(self._db) - qs = qs.filter(**query) + queryset = queryset.filter(**query) + # Since we just bypassed this class' get_queryset(), we must manage # the reverse relation manually. - for rel_obj in qs: + for rel_obj in queryset: instance = instances_dict[rel_obj_attr(rel_obj)] setattr(rel_obj, rel_field.name, instance) cache_name = rel_field.related_query_name() - return qs, rel_obj_attr, instance_attr, False, cache_name + return queryset, rel_obj_attr, instance_attr, False, cache_name def add(self, *objs): objs = list(objs) @@ -563,15 +572,15 @@ def create_many_related_manager(superclass, rel): qs = qs.using(self._db) return qs._next_is_sticky().filter(**self.core_filters) - def get_prefetch_queryset(self, instances): - instance = instances[0] - db = self._db or router.db_for_read(instance.__class__, instance=instance) + def get_prefetch_queryset(self, instances, queryset=None): + if queryset is None: + queryset = super(ManyRelatedManager, self).get_queryset() + + queryset._add_hints(instance=instances[0]) + queryset = queryset.using(queryset._db or self._db) + query = {'%s__in' % self.query_field_name: instances} - qs = super(ManyRelatedManager, self).get_queryset() - qs._add_hints(instance=instance) - if self._db: - qs = qs.using(db) - qs = qs._next_is_sticky().filter(**query) + queryset = queryset._next_is_sticky().filter(**query) # M2M: need to annotate the query in order to get the primary model # that the secondary model was actually related to. We know that @@ -582,12 +591,12 @@ def create_many_related_manager(superclass, rel): # dealing with PK values. fk = self.through._meta.get_field(self.source_field_name) join_table = self.through._meta.db_table - connection = connections[db] + connection = connections[queryset.db] qn = connection.ops.quote_name - qs = qs.extra(select=dict( + queryset = queryset.extra(select=dict( ('_prefetch_related_val_%s' % f.attname, '%s.%s' % (qn(join_table), qn(f.column))) for f in fk.local_related_fields)) - return (qs, + return (queryset, lambda result: tuple(getattr(result, '_prefetch_related_val_%s' % f.attname) for f in fk.local_related_fields), lambda inst: tuple(getattr(inst, f.attname) for f in fk.foreign_related_fields), False, diff --git a/django/db/models/query.py b/django/db/models/query.py index 1289e14755..909f07114b 100644 --- a/django/db/models/query.py +++ b/django/db/models/query.py @@ -1619,6 +1619,59 @@ class RawQuerySet(object): return self._model_fields +class Prefetch(object): + def __init__(self, lookup, queryset=None, to_attr=None): + # `prefetch_through` is the path we traverse to perform the prefetch. + self.prefetch_through = lookup + # `prefetch_to` is the path to the attribute that stores the result. + self.prefetch_to = lookup + if to_attr: + self.prefetch_to = LOOKUP_SEP.join(lookup.split(LOOKUP_SEP)[:-1] + [to_attr]) + + self.queryset = queryset + self.to_attr = to_attr + + def add_prefix(self, prefix): + self.prefetch_through = LOOKUP_SEP.join([prefix, self.prefetch_through]) + self.prefetch_to = LOOKUP_SEP.join([prefix, self.prefetch_to]) + + def get_current_prefetch_through(self, level): + return LOOKUP_SEP.join(self.prefetch_through.split(LOOKUP_SEP)[:level + 1]) + + def get_current_prefetch_to(self, level): + return LOOKUP_SEP.join(self.prefetch_to.split(LOOKUP_SEP)[:level + 1]) + + def get_current_to_attr(self, level): + parts = self.prefetch_to.split(LOOKUP_SEP) + to_attr = parts[level] + to_list = self.to_attr and level == len(parts) - 1 + return to_attr, to_list + + def get_current_queryset(self, level): + if self.get_current_prefetch_to(level) == self.prefetch_to: + return self.queryset + return None + + def __eq__(self, other): + if isinstance(other, Prefetch): + return self.prefetch_to == other.prefetch_to + return False + + +def normalize_prefetch_lookups(lookups, prefix=None): + """ + Helper function that normalize lookups into Prefetch objects. + """ + ret = [] + for lookup in lookups: + if not isinstance(lookup, Prefetch): + lookup = Prefetch(lookup) + if prefix: + lookup.add_prefix(prefix) + ret.append(lookup) + return ret + + def prefetch_related_objects(result_cache, related_lookups): """ Helper function for prefetch_related functionality @@ -1626,13 +1679,15 @@ def prefetch_related_objects(result_cache, related_lookups): Populates prefetched objects caches for a list of results from a QuerySet """ + if len(result_cache) == 0: return # nothing to do + related_lookups = normalize_prefetch_lookups(related_lookups) + # We need to be able to dynamically add to the list of prefetch_related # lookups that we look up (see below). So we need some book keeping to # ensure we don't do duplicate work. - done_lookups = set() # list of lookups like foo__bar__baz done_queries = {} # dictionary of things like 'foo__bar': [results] auto_lookups = [] # we add to this as we go through. @@ -1640,25 +1695,27 @@ def prefetch_related_objects(result_cache, related_lookups): all_lookups = itertools.chain(related_lookups, auto_lookups) for lookup in all_lookups: - if lookup in done_lookups: - # We've done exactly this already, skip the whole thing + if lookup.prefetch_to in done_queries: + if lookup.queryset: + raise ValueError("'%s' lookup was already seen with a different queryset. " + "You may need to adjust the ordering of your lookups." % lookup.prefetch_to) + continue - done_lookups.add(lookup) # Top level, the list of objects to decorate is the result cache # from the primary QuerySet. It won't be for deeper levels. obj_list = result_cache - attrs = lookup.split(LOOKUP_SEP) - for level, attr in enumerate(attrs): + through_attrs = lookup.prefetch_through.split(LOOKUP_SEP) + for level, through_attr in enumerate(through_attrs): # Prepare main instances if len(obj_list) == 0: break - current_lookup = LOOKUP_SEP.join(attrs[:level + 1]) - if current_lookup in done_queries: + prefetch_to = lookup.get_current_prefetch_to(level) + if prefetch_to in done_queries: # Skip any prefetching, and any object preparation - obj_list = done_queries[current_lookup] + obj_list = done_queries[prefetch_to] continue # Prepare objects: @@ -1685,34 +1742,40 @@ def prefetch_related_objects(result_cache, related_lookups): # We assume that objects retrieved are homogenous (which is the premise # of prefetch_related), so what applies to first object applies to all. first_obj = obj_list[0] - prefetcher, descriptor, attr_found, is_fetched = get_prefetcher(first_obj, attr) + prefetcher, descriptor, attr_found, is_fetched = get_prefetcher(first_obj, through_attr) if not attr_found: raise AttributeError("Cannot find '%s' on %s object, '%s' is an invalid " "parameter to prefetch_related()" % - (attr, first_obj.__class__.__name__, lookup)) + (through_attr, first_obj.__class__.__name__, lookup.prefetch_through)) - if level == len(attrs) - 1 and prefetcher is None: + if level == len(through_attrs) - 1 and prefetcher is None: # Last one, this *must* resolve to something that supports # prefetching, otherwise there is no point adding it and the # developer asking for it has made a mistake. raise ValueError("'%s' does not resolve to a item that supports " "prefetching - this is an invalid parameter to " - "prefetch_related()." % lookup) + "prefetch_related()." % lookup.prefetch_through) if prefetcher is not None and not is_fetched: - obj_list, additional_prl = prefetch_one_level(obj_list, prefetcher, attr) + obj_list, additional_lookups = prefetch_one_level(obj_list, prefetcher, lookup, level) # We need to ensure we don't keep adding lookups from the # same relationships to stop infinite recursion. So, if we # are already on an automatically added lookup, don't add # the new lookups from relationships we've seen already. - if not (lookup in auto_lookups and - descriptor in followed_descriptors): - for f in additional_prl: - new_prl = LOOKUP_SEP.join([current_lookup, f]) - auto_lookups.append(new_prl) - done_queries[current_lookup] = obj_list + if not (lookup in auto_lookups and descriptor in followed_descriptors): + done_queries[prefetch_to] = obj_list + auto_lookups.extend(normalize_prefetch_lookups(additional_lookups, prefetch_to)) followed_descriptors.add(descriptor) + elif isinstance(getattr(first_obj, through_attr), list): + # The current part of the lookup relates to a custom Prefetch. + # This means that obj.attr is a list of related objects, and + # thus we must turn the obj.attr lists into a single related + # object list. + new_list = [] + for obj in obj_list: + new_list.extend(getattr(obj, through_attr)) + obj_list = new_list else: # Either a singly related object that has already been fetched # (e.g. via select_related), or hopefully some other property @@ -1724,7 +1787,7 @@ def prefetch_related_objects(result_cache, related_lookups): new_obj_list = [] for obj in obj_list: try: - new_obj = getattr(obj, attr) + new_obj = getattr(obj, through_attr) except exceptions.ObjectDoesNotExist: continue if new_obj is None: @@ -1755,6 +1818,11 @@ def get_prefetcher(instance, attr): try: rel_obj = getattr(instance, attr) attr_found = True + # If we are following a lookup path which leads us through a previous + # fetch from a custom Prefetch then we might end up into a list + # instead of related qs. This means the objects are already fetched. + if isinstance(rel_obj, list): + is_fetched = True except AttributeError: pass else: @@ -1776,7 +1844,7 @@ def get_prefetcher(instance, attr): return prefetcher, rel_obj_descriptor, attr_found, is_fetched -def prefetch_one_level(instances, prefetcher, attname): +def prefetch_one_level(instances, prefetcher, lookup, level): """ Helper function for prefetch_related_objects @@ -1799,14 +1867,14 @@ def prefetch_one_level(instances, prefetcher, attname): # 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 =\ - prefetcher.get_prefetch_queryset(instances) + rel_qs, rel_obj_attr, instance_attr, single, cache_name = ( + prefetcher.get_prefetch_queryset(instances, lookup.get_current_queryset(level))) # We have to handle the possibility that the default manager itself added # prefetch_related lookups to the QuerySet we just got back. We don't want to # trigger the prefetch_related functionality by evaluating the query. # Rather, we need to merge in the prefetch_related lookups. - additional_prl = getattr(rel_qs, '_prefetch_related_lookups', []) - if additional_prl: + additional_lookups = getattr(rel_qs, '_prefetch_related_lookups', []) + if additional_lookups: # Don't need to clone because the manager should have given us a fresh # instance, so we access an internal instead of using public interface # for performance reasons. @@ -1826,12 +1894,15 @@ def prefetch_one_level(instances, prefetcher, attname): # Need to assign to single cache on instance setattr(obj, cache_name, vals[0] if vals else None) else: - # Multi, attribute represents a manager with an .all() method that - # returns a QuerySet - qs = getattr(obj, attname).all() - qs._result_cache = vals - # We don't want the individual qs doing prefetch_related now, since we - # have merged this into the current work. - qs._prefetch_done = True - obj._prefetched_objects_cache[cache_name] = qs - return all_related_objects, additional_prl + to_attr, to_list = lookup.get_current_to_attr(level) + if to_list: + setattr(obj, to_attr, vals) + else: + # Cache in the QuerySet.all(). + qs = getattr(obj, to_attr).all() + qs._result_cache = vals + # We don't want the individual qs doing prefetch_related now, + # since we have merged this into the current work. + qs._prefetch_done = True + obj._prefetched_objects_cache[cache_name] = qs + return all_related_objects, additional_lookups diff --git a/docs/ref/models/queries.txt b/docs/ref/models/queries.txt index d07156c0b9..1979c16d40 100644 --- a/docs/ref/models/queries.txt +++ b/docs/ref/models/queries.txt @@ -129,3 +129,32 @@ In general, ``Q() objects`` make it possible to define and reuse conditions. This permits the :ref:`construction of complex database queries ` using ``|`` (``OR``) and ``&`` (``AND``) operators; in particular, it is not otherwise possible to use ``OR`` in ``QuerySets``. + +``Prefetch()`` objects +====================== + +.. versionadded:: 1.7 + +.. class:: Prefetch(lookup, queryset=None, to_attr=None) + +The ``Prefetch()`` object can be used to control the operation of +:meth:`~django.db.models.query.QuerySet.prefetch_related()`. + +The ``lookup`` argument describes the relations to follow and works the same +as the string based lookups passed to +:meth:`~django.db.models.query.QuerySet.prefetch_related()`. + +The ``queryset`` argument supplies a base ``QuerySet`` for the given lookup. +This is useful to further filter down the prefetch operation, or to call +:meth:`~django.db.models.query.QuerySet.select_related()` from the prefetched +relation, hence reducing the number of queries even further. + +The ``to_attr`` argument sets the result of the prefetch operation to a custom +attribute. + +.. note:: + + When using ``to_attr`` the prefetched result is stored in a list. + This can provide a significant speed improvement over traditional + ``prefetch_related`` calls which store the cached result within a + ``QuerySet`` instance. diff --git a/docs/ref/models/querysets.txt b/docs/ref/models/querysets.txt index 170df975bd..d1b6caba8d 100644 --- a/docs/ref/models/querysets.txt +++ b/docs/ref/models/querysets.txt @@ -898,7 +898,7 @@ objects have already been fetched, and it will skip fetching them again. Chaining ``prefetch_related`` calls will accumulate the lookups that are prefetched. To clear any ``prefetch_related`` behavior, pass ``None`` as a -parameter:: +parameter: >>> non_prefetched = qs.prefetch_related(None) @@ -925,6 +925,91 @@ profile for your use case! Note that if you use ``iterator()`` to run the query, ``prefetch_related()`` calls will be ignored since these two optimizations do not make sense together. +.. versionadded:: 1.7 + +You can use the :class:`~django.db.models.Prefetch` object to further control +the prefetch operation. + +In its simplest form ``Prefetch`` is equivalent to the traditional string based +lookups: + + >>> Restaurant.objects.prefetch_related(Prefetch('pizzas__toppings')) + +You can provide a custom queryset with the optional ``queryset`` argument. +This can be used to change the default ordering of the queryset: + + >>> Restaurant.objects.prefetch_related( + ... Prefetch('pizzas__toppings', queryset=Toppings.objects.order_by('name'))) + +Or to call :meth:`~django.db.models.query.QuerySet.select_related()` when +applicable to reduce the number of queries even further: + + >>> Pizza.objects.prefetch_related( + ... Prefetch('restaurants', queryset=Restaurant.objects.select_related('best_pizza'))) + +You can also assign the prefetched result to a custom attribute with the optional +``to_attr`` argument. The result will be stored directly in a list. + +This allows prefetching the same relation multiple times with a different +``QuerySet``; for instance: + + >>> vegetarian_pizzas = Pizza.objects.filter(vegetarian=True) + >>> Restaurant.objects.prefetch_related( + ... Prefetch('pizzas', to_attr('menu')), + ... Prefetch('pizzas', queryset=vegetarian_pizzas to_attr='vegetarian_menu')) + +Lookups created with custom ``to_attr`` can still be traversed as usual by other +lookups: + + >>> vegetarian_pizzas = Pizza.objects.filter(vegetarian=True) + >>> Restaurant.objects.prefetch_related( + ... Prefetch('pizzas', queryset=vegetarian_pizzas to_attr='vegetarian_menu'), + ... 'vegetarian_menu__toppings') + +Using ``to_attr`` is recommended when filtering down the prefetch result as it is +less ambiguous than storing a filtered result in the related manager's cache: + + >>> queryset = Pizza.objects.filter(vegetarian=True) + >>> + >>> # Recommended: + >>> restaurants = Restaurant.objects.prefetch_related( + ... Prefetch('pizzas', to_attr='vegetarian_pizzas' queryset=queryset)) + >>> vegetarian_pizzas = restaurants[0].vegetarian_pizzas + >>> + >>> # Not recommended: + >>> restaurants = Restaurant.objects.prefetch_related( + ... Prefetch('pizzas', queryset=queryset)) + >>> vegetarian_pizzas = restaurants[0].pizzas.all() + +.. note:: + + The ordering of lookups matters. + + Take the following examples: + + >>> prefetch_related('pizzas__toppings', 'pizzas') + + This works even though it's unordered because ``'pizzas__toppings'`` + already contains all the needed information, therefore the second argument + ``'pizzas'`` is actually redundant. + + >>> prefetch_related('pizzas__toppings', Prefetch('pizzas', queryset=Pizza.objects.all())) + + This will raise a ``ValueError`` because of the attempt to redefine the + queryset of a previously seen lookup. Note that an implicit queryset was + created to traverse ``'pizzas'`` as part of the ``'pizzas__toppings'`` + lookup. + + >>> prefetch_related('pizza_list__toppings', Prefetch('pizzas', to_attr='pizza_list')) + + This will trigger an ``AttributeError`` because ``'pizza_list'`` doesn't exist yet + when ``'pizza_list__toppings'`` is being processed. + + This consideration is not limited to the use of ``Prefetch`` objects. Some + advanced techniques may require that the lookups be performed in a + specific order to avoid creating extra queries; therefore it's recommended + to always carefully order ``prefetch_related`` arguments. + extra ~~~~~ diff --git a/docs/releases/1.7.txt b/docs/releases/1.7.txt index ad3166e973..b5e124ec3e 100644 --- a/docs/releases/1.7.txt +++ b/docs/releases/1.7.txt @@ -98,6 +98,21 @@ Using a custom manager when traversing reverse relations It is now possible to :ref:`specify a custom manager ` when traversing a reverse relationship. +New ``Prefetch`` object for advanced ``prefetch_related`` operations. +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The new :class:`~django.db.models.Prefetch` object allows customizing +prefetch operations. + +You can specify the ``QuerySet`` used to traverse a given relation +or customize the storage location of prefetch results. + +This enables things like filtering prefetched relations, calling +:meth:`~django.db.models.query.QuerySet.select_related()` from a prefetched +relation, or prefetching the same relation multiple times with different +querysets. See :meth:`~django.db.models.query.QuerySet.prefetch_related()` +for more details. + Admin shortcuts support time zones ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/prefetch_related/models.py b/tests/prefetch_related/models.py index 5e4bf4a3e4..076ca9d9d2 100644 --- a/tests/prefetch_related/models.py +++ b/tests/prefetch_related/models.py @@ -137,6 +137,9 @@ class TaggedItem(models.Model): def __str__(self): return self.tag + class Meta: + ordering = ['id'] + class Bookmark(models.Model): url = models.URLField() @@ -146,6 +149,9 @@ class Bookmark(models.Model): object_id_field='favorite_fkey', related_name='favorite_bookmarks') + class Meta: + ordering = ['id'] + class Comment(models.Model): comment = models.TextField() @@ -155,12 +161,16 @@ class Comment(models.Model): object_pk = models.TextField() content_object = generic.GenericForeignKey(ct_field="content_type", fk_field="object_pk") + class Meta: + ordering = ['id'] + ## Models for lookup ordering tests class House(models.Model): address = models.CharField(max_length=255) + owner = models.ForeignKey('Person', null=True) class Meta: ordering = ['id'] diff --git a/tests/prefetch_related/tests.py b/tests/prefetch_related/tests.py index 15ba01c1c5..76704743e2 100644 --- a/tests/prefetch_related/tests.py +++ b/tests/prefetch_related/tests.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals from django.contrib.contenttypes.models import ContentType from django.db import connection +from django.db.models import Prefetch from django.test import TestCase from django.test.utils import override_settings from django.utils import six @@ -13,9 +14,7 @@ from .models import (Author, Book, Reader, Qualification, Teacher, Department, class PrefetchRelatedTests(TestCase): - def setUp(self): - self.book1 = Book.objects.create(title="Poems") self.book2 = Book.objects.create(title="Jane Eyre") self.book3 = Book.objects.create(title="Wuthering Heights") @@ -207,6 +206,292 @@ class PrefetchRelatedTests(TestCase): self.assertTrue("name" in str(cm.exception)) +class CustomPrefetchTests(TestCase): + @classmethod + def traverse_qs(cls, obj_iter, path): + """ + Helper method that returns a list containing a list of the objects in the + obj_iter. Then for each object in the obj_iter, the path will be + recursively travelled and the found objects are added to the return value. + """ + ret_val = [] + + if hasattr(obj_iter, 'all'): + obj_iter = obj_iter.all() + + try: + iter(obj_iter) + except TypeError: + obj_iter = [obj_iter] + + for obj in obj_iter: + rel_objs = [] + for part in path: + if not part: + continue + rel_objs.extend(cls.traverse_qs(getattr(obj, part[0]), [part[1:]])) + ret_val.append((obj, rel_objs)) + return ret_val + + def setUp(self): + self.person1 = Person.objects.create(name="Joe") + self.person2 = Person.objects.create(name="Mary") + self.house1 = House.objects.create(address="123 Main St", owner=self.person1) + self.house2 = House.objects.create(address="45 Side St", owner=self.person1) + self.house3 = House.objects.create(address="6 Downing St", owner=self.person2) + self.house4 = House.objects.create(address="7 Regents St", owner=self.person2) + self.room1_1 = Room.objects.create(name="Dining room", house=self.house1) + self.room1_2 = Room.objects.create(name="Lounge", house=self.house1) + self.room1_3 = Room.objects.create(name="Kitchen", house=self.house1) + self.room2_1 = Room.objects.create(name="Dining room", house=self.house2) + self.room2_2 = Room.objects.create(name="Lounge", house=self.house2) + self.room2_3 = Room.objects.create(name="Kitchen", house=self.house2) + self.room3_1 = Room.objects.create(name="Dining room", house=self.house3) + self.room3_2 = Room.objects.create(name="Lounge", house=self.house3) + self.room3_3 = Room.objects.create(name="Kitchen", house=self.house3) + self.room4_1 = Room.objects.create(name="Dining room", house=self.house4) + self.room4_2 = Room.objects.create(name="Lounge", house=self.house4) + self.room4_3 = Room.objects.create(name="Kitchen", house=self.house4) + self.person1.houses.add(self.house1, self.house2) + self.person2.houses.add(self.house3, self.house4) + + def test_traverse_qs(self): + qs = Person.objects.prefetch_related('houses') + related_objs_normal = [list(p.houses.all()) for p in qs], + related_objs_from_traverse = [[inner[0] for inner in o[1]] + for o in self.traverse_qs(qs, [['houses']])] + self.assertEqual(related_objs_normal, (related_objs_from_traverse,)) + + def test_ambiguous(self): + # Ambiguous. + with self.assertRaises(ValueError): + self.traverse_qs( + Person.objects.prefetch_related('houses__rooms', Prefetch('houses', queryset=House.objects.all())), + [['houses', 'rooms']] + ) + + with self.assertRaises(AttributeError): + self.traverse_qs( + Person.objects.prefetch_related('houses_list__rooms', Prefetch('houses', queryset=House.objects.all(), to_attr='houses_lst')), + [['houses', 'rooms']] + ) + + # Not ambiguous. + self.traverse_qs( + Person.objects.prefetch_related('houses__rooms', 'houses'), + [['houses', 'rooms']] + ) + + self.traverse_qs( + Person.objects.prefetch_related('houses__rooms', Prefetch('houses', queryset=House.objects.all(), to_attr='houses_lst')), + [['houses', 'rooms']] + ) + + def test_m2m(self): + # Control lookups. + with self.assertNumQueries(2): + lst1 = self.traverse_qs( + Person.objects.prefetch_related('houses'), + [['houses']] + ) + + # Test lookups. + with self.assertNumQueries(2): + lst2 = self.traverse_qs( + Person.objects.prefetch_related(Prefetch('houses')), + [['houses']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(2): + lst2 = self.traverse_qs( + Person.objects.prefetch_related(Prefetch('houses', to_attr='houses_lst')), + [['houses_lst']] + ) + self.assertEqual(lst1, lst2) + + def test_reverse_m2m(self): + # Control lookups. + with self.assertNumQueries(2): + lst1 = self.traverse_qs( + House.objects.prefetch_related('occupants'), + [['occupants']] + ) + + # Test lookups. + with self.assertNumQueries(2): + lst2 = self.traverse_qs( + House.objects.prefetch_related(Prefetch('occupants')), + [['occupants']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(2): + lst2 = self.traverse_qs( + House.objects.prefetch_related(Prefetch('occupants', to_attr='occupants_lst')), + [['occupants_lst']] + ) + self.assertEqual(lst1, lst2) + + def test_m2m_through_fk(self): + # Control lookups. + with self.assertNumQueries(3): + lst1 = self.traverse_qs( + Room.objects.prefetch_related('house__occupants'), + [['house', 'occupants']] + ) + + # Test lookups. + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Room.objects.prefetch_related(Prefetch('house__occupants')), + [['house', 'occupants']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Room.objects.prefetch_related(Prefetch('house__occupants', to_attr='occupants_lst')), + [['house', 'occupants_lst']] + ) + self.assertEqual(lst1, lst2) + + def test_m2m_through_gfk(self): + TaggedItem.objects.create(tag="houses", content_object=self.house1) + TaggedItem.objects.create(tag="houses", content_object=self.house2) + + # Control lookups. + with self.assertNumQueries(3): + lst1 = self.traverse_qs( + TaggedItem.objects.filter(tag='houses').prefetch_related('content_object__rooms'), + [['content_object', 'rooms']] + ) + + # Test lookups. + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + TaggedItem.objects.prefetch_related( + Prefetch('content_object'), + Prefetch('content_object__rooms', to_attr='rooms_lst') + ), + [['content_object', 'rooms_lst']] + ) + self.assertEqual(lst1, lst2) + + def test_o2m_through_m2m(self): + # Control lookups. + with self.assertNumQueries(3): + lst1 = self.traverse_qs( + Person.objects.prefetch_related('houses', 'houses__rooms'), + [['houses', 'rooms']] + ) + + # Test lookups. + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Person.objects.prefetch_related(Prefetch('houses'), 'houses__rooms'), + [['houses', 'rooms']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Person.objects.prefetch_related(Prefetch('houses'), Prefetch('houses__rooms')), + [['houses', 'rooms']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Person.objects.prefetch_related(Prefetch('houses', to_attr='houses_lst'), 'houses_lst__rooms'), + [['houses_lst', 'rooms']] + ) + self.assertEqual(lst1, lst2) + with self.assertNumQueries(3): + lst2 = self.traverse_qs( + Person.objects.prefetch_related( + Prefetch('houses', to_attr='houses_lst'), + Prefetch('houses_lst__rooms', to_attr='rooms_lst') + ), + [['houses_lst', 'rooms_lst']] + ) + self.assertEqual(lst1, lst2) + + def test_generic_rel(self): + bookmark = Bookmark.objects.create(url='http://www.djangoproject.com/') + TaggedItem.objects.create(content_object=bookmark, tag='django') + TaggedItem.objects.create(content_object=bookmark, favorite=bookmark, tag='python') + + # Control lookups. + with self.assertNumQueries(4): + lst1 = self.traverse_qs( + Bookmark.objects.prefetch_related('tags', 'tags__content_object', 'favorite_tags'), + [['tags', 'content_object'], ['favorite_tags']] + ) + + # Test lookups. + with self.assertNumQueries(4): + lst2 = self.traverse_qs( + Bookmark.objects.prefetch_related( + Prefetch('tags', to_attr='tags_lst'), + Prefetch('tags_lst__content_object'), + Prefetch('favorite_tags'), + ), + [['tags_lst', 'content_object'], ['favorite_tags']] + ) + self.assertEqual(lst1, lst2) + + def test_custom_qs(self): + # Test basic. + with self.assertNumQueries(2): + lst1 = list(Person.objects.prefetch_related('houses')) + with self.assertNumQueries(2): + lst2 = list(Person.objects.prefetch_related( + Prefetch('houses', queryset=House.objects.all(), to_attr='houses_lst'))) + self.assertEqual( + self.traverse_qs(lst1, [['houses']]), + self.traverse_qs(lst2, [['houses_lst']]) + ) + + # Test queryset filtering. + with self.assertNumQueries(2): + lst2 = list(Person.objects.prefetch_related( + Prefetch('houses', queryset=House.objects.filter(pk__in=[self.house1.pk, self.house3.pk]), to_attr='houses_lst'))) + self.assertEqual(len(lst2[0].houses_lst), 1) + self.assertEqual(lst2[0].houses_lst[0], self.house1) + self.assertEqual(len(lst2[1].houses_lst), 1) + self.assertEqual(lst2[1].houses_lst[0], self.house3) + + # Test flattened. + with self.assertNumQueries(3): + lst1 = list(Person.objects.prefetch_related('houses__rooms')) + with self.assertNumQueries(3): + lst2 = list(Person.objects.prefetch_related( + Prefetch('houses__rooms', queryset=Room.objects.all(), to_attr='rooms_lst'))) + self.assertEqual( + self.traverse_qs(lst1, [['houses', 'rooms']]), + self.traverse_qs(lst2, [['houses', 'rooms_lst']]) + ) + + # Test inner select_related. + with self.assertNumQueries(3): + lst1 = list(Person.objects.prefetch_related('houses__owner')) + with self.assertNumQueries(2): + lst2 = list(Person.objects.prefetch_related( + Prefetch('houses', queryset=House.objects.select_related('owner')))) + self.assertEqual( + self.traverse_qs(lst1, [['houses', 'owner']]), + self.traverse_qs(lst2, [['houses', 'owner']]) + ) + + # Test inner prefetch. + inner_rooms_qs = Room.objects.filter(pk__in=[self.room1_1.pk, self.room1_2.pk]) + houses_qs_prf = House.objects.prefetch_related( + Prefetch('rooms', queryset=inner_rooms_qs, to_attr='rooms_lst')) + with self.assertNumQueries(3): + lst2 = list(Person.objects.prefetch_related( + Prefetch('houses', queryset=houses_qs_prf.filter(pk=self.house1.pk), to_attr='houses_lst'))) + + self.assertEqual(len(lst2[0].houses_lst[0].rooms_lst), 2) + self.assertEqual(lst2[0].houses_lst[0].rooms_lst[0], self.room1_1) + self.assertEqual(lst2[0].houses_lst[0].rooms_lst[1], self.room1_2) + self.assertEqual(len(lst2[1].houses_lst), 0) + class DefaultManagerTests(TestCase): def setUp(self): @@ -627,6 +912,45 @@ class MultiDbTests(TestCase): self.assertEqual(ages, "50, 49") + def test_using_is_honored_custom_qs(self): + B = Book.objects.using('other') + A = Author.objects.using('other') + book1 = B.create(title="Poems") + book2 = B.create(title="Sense and Sensibility") + + A.create(name="Charlotte Bronte", first_book=book1) + A.create(name="Jane Austen", first_book=book2) + + # Implicit hinting + with self.assertNumQueries(2, using='other'): + prefetch = Prefetch('first_time_authors', queryset=Author.objects.all()) + books = "".join("%s (%s)\n" % + (b.title, ", ".join(a.name for a in b.first_time_authors.all())) + for b in B.prefetch_related(prefetch)) + self.assertEqual(books, + "Poems (Charlotte Bronte)\n" + "Sense and Sensibility (Jane Austen)\n") + + # Explicit using on the same db. + with self.assertNumQueries(2, using='other'): + prefetch = Prefetch('first_time_authors', queryset=Author.objects.using('other')) + books = "".join("%s (%s)\n" % + (b.title, ", ".join(a.name for a in b.first_time_authors.all())) + for b in B.prefetch_related(prefetch)) + self.assertEqual(books, + "Poems (Charlotte Bronte)\n" + "Sense and Sensibility (Jane Austen)\n") + + # Explicit using on a different db. + with self.assertNumQueries(1, using='default'), self.assertNumQueries(1, using='other'): + prefetch = Prefetch('first_time_authors', queryset=Author.objects.using('default')) + books = "".join("%s (%s)\n" % + (b.title, ", ".join(a.name for a in b.first_time_authors.all())) + for b in B.prefetch_related(prefetch)) + self.assertEqual(books, + "Poems ()\n" + "Sense and Sensibility ()\n") + class Ticket19607Tests(TestCase):