diff --git a/pymongo/collection.py b/pymongo/collection.py index b95bfb832..958f7ec21 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -2410,7 +2410,7 @@ class Collection(common.BaseObject, Generic[_DocumentType]): .. versionadded:: 4.5 """ if not isinstance(model, SearchIndexModel): - model = SearchIndexModel(model["definition"], model.get("name")) + model = SearchIndexModel(model["definition"], model.get("name"), model.get("type")) return self.create_search_indexes([model], session, comment, **kwargs)[0] def create_search_indexes( diff --git a/pymongo/operations.py b/pymongo/operations.py index a586fa508..0240b6e17 100644 --- a/pymongo/operations.py +++ b/pymongo/operations.py @@ -588,7 +588,12 @@ class SearchIndexModel: __slots__ = ("__document",) - def __init__(self, definition: Mapping[str, Any], name: Optional[str] = None) -> None: + def __init__( + self, + definition: Mapping[str, Any], + name: Optional[str] = None, + type: Optional[str] = "search", + ) -> None: """Create a Search Index instance. For use with :meth:`~pymongo.collection.Collection.create_search_index` and :meth:`~pymongo.collection.Collection.create_search_indexes`. @@ -604,6 +609,7 @@ class SearchIndexModel: self.__document = dict(name=name, definition=definition) else: self.__document = dict(definition=definition) + self.__document["type"] = type # type: ignore[assignment] @property def document(self) -> Mapping[str, Any]: diff --git a/test/index_management/createSearchIndex.json b/test/index_management/createSearchIndex.json index f9c4e44d3..31f4c3fdf 100644 --- a/test/index_management/createSearchIndex.json +++ b/test/index_management/createSearchIndex.json @@ -40,52 +40,6 @@ "tests": [ { "description": "no name provided for an index definition", - "operations": [ - { - "name": "createSearchIndex", - "object": "collection0", - "arguments": { - "model": { - "definition": { - "mappings": { - "dynamic": true - } - } - } - }, - "expectError": { - "isError": true, - "errorContains": "Atlas" - } - } - ], - "expectEvents": [ - { - "client": "client0", - "events": [ - { - "commandStartedEvent": { - "command": { - "createSearchIndexes": "collection0", - "indexes": [ - { - "definition": { - "mappings": { - "dynamic": true - } - } - } - ], - "$db": "database0" - } - } - } - ] - } - ] - }, - { - "description": "name provided for an index definition", "operations": [ { "name": "createSearchIndex", @@ -97,7 +51,7 @@ "dynamic": true } }, - "name": "test index" + "type": "search" } }, "expectError": { @@ -121,7 +75,117 @@ "dynamic": true } }, - "name": "test index" + "type": "search" + } + ], + "$db": "database0" + } + } + } + ] + } + ] + }, + { + "description": "name provided for an index definition", + "operations": [ + { + "name": "createSearchIndex", + "object": "collection0", + "arguments": { + "model": { + "definition": { + "mappings": { + "dynamic": true + } + }, + "name": "test index", + "type": "search" + } + }, + "expectError": { + "isError": true, + "errorContains": "Atlas" + } + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "createSearchIndexes": "collection0", + "indexes": [ + { + "definition": { + "mappings": { + "dynamic": true + } + }, + "name": "test index", + "type": "search" + } + ], + "$db": "database0" + } + } + } + ] + } + ] + }, + { + "description": "create a vector search index", + "operations": [ + { + "name": "createSearchIndex", + "object": "collection0", + "arguments": { + "model": { + "definition": { + "fields": [ + { + "type": "vector", + "path": "plot_embedding", + "numDimensions": 1536, + "similarity": "euclidean" + } + ] + }, + "name": "test index", + "type": "vectorSearch" + } + }, + "expectError": { + "isError": true, + "errorContains": "Atlas" + } + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "createSearchIndexes": "collection0", + "indexes": [ + { + "definition": { + "fields": [ + { + "type": "vector", + "path": "plot_embedding", + "numDimensions": 1536, + "similarity": "euclidean" + } + ] + }, + "name": "test index", + "type": "vectorSearch" } ], "$db": "database0" diff --git a/test/index_management/createSearchIndexes.json b/test/index_management/createSearchIndexes.json index 3cf56ce12..be2d02cfc 100644 --- a/test/index_management/createSearchIndexes.json +++ b/test/index_management/createSearchIndexes.json @@ -72,54 +72,6 @@ }, { "description": "no name provided for an index definition", - "operations": [ - { - "name": "createSearchIndexes", - "object": "collection0", - "arguments": { - "models": [ - { - "definition": { - "mappings": { - "dynamic": true - } - } - } - ] - }, - "expectError": { - "isError": true, - "errorContains": "Atlas" - } - } - ], - "expectEvents": [ - { - "client": "client0", - "events": [ - { - "commandStartedEvent": { - "command": { - "createSearchIndexes": "collection0", - "indexes": [ - { - "definition": { - "mappings": { - "dynamic": true - } - } - } - ], - "$db": "database0" - } - } - } - ] - } - ] - }, - { - "description": "name provided for an index definition", "operations": [ { "name": "createSearchIndexes", @@ -132,7 +84,7 @@ "dynamic": true } }, - "name": "test index" + "type": "search" } ] }, @@ -157,7 +109,121 @@ "dynamic": true } }, - "name": "test index" + "type": "search" + } + ], + "$db": "database0" + } + } + } + ] + } + ] + }, + { + "description": "name provided for an index definition", + "operations": [ + { + "name": "createSearchIndexes", + "object": "collection0", + "arguments": { + "models": [ + { + "definition": { + "mappings": { + "dynamic": true + } + }, + "name": "test index", + "type": "search" + } + ] + }, + "expectError": { + "isError": true, + "errorContains": "Atlas" + } + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "createSearchIndexes": "collection0", + "indexes": [ + { + "definition": { + "mappings": { + "dynamic": true + } + }, + "name": "test index", + "type": "search" + } + ], + "$db": "database0" + } + } + } + ] + } + ] + }, + { + "description": "create a vector search index", + "operations": [ + { + "name": "createSearchIndexes", + "object": "collection0", + "arguments": { + "models": [ + { + "definition": { + "fields": [ + { + "type": "vector", + "path": "plot_embedding", + "numDimensions": 1536, + "similarity": "euclidean" + } + ] + }, + "name": "test index", + "type": "vectorSearch" + } + ] + }, + "expectError": { + "isError": true, + "errorContains": "Atlas" + } + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "createSearchIndexes": "collection0", + "indexes": [ + { + "definition": { + "fields": [ + { + "type": "vector", + "path": "plot_embedding", + "numDimensions": 1536, + "similarity": "euclidean" + } + ] + }, + "name": "test index", + "type": "vectorSearch" } ], "$db": "database0" diff --git a/test/test_index_management.py b/test/test_index_management.py index 8f60b5d8d..b6d5c4775 100644 --- a/test/test_index_management.py +++ b/test/test_index_management.py @@ -217,6 +217,70 @@ class TestSearchIndexProse(unittest.TestCase): # Run a ``dropSearchIndex`` command and assert that no error is thrown. coll0.drop_search_index("foo") + def test_case_7(self): + """Driver handles index types.""" + + # Create a collection with the "create" command using a randomly generated name (referred to as ``coll0``). + coll0 = self.db[f"col{uuid.uuid4()}"] + coll0.insert_one({}) + + # Use these search and vector search definitions for indexes. + search_definition = {"mappings": {"dynamic": False}} + vector_search_definition = { + "fields": [ + { + "type": "vector", + "path": "plot_embedding", + "numDimensions": 1536, + "similarity": "euclidean", + }, + ] + } + + # Create a new search index on ``coll0`` that implicitly passes its type. + implicit_search_resp = coll0.create_search_index( + model={"name": _NAME + "-implicit", "definition": search_definition} + ) + + # Get the index definition. + resp = coll0.list_search_indexes(name=implicit_search_resp).next() + + # Assert that the index model contains the correct index type: ``"search"``. + self.assertEqual(resp["type"], "search") + + # Create a new search index on ``coll0`` that explicitly passes its type. + explicit_search_resp = coll0.create_search_index( + model={"name": _NAME + "-explicit", "type": "search", "definition": search_definition} + ) + + # Get the index definition. + resp = coll0.list_search_indexes(name=explicit_search_resp).next() + + # Assert that the index model contains the correct index type: ``"search"``. + self.assertEqual(resp["type"], "search") + + # Create a new vector search index on ``coll0`` that explicitly passes its type. + explicit_vector_resp = coll0.create_search_index( + model={ + "name": _NAME + "-vector", + "type": "vectorSearch", + "definition": vector_search_definition, + } + ) + + # Get the index definition. + resp = coll0.list_search_indexes(name=explicit_vector_resp).next() + + # Assert that the index model contains the correct index type: ``"vectorSearch"``. + self.assertEqual(resp["type"], "vectorSearch") + + # Catch the error raised when trying to create a vector search index without specifying the type + with self.assertRaises(OperationFailure) as e: + coll0.create_search_index( + model={"name": _NAME + "-error", "definition": vector_search_definition} + ) + self.assertIn("Attribute mappings missing.", e.exception.details["errmsg"]) + globals().update( generate_test_classes(