Repository: adamchainz/django-mysql

Issue:
Allow the smart iteration id-range-chunking code to be used with raw SQL
When you want to ID-slice a table but not using a queryset, SmartChunkedIterator doesn't seem to help. It can be done with a dumb:

``` py
for chunk in Model.objects.iter_smart_chunks():
    limits = chunk.aggregate(min_id=Min('id'), max_id=Max('id'))
```

And using `limits` as the dict of first and last, but that's issuing a query to get something that's in python.

We could change it to output the slices quite easily. I'm imagining a call like:

``` py
for id_start, id_end in User.objects.iter_smart_id_slices():
    cursor.execute("SELECT * FROM x WHERE id >= %s AND id < %s", (id_start, id_end))
```



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
@@ -116,6 +116,11 @@ def iter_smart_chunks(self, **kwargs):
             "You can't pass another queryset in through iter_smart_chunks!"
         return SmartChunkedIterator(queryset=self, **kwargs)
 
+    def iter_smart_pk_ranges(self, **kwargs):
+        assert 'queryset' not in kwargs, \
+            "You can't pass another queryset in through iter_smart_pk_ranges!"
+        return SmartPKRangeIterator(queryset=self, **kwargs)
+
     def pt_visual_explain(self, display=True):
         return pt_visual_explain(self, display)
 
@@ -180,6 +185,9 @@ def __iter__(self):
 
             with StopWatch() as timer, self.maybe_atomic(using=db_alias):
                 chunk = self.queryset.filter(pk__gte=start_pk, pk__lt=end_pk)
+                # Attach the start_pk, end_pk onto the chunk queryset so they
+                # can be read by SmartRangeIterator or other client code
+                chunk._smart_iterator_pks = (start_pk, end_pk)
                 yield chunk
                 self.update_progress(chunk=chunk, end_pk=end_pk)
 
@@ -337,6 +345,13 @@ def __iter__(self):
                 yield obj
 
 
+class SmartPKRangeIterator(SmartChunkedIterator):
+    def __iter__(self):
+        for chunk in super(SmartPKRangeIterator, self).__iter__():
+            start_pk, end_pk = chunk._smart_iterator_pks
+            yield start_pk, end_pk
+
+
 def approx_count(queryset):
     # Returns the approximate count or raises a ValueError if this queryset
     # cannot be approximately counted


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
@@ -312,6 +312,30 @@ def test_pk_hole(self):
         seen = [author.id for author in Author.objects.iter_smart()]
         assert seen == [first.id, last.id]
 
+    def test_iter_smart_pk_range(self):
+        seen = []
+        for start_pk, end_pk in Author.objects.iter_smart_pk_ranges():
+            seen.extend(
+                Author.objects.filter(id__gte=start_pk, id__lt=end_pk)
+                              .values_list('id', flat=True)
+            )
+        all_ids = list(Author.objects.order_by('id')
+                                     .values_list('id', flat=True))
+        assert seen == all_ids
+
+    def test_iter_smart_pk_range_with_raw(self):
+        seen = []
+        for start_pk, end_pk in Author.objects.iter_smart_pk_ranges():
+            authors = Author.objects.raw("""
+                SELECT id FROM {}
+                WHERE id >= %s AND id < %s
+            """.format(Author._meta.db_table), (start_pk, end_pk))
+            seen.extend(author.id for author in authors)
+
+        all_ids = list(Author.objects.order_by('id')
+                                     .values_list('id', flat=True))
+        assert seen == all_ids
+
     def test_reporting(self):
         with captured_stdout() as output:
             qs = Author.objects.all()
@@ -370,7 +394,7 @@ def test_reporting_on_uncounted_qs(self):
 
         assert lines[1] == 'Finished!'
 
-    def test_running_on_non_mysql_model(self):
+    def test_filter_and_delete(self):
         VanillaAuthor.objects.create(name="Alpha")
         VanillaAuthor.objects.create(name="pants")
         VanillaAuthor.objects.create(name="Beta")