implement SchemaEditor.add/remove_index()

And also creating indexes in create_model() and add_field().

Author: timgraham

django_mongodb/features.py

@@ -77,29 +77,13 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "backends.tests.ThreadTests.test_closing_non_shared_connections",
         "backends.tests.ThreadTests.test_default_connection_thread_local",
         # AddField
-        "schema.tests.SchemaTests.test_add_indexed_charfield",
         "schema.tests.SchemaTests.test_add_unique_charfield",
-        # Add/RemoveIndex
-        "migrations.test_operations.OperationTests.test_add_index",
-        "migrations.test_operations.OperationTests.test_alter_field_with_index",
-        "migrations.test_operations.OperationTests.test_remove_index",
-        "migrations.test_operations.OperationTests.test_rename_index",
-        "migrations.test_operations.OperationTests.test_rename_index_unknown_unnamed_index",
-        "migrations.test_operations.OperationTests.test_rename_index_unnamed_index",
-        "schema.tests.SchemaTests.test_add_remove_index",
-        "schema.tests.SchemaTests.test_composed_desc_index_with_fk",
-        "schema.tests.SchemaTests.test_composed_index_with_fk",
-        "schema.tests.SchemaTests.test_create_index_together",
-        "schema.tests.SchemaTests.test_order_index",
-        "schema.tests.SchemaTests.test_text_field_with_db_index",
         # AlterField
         "schema.tests.SchemaTests.test_alter_field_add_index_to_integerfield",
-        "schema.tests.SchemaTests.test_alter_field_fk_keeps_index",
         "schema.tests.SchemaTests.test_alter_field_fk_to_o2o",
         "schema.tests.SchemaTests.test_alter_field_o2o_keeps_unique",
         "schema.tests.SchemaTests.test_alter_field_o2o_to_fk",
         "schema.tests.SchemaTests.test_alter_int_pk_to_int_unique",
-        "schema.tests.SchemaTests.test_alter_not_unique_field_to_primary_key",
         # AlterField (db_index)
         "schema.tests.SchemaTests.test_alter_renames_index",
         "schema.tests.SchemaTests.test_indexes",
@@ -121,9 +105,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "schema.tests.SchemaTests.test_composed_constraint_with_fk",
         "schema.tests.SchemaTests.test_remove_ignored_unique_constraint_not_create_fk_index",
         "schema.tests.SchemaTests.test_unique_constraint",
-        # subclasses of BaseDatabaseIntrospection may require a get_constraints() method
-        "migrations.test_operations.OperationTests.test_add_func_unique_constraint",
-        "migrations.test_operations.OperationTests.test_remove_func_unique_constraint",
     }
     # $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
     _django_test_expected_failures_bitwise = {
@@ -227,6 +208,7 @@ def django_test_expected_failures(self):
             "get_or_create.tests.GetOrCreateThroughManyToMany.test_something",
             "get_or_create.tests.UpdateOrCreateTests.test_manual_primary_key_test",
             "get_or_create.tests.UpdateOrCreateTestsWithManualPKs.test_create_with_duplicate_primary_key",
+            "introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
             "model_fields.test_filefield.FileFieldTests.test_unique_when_same_filename",
             "one_to_one.tests.OneToOneTests.test_multiple_o2o",
             "queries.test_bulk_update.BulkUpdateTests.test_database_routing_batch_atomicity",
@@ -618,14 +600,8 @@ def django_test_expected_failures(self):
             "introspection.tests.IntrospectionTests.test_get_table_description_types",
             "introspection.tests.IntrospectionTests.test_smallautofield",
         },
-        "DatabaseIntrospection.get_constraints() not implemented.": {
-            "introspection.tests.IntrospectionTests.test_get_constraints",
-            "introspection.tests.IntrospectionTests.test_get_constraints_index_types",
-            "introspection.tests.IntrospectionTests.test_get_constraints_indexes_orders",
-            "introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders",
-            "introspection.tests.IntrospectionTests.test_get_primary_key_column",
-        },
         "MongoDB can't introspect primary key.": {
+            "introspection.tests.IntrospectionTests.test_get_primary_key_column",
             "schema.tests.SchemaTests.test_alter_primary_key_the_same_name",
             "schema.tests.SchemaTests.test_primary_key",
         },

django_mongodb/introspection.py

@@ -1,6 +1,31 @@
 from django.db.backends.base.introspection import BaseDatabaseIntrospection
+from django.db.models import Index
 
 
 class DatabaseIntrospection(BaseDatabaseIntrospection):
+    ORDER_DIR = {1: "ASC", -1: "DESC"}
+
     def table_names(self, cursor=None, include_views=False):
         return sorted([x["name"] for x in self.connection.database.list_collections()])
+
+    def get_constraints(self, cursor, table_name):
+        indexes = self.connection.database[table_name].index_information()
+        constraints = {}
+        for name, details in indexes.items():
+            # Remove underscore prefix from "_id" columns in primary key index.
+            if is_primary_key := name == "_id_":
+                name = "id"
+                details["key"] = [("id", 1)]
+            constraints[name] = {
+                "check": False,
+                "columns": [field for field, order in details["key"]],
+                "definition": None,
+                "foreign_key": None,
+                "index": True,
+                "orders": [self.ORDER_DIR[order] for field, order in details["key"]],
+                "primary_key": is_primary_key,
+                "type": Index.suffix,
+                "unique": details.get("unique", False),
+                "options": {},
+            }
+        return constraints

django_mongodb/schema.py

@@ -1,4 +1,6 @@
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.models import Index
+from pymongo.operations import IndexModel
 
 from .query import wrap_database_errors
 
@@ -7,11 +9,30 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     @wrap_database_errors
     def create_model(self, model):
         self.connection.database.create_collection(model._meta.db_table)
+        self._create_model_indexes(model)
         # Make implicit M2M tables.
         for field in model._meta.local_many_to_many:
             if field.remote_field.through._meta.auto_created:
                 self.create_model(field.remote_field.through)
 
+    def _create_model_indexes(self, model):
+        """
+        Create all indexes (field indexes, index_together, Meta.indexes) for
+        the specified model.
+        """
+        if not model._meta.managed or model._meta.proxy or model._meta.swapped:
+            return
+        # Field indexes
+        for field in model._meta.local_fields:
+            if self._field_should_be_indexed(model, field):
+                self._add_field_index(model, field)
+        # Meta.index_together (RemovedInDjango51Warning)
+        for field_names in model._meta.index_together:
+            self._add_composed_index(model, field_names)
+        # Meta.indexes
+        for index in model._meta.indexes:
+            self.add_index(model, index)
+
     def delete_model(self, model):
         # Delete implicit M2m tables.
         for field in model._meta.local_many_to_many:
@@ -29,6 +50,9 @@ def add_field(self, model, field):
             self.connection.database[model._meta.db_table].update_many(
                 {}, [{"$set": {column: self.effective_default(field)}}]
             )
+        # Add an index, if required.
+        if self._field_should_be_indexed(model, field):
+            self._add_field_index(model, field)
 
     def _alter_field(
         self,
@@ -67,14 +91,39 @@ def alter_index_together(self, model, old_index_together, new_index_together):
     def alter_unique_together(self, model, old_unique_together, new_unique_together):
         pass
 
-    def add_index(self, model, index):
-        pass
+    def add_index(self, model, index, field=None):
+        if index.contains_expressions:
+            return
+        index_orders = (
+            [(field.column, 1)]
+            if field
+            else [
+                (model._meta.get_field(field_name).column, 1 if order == "" else -1)
+                for field_name, order in index.fields_orders
+            ]
+        )
+        idx = IndexModel(
+            index_orders,
+            name=index.name,
+        )
+        self.connection.database[model._meta.db_table].create_indexes([idx])
 
-    def rename_index(self, model, old_index, new_index):
-        pass
+    def _add_composed_index(self, model, field_names):
+        """Add an index on the given list of field_names."""
+        idx = Index(fields=field_names)
+        idx.set_name_with_model(model)
+        self.add_index(model, idx)
+
+    def _add_field_index(self, model, field):
+        """Add an index for a field that has db_index=True."""
+        index = Index(fields=[field.name])
+        index.name = self._create_index_name(model._meta.db_table, [field.column])
+        self.add_index(model, index, field=field)
 
     def remove_index(self, model, index):
-        pass
+        if index.contains_expressions:
+            return
+        self.connection.database[model._meta.db_table].drop_index(index.name)
 
     def add_constraint(self, model, constraint):
         pass