Edits

Author: WaVEV

django_mongodb_backend/indexes.py

@@ -124,8 +124,8 @@ def create_mongodb_index(
     return IndexModel(index_orders, name=self.name, **kwargs)
 
 
-class AtlasSearchIndex(Index):
-    suffix = "atlas_search"
+class SearchIndex(Index):
+    suffix = "search"
 
     def __init__(self, *expressions, **kwargs):
         super().__init__(*expressions, **kwargs)
@@ -144,8 +144,8 @@ def create_mongodb_index(
         )
 
 
-class AtlasVectorSearchIndex(Index):
-    suffix = "atlas_vector_search"
+class VectorSearchIndex(Index):
+    suffix = "vector_search"
     ALLOWED_SIMILARITY_FUNCTIONS = ("euclidean", "cosine", "dotProduct")
 
     def __init__(self, *expressions, similarities="cosine", **kwargs):

django_mongodb_backend/introspection.py

@@ -2,7 +2,7 @@
 from django.db.models import Index
 from pymongo import ASCENDING, DESCENDING
 
-from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
+from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
 
 
 class DatabaseIntrospection(BaseDatabaseIntrospection):
@@ -33,18 +33,18 @@ def _get_index_info(self, table_name):
             }
         return constraints
 
-    def _get_atlas_index_info(self, table_name):
+    def _get_search_index_info(self, table_name):
         constraints = {}
         indexes = self.connection.get_collection(table_name).list_search_indexes()
         for details in indexes:
             if details["type"] == "vectorSearch":
                 columns = [field["path"] for field in details["latestDefinition"]["fields"]]
-                type_ = AtlasVectorSearchIndex.suffix
+                type_ = VectorSearchIndex.suffix
                 options = details
             else:
                 columns = list(details["latestDefinition"]["mappings"].get("fields", {}).keys())
                 options = details["latestDefinition"]["mappings"]
-                type_ = AtlasSearchIndex.suffix
+                type_ = SearchIndex.suffix
             constraints[details["name"]] = {
                 "check": False,
                 "columns": columns,
@@ -60,4 +60,4 @@ def _get_atlas_index_info(self, table_name):
         return constraints
 
     def get_constraints(self, cursor, table_name):
-        return {**self._get_index_info(table_name), **self._get_atlas_index_info(table_name)}
+        return {**self._get_index_info(table_name), **self._get_search_index_info(table_name)}

django_mongodb_backend/schema.py

@@ -2,7 +2,7 @@
 from django.db.models import Index, UniqueConstraint
 from pymongo.operations import IndexModel, SearchIndexModel
 
-from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
+from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
 
 from .fields import EmbeddedModelField
 from .query import wrap_database_errors
@@ -259,13 +259,6 @@ def alter_unique_together(
                 model, constraint, parent_model=parent_model, column_prefix=column_prefix
             )
 
-    def _add_index(self, index, model):
-        if isinstance(index, SearchIndexModel):
-            return self.get_collection(model._meta.db_table).create_search_index(index)
-        if isinstance(index, IndexModel):
-            return self.get_collection(model._meta.db_table).create_indexes([index])
-        raise ValueError(f"{type(index)} isn't a supported index type")
-
     @ignore_embedded_models
     def add_index(
         self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None
@@ -280,7 +273,12 @@ def add_index(
         )
         if idx:
             model = parent_model or model
-            self._add_index(idx, model)
+            if isinstance(idx, SearchIndexModel):
+                self.get_collection(model._meta.db_table).create_search_index(idx)
+            elif isinstance(idx, IndexModel):
+                self.get_collection(model._meta.db_table).create_indexes([idx])
+            else:
+                raise ValueError(f"{type(idx)} isn't a supported index type")
 
     def _add_composed_index(self, model, field_names, column_prefix="", parent_model=None):
         """Add an index on the given list of field_names."""
@@ -294,18 +292,16 @@ def _add_field_index(self, model, field, *, column_prefix=""):
         index.name = self._create_index_name(model._meta.db_table, [column_prefix + field.column])
         self.add_index(model, index, field=field, column_prefix=column_prefix)
 
-    def _remove_index(self, index, model):
-        if isinstance(index, AtlasSearchIndex | AtlasVectorSearchIndex):
-            return self.get_collection(model._meta.db_table).drop_search_index(index.name)
-        if isinstance(index, Index):
-            return self.get_collection(model._meta.db_table).drop_index(index.name)
-        raise ValueError(f"{type(index)} isn't a supported index type")
-
     @ignore_embedded_models
     def remove_index(self, model, index):
         if index.contains_expressions:
             return
-        self._remove_index(index, model)
+        if isinstance(index, SearchIndex | VectorSearchIndex):
+            self.get_collection(model._meta.db_table).drop_search_index(index.name)
+        elif isinstance(index, Index):
+            self.get_collection(model._meta.db_table).drop_index(index.name)
+        else:
+            raise ValueError(f"{type(index)} isn't a supported index type")
 
     def _remove_composed_index(
         self, model, field_names, constraint_kwargs, column_prefix="", parent_model=None

tests/indexes_/test_atlas_indexes.py

@@ -4,7 +4,7 @@
 from django.db import connection
 from django.test import TestCase
 
-from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
+from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
 
 from .models import Article, Data
 
@@ -34,7 +34,7 @@ def assertAddRemoveIndex(self, editor, model, index):
 
     def test_simple(self):
         with connection.schema_editor() as editor:
-            index = AtlasSearchIndex(
+            index = SearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
@@ -43,7 +43,7 @@ def test_simple(self):
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
-            index = AtlasSearchIndex(
+            index = SearchIndex(
                 name="recent_article_idx",
                 fields=["headline", "number", "body", "data", "embedded", "auto_now"],
             )
@@ -83,7 +83,7 @@ def test_multiple_fields(self):
             self.assertAddRemoveIndex(editor, Article, index)
 
     def test_field_not_exists(self):
-        index = AtlasSearchIndex(
+        index = SearchIndex(
             name="recent_article_idx",
             fields=["headline", "non_existing_name"],
         )
@@ -137,7 +137,7 @@ def assertAddRemoveIndex(self, editor, model, index):
 
     def test_simple(self):
         with connection.schema_editor() as editor:
-            index = AtlasSearchIndex(
+            index = SearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
@@ -146,7 +146,7 @@ def test_simple(self):
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
-            index = AtlasSearchIndex(
+            index = SearchIndex(
                 name="recent_article_idx",
                 fields=["headline", "number", "body", "data", "embedded", "auto_now"],
             )
@@ -186,7 +186,7 @@ def test_multiple_fields(self):
             self.assertAddRemoveIndex(editor, Article, index)
 
 
-class AtlasSearchIndexTests(TestCase):
+class SearchIndexTests(TestCase):
     # Schema editor is used to  create the index to test that it works.
     # available_apps = ["indexes"]
     available_apps = None  # could be removed?
@@ -209,9 +209,9 @@ def assertAddRemoveIndex(self, editor, model, index):
             ),
         )
 
-    def test_simple_atlas_vector_search(self):
+    def test_simple_vector_search(self):
         with connection.schema_editor() as editor:
-            index = AtlasVectorSearchIndex(
+            index = VectorSearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
@@ -220,7 +220,7 @@ def test_simple_atlas_vector_search(self):
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
-            index = AtlasVectorSearchIndex(
+            index = VectorSearchIndex(
                 name="recent_article_idx",
                 fields=["headline", "number", "body", "description_embedded"],
             )
@@ -255,7 +255,7 @@ def test_multiple_fields(self):
             self.assertAddRemoveIndex(editor, Article, index)
 
     def test_field_not_exists(self):
-        index = AtlasVectorSearchIndex(
+        index = VectorSearchIndex(
             name="recent_article_idx",
             fields=["headline", "non_existing_name", "title_embedded"],
         )
@@ -267,7 +267,7 @@ def test_field_not_exists(self):
                 editor.add_index(index=index, model=Article)
 
 
-class AtlasSearchIndexTestsWithData(TestCase):
+class SearchIndexTestsWithData(TestCase):
     available_apps = None  # could be removed?
 
     @classmethod
@@ -307,9 +307,9 @@ def assertAddRemoveIndex(self, editor, model, index):
             ),
         )
 
-    def test_simple_atlas_vector_search(self):
+    def test_simple_vector_search(self):
         with connection.schema_editor() as editor:
-            index = AtlasVectorSearchIndex(
+            index = VectorSearchIndex(
                 name="recent_article_idx",
                 fields=["number"],
             )
@@ -318,7 +318,7 @@ def test_simple_atlas_vector_search(self):
 
     def test_multiple_fields(self):
         with connection.schema_editor() as editor:
-            index = AtlasVectorSearchIndex(
+            index = VectorSearchIndex(
                 name="recent_article_idx",
                 fields=["headline", "number", "body", "description_embedded"],
             )
@@ -367,7 +367,7 @@ def test_corrupted_data(self):
                 number_list=[2] * 22,
                 name_list=[f"name_{2}"] * 7,
             )
-            index = AtlasVectorSearchIndex(
+            index = VectorSearchIndex(
                 name="recent_article_idx",
                 fields=["headline", "number", "body", "description_embedded"],
             )