Skip to content

Commit b978a65

Browse files
Jibolatimgraham
Jibola
authored and
timgraham
committed
INTPYTHON-729 Allow creating search indexes with field mappings
1 parent 2787e0d commit b978a65

File tree

2 files changed

+84
-108
lines changed

2 files changed

+84
-108
lines changed

‎django_mongodb_backend/indexes.py

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

112-
def __init__(self, *, fields=(), name=None):
112+
def __init__(self, *, fields=(), name=None, field_mappings=None):
113+
if field_mappings and not isinstance(field_mappings, dict):
114+
raise ValueError(
115+
"field_mappings must be a dictionary mapping field names to their "
116+
"Atlas Search field mappings."
117+
)
118+
self.field_mappings = field_mappings or {}
119+
120+
fields = list({*fields, *self.field_mappings.keys()})
113121
super().__init__(fields=fields, name=name)
114122

123+
def deconstruct(self):
124+
path, args, kwargs = super().deconstruct()
125+
if self.field_mappings is not None:
126+
kwargs["field_mappings"] = self.field_mappings
127+
return path, args, kwargs
128+
115129
def check(self, model, connection):
116130
errors = []
117131
if not connection.features.supports_atlas_search:
@@ -152,22 +166,38 @@ def get_pymongo_index_model(
152166
return None
153167
fields = {}
154168
for field_name, _ in self.fields_orders:
155-
field = model._meta.get_field(field_name)
156-
type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
157169
field_path = column_prefix + model._meta.get_field(field_name).column
158-
fields[field_path] = {"type": type_}
170+
if field_name in self.field_mappings:
171+
fields[field_path] = self.field_mappings[field_name].copy()
172+
else:
173+
# If no field mapping is provided, use the default search index data type.
174+
field = model._meta.get_field(field_name)
175+
type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
176+
fields[field_path] = {"type": type_}
159177
return SearchIndexModel(
160178
definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name
161179
)
162180

163181

182+
class DynamicSearchIndex(SearchIndex):
183+
suffix = "dsix"
184+
_error_id_prefix = "django_mongodb_backend.indexes.DynamicSearchIndex"
185+
186+
def get_pymongo_index_model(
187+
self, model, schema_editor, field=None, unique=False, column_prefix=""
188+
):
189+
if not schema_editor.connection.features.supports_atlas_search:
190+
return None
191+
return SearchIndexModel(definition={"mappings": {"dynamic": True}}, name=self.name)
192+
193+
164194
class VectorSearchIndex(SearchIndex):
165195
suffix = "vsi"
166196
_error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
167197
VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
168198
VALID_SIMILARITIES = frozenset(("cosine", "dotProduct", "euclidean"))
169199

170-
def __init__(self, *, fields=(), name=None, similarities):
200+
def __init__(self, *, fields=(), name=None, similarities=()):
171201
super().__init__(fields=fields, name=name)
172202
self.similarities = similarities
173203
self._multiple_similarities = isinstance(similarities, (tuple, list))

‎tests/atlas_search_/test_search.py

Lines changed: 49 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
from django.db.models.query import QuerySet
99
from django.db.utils import DatabaseError
1010
from django.test import TransactionTestCase, skipUnlessDBFeature
11-
from pymongo.operations import SearchIndexModel
1211

1312
from django_mongodb_backend.expressions import (
1413
CompoundExpression,
@@ -28,7 +27,7 @@
2827
SearchVector,
2928
SearchWildcard,
3029
)
31-
from django_mongodb_backend.schema import DatabaseSchemaEditor
30+
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
3231

3332
from .models import Article, Location, Writer
3433

@@ -75,22 +74,15 @@ class SearchUtilsMixin(TransactionTestCase):
7574
assertListEqual = _delayed_assertion(timeout=2)(TransactionTestCase.assertListEqual)
7675
assertQuerySetEqual = _delayed_assertion(timeout=2)(TransactionTestCase.assertQuerySetEqual)
7776

78-
@staticmethod
79-
def _get_collection(model):
80-
return connection.database.get_collection(model._meta.db_table)
81-
8277
@classmethod
83-
def create_search_index(cls, model, index_name, definition, type="search"):
84-
# TODO: create/delete indexes using DatabaseSchemaEditor when
85-
# SearchIndexes support mappings (INTPYTHON-729).
86-
collection = cls._get_collection(model)
87-
idx = SearchIndexModel(definition=definition, name=index_name, type=type)
88-
collection.create_search_index(idx)
89-
DatabaseSchemaEditor.wait_until_index_created(collection, index_name)
78+
def create_search_index(cls, model, index_name, definition, index_cls=SearchIndex):
79+
idx = index_cls(field_mappings=definition, name=index_name)
80+
with connection.schema_editor() as editor:
81+
editor.add_index(model, idx)
9082

9183
def drop_index():
92-
collection.drop_search_index(index_name)
93-
DatabaseSchemaEditor.wait_until_index_dropped(collection, index_name)
84+
with connection.schema_editor() as editor:
85+
editor.remove_index(model, idx)
9486

9587
cls.addClassCleanup(drop_index)
9688

@@ -101,12 +93,7 @@ def setUpClass(cls):
10193
cls.create_search_index(
10294
Article,
10395
"equals_headline_index",
104-
{
105-
"mappings": {
106-
"dynamic": False,
107-
"fields": {"headline": {"type": "token"}, "number": {"type": "number"}},
108-
}
109-
},
96+
{"headline": {"type": "token"}, "number": {"type": "number"}},
11097
)
11198

11299
def setUp(self):
@@ -167,32 +154,27 @@ def setUpClass(cls):
167154
Article,
168155
"autocomplete_headline_index",
169156
{
170-
"mappings": {
171-
"dynamic": False,
157+
"headline": {
158+
"type": "autocomplete",
159+
"analyzer": "lucene.standard",
160+
"tokenization": "edgeGram",
161+
"minGrams": 3,
162+
"maxGrams": 5,
163+
"foldDiacritics": False,
164+
},
165+
"writer": {
166+
"type": "document",
172167
"fields": {
173-
"headline": {
168+
"name": {
174169
"type": "autocomplete",
175170
"analyzer": "lucene.standard",
176171
"tokenization": "edgeGram",
177172
"minGrams": 3,
178173
"maxGrams": 5,
179174
"foldDiacritics": False,
180-
},
181-
"writer": {
182-
"type": "document",
183-
"fields": {
184-
"name": {
185-
"type": "autocomplete",
186-
"analyzer": "lucene.standard",
187-
"tokenization": "edgeGram",
188-
"minGrams": 3,
189-
"maxGrams": 5,
190-
"foldDiacritics": False,
191-
}
192-
},
193-
},
175+
}
194176
},
195-
}
177+
},
196178
},
197179
)
198180

@@ -253,7 +235,7 @@ def setUpClass(cls):
253235
cls.create_search_index(
254236
Article,
255237
"exists_body_index",
256-
{"mappings": {"dynamic": False, "fields": {"body": {"type": "token"}}}},
238+
{"body": {"type": "token"}},
257239
)
258240

259241
def setUp(self):
@@ -282,7 +264,7 @@ def setUpClass(cls):
282264
cls.create_search_index(
283265
Article,
284266
"in_headline_index",
285-
{"mappings": {"dynamic": False, "fields": {"headline": {"type": "token"}}}},
267+
{"headline": {"type": "token"}},
286268
)
287269

288270
def setUp(self):
@@ -316,7 +298,7 @@ def setUpClass(cls):
316298
cls.create_search_index(
317299
Article,
318300
"phrase_body_index",
319-
{"mappings": {"dynamic": False, "fields": {"body": {"type": "string"}}}},
301+
{"body": {"type": "string"}},
320302
)
321303

322304
def setUp(self):
@@ -356,13 +338,8 @@ def setUpClass(cls):
356338
Article,
357339
"query_string_index",
358340
{
359-
"mappings": {
360-
"dynamic": False,
361-
"fields": {
362-
"headline": {"type": "string"},
363-
"body": {"type": "string"},
364-
},
365-
}
341+
"headline": {"type": "string"},
342+
"body": {"type": "string"},
366343
},
367344
)
368345

@@ -416,7 +393,7 @@ def setUpClass(cls):
416393
cls.create_search_index(
417394
Article,
418395
"range_number_index",
419-
{"mappings": {"dynamic": False, "fields": {"number": {"type": "number"}}}},
396+
{"number": {"type": "number"}},
420397
)
421398
Article.objects.create(headline="x", number=5, body="z")
422399

@@ -453,12 +430,7 @@ def setUpClass(cls):
453430
cls.create_search_index(
454431
Article,
455432
"regex_headline_index",
456-
{
457-
"mappings": {
458-
"dynamic": False,
459-
"fields": {"headline": {"type": "string", "analyzer": "lucene.keyword"}},
460-
}
461-
},
433+
{"headline": {"type": "string", "analyzer": "lucene.keyword"}},
462434
)
463435

464436
def setUp(self):
@@ -498,7 +470,7 @@ def setUpClass(cls):
498470
cls.create_search_index(
499471
Article,
500472
"text_body_index",
501-
{"mappings": {"dynamic": False, "fields": {"body": {"type": "string"}}}},
473+
{"body": {"type": "string"}},
502474
)
503475

504476
def setUp(self):
@@ -560,12 +532,7 @@ def setUpClass(cls):
560532
cls.create_search_index(
561533
Article,
562534
"wildcard_headline_index",
563-
{
564-
"mappings": {
565-
"dynamic": False,
566-
"fields": {"headline": {"type": "string", "analyzer": "lucene.keyword"}},
567-
}
568-
},
535+
{"headline": {"type": "string", "analyzer": "lucene.keyword"}},
569536
)
570537

571538
def setUp(self):
@@ -603,12 +570,7 @@ def setUpClass(cls):
603570
cls.create_search_index(
604571
Article,
605572
"geoshape_location_index",
606-
{
607-
"mappings": {
608-
"dynamic": False,
609-
"fields": {"location": {"type": "geo", "indexShapes": True}},
610-
}
611-
},
573+
{"location": {"type": "geo", "indexShapes": True}},
612574
)
613575

614576
def setUp(self):
@@ -668,7 +630,7 @@ def setUpClass(cls):
668630
cls.create_search_index(
669631
Article,
670632
"geowithin_location_index",
671-
{"mappings": {"dynamic": False, "fields": {"location": {"type": "geo"}}}},
633+
{"location": {"type": "geo"}},
672634
)
673635

674636
def setUp(self):
@@ -743,12 +705,7 @@ def setUpClass(cls):
743705
cls.create_search_index(
744706
Article,
745707
"mlt_index",
746-
{
747-
"mappings": {
748-
"dynamic": False,
749-
"fields": {"body": {"type": "string"}, "headline": {"type": "string"}},
750-
}
751-
},
708+
{"body": {"type": "string"}, "headline": {"type": "string"}},
752709
)
753710
cls.article1 = Article.objects.create(
754711
headline="Space exploration", number=1, body="Webb telescope"
@@ -782,14 +739,9 @@ def setUpClass(cls):
782739
Article,
783740
"compound_index",
784741
{
785-
"mappings": {
786-
"dynamic": False,
787-
"fields": {
788-
"headline": [{"type": "token"}, {"type": "string"}],
789-
"body": {"type": "string"},
790-
"number": {"type": "number"},
791-
},
792-
}
742+
"headline": [{"type": "token"}, {"type": "string"}],
743+
"body": {"type": "string"},
744+
"number": {"type": "number"},
793745
},
794746
)
795747

@@ -962,26 +914,20 @@ def test_str_returns_expected_format(self):
962914
class SearchVectorTests(SearchUtilsMixin):
963915
@classmethod
964916
def setUpClass(cls):
965-
cls.create_search_index(
966-
Article,
967-
"vector_index",
968-
{
969-
"fields": [
970-
{
971-
"type": "vector",
972-
"path": "plot_embedding",
973-
"numDimensions": 3,
974-
"similarity": "cosine",
975-
"quantization": "scalar",
976-
},
977-
{
978-
"type": "filter",
979-
"path": "number",
980-
},
981-
]
982-
},
983-
type="vectorSearch",
917+
model = Article
918+
idx = VectorSearchIndex(
919+
fields=["plot_embedding", "number"],
920+
name="vector_index",
921+
similarities="cosine",
984922
)
923+
with connection.schema_editor() as editor:
924+
editor.add_index(model, idx)
925+
926+
def drop_index():
927+
with connection.schema_editor() as editor:
928+
editor.remove_index(model, idx)
929+
930+
cls.addClassCleanup(drop_index)
985931

986932
def setUp(self):
987933
self.mars = Article.objects.create(

0 commit comments

Comments
 (0)