Browse Source

Fixed #28817 -- Made QuerySet.iterator() use server-side cursors after values() and values_list().

Dražen Odobašić 7 years ago
parent
commit
d97f026a7a

+ 7 - 4
django/db/models/query.py

@@ -105,7 +105,7 @@ class ValuesIterable(BaseIterable):
         names = extra_names + field_names + annotation_names
 
         indexes = range(len(names))
-        for row in compiler.results_iter():
+        for row in compiler.results_iter(chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size):
             yield {names[i]: row[i] for i in indexes}
 
 
@@ -133,8 +133,11 @@ class ValuesListIterable(BaseIterable):
                 # Reorder according to fields.
                 index_map = {name: idx for idx, name in enumerate(names)}
                 rowfactory = operator.itemgetter(*[index_map[f] for f in fields])
-                return map(rowfactory, compiler.results_iter())
-        return compiler.results_iter(tuple_expected=True)
+                return map(
+                    rowfactory,
+                    compiler.results_iter(chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size)
+                )
+        return compiler.results_iter(tuple_expected=True, chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size)
 
 
 class NamedValuesListIterable(ValuesListIterable):
@@ -174,7 +177,7 @@ class FlatValuesListIterable(BaseIterable):
     def __iter__(self):
         queryset = self.queryset
         compiler = queryset.query.get_compiler(queryset.db)
-        return chain.from_iterable(compiler.results_iter())
+        return chain.from_iterable(compiler.results_iter(chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size))
 
 
 class QuerySet:

+ 3 - 2
django/db/models/sql/compiler.py

@@ -1002,10 +1002,11 @@ class SQLCompiler:
                 row[pos] = value
             yield row
 
-    def results_iter(self, results=None, tuple_expected=False):
+    def results_iter(self, results=None, tuple_expected=False, chunked_fetch=False,
+                     chunk_size=GET_ITERATOR_CHUNK_SIZE):
         """Return an iterator over the results from executing this query."""
         if results is None:
-            results = self.execute_sql(MULTI)
+            results = self.execute_sql(MULTI, chunked_fetch=chunked_fetch, chunk_size=chunk_size)
         fields = [s[0] for s in self.select[0:self.col_count]]
         converters = self.get_converters(fields)
         rows = chain.from_iterable(results)

+ 3 - 0
docs/releases/1.11.8.txt

@@ -18,3 +18,6 @@ Bugfixes
 
 * Fixed incorrect index name truncation when using a namespaced ``db_table``
   (:ticket:`28792`).
+
+* Made ``QuerySet.iterator()`` use server-side cursors on PostgreSQL after
+  ``values()`` and ``values_list()`` (:ticket:`28817`).

+ 14 - 1
tests/backends/postgresql/test_server_side_cursors.py

@@ -3,7 +3,7 @@ import unittest
 from collections import namedtuple
 from contextlib import contextmanager
 
-from django.db import connection
+from django.db import connection, models
 from django.test import TestCase
 
 from ..models import Person
@@ -53,6 +53,19 @@ class ServerSideCursorsPostgres(TestCase):
     def test_server_side_cursor(self):
         self.assertUsesCursor(Person.objects.iterator())
 
+    def test_values(self):
+        self.assertUsesCursor(Person.objects.values('first_name').iterator())
+
+    def test_values_list(self):
+        self.assertUsesCursor(Person.objects.values_list('first_name').iterator())
+
+    def test_values_list_flat(self):
+        self.assertUsesCursor(Person.objects.values_list('first_name', flat=True).iterator())
+
+    def test_values_list_fields_not_equal_to_names(self):
+        expr = models.Count('id')
+        self.assertUsesCursor(Person.objects.annotate(id__count=expr).values_list(expr, 'id__count').iterator())
+
     def test_server_side_cursor_many_cursors(self):
         persons = Person.objects.iterator()
         persons2 = Person.objects.iterator()