Add support for models with auto-generated integer PKs on Firestore

adamalton · Jan 8, 2024 · 2f5d764 · 2f5d764
1 parent 3f08bd9
commit 2f5d764
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/massmigration/backends/djangae.py b/massmigration/backends/djangae.py
@@ -5,6 +5,7 @@
 from djangae.processing import (
     datastore_key_ranges,
     firestore_name_key_ranges,
+    firestore_scattered_int_key_ranges,
     sequential_int_key_ranges,
     uuid_key_ranges,
 )
@@ -80,6 +81,8 @@ def _key_ranges_getter(self, queryset):
         elif engine == "gcloudc.db.backends.firestore":
             if isinstance(pk_field, AutoCharField):
                 return firestore_name_key_ranges
+            elif isinstance(pk_field, (models.AutoField, models.BigAutoField, models.SmallAutoField)):
+                return firestore_scattered_int_key_ranges
         else:  # SQL
             if isinstance(pk_field, models.IntegerField):
                 return sequential_int_key_ranges