Repository: adamchainz/django-mysql

Issue:
iter_smart and finding the min id on large tables
Sometimes when trying to find the first ID to start at on a large table, the query ends up being terribly long and blocking the database still, especially when nothing matches the given queryset. This could be done more intelligently, e.g. using the same slicing strategy to scan up the table. Unfortunately there is no example code in `pt-online-schema-change` to copy since it does not have this problem as it always copies all the rows and has no min/max ids to worry about.



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
@@ -119,8 +119,8 @@ def __str__(self):
 
 class SmartChunkedIterator(object):
     def __init__(self, queryset, atomically=True, status_thresholds=None,
-                 chunk_time=0.5, chunk_max=10000, report_progress=False,
-                 total=None):
+                 pk_range=None, chunk_time=0.5, chunk_max=10000,
+                 report_progress=False, total=None):
         self.queryset = self.sanitize_queryset(queryset)
 
         if atomically:
@@ -131,6 +131,7 @@ def __init__(self, queryset, atomically=True, status_thresholds=None,
             self.maybe_atomic = noop_context
 
         self.status_thresholds = status_thresholds
+        self.pk_range = pk_range
 
         self.rate = WeightedAverageRate(chunk_time)
         self.chunk_size = 2  # Small but will expand rapidly anyhow
@@ -187,8 +188,18 @@ def sanitize_queryset(self, queryset):
         return queryset.order_by('pk')
 
     def get_min_and_max(self):
-        min_qs = self.queryset.order_by('pk').values_list('pk', flat=True)
-        max_qs = self.queryset.order_by('-pk').values_list('pk', flat=True)
+        if isinstance(self.pk_range, tuple) and len(self.pk_range) == 2:
+            return self.pk_range
+        elif self.pk_range == 'all':
+            base_qs = self.queryset.model.objects.using(self.queryset.db).all()
+        elif self.pk_range is None:
+            base_qs = self.queryset
+        else:
+            raise ValueError("Unrecognized value for pk_range: {}"
+                             .format(self.pk_range))
+
+        min_qs = base_qs.order_by('pk').values_list('pk', flat=True)
+        max_qs = base_qs.order_by('-pk').values_list('pk', flat=True)
         try:
             min_pk = min_qs[0]
         except IndexError:


Test patch:
diff --git a/tests/django_mysql_tests/test_models.py b/tests/django_mysql_tests/test_models.py
--- a/tests/django_mysql_tests/test_models.py
+++ b/tests/django_mysql_tests/test_models.py
@@ -109,6 +109,33 @@ def test_objects_non_atomic(self):
                                      .values_list('id', flat=True))
         self.assertEqual(seen, all_ids)
 
+    def test_objects_pk_range_all(self):
+        seen = [author.id for author in
+                Author.objects.iter_smart(pk_range='all')]
+        all_ids = list(Author.objects.order_by('id')
+                                     .values_list('id', flat=True))
+        self.assertEqual(seen, all_ids)
+
+    def test_objects_pk_range_tuple(self):
+        seen = [author.id for author in
+                Author.objects.iter_smart(pk_range=(0, 0))]
+        self.assertEqual(seen, [])
+
+        min_id = Author.objects.earliest('id').id
+        max_id = Author.objects.order_by('id')[5].id
+
+        seen = [author.id for author in
+                Author.objects.iter_smart(pk_range=(min_id, max_id))]
+        cut_ids = list(Author.objects.order_by('id')
+                                     .filter(id__gte=min_id, id__lte=max_id)
+                                     .values_list('id', flat=True))
+        self.assertEqual(seen, cut_ids)
+
+    def test_objects_pk_range_bad(self):
+        with self.assertRaises(ValueError) as cm:
+            list(Author.objects.iter_smart(pk_range="My Bad Value"))
+        self.assertIn("Unrecognized value for pk_range", str(cm.exception))
+
     def test_objects_max_size(self):
         seen = [author.id for author in
                 Author.objects.iter_smart(chunk_max=1)]
@@ -144,6 +171,24 @@ def test_reporting(self):
 
         self.assertEqual(lines[1], 'Finished!')
 
+    def test_reporting_with_total(self):
+        with captured_stdout() as output:
+            qs = Author.objects.all()
+            for authors in qs.iter_smart_chunks(report_progress=True, total=4):
+                list(authors)  # fetch them
+
+        lines = output.getvalue().split('\n')
+
+        reports = lines[0].split('\r')
+        for report in reports:
+            self.assertRegexpMatches(
+                report,
+                r"AuthorSmartChunkedIterator processed \d+/4 objects "
+                r"\(\d+\.\d+%\) in \d+ chunks(; highest pk so far \d+)?"
+            )
+
+        self.assertEqual(lines[1], 'Finished!')
+
     def test_reporting_on_uncounted_qs(self):
         Author.objects.create(name="pants")