diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py index 99c1ef5f..78c9e9ab 100644 --- a/django_mongodb_backend/indexes.py +++ b/django_mongodb_backend/indexes.py @@ -109,9 +109,25 @@ class SearchIndex(Index): suffix = "six" _error_id_prefix = "django_mongodb_backend.indexes.SearchIndex" - def __init__(self, *, fields=(), name=None): + def __init__(self, *, fields=(), name=None, field_mappings=None): + if field_mappings and not isinstance(field_mappings, dict): + raise ValueError( + "field_mappings must be a dictionary mapping field names to their " + "Atlas Search field mappings." + ) + self.field_mappings = field_mappings + if field_mappings: + if fields: + raise ValueError("Cannot provide fields and fields_mappings") + fields = [*self.field_mappings.keys()] super().__init__(fields=fields, name=name) + def deconstruct(self): + path, args, kwargs = super().deconstruct() + if self.field_mappings is not None: + kwargs["field_mappings"] = self.field_mappings + return path, args, kwargs + def check(self, model, connection): errors = [] if not connection.features.supports_atlas_search: @@ -152,15 +168,30 @@ def get_pymongo_index_model( return None fields = {} for field_name, _ in self.fields_orders: - field = model._meta.get_field(field_name) - type_ = self.search_index_data_types(field.db_type(schema_editor.connection)) field_path = column_prefix + model._meta.get_field(field_name).column - fields[field_path] = {"type": type_} + if self.field_mappings: + fields[field_path] = self.field_mappings[field_name] + else: + field = model._meta.get_field(field_name) + type_ = self.search_index_data_types(field.db_type(schema_editor.connection)) + fields[field_path] = {"type": type_} return SearchIndexModel( definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name ) +class DynamicSearchIndex(SearchIndex): + suffix = "dsix" + _error_id_prefix = "django_mongodb_backend.indexes.DynamicSearchIndex" + + def get_pymongo_index_model( + self, model, schema_editor, field=None, unique=False, column_prefix="" + ): + if not schema_editor.connection.features.supports_atlas_search: + return None + return SearchIndexModel(definition={"mappings": {"dynamic": True}}, name=self.name) + + class VectorSearchIndex(SearchIndex): suffix = "vsi" _error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex" diff --git a/tests/atlas_search_/test_search.py b/tests/atlas_search_/test_search.py index b1102a74..192c231d 100644 --- a/tests/atlas_search_/test_search.py +++ b/tests/atlas_search_/test_search.py @@ -8,7 +8,6 @@ from django.db.models.query import QuerySet from django.db.utils import DatabaseError from django.test import TransactionTestCase, skipUnlessDBFeature -from pymongo.operations import SearchIndexModel from django_mongodb_backend.expressions import ( CompoundExpression, @@ -28,7 +27,7 @@ SearchVector, SearchWildcard, ) -from django_mongodb_backend.schema import DatabaseSchemaEditor +from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex from .models import Article, Location, Writer @@ -75,22 +74,15 @@ class SearchUtilsMixin(TransactionTestCase): assertListEqual = _delayed_assertion(timeout=2)(TransactionTestCase.assertListEqual) assertQuerySetEqual = _delayed_assertion(timeout=2)(TransactionTestCase.assertQuerySetEqual) - @staticmethod - def _get_collection(model): - return connection.database.get_collection(model._meta.db_table) - @classmethod - def create_search_index(cls, model, index_name, definition, type="search"): - # TODO: create/delete indexes using DatabaseSchemaEditor when - # SearchIndexes support mappings (INTPYTHON-729). - collection = cls._get_collection(model) - idx = SearchIndexModel(definition=definition, name=index_name, type=type) - collection.create_search_index(idx) - DatabaseSchemaEditor.wait_until_index_created(collection, index_name) + def create_search_index(cls, model, index_name, definition, index_cls=SearchIndex): + idx = index_cls(field_mappings=definition, name=index_name) + with connection.schema_editor() as editor: + editor.add_index(model, idx) def drop_index(): - collection.drop_search_index(index_name) - DatabaseSchemaEditor.wait_until_index_dropped(collection, index_name) + with connection.schema_editor() as editor: + editor.remove_index(model, idx) cls.addClassCleanup(drop_index) @@ -101,12 +93,7 @@ def setUpClass(cls): cls.create_search_index( Article, "equals_headline_index", - { - "mappings": { - "dynamic": False, - "fields": {"headline": {"type": "token"}, "number": {"type": "number"}}, - } - }, + {"headline": {"type": "token"}, "number": {"type": "number"}}, ) def setUp(self): @@ -167,32 +154,27 @@ def setUpClass(cls): Article, "autocomplete_headline_index", { - "mappings": { - "dynamic": False, + "headline": { + "type": "autocomplete", + "analyzer": "lucene.standard", + "tokenization": "edgeGram", + "minGrams": 3, + "maxGrams": 5, + "foldDiacritics": False, + }, + "writer": { + "type": "document", "fields": { - "headline": { + "name": { "type": "autocomplete", "analyzer": "lucene.standard", "tokenization": "edgeGram", "minGrams": 3, "maxGrams": 5, "foldDiacritics": False, - }, - "writer": { - "type": "document", - "fields": { - "name": { - "type": "autocomplete", - "analyzer": "lucene.standard", - "tokenization": "edgeGram", - "minGrams": 3, - "maxGrams": 5, - "foldDiacritics": False, - } - }, - }, + } }, - } + }, }, ) @@ -253,7 +235,7 @@ def setUpClass(cls): cls.create_search_index( Article, "exists_body_index", - {"mappings": {"dynamic": False, "fields": {"body": {"type": "token"}}}}, + {"body": {"type": "token"}}, ) def setUp(self): @@ -282,7 +264,7 @@ def setUpClass(cls): cls.create_search_index( Article, "in_headline_index", - {"mappings": {"dynamic": False, "fields": {"headline": {"type": "token"}}}}, + {"headline": {"type": "token"}}, ) def setUp(self): @@ -316,7 +298,7 @@ def setUpClass(cls): cls.create_search_index( Article, "phrase_body_index", - {"mappings": {"dynamic": False, "fields": {"body": {"type": "string"}}}}, + {"body": {"type": "string"}}, ) def setUp(self): @@ -356,13 +338,8 @@ def setUpClass(cls): Article, "query_string_index", { - "mappings": { - "dynamic": False, - "fields": { - "headline": {"type": "string"}, - "body": {"type": "string"}, - }, - } + "headline": {"type": "string"}, + "body": {"type": "string"}, }, ) @@ -416,7 +393,7 @@ def setUpClass(cls): cls.create_search_index( Article, "range_number_index", - {"mappings": {"dynamic": False, "fields": {"number": {"type": "number"}}}}, + {"number": {"type": "number"}}, ) Article.objects.create(headline="x", number=5, body="z") @@ -453,12 +430,7 @@ def setUpClass(cls): cls.create_search_index( Article, "regex_headline_index", - { - "mappings": { - "dynamic": False, - "fields": {"headline": {"type": "string", "analyzer": "lucene.keyword"}}, - } - }, + {"headline": {"type": "string", "analyzer": "lucene.keyword"}}, ) def setUp(self): @@ -498,7 +470,7 @@ def setUpClass(cls): cls.create_search_index( Article, "text_body_index", - {"mappings": {"dynamic": False, "fields": {"body": {"type": "string"}}}}, + {"body": {"type": "string"}}, ) def setUp(self): @@ -560,12 +532,7 @@ def setUpClass(cls): cls.create_search_index( Article, "wildcard_headline_index", - { - "mappings": { - "dynamic": False, - "fields": {"headline": {"type": "string", "analyzer": "lucene.keyword"}}, - } - }, + {"headline": {"type": "string", "analyzer": "lucene.keyword"}}, ) def setUp(self): @@ -603,12 +570,7 @@ def setUpClass(cls): cls.create_search_index( Article, "geoshape_location_index", - { - "mappings": { - "dynamic": False, - "fields": {"location": {"type": "geo", "indexShapes": True}}, - } - }, + {"location": {"type": "geo", "indexShapes": True}}, ) def setUp(self): @@ -668,7 +630,7 @@ def setUpClass(cls): cls.create_search_index( Article, "geowithin_location_index", - {"mappings": {"dynamic": False, "fields": {"location": {"type": "geo"}}}}, + {"location": {"type": "geo"}}, ) def setUp(self): @@ -743,12 +705,7 @@ def setUpClass(cls): cls.create_search_index( Article, "mlt_index", - { - "mappings": { - "dynamic": False, - "fields": {"body": {"type": "string"}, "headline": {"type": "string"}}, - } - }, + {"body": {"type": "string"}, "headline": {"type": "string"}}, ) cls.article1 = Article.objects.create( headline="Space exploration", number=1, body="Webb telescope" @@ -782,14 +739,9 @@ def setUpClass(cls): Article, "compound_index", { - "mappings": { - "dynamic": False, - "fields": { - "headline": [{"type": "token"}, {"type": "string"}], - "body": {"type": "string"}, - "number": {"type": "number"}, - }, - } + "headline": [{"type": "token"}, {"type": "string"}], + "body": {"type": "string"}, + "number": {"type": "number"}, }, ) @@ -962,26 +914,20 @@ def test_str_returns_expected_format(self): class SearchVectorTests(SearchUtilsMixin): @classmethod def setUpClass(cls): - cls.create_search_index( - Article, - "vector_index", - { - "fields": [ - { - "type": "vector", - "path": "plot_embedding", - "numDimensions": 3, - "similarity": "cosine", - "quantization": "scalar", - }, - { - "type": "filter", - "path": "number", - }, - ] - }, - type="vectorSearch", + model = Article + idx = VectorSearchIndex( + fields=["plot_embedding", "number"], + name="vector_index", + similarities="cosine", ) + with connection.schema_editor() as editor: + editor.add_index(model, idx) + + def drop_index(): + with connection.schema_editor() as editor: + editor.remove_index(model, idx) + + cls.addClassCleanup(drop_index) def setUp(self): self.mars = Article.objects.create(