implement SchemaEditor.add/remove_index() and creating indexes in create_model()

Author: timgraham

django_mongodb/features.py

@@ -84,29 +84,14 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "backends.tests.ThreadTests.test_pass_connection_between_threads",
         "backends.tests.ThreadTests.test_closing_non_shared_connections",
         "backends.tests.ThreadTests.test_default_connection_thread_local",
-        # 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_auto_field_to_integer_field",
         "schema.tests.SchemaTests.test_alter_field_add_index_to_integerfield",
         "schema.tests.SchemaTests.test_alter_field_default_dropped",
-        "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",
         "schema.tests.SchemaTests.test_alter_null_to_not_null",
         "schema.tests.SchemaTests.test_alter_primary_key_the_same_name",
         "schema.tests.SchemaTests.test_autofield_to_o2o",
@@ -124,21 +109,17 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "migrations.test_operations.OperationTests.test_alter_unique_together",
         "schema.tests.SchemaTests.test_unique_together",
         # add/remove_constraint
+        "introspection.tests.IntrospectionTests.test_get_constraints",
         "migrations.test_operations.OperationTests.test_add_partial_unique_constraint",
         "migrations.test_operations.OperationTests.test_create_model_with_partial_unique_constraint",
         "migrations.test_operations.OperationTests.test_remove_partial_unique_constraint",
         "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",
-        # constraints not fully implemented.
-        "introspection.tests.IntrospectionTests.test_get_constraints",
         # pymongo.errors.OperationFailure: Can't rename a collection to itself
         "migrations.test_operations.OperationTests.test_alter_model_table_noop",
         "migrations.test_operations.OperationTests.test_rename_model_no_relations_with_db_table_noop",
         "migrations.test_operations.OperationTests.test_rename_model_with_db_table_rename_m2m",
-        # 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 = {
@@ -241,6 +222,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",
@@ -624,14 +606,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_primary_key",
         },
         "Known issue querying JSONField.": {

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,14 +1,39 @@
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.models import Index
+from pymongo.operations import IndexModel
 
 
 class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     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):
+                index = Index(fields=[field.name])
+                index.set_name_with_model(model)
+                self.add_index(model, index)
+        # Meta.index_together (RemovedInDjango51Warning)
+        for field_names in model._meta.index_together:
+            index = Index(fields=field_names)
+            index.set_name_with_model(model)
+            self.add_index(model, index)
+        # 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:
@@ -60,13 +85,21 @@ def alter_unique_together(self, model, old_unique_together, new_unique_together)
         pass
 
     def add_index(self, model, index):
-        pass
-
-    def rename_index(self, model, old_index, new_index):
-        pass
+        if index.contains_expressions:
+            return
+        idx = IndexModel(
+            [
+                (model._meta.get_field(field_name).column, 1 if order == "" else -1)
+                for field_name, order in index.fields_orders
+            ],
+            name=index.name,
+        )
+        self.connection.database[model._meta.db_table].create_indexes([idx])
 
     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