INTPYTHON-730 Make search index creation and deletion synchronous

Author: timgraham

django_mongodb_backend/schema.py

@@ -1,3 +1,5 @@
+from time import monotonic, sleep
+
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
 from pymongo.operations import SearchIndexModel
@@ -269,10 +271,12 @@ def add_index(
         )
         if idx:
             model = parent_model or model
+            collection = self.get_collection(model._meta.db_table)
             if isinstance(idx, SearchIndexModel):
-                self.get_collection(model._meta.db_table).create_search_index(idx)
+                collection.create_search_index(idx)
+                self.wait_until_index_created(collection, index.name)
             else:
-                self.get_collection(model._meta.db_table).create_indexes([idx])
+                collection.create_indexes([idx])
 
     def _add_composed_index(self, model, field_names, column_prefix="", parent_model=None):
         """Add an index on the given list of field_names."""
@@ -290,12 +294,14 @@ def _add_field_index(self, model, field, *, column_prefix=""):
     def remove_index(self, model, index):
         if index.contains_expressions:
             return
+        collection = self.get_collection(model._meta.db_table)
         if isinstance(index, SearchIndex):
             # Drop the index if it's supported.
             if self.connection.features.supports_atlas_search:
-                self.get_collection(model._meta.db_table).drop_search_index(index.name)
+                collection.drop_search_index(index.name)
+                self.wait_until_index_dropped(collection, index.name)
         else:
-            self.get_collection(model._meta.db_table).drop_index(index.name)
+            collection.drop_index(index.name)
 
     def _remove_composed_index(
         self, model, field_names, constraint_kwargs, column_prefix="", parent_model=None
@@ -420,6 +426,32 @@ def _field_should_have_unique(self, field):
         # The _id column is automatically unique.
         return db_type and field.unique and field.column != "_id"
 
+    @staticmethod
+    def wait_until_index_created(collection, index_name, timeout=60 * 60, interval=0.5):
+        """
+        Wait up to an hour until an index is created. Index creation time
+        depends on the size of the collection being indexed.
+        """
+        start = monotonic()
+        while monotonic() - start < timeout:
+            indexes = list(collection.list_search_indexes())
+            for idx in indexes:
+                if idx["name"] == index_name and idx["status"] == "READY":
+                    return True
+            sleep(interval)
+        raise TimeoutError(f"Index {index_name} not ready after {timeout} seconds.")
+
+    @staticmethod
+    def wait_until_index_dropped(collection, index_name, timeout=60, interval=0.5):
+        """Wait up to 60 seconds until an index is dropped."""
+        start = monotonic()
+        while monotonic() - start < timeout:
+            indexes = list(collection.list_search_indexes())
+            if all(idx["name"] != index_name for idx in indexes):
+                return True
+            sleep(interval)
+        raise TimeoutError(f"Index {index_name} not dropped after {timeout} seconds.")
+
 
 # GISSchemaEditor extends some SchemaEditor methods.
 class DatabaseSchemaEditor(GISSchemaEditor, BaseSchemaEditor):

docs/ref/models/indexes.rst

@@ -9,8 +9,22 @@ Some MongoDB-specific :doc:`indexes <django:ref/models/indexes>`, for use on a
 model's :attr:`Meta.indexes <django.db.models.Options.indexes>` option, are
 available in ``django_mongodb_backend.indexes``.
 
+Search indexes
+==============
+
+MongoDB creates these indexes asynchronously, however, Django's
+:class:`~django.db.migrations.operations.AddIndex` and
+:class:`~django.db.migrations.operations.RemoveIndex` operations will wait
+until the index is created or deleted so that the database state is
+consistent in the operations that follow. Adding indexes may take seconds or
+minutes, depending on the size of the collection.
+
+.. versionchanged:: 5.2.1
+
+    The aforementioned waiting was added.
+
 ``SearchIndex``
-===============
+---------------
 
 .. class:: SearchIndex(fields=(), name=None)
 
@@ -30,7 +44,7 @@ available in ``django_mongodb_backend.indexes``.
     your model has multiple indexes).
 
 ``VectorSearchIndex``
-=====================
+---------------------
 
 .. class:: VectorSearchIndex(*, fields=(), name=None, similarities)
 

docs/releases/5.2.x.rst

@@ -15,7 +15,9 @@ New features
 Bug fixes
 ---------
 
-- ...
+- Migrations operations that add or delete search indexes now wait until the
+  operation is completed on the server to prevent conflicts when running
+  multiple operations sequentially.
 
 5.2.0
 =====

tests/atlas_search_/test_search.py

@@ -28,21 +28,11 @@
     SearchVector,
     SearchWildcard,
 )
+from django_mongodb_backend.schema import DatabaseSchemaEditor
 
 from .models import Article, Location, Writer
 
 
-def wait_until_index_ready(collection, index_name, timeout: float = 5, interval: float = 0.5):
-    start = monotonic()
-    while monotonic() - start < timeout:
-        indexes = list(collection.list_search_indexes())
-        for idx in indexes:
-            if idx["name"] == index_name and idx["status"] == "READY":
-                return True
-        sleep(interval)
-    raise TimeoutError(f"Index {index_name} not ready after {timeout} seconds")
-
-
 def _delayed_assertion(timeout: float = 4, interval: float = 0.5):
     def decorator(assert_func):
         @wraps(assert_func)
@@ -91,13 +81,16 @@ def _get_collection(model):
 
     @classmethod
     def create_search_index(cls, model, index_name, definition, type="search"):
+        # TODO: create/delete indexes using DatabaseSchemaEditor when
+        # SearchIndexes support mappings (INTPYTHON-729).
         collection = cls._get_collection(model)
         idx = SearchIndexModel(definition=definition, name=index_name, type=type)
         collection.create_search_index(idx)
-        wait_until_index_ready(collection, index_name)
+        DatabaseSchemaEditor.wait_until_index_created(collection, index_name)
 
         def drop_index():
             collection.drop_search_index(index_name)
+            DatabaseSchemaEditor.wait_until_index_dropped(collection, index_name)
 
         cls.addClassCleanup(drop_index)
 

tests/indexes_/test_base.py

@@ -3,8 +3,7 @@
 
 class SchemaAssertionMixin:
     def assertAddRemoveIndex(self, editor, model, index):
-        with self.assertNumQueries(1):
-            editor.add_index(index=index, model=model)
+        editor.add_index(index=index, model=model)
         try:
             self.assertIn(
                 index.name,

tests/indexes_/test_search_indexes.py

@@ -237,7 +237,7 @@ def test_multiple_fields(self):
                 },
                 "latestVersion": 0,
                 "name": "recent_test_idx",
-                "queryable": False,
+                "queryable": True,
                 "type": "vectorSearch",
             }
             self.assertCountEqual(index_info[index.name]["columns"], index.fields)
@@ -280,7 +280,7 @@ def test_similarities_list(self):
                 },
                 "latestVersion": 0,
                 "name": "recent_test_idx",
-                "queryable": False,
+                "queryable": True,
                 "type": "vectorSearch",
             }
             self.assertCountEqual(index_info[index.name]["columns"], index.fields)