Code

Ticket #17668: 17668.diff

File 17668.diff, 3.6 KB (added by akaariai, 2 years ago)
Line 
1diff --git a/django/db/models/query.py b/django/db/models/query.py
2index 41c24c7..6ae0a52 100644
3--- a/django/db/models/query.py
4+++ b/django/db/models/query.py
5@@ -82,7 +82,7 @@ class QuerySet(object):
6             if self._iter:
7                 self._result_cache = list(self._iter)
8             else:
9-                self._result_cache = list(self.iterator())
10+                self._result_cache = list(self.iterator(_check_prefetch=False))
11         elif self._iter:
12             self._result_cache.extend(self._iter)
13         if self._prefetch_related_lookups and not self._prefetch_done:
14@@ -228,11 +228,18 @@ class QuerySet(object):
15     # METHODS THAT DO DATABASE QUERIES #
16     ####################################
17 
18-    def iterator(self):
19+    def iterator(self, _check_prefetch=True):
20         """
21         An iterator over the results from applying this QuerySet to the
22         database.
23+
24+        Calling this method with _check_prefetch=True will cause an error
25+        if there are prefetches defined for the qs. Internal methods can
26+        skip prefetch_related checks by defining _check_prefetch=False.
27         """
28+        if _check_prefetch and self._prefetch_related_lookups:
29+            raise ValueError("Using iterator() when prefetch_related is used "
30+                             "is not supported")
31         fill_cache = False
32         if connections[self.db].features.supports_select_related:
33             fill_cache = self.query.select_related
34@@ -485,7 +492,10 @@ class QuerySet(object):
35         qs = self._clone()
36         qs.query.add_filter(('pk__in', id_list))
37         qs.query.clear_ordering(force_empty=True)
38-        return dict([(obj._get_pk_val(), obj) for obj in qs.iterator()])
39+        # Populate the query, making sure prefetch_related etc. are handled
40+        # properly.
41+        objs = list(self)
42+        return dict([(obj._get_pk_val(), obj) for obj in objs])
43 
44     def delete(self):
45         """
46diff --git a/docs/ref/models/querysets.txt b/docs/ref/models/querysets.txt
47index 103cae1..725f38d 100644
48--- a/docs/ref/models/querysets.txt
49+++ b/docs/ref/models/querysets.txt
50@@ -1430,6 +1430,9 @@ performance and a significant reduction in memory.
51 Note that using ``iterator()`` on a ``QuerySet`` which has already been
52 evaluated will force it to evaluate again, repeating the query.
53 
54+It is an error to use ``iterator()`` when ``prefetch_related`` is also used
55+in the same ``QuerySet``.
56+
57 latest
58 ~~~~~~
59 
60diff --git a/tests/modeltests/prefetch_related/tests.py b/tests/modeltests/prefetch_related/tests.py
61index 4c51a83..f457fc9 100644
62--- a/tests/modeltests/prefetch_related/tests.py
63+++ b/tests/modeltests/prefetch_related/tests.py
64@@ -470,3 +470,26 @@ class NullableTest(TestCase):
65                         for e in qs2]
66 
67         self.assertEqual(co_serfs, co_serfs2)
68+
69+class TestIterator(TestCase):
70+
71+    def test_iterator_raises_error(self):
72+        """
73+        Using an iterator when prefetch_related is defined is an error.
74+        """
75+        def raises():
76+            list(Employee.objects.prefetch_related('boss').iterator())
77+        self.assertRaises(ValueError, raises)
78+
79+    def test_in_bulk(self):
80+        """
81+        In-bulk does correctly prefetch objects by not using .iterator()
82+        directly.
83+        """
84+        boss = Employee.objects.create(name="Peter")
85+        boss = Employee.objects.create(name="Jack")
86+        with self.assertNumQueries(2):
87+            # Check that prefetch is done and it does not cause any errors.
88+            bulk = Employee.objects.prefetch_related('serfs').in_bulk([boss.pk])
89+            for b in bulk.values():
90+                list(b.serfs.all())