Repository: adamchainz/django-mysql

Issue:
SQL_CALC_FOUND_ROWS
With the work for query hints in place from #114 , it wouldn't be hard to add in [SQL_COUNT_FOUND_ROWS](https://mariadb.com/kb/en/mariadb/select/#sql_calc_found_rows). This could then be used in conjunction with the database function [FOUND_ROWS()](https://mariadb.com/kb/en/mariadb/found_rows/) to override Django's Paginator class to count the rows once - if you check [its source](https://github.com/django/django/blob/master/django/core/paginator.py) it doesn't seem too hard, as only `_get_count` needs overriding to return the count found when taking the slice of `object_list` in `_get_page`.



Discussion:

Resolving patch:
diff --git a/django_mysql/models/query.py b/django_mysql/models/query.py
--- a/django_mysql/models/query.py
+++ b/django_mysql/models/query.py
@@ -31,6 +31,10 @@ def __init__(self, *args, **kwargs):
     def _clone(self, *args, **kwargs):
         clone = super(QuerySetMixin, self)._clone(*args, **kwargs)
         clone._count_tries_approx = copy(self._count_tries_approx)
+        if hasattr(self, '_found_rows'):
+            # If it's a number, don't copy it - the clone has a fresh result
+            # cache
+            clone._found_rows = None
         return clone
 
     # approx_count features
@@ -106,6 +110,32 @@ def sql_cache(self):
     def sql_no_cache(self):
         return self.extra(where=["/*QueryRewrite':SQL_NO_CACHE*/1"])
 
+    def sql_calc_found_rows(self):
+        qs = self.extra(where=["/*QueryRewrite':SQL_CALC_FOUND_ROWS*/1"])
+        qs._found_rows = None
+        return qs
+
+    @property
+    def found_rows(self):
+        if not hasattr(self, '_found_rows'):
+            raise ValueError(
+                "found_rows can only be used if you call sql_calc_found_rows()"
+            )
+        if self._found_rows is None:
+            raise RuntimeError(
+                "A QuerySet with sql_calc_found_rows must be iterated before "
+                "found_rows can be accessed"
+            )
+        return self._found_rows
+
+    def iterator(self):
+        for row in super(QuerySetMixin, self).iterator():
+            yield row
+        if getattr(self, '_found_rows', 0) is None:
+            with connections[self.db].cursor() as cursor:
+                cursor.execute("SELECT FOUND_ROWS()")
+                self._found_rows = cursor.fetchone()[0]
+
     def use_index(self, *index_names, **kwargs):
         kwargs['hint'] = 'USE'
         return self._index_hint(*index_names, **kwargs)


Test patch:
diff --git a/tests/testapp/test_models.py b/tests/testapp/test_models.py
--- a/tests/testapp/test_models.py
+++ b/tests/testapp/test_models.py
@@ -1,5 +1,6 @@
 # -*- coding:utf-8 -*-
 import mock
+import pickle
 import re
 from unittest import skipUnless
 
@@ -323,6 +324,64 @@ def test_force_index_table_name_doesnt_exist_ignored(self):
         assert ' FORCE INDEX ' not in cap.query
 
 
+class FoundRowsTests(TestCase):
+    def setUp(self):
+        super(FoundRowsTests, self).setUp()
+        Author.objects.bulk_create([Author() for i in range(10)])
+
+    def test_found_rows_requires_sql_calc_found_rows(self):
+        authors = Author.objects.all()
+        with pytest.raises(ValueError) as excinfo:
+            authors.found_rows
+        assert "can only be used if" in str(excinfo.value)
+
+    def test_found_rows_requires_iteration(self):
+        authors = Author.objects.sql_calc_found_rows()[:5]
+        with self.assertRaises(RuntimeError):
+            authors.found_rows
+
+    def test_it_working(self):
+        with self.assertNumQueries(2):
+            authors = Author.objects.sql_calc_found_rows()[:5]
+            list(authors)
+            assert authors.found_rows == 10
+
+    def test_no_repeats(self):
+        with self.assertNumQueries(2):
+            authors = (
+                Author.objects.sql_calc_found_rows()
+                              .sql_calc_found_rows()[:5]
+            )
+            list(authors)
+            assert authors.found_rows == 10
+
+    def test_it_working_prefetching(self):
+        with self.assertNumQueries(3):
+            authors = (
+                Author.objects.sql_calc_found_rows()
+                              .prefetch_related('tutor')[:5]
+            )
+            list(authors)
+            assert authors.found_rows == 10
+
+    def test_pickleability(self):
+        authors = Author.objects.sql_calc_found_rows()
+        pickle.dumps(authors)
+
+    def test_cloning_doesnt_copy_found_rows_number(self):
+        authors = Author.objects.sql_calc_found_rows()
+        list(authors)
+        assert authors.found_rows == 10
+
+        authors2 = authors.filter(id__gte=1)
+        with pytest.raises(RuntimeError):
+            authors2.found_rows
+
+        with self.assertNumQueries(2):
+            list(authors2)
+            assert authors2.found_rows == 10
+
+
 class SmartIteratorTests(TestCase):
 
     def setUp(self):