Coverage for mongodb_odm / utils / apply_indexes.py: 100%
182 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-07 08:37 +0000
1import logging 1fabcde
2from collections.abc import MutableMapping 1fabcde
3from typing import Any, Optional, Union 1fabcde
4from typing import cast as type_cast 1fabcde
6from bson import SON 1fabcde
7from mongodb_odm.connection import db 1fabcde
8from mongodb_odm.exceptions import ConnectionError, InvalidConnection 1fabcde
9from mongodb_odm.models import INHERITANCE_FIELD_NAME, Document 1fabcde
10from mongodb_odm.types import DICT_TYPE 1fabcde
11from pydantic import BaseModel 1fabcde
12from pymongo import ASCENDING, TEXT, IndexModel 1fabcde
13from pymongo.asynchronous.collection import AsyncCollection 1fabcde
14from pymongo.asynchronous.command_cursor import AsyncCommandCursor 1fabcde
15from pymongo.collection import Collection 1fabcde
16from pymongo.command_cursor import CommandCursor 1fabcde
18logger = logging.getLogger(__name__) 1fabcde
21class IndexOperation(BaseModel): 1fabcde
22 collection_name: str 1fabcde
23 model_indexes: list[Any] 1fabcde
24 database_name: Optional[str] = None 1fabcde
27def _get_dict_db_indexes(database_indexes: Any) -> list[DICT_TYPE]: 1fabcde
28 dict_db_indexes: list[DICT_TYPE] = [] 1fabcde
30 for index in database_indexes: 1fabcde
31 old_index = index.to_dict() 1fabcde
32 # Skip "_id" index since it's create by mongodb system
33 if "_id" in old_index["key"]: 1fabcde
34 continue 1fabcde
36 old_index.pop("v", None) 1fabcde
37 dict_db_indexes.append(old_index) 1fabcde
39 return dict_db_indexes 1fabcde
42def _get_dict_model_indexes( 1fabcde
43 model_indexes: list[IndexModel],
44) -> tuple[list[Union[DICT_TYPE, None]], dict[str, IndexModel]]:
45 dict_model_indexes: list[Union[DICT_TYPE, None]] = [] 1fabcde
46 model_indexes_map: dict[str, IndexModel] = {} 1fabcde
48 for index in model_indexes: 1fabcde
49 new_index: Any = index.document 1fabcde
50 # Replace SON object with dict
51 if isinstance(new_index["key"], SON): 1fabcde
52 new_index["key"] = new_index["key"].to_dict() 1fabcde
53 if isinstance(new_index["key"], dict): 1fabcde
54 new_index["key"] = new_index["key"] 1fabcde
55 else:
56 continue 1abcde
57 dict_model_indexes.append(new_index) 1fabcde
58 # Store index object for future use
59 model_indexes_map[new_index["name"]] = index 1fabcde
61 return dict_model_indexes, model_indexes_map 1fabcde
64def _check_exists_in_model_indexes( 1fabcde
65 db_dict_index: DICT_TYPE, dict_model_indexes: list[Union[DICT_TYPE, None]]
66) -> bool:
67 def mark_as_exist(position: int) -> bool: 1fabcde
68 """
69 Since the index already exists in the model indexes,
70 we will mark it as None.
71 """
72 dict_model_indexes[position] = None 1fabcde
74 return True 1fabcde
76 # Iterate over the indexes that are currently assigned in the model
77 for idx, index_obj in enumerate(dict_model_indexes): 1fabcde
78 if not index_obj or not isinstance(index_obj, dict): 1fabcde
79 continue 1abcde
81 if db_dict_index == index_obj: 1fabcde
82 mark_as_exist(idx) 1fabcde
83 return True 1fabcde
85 if db_dict_index.get("name") == index_obj.get("name"): 1fabcde
86 # This condition will handle some of the special indexes like text-based indexes.
87 db_value, new_value = db_dict_index, index_obj 1fabcde
88 if TEXT in db_value["key"].values() and TEXT in new_value["key"].values(): 1fabcde
89 """ 1abcde
90 based on condition this is a text-based index.
91 Check that if value is changed.
92 If not changed then assign null otherwise not.
93 NOTE:
94 The data check may not be sufficient enough.
95 Add conditions based on the issue.
96 """
98 new_keys = new_value["key"].keys() 1fabcde
99 default_language = new_value.get("default_language", "english") 1fabcde
100 new_weight = new_value.get("weights") or dict.fromkeys(new_keys, 1) 1fabcde
102 if ( 1f
103 db_value["weights"].keys() == new_keys
104 and new_weight == db_value["weights"]
105 and db_value["default_language"] == default_language
106 ):
107 """Check all key, weights, and default_language match with existing values""" 1abcde
108 mark_as_exist(idx) 1fabcde
109 return True 1fabcde
111 return False 1fabcde
114def _get_calculated_indexes_for_collection( 1fabcde
115 database_indexes: Any,
116 model_indexes: list[IndexModel],
117) -> tuple[list[IndexModel], list[DICT_TYPE]]:
118 dict_db_indexes = _get_dict_db_indexes(database_indexes) 1fabcde
119 dict_model_indexes, model_indexes_map = _get_dict_model_indexes(model_indexes) 1fabcde
121 delete_db_indexes: list[DICT_TYPE] = [] 1fabcde
123 # Iterate over indexes that are already created for a collection
124 for db_dict_index in dict_db_indexes: 1fabcde
125 is_matched = _check_exists_in_model_indexes(db_dict_index, dict_model_indexes) 1fabcde
127 if not is_matched: 1fabcde
128 # If the index does not exist in model indexes then delete it.
129 delete_db_indexes.append(db_dict_index) 1fabcde
131 new_indexes = [ 1fabcde
132 model_indexes_map[new_index["name"]]
133 for new_index in dict_model_indexes
134 if new_index
135 ]
137 return new_indexes, delete_db_indexes 1fabcde
140def _get_model_indexes(model: type[Document]) -> list[IndexModel]: 1fabcde
141 # Get define indexes for a model
142 if hasattr(model.ODMConfig, "indexes"): 1fabcde
143 return list(model.ODMConfig.indexes) 1fabcde
144 return [] 1fabcde
147def _get_all_indexes() -> list[IndexOperation]: 1fabcde
148 """
149 First imports all child models of Document since it's the abstract parent model.
150 Then retrieve all the child modules and will try to get indexes inside the ODMConfig class.
151 """
152 operations: list[IndexOperation] = [] 1fabcde
154 def get_operation_obj(model: type[Document]) -> IndexOperation: 1fabcde
155 return IndexOperation( 1fabcde
156 collection_name=model._get_collection_name(),
157 model_indexes=_get_model_indexes(model),
158 database_name=model._database_name(),
159 )
161 for model in Document.__subclasses__(): 1fabcde
162 obj = get_operation_obj(model) 1fabcde
163 if ( 1fab
164 hasattr(model.ODMConfig, "allow_inheritance")
165 and model.ODMConfig.allow_inheritance is True
166 ):
167 """If a model has child model""" 1abcde
168 if model.ODMConfig.index_inheritance_field is True: 1fabcde
169 """ 1abcde
170 If index_inheritance_field is true then create an index '_cls'
171 that will store the name of the child collection name.
172 """
173 obj.model_indexes.append( 1fabcde
174 IndexModel([(INHERITANCE_FIELD_NAME, ASCENDING)])
175 )
176 for child_model in model.__subclasses__(): 1fabcde
177 """Get all indexes that are defined in child model""" 1abcde
178 obj.model_indexes += _get_model_indexes(child_model) 1fabcde
180 if obj and obj.model_indexes: 1fabcde
181 operations.append(obj) 1fabcde
183 return operations 1fabcde
186def _get_collection( 1fabcde
187 operation: IndexOperation,
188) -> Union[Collection[Any], AsyncCollection[Any]]:
189 return db(operation.database_name)[operation.collection_name] 1fabcde
192def _get_created_and_deleted_indexes_count( 1fabcde
193 operation: IndexOperation,
194 new_indexes: list[IndexModel],
195 delete_db_indexes: list[DICT_TYPE],
196) -> tuple[int, int]:
197 ne, de = len(new_indexes), len(delete_db_indexes) 1fabcde
198 if ne > 0 or de > 0: 1fabcde
199 logger.info( 1fabcde
200 f'Applied for "{operation.collection_name}": {de} deleted, {ne} added'
201 )
203 return ne, de 1fabcde
206def _get_collection_and_indexes( 1fabcde
207 operation: IndexOperation,
208) -> tuple[Union[Collection[Any], AsyncCollection[Any]], list[IndexModel]]:
209 try: 1fabcde
210 collection = _get_collection(operation) 1fabcde
211 model_indexes = operation.model_indexes 1fabcde
212 except Exception as e: 1fabcde
213 raise InvalidConnection( 1fabcde
214 f"Invalid database:'{operation.database_name}' and collection:{operation.collection_name}"
215 ) from e
217 return collection, model_indexes 1fabcde
220def _sync_get_database_indexes( 1fabcde
221 collection: Union[Collection[Any], AsyncCollection[Any]],
222) -> CommandCursor[MutableMapping[str, Any]]:
223 if isinstance(collection, Collection): 1fabcde
224 return collection.list_indexes() 1fabcde
226 if isinstance(collection, AsyncCollection): 1fabcde
227 raise ConnectionError("Use synchronous collection for indexes.") 1fabcde
230def _sync_apply_indexes_to_db( 1fabcde
231 operation: IndexOperation,
232 new_indexes: list[IndexModel],
233 delete_db_indexes: list[DICT_TYPE],
234) -> tuple[int, int]:
235 _collection = _get_collection(operation) 1fabcde
236 collection = type_cast(Collection[Any], _collection) 1fabcde
238 for db_index in delete_db_indexes: 1fabcde
239 # If the DB index does not exist in new_indexes then drop that index.
240 if db_index is not None: 1fabcde
241 collection.drop_index(db_index["name"]) 1fabcde
243 if len(new_indexes) > 0: 1fabcde
244 try: 1fabcde
245 collection.create_indexes(new_indexes) 1fabcde
246 except Exception as e: 1fabcde
247 logger.error(f'\nProblem arise at "{operation.collection_name}": {e}\n') 1fabcde
248 raise e 1fabcde
250 return _get_created_and_deleted_indexes_count( 1fabcde
251 operation, new_indexes, delete_db_indexes
252 )
255def _sync_apply_indexes_for_a_collection(operation: IndexOperation) -> tuple[int, int]: 1fabcde
256 collection, model_indexes = _get_collection_and_indexes(operation) 1fabcde
257 database_indexes = _sync_get_database_indexes(collection) 1fabcde
258 new_indexes, delete_db_indexes = _get_calculated_indexes_for_collection( 1fabcde
259 database_indexes, model_indexes
260 )
261 ne, de = _sync_apply_indexes_to_db(operation, new_indexes, delete_db_indexes) 1fabcde
263 return ne, de 1fabcde
266async def _async_get_database_indexes( 1fabcde
267 collection: Union[Collection[Any], AsyncCollection[Any]],
268) -> AsyncCommandCursor[MutableMapping[str, Any]]:
269 if isinstance(collection, AsyncCollection): 1fabcde
270 return await collection.list_indexes() 1fabcde
272 if isinstance(collection, Collection): 1fabcde
273 raise ConnectionError("Use asynchronous collection for indexes.") 1fabcde
276async def _async_apply_indexes_to_db( 1fabcde
277 operation: IndexOperation,
278 new_indexes: list[IndexModel],
279 delete_db_indexes: list[DICT_TYPE],
280) -> tuple[int, int]:
281 _collection = _get_collection(operation) 1fabcde
282 collection = type_cast(AsyncCollection[Any], _collection) 1fabcde
284 for db_index in delete_db_indexes: 1fabcde
285 # If the DB index does not exist in new_indexes then drop that index.
286 if db_index is not None: 1fabcde
287 await collection.drop_index(db_index["name"]) 1fabcde
289 if len(new_indexes) > 0: 1fabcde
290 try: 1fabcde
291 await collection.create_indexes(new_indexes) 1fabcde
292 except Exception as e: 1fabcde
293 logger.error(f'\nProblem arise at "{operation.collection_name}": {e}\n') 1fabcde
294 raise e 1fabcde
296 return _get_created_and_deleted_indexes_count( 1fabcde
297 operation, new_indexes, delete_db_indexes
298 )
301async def _async_apply_indexes_for_a_collection( 1fabcde
302 operation: IndexOperation,
303) -> tuple[int, int]:
304 collection, model_indexes = _get_collection_and_indexes(operation) 1fabcde
305 database_indexes = [ 1fabcde
306 obj async for obj in await _async_get_database_indexes(collection)
307 ]
308 new_indexes, delete_db_indexes = _get_calculated_indexes_for_collection( 1fabcde
309 database_indexes, model_indexes
310 )
311 ne, de = await _async_apply_indexes_to_db(operation, new_indexes, delete_db_indexes) 1fabcde
313 return ne, de 1fabcde
316def log_final_results(new_index: int, delete_index: int) -> None: 1fabcde
317 if delete_index: 1fabcde
318 logger.info(f"{delete_index}, index deleted.") 1fabcde
319 if new_index: 1fabcde
320 logger.info(f"{new_index}, index created.") 1fabcde
321 if [new_index, delete_index] == [0, 0]: 1fabcde
322 logger.info("No change detected.") 1fabcde
325def apply_indexes() -> None: 1fabcde
326 """First get all indexes from all model."""
327 operations = _get_all_indexes() 1fabcde
329 """Then execute each indexes operation for each model.""" 1abcde
330 new_index, delete_index = 0, 0 1fabcde
331 for operation in operations: 1fabcde
332 ne, de = _sync_apply_indexes_for_a_collection(operation) 1fabcde
333 new_index += ne 1fabcde
334 delete_index += de 1fabcde
336 log_final_results(new_index, delete_index) 1fabcde
339async def async_apply_indexes() -> None: 1fabcde
340 """First get all indexes from all model."""
341 operations = _get_all_indexes() 1fabcde
343 """Then execute each indexes operation for each model.""" 1abcde
344 new_index, delete_index = 0, 0 1fabcde
345 for operation in operations: 1fabcde
346 ne, de = await _async_apply_indexes_for_a_collection(operation) 1fabcde
347 new_index += ne 1fabcde
348 delete_index += de 1fabcde
350 log_final_results(new_index, delete_index) 1fabcde