Factor auto field logic to its own method, detect engine using writable route

Author: macro1

simple_history/models.py

@@ -6,7 +6,7 @@
     from django.apps import apps  # Django >= 1.7
 except ImportError:
     apps = None
-from django.db import models
+from django.db import models, router
 from django.db.models.fields.related import RelatedField
 from django.db.models.related import RelatedObject
 from django.conf import settings
@@ -53,8 +53,6 @@ def python_2_unicode_compatible(klass):
 
 registered_models = {}
 
-nonrel_dbs = ('django_mongodb_engine')
-
 
 class HistoricalRecords(object):
     thread = threading.local()
@@ -267,11 +265,7 @@ def get_field(self, other, cls):
         if isinstance(to_field, models.OneToOneField):
             field = self.get_one_to_one_field(to_field, other)
         elif isinstance(to_field, models.AutoField):
-            # Check if AutoField is string for django-non-rel support
-            if settings.DATABASES['default']['ENGINE'] in nonrel_dbs:
-                field.__class__ = models.TextField
-            else:
-                field.__class__ = models.IntegerField
+            field.__class__ = convert_auto_field(to_field)
         else:
             field.__class__ = to_field.__class__
             excluded_prefixes = ("_", "__")
@@ -323,13 +317,7 @@ def transform_field(field):
     """Customize field appropriately for use in historical model"""
     field.name = field.attname
     if isinstance(field, models.AutoField):
-        # The historical model gets its own AutoField, so any
-        # existing one must be replaced with an IntegerField.
-        if settings.DATABASES['default']['ENGINE'] in nonrel_dbs:
-            # Check if AutoField is string for django-non-rel support
-            field.__class__ = models.TextField
-        else:
-            field.__class__ = models.IntegerField
+        field.__class__ = convert_auto_field(field)
 
     elif isinstance(field, models.FileField):
         # Don't copy file, just path.
@@ -348,6 +336,19 @@ def transform_field(field):
         field.serialize = True
 
 
+def convert_auto_field(field):
+    """Convert AutoField to a non-incrementing type
+
+    The historical model gets its own AutoField, so any existing one
+    must be replaced with an IntegerField.
+    """
+    connection = router.db_for_write(field.model)
+    if settings.DATABASES[connection]['ENGINE'] in ('django_mongodb_engine',):
+        # Check if AutoField is string for django-non-rel support
+        return models.TextField
+    return models.IntegerField
+
+
 class HistoricalObjectDescriptor(object):
     def __init__(self, model):
         self.model = model

simple_history/tests/tests/test_models.py

@@ -7,11 +7,12 @@
     User = get_user_model()
 except ImportError:  # django 1.4 compatibility
     from django.contrib.auth.models import User
+from django.db import models
 from django.db.models.loading import get_model
 from django.test import TestCase
 from django.core.files.base import ContentFile
 
-from simple_history.models import HistoricalRecords
+from simple_history.models import HistoricalRecords, convert_auto_field
 from simple_history import register
 from ..models import (
     AdminProfile, Bookcase, MultiOneToOne, Poll, Choice, Restaurant, Person,
@@ -469,3 +470,32 @@ def test_import_related(self):
     def test_string_related(self):
         field_object = HistoricalState._meta.get_field_by_name('library_id')[0]
         self.assertEqual(field_object.related.model, State)
+
+
+class TestConvertAutoField(TestCase):
+    """Check what AutoFields get converted to."""
+
+    def setUp(self):
+        for field in Poll._meta.fields:
+            if isinstance(field, models.AutoField):
+                self.field = field
+                break
+
+    def test_relational(self):
+        """Relational test
+
+        Default Django ORM uses an integer-based auto field.
+        """
+        with self.settings(DATABASES={'default': {
+                'ENGINE': 'django.db.backends.postgresql_psycopg2'}}):
+            assert convert_auto_field(self.field) == models.IntegerField
+
+    def test_non_relational(self):
+        """Non-relational test
+
+        MongoDB uses a string-based auto field. We need to make sure
+        the converted field type is string.
+        """
+        with self.settings(DATABASES={'default': {
+                'ENGINE': 'django_mongodb_engine'}}):
+            assert convert_auto_field(self.field) == models.TextField