Skip to content

Commit 0f28afe

Browse files
timgrahamtimgraham
committed
more edits
1 parent 252f364 commit 0f28afe

File tree

8 files changed

+123
-61
lines changed

8 files changed

+123
-61
lines changed

‎.github/workflows/test-python-atlas.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Python Atlas Tests
1+
name: Python Tests on Atlas
22

33
on:
44
pull_request:

‎django_mongodb_backend/features.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -634,8 +634,9 @@ def supports_atlas_search(self):
634634
# An existing collection must be used on MongoDB 6, otherwise
635635
# the operation will not error when unsupported.
636636
self.connection.get_collection("django_migrations").list_search_indexes()
637-
except OperationFailure:
638-
# Error: $listSearchIndexes stage is only allowed on MongoDB Atlas
639-
return False
637+
except OperationFailure as exc:
638+
if "$listSearchIndexes stage is only allowed on MongoDB Atlas" in str(exc):
639+
return False
640+
raise
640641
else:
641642
return True

‎django_mongodb_backend/indexes.py

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,9 @@ class SearchIndex(Index):
109109
suffix = "six"
110110
_error_id_prefix = "django_mongodb_backend.indexes.SearchIndex"
111111

112+
def __init__(self, *, fields=(), name=None):
113+
super().__init__(fields=fields, name=name)
114+
112115
def check(self, model, connection):
113116
errors = []
114117
if not connection.features.supports_atlas_search:
@@ -161,10 +164,17 @@ def get_pymongo_index_model(
161164
class VectorSearchIndex(SearchIndex):
162165
suffix = "vsi"
163166
VALID_SIMILARITIES = frozenset(("euclidean", "cosine", "dotProduct"))
167+
VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
164168
_error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
165169

166-
def __init__(self, *expressions, fields=(), similarities="cosine", name=None, **kwargs):
167-
super().__init__(*expressions, fields=fields, name=name, **kwargs)
170+
def __init__(
171+
self,
172+
*,
173+
fields=(),
174+
similarities="cosine",
175+
name=None,
176+
):
177+
super().__init__(fields=fields, name=name)
168178
self.similarities = similarities
169179
self._multiple_similarities = isinstance(similarities, tuple | list)
170180
for func in similarities if self._multiple_similarities else (similarities,):
@@ -209,27 +219,24 @@ def check(self, model, connection):
209219
)
210220
else:
211221
search_type = self.search_index_data_types(field.db_type(connection))
212-
# filter - for fields that contain boolean, date, objectId,
213-
# numeric, string, or UUID values. Reference:
214-
# https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#atlas-vector-search-index-fields
215-
if search_type not in {"number", "string", "boolean", "objectId", "uuid", "date"}:
222+
# Validate allowed search types.
223+
if search_type not in self.VALID_FIELD_TYPES:
216224
errors.append(
217225
Error(
218-
"VectorSearchIndex does not support "
219-
f"{field.get_internal_type()} '{field_name}'.",
226+
"VectorSearchIndex does not support field "
227+
f"'{field_name}' ({field.get_internal_type()}).",
220228
obj=model,
221229
id=f"{self._error_id_prefix}.E004",
230+
hint=f"Allowed types are {', '.join(sorted(self.VALID_FIELD_TYPES))}.",
222231
)
223232
)
224233
if self._multiple_similarities and expected_similarities != len(self.similarities):
225-
similarity_function_text = (
226-
"similarity" if expected_similarities == 1 else "similarities"
227-
)
228234
errors.append(
229235
Error(
230-
f"VectorSearchIndex requires the same number of similarities and "
231-
f"vector fields; expected {expected_similarities} "
232-
f"{similarity_function_text} but got {len(self.similarities)}.",
236+
f"VectorSearchIndex requires the same number of similarities "
237+
f"and vector fields; {model._meta.object_name} has "
238+
f"{expected_similarities} ArrayField(s) but similarities "
239+
f"has {len(self.similarities)} element(s).",
233240
obj=model,
234241
id=f"{self._error_id_prefix}.E005",
235242
)

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ Some MongoDB-specific indexes are available in
1818
Creates a basic :doc:`search index <atlas:atlas-search/index-definitions>` on
1919
the given field(s).
2020

21+
Some fields such as :class:`~django.db.models.DecimalField` aren't
22+
supported. See the :ref:`Atlas documentation <atlas:bson-data-chart>` for a
23+
complete list of unsupported data types.
24+
2125
If ``name`` isn't provided, one will be generated automatically. If you need
2226
to reference the name in your search query and don't provide your own name,
2327
you can lookup the generated one using ``Model._meta.indexes[0].name``
@@ -36,10 +40,10 @@ A subclass of :class:`SearchIndex` that creates a :doc:`vector search index
3640

3741
Each index should references at least one vector field: an :class:`.ArrayField`
3842
with a :attr:`~.ArrayField.base_field` of :class:`~django.db.models.FloatField`
39-
or :class:`~django.db.models.DecimalField`.
43+
or :class:`~django.db.models.IntegerField`.
4044

4145
It may also have other fields to filter on, provided the field stores
42-
``boolean``, ``date``, ``objectId``, ``numeric``, ``string``, or ``UUID``.
46+
``boolean``, ``date``, ``objectId``, ``numeric``, ``string``, or ``uuid``.
4347

4448
Available values for ``similarities`` are ``"euclidean"``, ``"cosine"``, and
4549
``"dotProduct"``. You can provide this value either a string, in which case

‎docs/source/releases/5.2.x.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,8 @@ Initial release from the state of :ref:`django-mongodb-backend 5.1.0 beta 2
1212

1313
Regarding new features in Django 5.2,
1414
:class:`~django.db.models.CompositePrimaryKey` isn't supported.
15+
16+
Other changes since Django MongoDB Backend 5.1.x:
17+
18+
- Added support for :class:`.SearchIndex` and :class:`.VectorSearchIndex`.
19+
- The minimum supported version of ``pymongo`` is increased from 4.6 to 4.7.

‎tests/indexes_/models.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,15 @@ class Data(EmbeddedModel):
1515

1616

1717
class SearchIndexTestModel(models.Model):
18+
big_integer = models.BigIntegerField()
19+
binary = models.BinaryField()
1820
boolean = models.BooleanField()
19-
date = models.DateTimeField(auto_now=True)
20-
embedded = EmbeddedModelField(Data)
21+
char = models.CharField(max_length=100)
22+
datetime = models.DateTimeField(auto_now=True)
23+
embedded_model = EmbeddedModelField(Data)
2124
float = models.FloatField()
22-
json_data = models.JSONField()
23-
number = models.IntegerField()
25+
integer = models.IntegerField()
26+
json = models.JSONField()
2427
object_id = ObjectIdField()
25-
text = models.CharField(max_length=100)
2628
vector_float = ArrayField(models.FloatField(), size=10)
2729
vector_integer = ArrayField(models.IntegerField(), size=10)

‎tests/indexes_/test_checks.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,10 @@ class Meta:
117117
errors,
118118
[
119119
checks.Error(
120-
"VectorSearchIndex does not support JSONField 'data'.",
120+
"VectorSearchIndex does not support field 'data' (JSONField).",
121121
id="django_mongodb_backend.indexes.VectorSearchIndex.E004",
122122
obj=Article,
123+
hint="Allowed types are boolean, date, number, objectId, string, uuid.",
123124
)
124125
],
125126
)
@@ -142,7 +143,8 @@ class Meta:
142143
[
143144
checks.Error(
144145
"VectorSearchIndex requires the same number of similarities "
145-
"and vector fields; expected 1 similarity but got 2.",
146+
"and vector fields; Article has 1 ArrayField(s) but similarities "
147+
"has 2 element(s).",
146148
id="django_mongodb_backend.indexes.VectorSearchIndex.E005",
147149
obj=Article,
148150
),
@@ -168,7 +170,8 @@ class Meta:
168170
[
169171
checks.Error(
170172
"VectorSearchIndex requires the same number of similarities "
171-
"and vector fields; expected 2 similarities but got 1.",
173+
"and vector fields; Article has 2 ArrayField(s) but similarities "
174+
"has 1 element(s).",
172175
id="django_mongodb_backend.indexes.VectorSearchIndex.E005",
173176
obj=Article,
174177
),

‎tests/indexes_/test_search_indexes.py

Lines changed: 73 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,33 @@ def test_vector_index_not_created(self):
3838
)
3939

4040

41+
class SearchIndexTests(SimpleTestCase):
42+
def test_no_init_args(self):
43+
"""All arguments must be kwargs."""
44+
msg = "SearchIndex.__init__() takes 1 positional argument but 2 were given"
45+
with self.assertRaisesMessage(TypeError, msg):
46+
SearchIndex("foo")
47+
48+
def test_no_extra_kargs(self):
49+
"""Unused wargs that appear on Index aren't accepted."""
50+
msg = "SearchIndex.__init__() got an unexpected keyword argument 'condition'"
51+
with self.assertRaisesMessage(TypeError, msg):
52+
SearchIndex(condition="")
53+
54+
4155
class VectorSearchIndexTests(SimpleTestCase):
56+
def test_no_init_args(self):
57+
"""All arguments must be kwargs."""
58+
msg = "VectorSearchIndex.__init__() takes 1 positional argument but 2 were given"
59+
with self.assertRaisesMessage(TypeError, msg):
60+
VectorSearchIndex("foo")
61+
62+
def test_no_extra_kargs(self):
63+
"""Unused wargs that appear on Index aren't accepted."""
64+
msg = "VectorSearchIndex.__init__() got an unexpected keyword argument 'condition'"
65+
with self.assertRaisesMessage(TypeError, msg):
66+
VectorSearchIndex(condition="")
67+
4268
def test_deconstruct(self):
4369
index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
4470
name, args, kwargs = index.deconstruct()
@@ -48,15 +74,15 @@ def test_deconstruct(self):
4874
def test_deconstruct_with_similarities(self):
4975
index = VectorSearchIndex(
5076
name="recent_test_idx",
51-
fields=["number", "text"],
77+
fields=["number", "char"],
5278
similarities=["cosine", "dotProduct"],
5379
)
5480
path, args, kwargs = index.deconstruct()
5581
self.assertEqual(
5682
kwargs,
5783
{
5884
"name": "recent_test_idx",
59-
"fields": ["number", "text"],
85+
"fields": ["number", "char"],
6086
"similarities": ["cosine", "dotProduct"],
6187
},
6288
)
@@ -87,23 +113,25 @@ class SearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
87113
def test_simple(self):
88114
index = SearchIndex(
89115
name="recent_test_idx",
90-
fields=["text"],
116+
fields=["char"],
91117
)
92118
with connection.schema_editor() as editor:
93119
self.assertAddRemoveIndex(editor, index=index, model=SearchIndexTestModel)
94120

95-
def test_all_fields(self):
121+
def test_valid_fields(self):
96122
index = SearchIndex(
97123
name="recent_test_idx",
98124
fields=[
125+
"big_integer",
126+
"binary",
127+
"char",
99128
"boolean",
100-
"date",
101-
"embedded",
129+
"datetime",
130+
"embedded_model",
102131
"float",
103-
"json_data",
104-
"number",
132+
"integer",
133+
"json",
105134
"object_id",
106-
"text",
107135
"vector_integer",
108136
"vector_float",
109137
],
@@ -118,29 +146,41 @@ def test_all_fields(self):
118146
expected_options = {
119147
"dynamic": False,
120148
"fields": {
149+
"big_integer": {
150+
"indexDoubles": True,
151+
"indexIntegers": True,
152+
"representation": "double",
153+
"type": "number",
154+
},
155+
"binary": {
156+
"indexOptions": "offsets",
157+
"norms": "include",
158+
"store": True,
159+
"type": "string",
160+
},
121161
"boolean": {"type": "boolean"},
122-
"date": {"type": "date"},
123-
"embedded": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
162+
"char": {
163+
"indexOptions": "offsets",
164+
"norms": "include",
165+
"store": True,
166+
"type": "string",
167+
},
168+
"datetime": {"type": "date"},
169+
"embedded_model": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
124170
"float": {
125171
"indexDoubles": True,
126172
"indexIntegers": True,
127173
"representation": "double",
128174
"type": "number",
129175
},
130-
"json_data": {"dynamic": False, "fields": {}, "type": "document"},
131-
"number": {
176+
"integer": {
132177
"indexDoubles": True,
133178
"indexIntegers": True,
134179
"representation": "double",
135180
"type": "number",
136181
},
182+
"json": {"dynamic": False, "fields": {}, "type": "document"},
137183
"object_id": {"type": "objectId"},
138-
"text": {
139-
"indexOptions": "offsets",
140-
"norms": "include",
141-
"store": True,
142-
"type": "string",
143-
},
144184
"vector_float": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
145185
"vector_integer": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
146186
},
@@ -155,22 +195,22 @@ def test_all_fields(self):
155195
@skipUnlessDBFeature("supports_atlas_search")
156196
class VectorSearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
157197
def test_simple(self):
158-
index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
198+
index = VectorSearchIndex(name="recent_test_idx", fields=["integer"])
159199
with connection.schema_editor() as editor:
160200
self.assertAddRemoveIndex(editor, index=index, model=SearchIndexTestModel)
161201

162202
def test_multiple_fields(self):
163203
index = VectorSearchIndex(
164204
name="recent_test_idx",
165205
fields=[
166-
"text",
206+
"boolean",
207+
"char",
208+
"datetime",
209+
"embedded_model",
210+
"integer",
167211
"object_id",
168-
"number",
169-
"embedded",
170-
"vector_integer",
171212
"vector_float",
172-
"boolean",
173-
"date",
213+
"vector_integer",
174214
],
175215
)
176216
with connection.schema_editor() as editor:
@@ -183,24 +223,24 @@ def test_multiple_fields(self):
183223
expected_options = {
184224
"latestDefinition": {
185225
"fields": [
186-
{"path": "text", "type": "filter"},
226+
{"path": "boolean", "type": "filter"},
227+
{"path": "char", "type": "filter"},
228+
{"path": "datetime", "type": "filter"},
229+
{"path": "embedded_model", "type": "filter"},
230+
{"path": "integer", "type": "filter"},
187231
{"path": "object_id", "type": "filter"},
188-
{"path": "number", "type": "filter"},
189-
{"path": "embedded", "type": "filter"},
190232
{
191233
"numDimensions": 10,
192-
"path": "vector_integer",
234+
"path": "vector_float",
193235
"similarity": "cosine",
194236
"type": "vector",
195237
},
196238
{
197239
"numDimensions": 10,
198-
"path": "vector_float",
240+
"path": "vector_integer",
199241
"similarity": "cosine",
200242
"type": "vector",
201243
},
202-
{"path": "boolean", "type": "filter"},
203-
{"path": "date", "type": "filter"},
204244
]
205245
},
206246
"latestVersion": 0,

0 commit comments

Comments
 (0)