fixup! Implement index definition and creation

art049 · Sep 7, 2022 · eeb0596 · eeb0596
1 parent 3e3d41e
commit eeb0596
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 9 deletions.
diff --git a/odmantic/config.py b/odmantic/config.py
@@ -3,6 +3,7 @@
 import json
 from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Optional, Type, Union
 
+import pymongo
 from pydantic.main import BaseConfig
 from pydantic.typing import AnyCallable
 
@@ -30,7 +31,7 @@ class BaseODMConfig:
     parse_doc_with_default_factories: bool = False
 
     @staticmethod
-    def indexes() -> Iterable[ODMIndex.Index]:
+    def indexes() -> Iterable[Union[ODMIndex.Index, pymongo.IndexModel]]:
         return []
 
     # Inherited from pydantic

diff --git a/odmantic/engine.py b/odmantic/engine.py
@@ -28,6 +28,7 @@
 
 from odmantic.exceptions import DocumentNotFoundError, DuplicateKeyError
 from odmantic.field import FieldProxy, ODMReference
+from odmantic.index import ODMBaseIndex
 from odmantic.model import Model
 from odmantic.query import QueryExpression, SortExpression, and_
 from odmantic.session import (
@@ -366,7 +367,11 @@ async def configure_database(
         for model in models:
             collection = self.get_collection(model)
             for index in model.__indexes__():
-                pymongo_index = index.get_pymongo_index()
+                pymongo_index = (
+                    index.get_pymongo_index()
+                    if isinstance(index, ODMBaseIndex)
+                    else index
+                )
                 try:
                     await collection.create_indexes(
                         [pymongo_index], session=driver_session
@@ -377,7 +382,7 @@ async def configure_database(
                         86,  # aka IndexKeySpecsConflict for MongoDB > 5
                     ):
                         await collection.drop_index(
-                            index.get_index_specifier(), session=driver_session
+                            pymongo_index.document["name"], session=driver_session
                         )
                         await collection.create_indexes(
                             [pymongo_index], session=driver_session
@@ -776,7 +781,11 @@ def configure_database(
         for model in models:
             collection = self.get_collection(model)
             for index in model.__indexes__():
-                pymongo_index = index.get_pymongo_index()
+                pymongo_index = (
+                    index.get_pymongo_index()
+                    if isinstance(index, ODMBaseIndex)
+                    else index
+                )
                 try:
                     collection.create_indexes([pymongo_index], session=driver_session)
                 except pymongo.errors.OperationFailure as exc:
@@ -785,7 +794,7 @@ def configure_database(
                         86,  # aka IndexKeySpecsConflict for MongoDB > 5
                     ):
                         collection.drop_index(
-                            index.get_index_specifier(), session=driver_session
+                            pymongo_index.document["name"], session=driver_session
                         )
                         collection.create_indexes(
                             [pymongo_index], session=driver_session

diff --git a/odmantic/model.py b/odmantic/model.py
@@ -28,6 +28,7 @@
 
 import bson
 import pydantic
+import pymongo
 from pydantic.error_wrappers import ErrorWrapper, ValidationError
 from pydantic.fields import Field as PDField
 from pydantic.fields import FieldInfo as PDFieldInfo
@@ -61,7 +62,7 @@
     ODMFieldInfo,
     ODMReference,
 )
-from odmantic.index import ODMBaseIndex, ODMSingleFieldIndex
+from odmantic.index import Index, ODMBaseIndex, ODMSingleFieldIndex
 from odmantic.reference import ODMReferenceInfo
 from odmantic.utils import (
     is_dunder,
@@ -780,8 +781,8 @@ def __setattr__(self, name: str, value: Any) -> None:
         super().__setattr__(name, value)
 
     @classmethod
-    def __indexes__(cls) -> Tuple[ODMBaseIndex, ...]:
-        indexes: List[ODMBaseIndex] = []
+    def __indexes__(cls) -> Tuple[Union[ODMBaseIndex, pymongo.IndexModel], ...]:
+        indexes: List[Union[ODMBaseIndex, pymongo.IndexModel]] = []
         for field in cls.__odm_fields__.values():
             if isinstance(field, ODMBaseIndexableField) and (
                 field.index or field.unique
@@ -794,7 +795,7 @@ def __indexes__(cls) -> Tuple[ODMBaseIndex, ...]:
                 )
 
         for index in cast(BaseODMConfig, cls.Config).indexes():
-            indexes.append(index.to_odm_index())
+            indexes.append(index.to_odm_index() if isinstance(index, Index) else index)
         return tuple(indexes)
 
     def update(

diff --git a/tests/integration/test_index.py b/tests/integration/test_index.py
@@ -1,3 +1,4 @@
+import pymongo
 import pytest
 
 from odmantic.engine import AIOEngine, SyncEngine
@@ -310,3 +311,37 @@ class Config:
         )
         is not None
     )
+
+
+async def test_custom_text_index(aio_engine: AIOEngine):
+    class Post(Model):
+        title: str
+        content: str
+
+        class Config:
+            @staticmethod
+            def indexes():
+                yield pymongo.IndexModel(
+                    [("title", pymongo.TEXT), ("content", pymongo.TEXT)]
+                )
+
+    await aio_engine.configure_database([Post])
+    await aio_engine.save(Post(title="My post on python", content="It's awesome!"))
+    assert await aio_engine.find_one(Post, {"$text": {"$search": "python"}}) is not None
+
+
+async def test_sync_custom_text_index(sync_engine: SyncEngine):
+    class Post(Model):
+        title: str
+        content: str
+
+        class Config:
+            @staticmethod
+            def indexes():
+                yield pymongo.IndexModel(
+                    [("title", pymongo.TEXT), ("content", pymongo.TEXT)]
+                )
+
+    sync_engine.configure_database([Post])
+    sync_engine.save(Post(title="My post on python", content="It's awesome!"))
+    assert sync_engine.find_one(Post, {"$text": {"$search": "python"}}) is not None