create spatial indexes on embedded GIS fields

Author: timgraham

django_mongodb_backend/gis/schema.py

@@ -9,8 +9,8 @@ def _field_should_be_indexed(self, model, field):
         return super()._field_should_be_indexed(model, field)
 
     def _add_field_index(self, model, field, *, column_prefix=""):
-        if hasattr(field, "geodetic"):
-            self._add_spatial_index(model, field)
+        if getattr(field, "spatial_index", False):
+            self._add_spatial_index(model, field, column_prefix)
         else:
             super()._add_field_index(model, field, column_prefix=column_prefix)
 
@@ -42,15 +42,15 @@ def _alter_field(
         elif old_field_spatial_index and not new_field_spatial_index:
             self._delete_spatial_index(model, new_field)
 
-    def _add_spatial_index(self, model, field):
-        index_name = self._create_spatial_index_name(model, field)
+    def _add_spatial_index(self, model, field, column_prefix=""):
+        index_name = self._create_spatial_index_name(model, field, column_prefix)
         self.get_collection(model._meta.db_table).create_indexes(
-            [IndexModel([(field.column, GEOSPHERE)], name=index_name)]
+            [IndexModel([(column_prefix + field.column, GEOSPHERE)], name=index_name)]
         )
 
     def _delete_spatial_index(self, model, field):
         index_name = self._create_spatial_index_name(model, field)
         self.get_collection(model._meta.db_table).drop_index(index_name)
 
-    def _create_spatial_index_name(self, model, field):
-        return f"{model._meta.db_table}_{field.column}_id"
+    def _create_spatial_index_name(self, model, field, column_prefix=""):
+        return f"{model._meta.db_table}_{column_prefix}{field.column}_id"

django_mongodb_backend/schema.py

@@ -28,7 +28,7 @@ def wrapper(self, model, *args, **kwargs):
     return wrapper
 
 
-class DatabaseSchemaEditor(GISSchemaEditor, BaseDatabaseSchemaEditor):
+class BaseSchemaEditor(BaseDatabaseSchemaEditor):
     def get_collection(self, name):
         if self.collect_sql:
             return OperationCollector(self.collected_sql, collection=self.connection.database[name])
@@ -419,3 +419,8 @@ def _field_should_have_unique(self, field):
         db_type = field.db_type(self.connection)
         # The _id column is automatically unique.
         return db_type and field.unique and field.column != "_id"
+
+
+# GISSchemaEditor extends some SchemaEditor methods.
+class DatabaseSchemaEditor(GISSchemaEditor, BaseSchemaEditor):
+    pass

tests/schema_/test_embedded_model.py

@@ -1,7 +1,7 @@
 import itertools
 
 from django.db import connection, models
-from django.test import TransactionTestCase
+from django.test import TransactionTestCase, skipUnlessDBFeature
 from django.test.utils import isolate_apps
 
 from django_mongodb_backend.fields import EmbeddedModelField
@@ -552,3 +552,41 @@ def test_embedded_alter_field_ignored(self):
         new_field.set_attributes_from_name("new")
         with connection.schema_editor() as editor, self.assertNumQueries(0):
             editor.alter_field(Author, old_field, new_field)
+
+
+@skipUnlessDBFeature("gis_enabled")
+class GISTests(TestMixin, TransactionTestCase):
+    @isolate_apps("schema_")
+    def test_create_model(self):
+        """
+        Spatial indexes for embedded GIS fields are created when the collections are
+        created.
+        """
+        from django.contrib.gis.db.models import PointField
+
+        class Place(EmbeddedModel):
+            name = models.CharField(max_length=10)
+            location = PointField()
+
+            class Meta:
+                app_label = "schema_"
+
+        class Author(models.Model):
+            birthplace = EmbeddedModelField(Place)
+
+            class Meta:
+                app_label = "schema_"
+
+        with connection.schema_editor() as editor:
+            # Create the table
+            editor.create_model(Author)
+            self.assertTableExists(Author)
+            # The embedded GEO indexes is created.
+            constraint_name = "schema__author_birthplace.location_id"
+            self.assertEqual(
+                self.get_constraints_for_columns(Author, ["birthplace.location"]),
+                [constraint_name],
+            )
+            self.assertIndexOrder(Author._meta.db_table, constraint_name, ["GEO"])
+            editor.delete_model(Author)
+        self.assertTableNotExists(Author)