Fixed #26226 -- Made related managers honor the queryset used for prefetching their results.

Thanks Loïc for the suggested improvements and Tim for the review.

6 files changed, 133 insertions, 22 deletions

diff --git a/django/contrib/contenttypes/fields.py b/django/contrib/contenttypes/fields.py
index a51ecdd48e..f9d5d7cfd2 100644
--- a/django/contrib/contenttypes/fields.py
+++ b/django/contrib/contenttypes/fields.py
@@ -521,12 +521,19 @@ def create_generic_related_manager(superclass, rel):
         def __str__(self):
             return repr(self)
 
+        def _apply_rel_filters(self, queryset):
+            """
+            Filter the queryset for the instance this manager is bound to.
+            """
+            db = self._db or router.db_for_read(self.model, instance=self.instance)
+            return queryset.using(db).filter(**self.core_filters)
+
         def get_queryset(self):
             try:
                 return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
             except (AttributeError, KeyError):
-                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)
+                queryset = super(GenericRelatedObjectManager, self).get_queryset()
+                return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
             if queryset is None:
diff --git a/django/db/models/fields/related_descriptors.py b/django/db/models/fields/related_descriptors.py
index eae3a9628d..b6349e5c49 100644
--- a/django/db/models/fields/related_descriptors.py
+++ b/django/db/models/fields/related_descriptors.py
@@ -502,23 +502,29 @@ def create_reverse_many_to_one_manager(superclass, rel):
             return manager_class(self.instance)
         do_not_call_in_templates = True
 
+        def _apply_rel_filters(self, queryset):
+            """
+            Filter the queryset for the instance this manager is bound to.
+            """
+            db = self._db or router.db_for_read(self.model, instance=self.instance)
+            empty_strings_as_null = connections[db].features.interprets_empty_strings_as_nulls
+            queryset._add_hints(instance=self.instance)
+            if self._db:
+                queryset = queryset.using(self._db)
+            queryset = queryset.filter(**self.core_filters)
+            for field in self.field.foreign_related_fields:
+                val = getattr(self.instance, field.attname)
+                if val is None or (val == '' and empty_strings_as_null):
+                    return queryset.none()
+            queryset._known_related_objects = {self.field: {self.instance.pk: self.instance}}
+            return queryset
+
         def get_queryset(self):
             try:
                 return self.instance._prefetched_objects_cache[self.field.related_query_name()]
             except (AttributeError, KeyError):
-                db = self._db or router.db_for_read(self.model, instance=self.instance)
-                empty_strings_as_null = connections[db].features.interprets_empty_strings_as_nulls
-                qs = super(RelatedManager, self).get_queryset()
-                qs._add_hints(instance=self.instance)
-                if self._db:
-                    qs = qs.using(self._db)
-                qs = qs.filter(**self.core_filters)
-                for field in self.field.foreign_related_fields:
-                    val = getattr(self.instance, field.attname)
-                    if val is None or (val == '' and empty_strings_as_null):
-                        return qs.none()
-                qs._known_related_objects = {self.field: {self.instance.pk: self.instance}}
-                return qs
+                queryset = super(RelatedManager, self).get_queryset()
+                return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
             if queryset is None:
@@ -776,15 +782,21 @@ def create_forward_many_to_many_manager(superclass, rel, reverse):
                 filters |= symmetrical_filters
             return filters
 
+        def _apply_rel_filters(self, queryset):
+            """
+            Filter the queryset for the instance this manager is bound to.
+            """
+            queryset._add_hints(instance=self.instance)
+            if self._db:
+                queryset = queryset.using(self._db)
+            return queryset._next_is_sticky().filter(**self.core_filters)
+
         def get_queryset(self):
             try:
                 return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
             except (AttributeError, KeyError):
-                qs = super(ManyRelatedManager, self).get_queryset()
-                qs._add_hints(instance=self.instance)
-                if self._db:
-                    qs = qs.using(self._db)
-                return qs._next_is_sticky().filter(**self.core_filters)
+                queryset = super(ManyRelatedManager, self).get_queryset()
+                return self._apply_rel_filters(queryset)
 
         def get_prefetch_queryset(self, instances, queryset=None):
             if queryset is None:
diff --git a/django/db/models/query.py b/django/db/models/query.py
index 71e9ba9bda..85567e88ac 100644
--- a/django/db/models/query.py
+++ b/django/db/models/query.py
@@ -23,6 +23,7 @@ from django.db.models.query_utils import (
 )
 from django.db.models.sql.constants import CURSOR
 from django.utils import six, timezone
+from django.utils.deprecation import RemovedInDjango20Warning
 from django.utils.functional import partition
 from django.utils.version import get_version
 
@@ -1608,6 +1609,9 @@ def prefetch_one_level(instances, prefetcher, lookup, level):
             msg = 'to_attr={} conflicts with a field on the {} model.'
             raise ValueError(msg.format(to_attr, model.__name__))
 
+    # Whether or not we're prefetching the last part of the lookup.
+    leaf = len(lookup.prefetch_through.split(LOOKUP_SEP)) - 1 == level
+
     for obj in instances:
         instance_attr_val = instance_attr(obj)
         vals = rel_obj_cache.get(instance_attr_val, [])
@@ -1621,8 +1625,23 @@ def prefetch_one_level(instances, prefetcher, lookup, level):
                 setattr(obj, to_attr, vals)
                 obj._prefetched_objects_cache[cache_name] = vals
             else:
-                # Cache in the QuerySet.all().
-                qs = getattr(obj, to_attr).all()
+                manager = getattr(obj, to_attr)
+                if leaf and lookup.queryset is not None:
+                    try:
+                        apply_rel_filter = manager._apply_rel_filters
+                    except AttributeError:
+                        warnings.warn(
+                            "The `%s.%s` class must implement a `_apply_rel_filters()` "
+                            "method that accepts a `QuerySet` as its single "
+                            "argument and returns an appropriately filtered version "
+                            "of it." % (manager.__class__.__module__, manager.__class__.__name__),
+                            RemovedInDjango20Warning,
+                        )
+                        qs = manager.get_queryset()
+                    else:
+                        qs = apply_rel_filter(lookup.queryset)
+                else:
+                    qs = manager.get_queryset()
                 qs._result_cache = vals
                 # We don't want the individual qs doing prefetch_related now,
                 # since we have merged this into the current work.
diff --git a/docs/internals/deprecation.txt b/docs/internals/deprecation.txt
index 2da4b6d6bb..e4a74abdf4 100644
--- a/docs/internals/deprecation.txt
+++ b/docs/internals/deprecation.txt
@@ -144,6 +144,9 @@ details on these changes.
 * The ``__search`` query lookup and the
   ``DatabaseOperations.fulltext_search_sql()`` method will be removed.
 
+* The shim for supporting custom related manager classes without a
+  ``_apply_rel_filters()`` method will be removed.
+
 .. _deprecation-removed-in-1.10:
 
 1.10
diff --git a/docs/releases/1.10.txt b/docs/releases/1.10.txt
index 1f2ddd9457..d935ab8c8f 100644
--- a/docs/releases/1.10.txt
+++ b/docs/releases/1.10.txt
@@ -762,6 +762,17 @@ features, is deprecated. Replace it with a custom lookup::
     models.CharField.register_lookup(Search)
     models.TextField.register_lookup(Search)
 
+Custom manager classes available through ``prefetch_related`` must define a ``_apply_rel_filters()`` method
+-----------------------------------------------------------------------------------------------------------
+
+If you defined a custom manager class available through
+:meth:`~django.db.models.query.QuerySet.prefetch_related` you must make sure
+it defines a ``_apply_rel_filters()`` method.
+
+This method must accept a :class:`~django.db.models.query.QuerySet` instance
+as its single argument and return a filtered version of the queryset for the
+model instance the manager is bound to.
+
 Miscellaneous
 -------------
 
diff --git a/tests/prefetch_related/tests.py b/tests/prefetch_related/tests.py
index d8c1458ce2..3c4a3a735e 100644
--- a/tests/prefetch_related/tests.py
+++ b/tests/prefetch_related/tests.py
@@ -1,5 +1,7 @@
 from __future__ import unicode_literals
 
+import warnings
+
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import connection
@@ -693,6 +695,16 @@ class CustomPrefetchTests(TestCase):
         ).first()
         self.assertIsNone(room.main_room_of_attr)
 
+        # The custom queryset filters should be applied to the queryset
+        # instance returned by the manager.
+        person = Person.objects.prefetch_related(
+            Prefetch('houses', queryset=House.objects.filter(name='House 1')),
+        ).get(pk=self.person1.pk)
+        self.assertEqual(
+            list(person.houses.all()),
+            list(person.houses.all().all()),
+        )
+
     def test_nested_prefetch_related_are_not_overwritten(self):
         # Regression test for #24873
         houses_2 = House.objects.prefetch_related(Prefetch('rooms'))
@@ -704,6 +716,30 @@ class CustomPrefetchTests(TestCase):
             self.room2_1
         )
 
+    def test_apply_rel_filters_deprecation_shim(self):
+        # Simulate a missing `_apply_rel_filters` method.
+        del Person.houses.related_manager_cls._apply_rel_filters
+        # Also remove `get_queryset` as it rely on `_apply_rel_filters`.
+        del Person.houses.related_manager_cls.get_queryset
+        try:
+            with warnings.catch_warnings(record=True) as warns:
+                warnings.simplefilter('always')
+                list(Person.objects.prefetch_related(
+                    Prefetch('houses', queryset=House.objects.filter(name='House 1'))
+                ))
+        finally:
+            # Deleting `related_manager_cls` will force the creation of a new
+            # class since it's a `cached_property`.
+            del Person.houses.related_manager_cls
+        msg = (
+            'The `django.db.models.fields.related_descriptors.ManyRelatedManager` class '
+            'must implement a `_apply_rel_filters()` method that accepts a `QuerySet` as '
+            'its single argument and returns an appropriately filtered version of it.'
+        )
+        self.assertEqual(len(warns), 2)  # Once person.
+        self.assertEqual(str(warns[0].message), msg)
+        self.assertEqual(str(warns[0].message), msg)
+
 
 class DefaultManagerTests(TestCase):
 
@@ -849,6 +885,10 @@ class GenericRelationTests(TestCase):
         with self.assertNumQueries(0):
             self.assertEqual(list(bookmark.tags.all()), [django_tag])
 
+        # The custom queryset filters should be applied to the queryset
+        # instance returned by the manager.
+        self.assertEqual(list(bookmark.tags.all()), list(bookmark.tags.all().all()))
+
 
 class MultiTableInheritanceTest(TestCase):
 
@@ -1296,6 +1336,25 @@ class Ticket25546Tests(TestCase):
             self.assertListEqual(list(book1.first_time_authors.all()[1].addresses.all()), [])
             self.assertListEqual(list(book2.first_time_authors.all()[0].addresses.all()), [self.author2_address1])
 
+        self.assertEqual(
+            list(book1.first_time_authors.all()), list(book1.first_time_authors.all().all())
+        )
+        self.assertEqual(
+            list(book2.first_time_authors.all()), list(book2.first_time_authors.all().all())
+        )
+        self.assertEqual(
+            list(book1.first_time_authors.all()[0].addresses.all()),
+            list(book1.first_time_authors.all()[0].addresses.all().all())
+        )
+        self.assertEqual(
+            list(book1.first_time_authors.all()[1].addresses.all()),
+            list(book1.first_time_authors.all()[1].addresses.all().all())
+        )
+        self.assertEqual(
+            list(book2.first_time_authors.all()[0].addresses.all()),
+            list(book2.first_time_authors.all()[0].addresses.all().all())
+        )
+
     def test_prefetch_with_to_attr(self):
         with self.assertNumQueries(3):
             books = Book.objects.filter(