Ticket #16937: prefetch_1.3.1.diff

django/db/models/fields/related.py

@@ -422 +422 @@
                 db = self._db or router.db_for_read(rel_model, instance=instance)
                 return superclass.get_query_set(self).using(db).filter(**(self.core_filters))
 
+            def get_prefetch_query_set(self, instances):
+                """
+                Return a queryset that does the bulk lookup needed
+                by prefetch_related functionality.
+                """
+                if not instances:
+                    return self.model.objects.none()
+
+                db = self._db or router.db_for_read(self.model, instance=instances[0])
+                query = {'%s__%s__in' % (rel_field.name, attname):
+                             [getattr(obj, attname) for obj in instances]}
+                return super(RelatedManager, self).get_query_set().using(db).filter(**query)
+
+            def select_matching_instances(self, obj, related_objects):
+                field_val = getattr(obj, attname)
+                other_attname = rel_field.get_attname()
+                return [rel_obj for rel_obj in related_objects
+                        if getattr(rel_obj, other_attname) == field_val]
+
+            def all(self):
+                try:
+                    return self.instance._prefetched_objects_cache[rel_field.related_query_name()]
+                except (AttributeError, KeyError):
+                    return super(RelatedManager, self).all()
+
             def add(self, *objs):
                 for obj in objs:
                     if not isinstance(obj, self.model):
@@ -476 +501 @@
     and adds behavior for many-to-many related objects."""
     through = rel.through
     class ManyRelatedManager(superclass):
-        def __init__(self, model=None, core_filters=None, instance=None, symmetrical=None,
-                join_table=None, source_field_name=None, target_field_name=None,
-                reverse=False):
+        def __init__(self, model=None, query_field_name=None, instance=None, symmetrical=None,
+                      source_field_name=None, target_field_name=None, reverse=False,
+                join_table=None, prefetch_cache_name=None):
             super(ManyRelatedManager, self).__init__()
-            self.core_filters = core_filters
+            self.query_field_name = query_field_name
+            self.core_filters = {'%s__pk' % query_field_name: instance._get_pk_val()}
             self.model = model
             self.symmetrical = symmetrical
             self.instance = instance
@@ -489 +515 @@
             self.through = through
             self._pk_val = self.instance.pk
             self.reverse = reverse
+            self.prefetch_cache_name = prefetch_cache_name
             if self._pk_val is None:
                 raise ValueError("%r instance needs to have a primary key value before a many-to-many relationship can be used." % instance.__class__.__name__)
 
@@ -496 +523 @@
             db = self._db or router.db_for_read(self.instance.__class__, instance=self.instance)
             return superclass.get_query_set(self).using(db)._next_is_sticky().filter(**(self.core_filters))
 
+        def get_prefetch_query_set(self, instances):
+            if not instances:
+                return self.model.objects.none()
+
+            from django.db import connections
+
+            db = self._db or router.db_for_read(self.model, instance=instances[0])
+            query = {'%s__pk__in' % self.query_field_name:
+                         [obj._get_pk_val() for obj in instances]}
+            qs = super(ManyRelatedManager, self).get_query_set().using(db)._next_is_sticky().filter(**query)
+
+            # M2M: need to annotate the query in order to get the PK of the
+            # primary model that the secondary model was actually related to.
+
+            # We know that there will already be a join on the join table, so we
+            # can just add the select.
+            join_table = self.through._meta.db_table
+            pk_col = "%s_id" % self.source_field_name
+            connection = connections[db]
+            qn = connection.ops.quote_name
+            qs = qs.extra(select={'_prefetch_related_pk':
+                                      '%s.%s' % (qn(join_table), qn(pk_col))})
+            return qs
+
+        def select_matching_instances(self, obj, related_objects):
+            pk_val = obj._get_pk_val()
+            return [rel_obj for rel_obj in related_objects
+                    if rel_obj._prefetch_related_pk == pk_val]
+
+        def all(self):
+            try:
+                return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
+            except (AttributeError, KeyError):
+                return super(ManyRelatedManager, self).all()
+
         # If the ManyToMany relation has an intermediary model,
         # the add and remove methods do not exist.
         if rel.through._meta.auto_created:
@@ -671 +733 @@
 
         manager = RelatedManager(
             model=rel_model,
-            core_filters={'%s__pk' % self.related.field.name: instance._get_pk_val()},
+            query_field_name=self.related.field.name,
+            prefetch_cache_name=self.related.field.related_query_name(),
             instance=instance,
             symmetrical=False,
             source_field_name=self.related.field.m2m_reverse_field_name(),
@@ -723 +786 @@
 
         manager = RelatedManager(
             model=rel_model,
-            core_filters={'%s__pk' % self.field.related_query_name(): instance._get_pk_val()},
+            query_field_name=self.field.related_query_name(),
+            prefetch_cache_name=self.field.name,
             instance=instance,
             symmetrical=self.field.rel.symmetrical,
             source_field_name=self.field.m2m_field_name(),

django/db/models/manager.py

@@ -167 +167 @@
     def select_related(self, *args, **kwargs):
         return self.get_query_set().select_related(*args, **kwargs)
 
+    def prefetch_related(self, *args, **kwargs):
+        return self.get_query_set().prefetch_related(*args, **kwargs)
+
     def values(self, *args, **kwargs):
         return self.get_query_set().values(*args, **kwargs)
 

django/db/models/query.py

@@ -37 +37 @@
         self._iter = None
         self._sticky_filter = False
         self._for_write = False
+        self._prefetch_related = set()
+        self._prefetch_done = False
 
     ########################
     # PYTHON MAGIC METHODS #
@@ -82 +84 @@
                 self._result_cache = list(self.iterator())
         elif self._iter:
             self._result_cache.extend(self._iter)
+        if self._prefetch_related and not self._prefetch_done:
+            self._prefetch_related_objects()
         return len(self._result_cache)
 
     def __iter__(self):
+        if self._prefetch_related:
+            # We need all the results in order to be able to do the prefetch
+            # in one go. To minimize code duplication, we use the __len__
+            # code path which also forces this, and also does the prefetch
+            len(self)
+
         if self._result_cache is None:
             self._iter = self.iterator()
             self._result_cache = []
@@ -107 +117 @@
                 self._fill_cache()
 
     def __nonzero__(self):
+        if self._prefetch_related:
+            # We need all the results in order to be able to do the prefetch
+            # in one go. To minimize code duplication, we use the __len__
+            # code path which also forces this, and also does the prefetch
+            len(self)
+
         if self._result_cache is not None:
             return bool(self._result_cache)
         try:
@@ -496 +512 @@
             return self.query.has_results(using=self.db)
         return bool(self._result_cache)
 
+    def _prefetch_related_objects(self):
+        # This method can only be called once the result cache has been filled.
+        prefetch_related_objects(self._result_cache, self._prefetch_related)
+        self._prefetch_done = True
+
     ##################################################
     # PUBLIC METHODS THAT RETURN A QUERYSET SUBCLASS #
     ##################################################
@@ -607 +628 @@
             obj.query.max_depth = depth
         return obj
 
+    def prefetch_related(self, *fields):
+        """
+        Returns a new QuerySet instance that will prefetch Many-To-One
+        and Many-To-Many related objects when the QuerySet is evaluated.
+
+        The fields specified must be attributes that return a RelatedManager of
+        some kind when used on instances of the evaluated QuerySet.
+
+        These RelatedManagers will be modified so that their 'all()' method will
+        return a QuerySet whose cache is already filled with objects that were
+        looked up in a single batch, rather than one query per object in the
+        current QuerySet.
+
+        When prefetch_related() is called more than once, the list of fields to
+        prefetch is added to. Call prefetch_related() with no arguments to
+        clears the list.
+        """
+        if fields:
+            new_fields = self._prefetch_related.union(set(fields))
+        else:
+            new_fields = set()
+        return self._clone(_prefetch_related=new_fields)
+
     def dup_select_related(self, other):
         """
         Copies the related selection status from the QuerySet 'other' to the
@@ -756 +800 @@
             query.filter_is_sticky = True
         c = klass(model=self.model, query=query, using=self._db)
         c._for_write = self._for_write
+        c._prefetch_related = self._prefetch_related
         c.__dict__.update(kwargs)
         if setup and hasattr(c, '_setup_query'):
             c._setup_query()
@@ -1434 +1479 @@
     query = sql.InsertQuery(model)
     query.insert_values(values, raw_values)
     return query.get_compiler(using=using).execute_sql(return_id)
+
+
+def prefetch_related_objects(result_cache, fields):
+    """
+    Populates prefetched objects caches for a list of results
+    from a QuerySet
+    """
+    from django.db.models.sql.constants import LOOKUP_SEP
+
+    if len(result_cache) == 0:
+        return # nothing to do
+
+    model = result_cache[0].__class__
+
+    # We need to be able to dynamically add to the list of prefetch_related
+    # fields that we look up (see below).  So we need some book keeping to
+    # ensure we don't do duplicate work.
+    done_fields = set() # list of fields like foo__bar__baz
+    done_lookups = {}   # dictionary of things like 'foo__bar': [results]
+    fields = list(fields)
+
+    # We may expand fields, so need a loop that allows for that
+    i = 0
+    while i < len(fields):
+        # 'field' can span several relationships, and so represent multiple
+        # lookups.
+        field = fields[i]
+
+        if field in done_fields:
+            # We've done exactly this already, skip the whole thing
+            i += 1
+            continue
+        done_fields.add(field)
+
+        # Top level, the list of objects to decorate is the the result cache
+        # from the primary QuerySet. It won't be for deeper levels.
+        obj_list = result_cache
+
+        attrs = field.split(LOOKUP_SEP)
+        for level, attr in enumerate(attrs):
+            # Prepare main instances
+            if len(obj_list) == 0:
+                break
+
+            good_objects = True
+            for obj in obj_list:
+                if not hasattr(obj, '_prefetched_objects_cache'):
+                    try:
+                        obj._prefetched_objects_cache = {}
+                    except AttributeError:
+                        # Must be in a QuerySet subclass that is not returning
+                        # Model instances, either in Django or 3rd
+                        # party. prefetch_related() doesn't make sense, so quit
+                        # now.
+                        good_objects = False
+                        break
+            if not good_objects:
+                break
+
+            # Descend down tree
+            try:
+                rel_obj = getattr(obj_list[0], attr)
+            except AttributeError:
+                raise AttributeError("Cannot find '%s' on %s object, '%s' is an invalid "
+                                     "parameter to prefetch_related()" %
+                                     (attr, obj_list[0].__class__.__name__, field))
+
+            can_prefetch = hasattr(rel_obj, 'get_prefetch_query_set')
+            if level == len(attrs) - 1 and not can_prefetch:
+                # Last one, this *must* resolve to a related manager.
+                raise ValueError("'%s' does not resolve to a supported 'many related"
+                                 " manager' for model %s - this is an invalid"
+                                 " parameter to prefetch_related()."
+                                 % (field, model.__name__))
+
+            if can_prefetch:
+                # Check we didn't do this already
+                lookup = LOOKUP_SEP.join(attrs[0:level+1])
+                if lookup in done_lookups:
+                    obj_list = done_lookups[lookup]
+                else:
+                    relmanager = rel_obj
+                    obj_list, additional_prf = _prefetch_one_level(obj_list, relmanager, attr)
+                    for f in additional_prf:
+                        new_prf = LOOKUP_SEP.join([lookup, f])
+                        fields.append(new_prf)
+                    done_lookups[lookup] = obj_list
+            else:
+                # Assume we've got some singly related object. We replace
+                # the current list of parent objects with that list.
+                obj_list = [getattr(obj, attr) for obj in obj_list]
+
+        i += 1
+
+
+def _prefetch_one_level(instances, relmanager, attname):
+    """
+    Runs prefetches on all instances using the manager relmanager,
+    assigning results to queryset against instance.attname.
+
+    The prefetched objects are returned, along with any additional
+    prefetches that must be done due to prefetch_related fields
+    found from default managers.
+    """
+    mainqs = relmanager.get_prefetch_query_set(instances)
+    # We have to handle the possibility that the default manager itself added
+    # prefetch_related fields 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 fields.
+    additional_prf = list(getattr(mainqs, '_prefetch_related', []))
+    if additional_prf:
+        mainqs = mainqs.prefetch_related()
+    all_related_objects = list(mainqs)
+    for obj in instances:
+        qs = getattr(obj, attname).all()
+        qs._result_cache = relmanager.select_matching_instances(obj, all_related_objects)
+        # 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[attname] = qs
+    return all_related_objects, additional_prf

docs/ref/models/querysets.txt

@@ -689 +689 @@
 ``OneToOneFields`` will not be traversed in the reverse direction if you
 are performing a depth-based ``select_related``.
 
+prefetch_related
+~~~~~~~~~~~~~~~~
+
+.. method:: prefetch_related(*fields)
+
+.. versionadded:: 1.4
+
+Returns a ``QuerySet`` that will automatically retrieve, in a single batch,
+related many-to-many and many-to-one objects for the specified fields.
+
+This is similar to ``select_related`` for the 'many related objects' case, with
+the following important differences:
+
+* ``prefetch_related`` causes a separate query to be issued for each set of
+  related objects that you request. This query is done as soon as the QuerySet
+  begins to be evaluated.
+
+  This is in contrast to ``select_related``, which modifies the original query
+  with joins in order to get the related objects in the same query as the main
+  objects.
+
+* It is 'single depth' only, and doesn't support join syntax.
+
+The fields that must be supplied to this method can be any attributes on the
+model instances which represent related queries that return multiple
+objects. This includes attributes representing the 'many' side of ``ForeignKey``
+relationships and ``ManyToManyField`` attributes.
+
+For example, suppose you have these models::
+
+    class Topping(models.Model):
+        name = models.CharField(max_length=30)
+
+    class Pizza(models.Model):
+        name = models.CharField(max_length=50)
+        toppings = models.ManyToManyField(Topping)
+
+        def __unicode__(self):
+            return u"%s (%s)" % (self.name, u", ".join([topping.name
+                                                        for topping in self.toppings.all()]))
+
+and run this code::
+
+    >>> Pizza.objects.all()
+    [u"Hawaiian (ham, pineaapple)", u"Seafood (prawns, smoked salmon)"...
+
+The problem with this code is that it will run a query on the Toppings table for
+**every** item in the Pizza ``QuerySet``.  Using ``prefetch_related``, this can
+be reduced to two:
+
+    >>> pizzas = Pizza.objects.all().prefetch_related('toppings')
+
+All the relevant toppings will be fetched in a single query, and used to make
+``QuerySets`` that have a pre-filled cache of the relevant results. These
+``QuerySets`` are then used in the ``self.toppings.all()`` calls.
+
+Please note that use of ``prefetch_related`` will mean that the additional
+queries run will **always** be executed - even if you never use the related
+objects - and it always fully populates the result cache on the primary
+``QuerySet`` (which can sometimes be avoided in other cases).
+
+Remember that, as always with QuerySets, any subsequent chained methods will
+ignore previously cached results, and retrieve data in a fresh database
+query. So, if you write the following:
+
+    >>> pizzas = Pizza.objects.prefetch_related('toppings')
+    >>> [list(pizza.topppings.filter(spicy=True) for pizza in pizzas]
+
+...then the fact that `pizza.toppings.all()` has been prefetched will not help
+you - in fact it hurts preformance, since you have done a database query that
+you haven't used. So use this feature with caution!
+
+Chaining ``prefetch_related`` calls will accumulate the fields that should have
+this behaviour applied. To clear any ``prefetch_related`` behaviour, call the
+method with no arguments.
+
+
 extra
 ~~~~~
 

tests/modeltests/prefetch_related/models.py

@@ +1 @@
+from django.db import models
+
+
+class Author(models.Model):
+    name = models.CharField(max_length=50)
+    first_book = models.ForeignKey('Book', related_name='first_time_authors')
+
+    def __unicode__(self):
+        return self.name
+
+class Book(models.Model):
+    title = models.CharField(max_length=255)
+
+    authors = models.ManyToManyField(Author, related_name='books')
+
+    def __unicode__(self):
+        return self.title

tests/modeltests/prefetch_related/tests.py

@@ +1 @@
+from django.test import TestCase
+
+from models import Author, Book
+
+
+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")
+
+        self.author1 = Author.objects.create(name="Charlotte",
+                                             first_book=self.book1)
+        self.author2 = Author.objects.create(name="Anne",
+                                             first_book=self.book1)
+        self.author3 = Author.objects.create(name="Emily",
+                                             first_book=self.book1)
+
+        self.book1.authors.add(self.author1)
+        self.book1.authors.add(self.author2)
+        self.book1.authors.add(self.author3)
+        self.book2.authors.add(self.author1)
+        self.book3.authors.add(self.author3)
+
+    def test_m2m_forward(self):
+        with self.assertNumQueries(2):
+            lists = [list(b.authors.all()) for b in Book.objects.prefetch_related('authors')]
+
+        normal_lists = [list(b.authors.all()) for b in Book.objects.all()]
+        self.assertEqual(lists, normal_lists)
+
+
+    def test_m2m_reverse(self):
+        with self.assertNumQueries(2):
+            lists = [list(a.books.all()) for a in Author.objects.prefetch_related('books')]
+
+        normal_lists = [list(a.books.all()) for a in Author.objects.all()]
+        self.assertEqual(lists, normal_lists)
+
+    def test_foreignkey_reverse(self):
+        with self.assertNumQueries(2):
+            lists = [list(b.first_time_authors.all())
+                     for b in Book.objects.prefetch_related('first_time_authors')]
+
+        self.assertQuerysetEqual(self.book2.authors.all(), [u"<Author: Charlotte>"])
+
+    def test_survives_clone(self):
+        with self.assertNumQueries(2):
+            lists = [list(b.first_time_authors.all())
+                     for b in Book.objects.prefetch_related('first_time_authors').exclude(id=1000)]
+
+    def test_len(self):
+        with self.assertNumQueries(2):
+            qs = Book.objects.prefetch_related('first_time_authors')
+            length = len(qs)
+            lists = [list(b.first_time_authors.all())
+                     for b in qs]
+
+    def test_bool(self):
+        with self.assertNumQueries(2):
+            qs = Book.objects.prefetch_related('first_time_authors')
+            x = bool(qs)
+            lists = [list(b.first_time_authors.all())
+                     for b in qs]
+
+    def test_clear(self):
+        with self.assertNumQueries(4):
+            with_prefetch = Author.objects.prefetch_related('books')
+            without_prefetch = with_prefetch.prefetch_related()
+            lists = [list(a.books.all()) for a in without_prefetch]