Repository: adamchainz/django-mysql

Issue:
Smart Iteration should allow you to specify the start chunk size
Add another argument to `SmartChunkedIterator` to override the default of 2.



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
@@ -141,7 +141,7 @@ def __str__(self):
 
 class SmartChunkedIterator(object):
     def __init__(self, queryset, atomically=True, status_thresholds=None,
-                 pk_range=None, chunk_time=0.5, chunk_max=10000,
+                 pk_range=None, chunk_time=0.5, chunk_size=2, chunk_max=10000,
                  report_progress=False, total=None):
         self.queryset = self.sanitize_queryset(queryset)
 
@@ -156,8 +156,8 @@ def __init__(self, queryset, atomically=True, status_thresholds=None,
         self.pk_range = pk_range
 
         self.rate = WeightedAverageRate(chunk_time)
-        self.chunk_size = 2  # Small but will expand rapidly anyhow
         self.chunk_max = chunk_max
+        self.chunk_size = min(chunk_size, chunk_max)
 
         self.report_progress = report_progress
         self.total = total


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
@@ -206,6 +206,11 @@ def setUp(self):
         super(SmartIteratorTests, self).setUp()
         Author.objects.bulk_create([Author() for i in range(10)])
 
+    def tearDown(self):
+        super(SmartIteratorTests, self).tearDown()
+        Author.objects.all().delete()
+        VanillaAuthor.objects.all().delete()
+
     def test_bad_querysets(self):
         with pytest.raises(ValueError) as excinfo:
             Author.objects.all().order_by('name').iter_smart_chunks()
@@ -288,9 +293,45 @@ def fail_second_slice(*args, **kwargs):
             seen = [author.id for author in Author.objects.iter_smart()]
         assert seen == []
 
+    def test_objects_chunk_size(self):
+        smart = iter(Author.objects.iter_smart_chunks(chunk_size=3))
+        chunk = next(smart)
+        assert len(list(chunk)) == 3
+
+    def test_objects_chunk_size_1(self):
+        smart = iter(Author.objects.iter_smart_chunks(chunk_size=1))
+        chunk = next(smart)
+        assert len(list(chunk)) == 1
+
     def test_objects_max_size(self):
-        seen = [author.id for author in
-                Author.objects.iter_smart(chunk_max=1)]
+        seen = []
+        for chunk in Author.objects.iter_smart_chunks(chunk_max=3):
+            ids = [author.id for author in chunk]
+            assert len(ids) <= 3
+            seen.extend(ids)
+        all_ids = list(Author.objects.order_by('id')
+                                     .values_list('id', flat=True))
+        assert seen == all_ids
+
+    def test_objects_max_size_1(self):
+        seen = []
+        for chunk in Author.objects.iter_smart_chunks(chunk_max=1):
+            ids = [author.id for author in chunk]
+            assert len(ids) == 1
+            seen.extend(ids)
+        all_ids = list(Author.objects.order_by('id')
+                                     .values_list('id', flat=True))
+        assert seen == all_ids
+
+    def test_objects_max_size_bounds_chunk_size(self):
+        smart = iter(Author.objects.iter_smart_chunks(chunk_max=5,
+                                                      chunk_size=1000))
+
+        seen = []
+        for chunk in smart:
+            ids = [author.id for author in chunk]
+            assert len(ids) <= 5
+            seen.extend(ids)
         all_ids = list(Author.objects.order_by('id')
                                      .values_list('id', flat=True))
         assert seen == all_ids