Skip to content

Commit 86489fa

Browse files
timgrahamtimgraham
committed
edits
1 parent f4a5b9a commit 86489fa

File tree

6 files changed

+145
-153
lines changed

6 files changed

+145
-153
lines changed

‎django_mongodb_backend/features.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -618,7 +618,7 @@ def supports_atlas_search(self):
618618
"""Does the server support Atlas search queries and search indexes?"""
619619
try:
620620
# An existing collection must be used on MongoDB 6, otherwise
621-
# the operation will not error.
621+
# the operation will not error when unsupported.
622622
self.connection.get_collection("django_migrations").list_search_indexes()
623623
except OperationFailure:
624624
# Error: $listSearchIndexes stage is only allowed on MongoDB Atlas

‎django_mongodb_backend/indexes.py

+27-31
Original file line numberDiff line numberDiff line change
@@ -132,11 +132,11 @@ def search_index_data_types(self, field, db_type):
132132
"""
133133
if field.get_internal_type() == "UUIDField":
134134
return "uuid"
135-
if field.get_internal_type() in ("ObjectIdAutoField", "ObjectIdField"):
135+
if field.get_internal_type() in {"ObjectIdAutoField", "ObjectIdField"}:
136136
return "ObjectId"
137137
if field.get_internal_type() == "EmbeddedModelField":
138138
return "embeddedDocuments"
139-
if db_type in ("int", "long"):
139+
if db_type in {"int", "long"}:
140140
return "number"
141141
if db_type == "binData":
142142
return "string"
@@ -164,26 +164,24 @@ def get_pymongo_index_model(
164164

165165
class VectorSearchIndex(SearchIndex):
166166
suffix = "vsi"
167-
ALLOWED_SIMILARITY_FUNCTIONS = frozenset(("euclidean", "cosine", "dotProduct"))
167+
VALID_SIMILARITIES = frozenset(("euclidean", "cosine", "dotProduct"))
168168
_error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
169169

170170
def __init__(self, *expressions, fields=(), similarities="cosine", name=None, **kwargs):
171171
super().__init__(*expressions, fields=fields, name=name, **kwargs)
172172
self.similarities = similarities
173-
for func in similarities if isinstance(similarities, list) else (similarities,):
174-
if func not in self.ALLOWED_SIMILARITY_FUNCTIONS:
173+
self._multiple_similarities = isinstance(similarities, tuple | list)
174+
for func in similarities if self._multiple_similarities else (similarities,):
175+
if func not in self.VALID_SIMILARITIES:
175176
raise ValueError(
176-
f"{func} isn't a valid similarity function, options "
177-
f"are {', '.join(sorted(self.ALLOWED_SIMILARITY_FUNCTIONS))}"
177+
f"'{func}' isn't a valid similarity function "
178+
f"({', '.join(sorted(self.VALID_SIMILARITIES))})."
178179
)
179-
viewed = set()
180+
seen_fields = set()
180181
for field_name, _ in self.fields_orders:
181-
if field_name in viewed:
182-
raise ValueError(
183-
f"Field '{field_name}' is defined more than once. Vector and filter "
184-
"fields must use distinct field names.",
185-
)
186-
viewed.add(field_name)
182+
if field_name in seen_fields:
183+
raise ValueError(f"Field '{field_name}' is duplicated in fields.")
184+
seen_fields.add(field_name)
187185

188186
def check(self, model, connection):
189187
errors = super().check(model, connection)
@@ -197,20 +195,20 @@ def check(self, model, connection):
197195
except (ValueError, TypeError):
198196
errors.append(
199197
Error(
200-
f"Atlas vector search requires size on {field_name}.",
198+
f"VectorSearchIndex requires 'size' on field '{field_name}'.",
201199
obj=model,
202-
id=f"{self._error_id_prefix}.E001",
200+
id=f"{self._error_id_prefix}.E002",
203201
)
204202
)
205203
if not isinstance(field_.base_field, FloatField | DecimalField):
206204
errors.append(
207205
Error(
208-
"An Atlas vector search index requires the base "
209-
"field of ArrayField Model.field_name "
210-
"to be FloatField or DecimalField but "
211-
f"is {field_.base_field.get_internal_type()}.",
206+
"VectorSearchIndex requires the base field of "
207+
f"ArrayField '{field_.name}' to be FloatField or "
208+
"DecimalField but is "
209+
f"{field_.base_field.get_internal_type()}.",
212210
obj=model,
213-
id=f"{self._error_id_prefix}.E002",
211+
id=f"{self._error_id_prefix}.E003",
214212
)
215213
)
216214
else:
@@ -223,24 +221,22 @@ def check(self, model, connection):
223221
errors.append(
224222
Error(
225223
"VectorSearchIndex does not support "
226-
f"'{field_.get_internal_type()}' {field_name}.",
224+
f"{field_.get_internal_type()} '{field_name}'.",
227225
obj=model,
228-
id=f"{self._error_id_prefix}.E003",
226+
id=f"{self._error_id_prefix}.E004",
229227
)
230228
)
231-
if isinstance(self.similarities, list) and expected_similarities != len(self.similarities):
229+
if self._multiple_similarities and expected_similarities != len(self.similarities):
232230
similarity_function_text = (
233-
"similarities functions" if expected_similarities != 1 else "similarity function"
231+
"similarity" if expected_similarities == 1 else "similarities"
234232
)
235233
errors.append(
236234
Error(
237-
f"An Atlas vector search index requires the same number of similarities and "
238-
f"vector fields, but {expected_similarities} "
239-
f"{similarity_function_text} were expected and "
240-
f"{len(self.similarities)} {'were' if len(self.similarities) != 1 else 'was'} "
241-
"provided.",
235+
f"VectorSearchIndex requires the same number of similarities and "
236+
f"vector fields; expected {expected_similarities} "
237+
f"{similarity_function_text} but got {len(self.similarities)}.",
242238
obj=model,
243-
id=f"{self._error_id_prefix}.E004",
239+
id=f"{self._error_id_prefix}.E005",
244240
)
245241
)
246242
return errors

‎django_mongodb_backend/schema.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,7 @@ def remove_index(self, model, index):
290290
if index.contains_expressions:
291291
return
292292
if isinstance(index, SearchIndex):
293-
# Drop the index if it is supported.
293+
# Drop the index if it's supported.
294294
if self.connection.features.supports_atlas_search:
295295
self.get_collection(model._meta.db_table).drop_search_index(index.name)
296296
else:

‎docs/source/ref/models/indexes.rst

+17-7
Original file line numberDiff line numberDiff line change
@@ -13,25 +13,35 @@ Some MongoDB-specific indexes are available in
1313

1414
.. class:: SearchIndex(fields=(), name=None)
1515

16+
.. versionadded:: 5.2.0b0
17+
1618
Creates a basic :doc:`search index <atlas:atlas-search/index-definitions>` on
1719
the given field(s).
1820

1921
If ``name`` isn't provided, one will be generated automatically. If you need
2022
to reference the name in your search query and don't provide your own name,
21-
you can lookup the generated one using: ``Model._meta.indexes[0].name``
22-
(substiting a different index as needed if your model has multiple indexes).
23+
you can lookup the generated one using ``Model._meta.indexes[0].name``
24+
(substiting the name of your model as well as a different list index if your
25+
model has multiple indexes).
2326

2427
``VectorSearchIndex``
2528
=====================
2629

2730
.. class:: VectorSearchIndex(fields=(), similarities="cosine", name=None)
2831

32+
.. versionadded:: 5.2.0b0
33+
2934
A subclass of :class:`SearchIndex` that creates a :doc:`vector search index
3035
<atlas:atlas-vector-search/vector-search-type>` on the given field(s).
3136

32-
Available values for ``similarities`` are ``"euclidean"``, ``"cosine"``, and
33-
``"dotProduct"``. You can provide either a string value, in which case that
34-
value will be applied to all fields, or a list or tuple of values of the same
35-
length as list/tuple of fields with a similarity value for each field.
37+
Each index should references at least one vector field: an :class:`.ArrayField`
38+
with a :attr:`~.ArrayField.base_field` of :class:`~django.db.models.FloatField`
39+
or :class:`~django.db.models.DecimalField`.
3640

37-
<document restrictions on arrayfield, etc.>
41+
It may also have other fields to filter on, provided the field stores
42+
``boolean``, ``date``, ``objectId``, ``numeric``, ``string``, or ``UUID``.
43+
44+
Available values for ``similarities`` are ``"euclidean"``, ``"cosine"``, and
45+
``"dotProduct"``. You can provide this value either a string, in which case
46+
that value will be applied to all vector fields, or a list or tuple of values
47+
with a similarity corresponding to each vector field.

‎tests/indexes_/test_checks.py

+20-25
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,18 @@
1+
from unittest import mock
2+
13
from django.core import checks
2-
from django.db import models
3-
from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
4-
from django.test.utils import (
5-
isolate_apps,
6-
override_system_checks,
7-
)
4+
from django.db import connection, models
5+
from django.test import TestCase
6+
from django.test.utils import isolate_apps, override_system_checks
87

98
from django_mongodb_backend.checks import check_indexes
109
from django_mongodb_backend.fields import ArrayField
1110
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
1211

1312

14-
@skipIfDBFeature("supports_atlas_search")
1513
@isolate_apps("indexes_", attr_name="apps")
1614
@override_system_checks([check_indexes])
15+
@mock.patch.object(connection.features, "supports_atlas_search", False)
1716
class UnsupportedSearchIndexesTests(TestCase):
1817
def test_search_requires_atlas_search_support(self):
1918
class Article(models.Model):
@@ -62,9 +61,9 @@ class Meta:
6261
)
6362

6463

65-
@skipUnlessDBFeature("supports_atlas_search")
6664
@isolate_apps("indexes_", attr_name="apps")
6765
@override_system_checks([check_indexes])
66+
@mock.patch.object(connection.features, "supports_atlas_search", True)
6867
class InvalidVectorSearchIndexesTests(TestCase):
6968
def test_requires_size(self):
7069
class Article(models.Model):
@@ -78,8 +77,8 @@ class Meta:
7877
errors,
7978
[
8079
checks.Error(
81-
"Atlas vector search requires size on title_embedded.",
82-
id="django_mongodb_backend.indexes.VectorSearchIndex.E001",
80+
"VectorSearchIndex requires 'size' on field 'title_embedded'.",
81+
id="django_mongodb_backend.indexes.VectorSearchIndex.E002",
8382
obj=Article,
8483
)
8584
],
@@ -97,9 +96,9 @@ class Meta:
9796
errors,
9897
[
9998
checks.Error(
100-
"An Atlas vector search index requires the base field of ArrayField "
101-
"Model.field_name to be FloatField or DecimalField but is CharField.",
102-
id="django_mongodb_backend.indexes.VectorSearchIndex.E002",
99+
"VectorSearchIndex requires the base field of ArrayField "
100+
"'title_embedded' to be FloatField or DecimalField but is CharField.",
101+
id="django_mongodb_backend.indexes.VectorSearchIndex.E003",
103102
obj=Article,
104103
)
105104
],
@@ -110,17 +109,15 @@ class Article(models.Model):
110109
data = models.JSONField()
111110

112111
class Meta:
113-
indexes = [
114-
VectorSearchIndex(fields=["data"]),
115-
]
112+
indexes = [VectorSearchIndex(fields=["data"])]
116113

117114
errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
118115
self.assertEqual(
119116
errors,
120117
[
121118
checks.Error(
122-
"VectorSearchIndex does not support 'JSONField' data.",
123-
id="django_mongodb_backend.indexes.VectorSearchIndex.E003",
119+
"VectorSearchIndex does not support JSONField 'data'.",
120+
id="django_mongodb_backend.indexes.VectorSearchIndex.E004",
124121
obj=Article,
125122
)
126123
],
@@ -143,9 +140,8 @@ class Meta:
143140
errors,
144141
[
145142
checks.Error(
146-
"An Atlas vector search index requires the same number of similarities "
147-
"and vector fields, but 1 similarity function were expected and 2 "
148-
"were provided.",
143+
"VectorSearchIndex requires the same number of similarities "
144+
"and vector fields; expected 1 similarity but got 2.",
149145
id="django_mongodb_backend.indexes.VectorSearchIndex.E004",
150146
obj=Article,
151147
),
@@ -170,10 +166,9 @@ class Meta:
170166
errors,
171167
[
172168
checks.Error(
173-
"An Atlas vector search index requires the same number of similarities "
174-
"and vector fields, but 2 similarities functions were expected and 1 "
175-
"was provided.",
176-
id="django_mongodb_backend.indexes.VectorSearchIndex.E004",
169+
"VectorSearchIndex requires the same number of similarities "
170+
"and vector fields; expected 2 similarities but got 1.",
171+
id="django_mongodb_backend.indexes.VectorSearchIndex.E005",
177172
obj=Article,
178173
),
179174
],

0 commit comments

Comments
 (0)