Skip to content

Commit dd02d80

Browse files
WaVEVWaVEV
committed
Edits
1 parent 7da1fdf commit dd02d80

File tree

4 files changed

+38
-42
lines changed

4 files changed

+38
-42
lines changed

‎django_mongodb_backend/indexes.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -124,8 +124,8 @@ def create_mongodb_index(
124124
return IndexModel(index_orders, name=self.name, **kwargs)
125125

126126

127-
class AtlasSearchIndex(Index):
128-
suffix = "atlas_search"
127+
class SearchIndex(Index):
128+
suffix = "search"
129129

130130
def __init__(self, *expressions, **kwargs):
131131
super().__init__(*expressions, **kwargs)
@@ -144,8 +144,8 @@ def create_mongodb_index(
144144
)
145145

146146

147-
class AtlasVectorSearchIndex(Index):
148-
suffix = "atlas_vector_search"
147+
class VectorSearchIndex(Index):
148+
suffix = "vector_search"
149149
ALLOWED_SIMILARITY_FUNCTIONS = ("euclidean", "cosine", "dotProduct")
150150

151151
def __init__(self, *expressions, similarities="cosine", **kwargs):

‎django_mongodb_backend/introspection.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from django.db.models import Index
33
from pymongo import ASCENDING, DESCENDING
44

5-
from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
5+
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
66

77

88
class DatabaseIntrospection(BaseDatabaseIntrospection):
@@ -33,18 +33,18 @@ def _get_index_info(self, table_name):
3333
}
3434
return constraints
3535

36-
def _get_atlas_index_info(self, table_name):
36+
def _get_search_index_info(self, table_name):
3737
constraints = {}
3838
indexes = self.connection.get_collection(table_name).list_search_indexes()
3939
for details in indexes:
4040
if details["type"] == "vectorSearch":
4141
columns = [field["path"] for field in details["latestDefinition"]["fields"]]
42-
type_ = AtlasVectorSearchIndex.suffix
42+
type_ = VectorSearchIndex.suffix
4343
options = details
4444
else:
4545
columns = list(details["latestDefinition"]["mappings"].get("fields", {}).keys())
4646
options = details["latestDefinition"]["mappings"]
47-
type_ = AtlasSearchIndex.suffix
47+
type_ = SearchIndex.suffix
4848
constraints[details["name"]] = {
4949
"check": False,
5050
"columns": columns,
@@ -60,4 +60,4 @@ def _get_atlas_index_info(self, table_name):
6060
return constraints
6161

6262
def get_constraints(self, cursor, table_name):
63-
return {**self._get_index_info(table_name), **self._get_atlas_index_info(table_name)}
63+
return {**self._get_index_info(table_name), **self._get_search_index_info(table_name)}

‎django_mongodb_backend/schema.py

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from django.db.models import Index, UniqueConstraint
33
from pymongo.operations import IndexModel, SearchIndexModel
44

5-
from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
5+
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
66

77
from .fields import EmbeddedModelField
88
from .query import wrap_database_errors
@@ -259,13 +259,6 @@ def alter_unique_together(
259259
model, constraint, parent_model=parent_model, column_prefix=column_prefix
260260
)
261261

262-
def _add_index(self, index, model):
263-
if isinstance(index, SearchIndexModel):
264-
return self.get_collection(model._meta.db_table).create_search_index(index)
265-
if isinstance(index, IndexModel):
266-
return self.get_collection(model._meta.db_table).create_indexes([index])
267-
raise ValueError(f"{type(index)} isn't a supported index type")
268-
269262
@ignore_embedded_models
270263
def add_index(
271264
self, model, index, *, field=None, unique=False, column_prefix="", parent_model=None
@@ -280,7 +273,12 @@ def add_index(
280273
)
281274
if idx:
282275
model = parent_model or model
283-
self._add_index(idx, model)
276+
if isinstance(idx, SearchIndexModel):
277+
self.get_collection(model._meta.db_table).create_search_index(idx)
278+
elif isinstance(idx, IndexModel):
279+
self.get_collection(model._meta.db_table).create_indexes([idx])
280+
else:
281+
raise ValueError(f"{type(idx)} isn't a supported index type")
284282

285283
def _add_composed_index(self, model, field_names, column_prefix="", parent_model=None):
286284
"""Add an index on the given list of field_names."""
@@ -294,18 +292,16 @@ def _add_field_index(self, model, field, *, column_prefix=""):
294292
index.name = self._create_index_name(model._meta.db_table, [column_prefix + field.column])
295293
self.add_index(model, index, field=field, column_prefix=column_prefix)
296294

297-
def _remove_index(self, index, model):
298-
if isinstance(index, AtlasSearchIndex | AtlasVectorSearchIndex):
299-
return self.get_collection(model._meta.db_table).drop_search_index(index.name)
300-
if isinstance(index, Index):
301-
return self.get_collection(model._meta.db_table).drop_index(index.name)
302-
raise ValueError(f"{type(index)} isn't a supported index type")
303-
304295
@ignore_embedded_models
305296
def remove_index(self, model, index):
306297
if index.contains_expressions:
307298
return
308-
self._remove_index(index, model)
299+
if isinstance(index, SearchIndex | VectorSearchIndex):
300+
self.get_collection(model._meta.db_table).drop_search_index(index.name)
301+
elif isinstance(index, Index):
302+
self.get_collection(model._meta.db_table).drop_index(index.name)
303+
else:
304+
raise ValueError(f"{type(index)} isn't a supported index type")
309305

310306
def _remove_composed_index(
311307
self, model, field_names, constraint_kwargs, column_prefix="", parent_model=None

‎tests/indexes_/test_atlas_indexes.py

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from django.db import connection
55
from django.test import TestCase
66

7-
from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
7+
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
88

99
from .models import Article, Data
1010

@@ -34,7 +34,7 @@ def assertAddRemoveIndex(self, editor, model, index):
3434

3535
def test_simple(self):
3636
with connection.schema_editor() as editor:
37-
index = AtlasSearchIndex(
37+
index = SearchIndex(
3838
name="recent_article_idx",
3939
fields=["number"],
4040
)
@@ -43,7 +43,7 @@ def test_simple(self):
4343

4444
def test_multiple_fields(self):
4545
with connection.schema_editor() as editor:
46-
index = AtlasSearchIndex(
46+
index = SearchIndex(
4747
name="recent_article_idx",
4848
fields=["headline", "number", "body", "data", "embedded", "auto_now"],
4949
)
@@ -83,7 +83,7 @@ def test_multiple_fields(self):
8383
self.assertAddRemoveIndex(editor, Article, index)
8484

8585
def test_field_not_exists(self):
86-
index = AtlasSearchIndex(
86+
index = SearchIndex(
8787
name="recent_article_idx",
8888
fields=["headline", "non_existing_name"],
8989
)
@@ -137,7 +137,7 @@ def assertAddRemoveIndex(self, editor, model, index):
137137

138138
def test_simple(self):
139139
with connection.schema_editor() as editor:
140-
index = AtlasSearchIndex(
140+
index = SearchIndex(
141141
name="recent_article_idx",
142142
fields=["number"],
143143
)
@@ -146,7 +146,7 @@ def test_simple(self):
146146

147147
def test_multiple_fields(self):
148148
with connection.schema_editor() as editor:
149-
index = AtlasSearchIndex(
149+
index = SearchIndex(
150150
name="recent_article_idx",
151151
fields=["headline", "number", "body", "data", "embedded", "auto_now"],
152152
)
@@ -186,7 +186,7 @@ def test_multiple_fields(self):
186186
self.assertAddRemoveIndex(editor, Article, index)
187187

188188

189-
class AtlasSearchIndexTests(TestCase):
189+
class SearchIndexTests(TestCase):
190190
# Schema editor is used to create the index to test that it works.
191191
# available_apps = ["indexes"]
192192
available_apps = None # could be removed?
@@ -209,9 +209,9 @@ def assertAddRemoveIndex(self, editor, model, index):
209209
),
210210
)
211211

212-
def test_simple_atlas_vector_search(self):
212+
def test_simple_vector_search(self):
213213
with connection.schema_editor() as editor:
214-
index = AtlasVectorSearchIndex(
214+
index = VectorSearchIndex(
215215
name="recent_article_idx",
216216
fields=["number"],
217217
)
@@ -220,7 +220,7 @@ def test_simple_atlas_vector_search(self):
220220

221221
def test_multiple_fields(self):
222222
with connection.schema_editor() as editor:
223-
index = AtlasVectorSearchIndex(
223+
index = VectorSearchIndex(
224224
name="recent_article_idx",
225225
fields=["headline", "number", "body", "description_embedded"],
226226
)
@@ -255,7 +255,7 @@ def test_multiple_fields(self):
255255
self.assertAddRemoveIndex(editor, Article, index)
256256

257257
def test_field_not_exists(self):
258-
index = AtlasVectorSearchIndex(
258+
index = VectorSearchIndex(
259259
name="recent_article_idx",
260260
fields=["headline", "non_existing_name", "title_embedded"],
261261
)
@@ -267,7 +267,7 @@ def test_field_not_exists(self):
267267
editor.add_index(index=index, model=Article)
268268

269269

270-
class AtlasSearchIndexTestsWithData(TestCase):
270+
class SearchIndexTestsWithData(TestCase):
271271
available_apps = None # could be removed?
272272

273273
@classmethod
@@ -307,9 +307,9 @@ def assertAddRemoveIndex(self, editor, model, index):
307307
),
308308
)
309309

310-
def test_simple_atlas_vector_search(self):
310+
def test_simple_vector_search(self):
311311
with connection.schema_editor() as editor:
312-
index = AtlasVectorSearchIndex(
312+
index = VectorSearchIndex(
313313
name="recent_article_idx",
314314
fields=["number"],
315315
)
@@ -318,7 +318,7 @@ def test_simple_atlas_vector_search(self):
318318

319319
def test_multiple_fields(self):
320320
with connection.schema_editor() as editor:
321-
index = AtlasVectorSearchIndex(
321+
index = VectorSearchIndex(
322322
name="recent_article_idx",
323323
fields=["headline", "number", "body", "description_embedded"],
324324
)
@@ -367,7 +367,7 @@ def test_corrupted_data(self):
367367
number_list=[2] * 22,
368368
name_list=[f"name_{2}"] * 7,
369369
)
370-
index = AtlasVectorSearchIndex(
370+
index = VectorSearchIndex(
371371
name="recent_article_idx",
372372
fields=["headline", "number", "body", "description_embedded"],
373373
)

0 commit comments

Comments
 (0)