Ticket #16937: prefetch.diff

django/db/models/fields/related.py

@@ -421 +421 @@
                 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 all(self):
+                try:
+                    return 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):
@@ -474 +480 @@
     through = rel.through
     class ManyRelatedManager(superclass):
         def __init__(self, model=None, core_filters=None, instance=None, symmetrical=None,
-                     source_field_name=None, target_field_name=None, reverse=False):
+                     source_field_name=None, target_field_name=None, reverse=False,
+                     prefetch_cache_name=None):
             super(ManyRelatedManager, self).__init__()
             self.core_filters = core_filters
             self.model = model
@@ -485 +492 @@
             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__)
 
@@ -492 +500 @@
             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 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:
@@ -669 +683 @@
         manager = RelatedManager(
             model=rel_model,
             core_filters={'%s__pk' % self.related.field.name: instance._get_pk_val()},
+            prefetch_cache_name=self.related.field.related_query_name(),
             instance=instance,
             symmetrical=False,
             source_field_name=self.related.field.m2m_reverse_field_name(),
@@ -721 +736 @@
         manager = RelatedManager(
             model=rel_model,
             core_filters={'%s__pk' % self.field.related_query_name(): instance._get_pk_val()},
+            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

@@ -172 +172 @@
     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

@@ -36 +36 @@
         self._iter = None
         self._sticky_filter = False
         self._for_write = False
+        self._prefetch_related = set()
 
     ########################
     # PYTHON MAGIC METHODS #
@@ -81 +82 @@
                 self._result_cache = list(self.iterator())
         elif self._iter:
             self._result_cache.extend(self._iter)
+        if self._prefetch_related:
+            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 = []
@@ -106 +115 @@
                 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:
@@ -526 +541 @@
             return self.query.has_results(using=self.db)
         return bool(self._result_cache)
 
+    def _prefetch_related_objects(self):
+        from django.db import connections
+        # Here we access some internals of how related managers work in order to
+        # fix a common efficiency problem with M2M and M2O lookups.
+
+        # This method can only be called once the result cache has been filled.
+        obj = None
+        for obj in self._result_cache:
+            if not hasattr(obj, '_prefetched_objects_cache'):
+                obj._prefetched_objects_cache = {}
+
+        if obj is None:
+            return # nothing to do
+
+
+        for attname in self._prefetch_related:
+            field, model, direct, m2m = self.model._meta.get_field_by_name(attname)
+
+            # A new related manager is created every time the attribute is
+            # accessed, so we can alter the following relmanager with safety
+            relmanager = getattr(obj, attname)
+
+            if m2m:
+                # Use a modified RelatedManager to build the query
+
+                # Collect values to filter on
+                pks = [obj._get_pk_val() for obj in self._result_cache]
+
+                if direct:
+                    field_name = field.related_query_name()
+                else:
+                    field_name = field.field.name
+                relmanager.core_filters = {'%s__pk__in' % field_name: pks}
+
+                # M2M: need to annotate the query in order to get the PK of the
+                # primary model that the secondary model was actually related
+                # to. For ForeignKeys, this information is in the secondary
+                # table already.
+
+                # We know that there will already be a join on the join table,
+                # so we can just add the select.
+                join_table = relmanager.through._meta.db_table
+                pk_col = "%s_id" % relmanager.source_field_name
+                connection = connections[self.db]
+                qn = connection.ops.quote_name
+                qs = relmanager.extra(select={'_prefetch_related_pk':'%s.%s' % (qn(join_table), qn(pk_col))})
+
+                # Execute query
+                all_related_objects = list(qs)
+
+                # Now decorate _result_cache with relevant prefetched data
+                for obj in self._result_cache:
+                    qs = getattr(obj, attname).all()
+                    pk_val = obj._get_pk_val()
+                    qs._result_cache = [rel_obj for rel_obj in all_related_objects
+                                        if rel_obj._prefetch_related_pk == pk_val]
+                    obj._prefetched_objects_cache[attname] = qs
+
+            else:
+                # Reverse ForeignKey
+
+                # Use a modified RelatedManager to build the query
+                rel_field = field.field
+                rel_attname = rel_field.rel.get_related_field().name
+                rel_vals = [getattr(obj, rel_attname) for obj in self._result_cache]
+                relmanager.core_filters = {'%s__%s__in' % (rel_field.name,
+                                                           rel_attname):
+                                               rel_vals}
+
+                # Execute query
+                all_related_objects = list(relmanager.all())
+
+                # Now decorate _result_cache with relevant prefetched data
+                for obj in self._result_cache:
+                    qs = getattr(obj, attname).all()
+                    field_val = getattr(obj, rel_attname)
+                    other_attname = rel_field.get_attname()
+                    qs._result_cache = [rel_obj for rel_obj in all_related_objects
+                                        if getattr(rel_obj, other_attname) == field_val]
+                    obj._prefetched_objects_cache[attname] = qs
+
+
     ##################################################
     # PUBLIC METHODS THAT RETURN A QUERYSET SUBCLASS #
     ##################################################
@@ -649 +746 @@
             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
@@ -798 +918 @@
             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()
@@ -941 +1062 @@
         c.aggregate_names = self.aggregate_names
         if setup and hasattr(c, '_setup_query'):
             c._setup_query()
+        # prefetch_related doesn't make sense for a ValuesQuerySet
+        c._prefetch_related = None
+
         return c
 
     def _merge_sanity_check(self, other):
@@ -1051 +1175 @@
         c._kind = self._kind
         if setup and hasattr(c, '_setup_query'):
             c._setup_query()
+        # prefetch_related doesn't make sense for a DateQuerySet
+        self._prefetch_related = None
         return c
 
 

docs/ref/models/querysets.txt

@@ -690 +690 @@
 A :class:`~django.db.models.OneToOneField` is not traversed in the reverse
 direction if you are performing a depth-based ``select_related()`` call.
 
+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 'many related objects' case, with the
+following important differences:
+
+* ``prefetch_related`` issues a separate query to get the related objects.
+  This is done as soon as the QuerySet begins to be evaluated.
+
+* 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 a
+``QuerySet`` that has a pre-filled cache of the relevant results. This
+``QuerySet`` is then used in the ``self.toppings.all()`` call.
+
+Please note the following:
+
+* use of ``prefetch_related`` will mean that it **always** runs the two
+  queries - even if you never use the related objects - and it always
+  fully populates the result cache on the primary ``QuerySet``.
+
+* **Only** the 'all()' method of the related manager is affected. If you want to
+  do additional filtering of the related objects using ``filter``, for example,
+  this will create a new ``QuerySet``, ignoring the pre-filled ``QuerySet``, and
+  will do that filtering in the database as always.
+
+Therefore you should only use ``prefetch_related`` when you are sure that the
+related objects are needed (or will be needed in the normal case), since you are
+adding overhead of issuing an additional database query creating (many) model
+instances. Where appropriate you could use this in a manager or default manager.
+
+Chaining ``prefetch_related`` will accumulate the fields that should have this
+behaviour applied. To clear any ``prefetch_related`` behaviour, call the method
+with no arguments.
+
+
 extra
 ~~~~~
 

new file 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')
+
+
+class Book(models.Model):
+    title = models.CharField(max_length=255)
+
+    authors = models.ManyToManyField(Author, related_name='books')
+

new file 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')]
+
+    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]