Repository: adamchainz/django-mysql

Issue:
Smart Iteration can't iterate a model whose primary key is an integer ForeignKey
You get the error:

```
... in sanitize_queryset(self, queryset)
    205             raise ValueError(
    206                 "You can't use %s on a model with a non-integer primary key." %
--> 207                 self.__class__.__name__
    208             )
    209 

ValueError: You can't use SmartChunkedIterator on a model with a non-integer primary key.
```

`sanitize_queryset` needs extending.



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
@@ -209,7 +209,9 @@ def sanitize_queryset(self, queryset):
             )
 
         pk = queryset.model._meta.pk
-        if not isinstance(pk, (models.IntegerField, models.AutoField)):
+        if not isinstance(pk, self.ALLOWED_PK_FIELD_CLASSES):
+            # If your custom field class should be allowed, just add it to
+            # ALLOWED_PK_FIELD_CLASSES
             raise ValueError(
                 "You can't use %s on a model with a non-integer primary key." %
                 self.__class__.__name__
@@ -217,6 +219,12 @@ def sanitize_queryset(self, queryset):
 
         return queryset.order_by('pk')
 
+    ALLOWED_PK_FIELD_CLASSES = (
+        models.IntegerField,  # Also covers e.g. PositiveIntegerField
+        models.AutoField,  # Is an integer field but doesn't subclass it :(
+        models.ForeignKey  # Should always point to an integer
+    )
+
     def get_min_and_max(self):
         if isinstance(self.pk_range, tuple) and len(self.pk_range) == 2:
             return self.pk_range


Test patch:
diff --git a/tests/django_mysql_tests/models.py b/tests/django_mysql_tests/models.py
--- a/tests/django_mysql_tests/models.py
+++ b/tests/django_mysql_tests/models.py
@@ -94,6 +94,11 @@ def __unicode__(self):
         return "{} {}".format(self.id, self.name)
 
 
+class AuthorExtra(Model):
+    author = ForeignKey(Author, primary_key=True)
+    legs = IntegerField(default=2)
+
+
 class NameAuthor(Model):
     name = CharField(max_length=32, primary_key=True)
 
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
@@ -11,7 +11,9 @@
 
 from django_mysql.models import ApproximateInt, SmartIterator
 from django_mysql.utils import have_program
-from django_mysql_tests.models import Author, Book, NameAuthor, VanillaAuthor
+from django_mysql_tests.models import (
+    Author, AuthorExtra, Book, NameAuthor, VanillaAuthor
+)
 from django_mysql_tests.utils import CaptureLastQuery, captured_stdout
 
 
@@ -377,6 +379,16 @@ def test_iter_smart_pk_range_with_raw(self):
                                      .values_list('id', flat=True))
         assert seen == all_ids
 
+    def test_iter_smart_fk_primary_key(self):
+        author, author2 = Author.objects.all()[:2]
+        AuthorExtra.objects.create(author=author, legs=2)
+        AuthorExtra.objects.create(author=author2, legs=1)
+
+        seen_author_ids = []
+        for extra in AuthorExtra.objects.iter_smart():
+            seen_author_ids.append(extra.author_id)
+        assert seen_author_ids == [author.id, author2.id]
+
     def test_reporting(self):
         with captured_stdout() as output:
             qs = Author.objects.all()