add index creation/deletion support to SchemaEditor.alter_field()

mongodb-labs · Sep 18, 2024 · 8520d2c · 8520d2c
1 parent 0c9e3d1
commit 8520d2c
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 4 deletions.
diff --git a/django_mongodb/features.py b/django_mongodb/features.py
@@ -86,7 +86,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "backends.tests.ThreadTests.test_default_connection_thread_local",
         # 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_to_o2o",
         "schema.tests.SchemaTests.test_alter_field_o2o_keeps_unique",
@@ -95,11 +94,9 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "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",
-        # AlterField (db_index)
-        "schema.tests.SchemaTests.test_indexes",
         "schema.tests.SchemaTests.test_remove_constraints_capital_letters",
-        "schema.tests.SchemaTests.test_remove_db_index_doesnt_remove_custom_indexes",
         # AlterField (unique)
+        "schema.tests.SchemaTests.test_indexes",
         "schema.tests.SchemaTests.test_unique",
         "schema.tests.SchemaTests.test_unique_and_reverse_m2m",
         # alter_unique_together

diff --git a/django_mongodb/schema.py b/django_mongodb/schema.py
@@ -68,11 +68,59 @@ def _alter_field(
         new_db_params,
         strict=False,
     ):
+        # Removed an index? (no strict check, as multiple indexes are possible)
+        # Remove indexes if db_index switched to False or a unique constraint
+        # will now be used in lieu of an index. The following lines from the
+        # truth table show all True cases; the rest are False:
+        #
+        # old_field.db_index | old_field.unique | new_field.db_index | new_field.unique
+        # ------------------------------------------------------------------------------
+        # True               | False            | False              | False
+        # True               | False            | False              | True
+        # True               | False            | True               | True
+        if (
+            old_field.db_index
+            and not old_field.unique
+            and (not new_field.db_index or new_field.unique)
+        ):
+            # Find the index for this field
+            meta_index_names = {index.name for index in model._meta.indexes}
+            # Retrieve only BTREE indexes since this is what's created with
+            # db_index=True.
+            index_names = self._constraint_names(
+                model,
+                [old_field.column],
+                index=True,
+                type_=Index.suffix,
+                exclude=meta_index_names,
+            )
+            for index_name in index_names:
+                # The only way to check if an index was created with
+                # db_index=True or with Index(['field'], name='foo')
+                # is to look at its name (refs #28053).
+                self.connection.database[model._meta.db_table].drop_index(index_name)
         # Have they renamed the column?
         if old_field.column != new_field.column:
             self.connection.database[model._meta.db_table].update_many(
                 {}, {"$rename": {old_field.column: new_field.column}}
             )
+        # Added an index? Add an index if db_index switched to True or a unique
+        # constraint will no longer be used in lieu of an index. The following
+        # lines from the truth table show all True cases; the rest are False:
+        #
+        # old_field.db_index | old_field.unique | new_field.db_index | new_field.unique
+        # ------------------------------------------------------------------------------
+        # False              | False            | True               | False
+        # False              | True             | True               | False
+        # True               | True             | True               | False
+        if (
+            (not old_field.db_index or old_field.unique)
+            and new_field.db_index
+            and not new_field.unique
+        ):
+            index = Index(fields=[new_field.name])
+            index.name = self._create_index_name(model._meta.db_table, [new_field.column])
+            self.add_index(model, index)
 
     def remove_field(self, model, field):
         # Remove implicit M2M tables.