Coverage for mongodb_odm / models.py: 99%

1from collections.abc import AsyncIterator, Iterator, Sequence 6 ctx1fabcde

2from datetime import datetime 6 ctx1fabcde

3from typing import ( 6 ctx1fabcde

4 Any, 

5 Optional, 

6 Union, 

7 cast, 

8) 

9 

10from mongodb_odm.connection import db, get_client 6 ctx1fabcde

11from mongodb_odm.data_conversion import dict2obj 6 ctx1fabcde

12from mongodb_odm.exceptions import InvalidConfiguration, ObjectDoesNotExist 6 ctx1fabcde

13from mongodb_odm.fields import Field 6 ctx1fabcde

14from mongodb_odm.types import DICT_TYPE, SORT_TYPE, ODMObjectId, WriteOp 6 ctx1fabcde

15from mongodb_odm.utils._internal_models import CollectionConfig, RelationalFieldInfo 6 ctx1fabcde

16from mongodb_odm.utils.utils import ( 6 ctx1fabcde

17 convert_model_to_collection, 

18 get_database_name, 

19 get_model_fields, 

20 get_relationship_fields_info, 

21) 

22from mongodb_odm.utils.validation import validate_filter_dict 6 ctx1fabcde

23from pydantic import BaseModel, PrivateAttr 6 ctx1fabcde

24from pydantic._internal._model_construction import ModelMetaclass 6 ctx1fabcde

25from pymongo import AsyncMongoClient, IndexModel, MongoClient 6 ctx1fabcde

26from pymongo.asynchronous.client_session import AsyncClientSession 6 ctx1fabcde

27from pymongo.asynchronous.collection import AsyncCollection 6 ctx1fabcde

28from pymongo.asynchronous.cursor import AsyncCursor 6 ctx1fabcde

29from pymongo.client_session import ClientSession 6 ctx1fabcde

30from pymongo.collection import Collection 6 ctx1fabcde

31from pymongo.cursor import Cursor 6 ctx1fabcde

32from pymongo.results import BulkWriteResult, DeleteResult, UpdateResult 6 ctx1fabcde

33from typing_extensions import Self 6 ctx1fabcde

34 

35INHERITANCE_FIELD_NAME = "_cls" 6 ctx1fabcde

36 

37RELATION_TYPE = dict[str, RelationalFieldInfo] 6 ctx1fabcde

38 

39_cashed_collection: dict[Any, CollectionConfig] = {} 6 ctx1fabcde

40_cashed_field_info: dict[str, RELATION_TYPE] = {} 6 ctx1fabcde

41 

42 

43def _clear_cache() -> None: 6 ctx1fabcde

44 global _cashed_collection, _cashed_field_info 

45 for key in list(_cashed_collection.keys()): 6 ctx1fabcde

46 del _cashed_collection[key] 6 ctx1fabcde

47 

48 for key in list(_cashed_field_info.keys()): 6 ctx1fabcde

49 del _cashed_field_info[key] 6 ctx1fabcde

50 

51 

52class ODMMeta(ModelMetaclass): 6 ctx1fabcde

53 def __getattr__(cls, name: str) -> str: 6 ctx1fabcde

54 # Avoid recursion for internal pydantic attributes 

55 if name.startswith("__pydantic"): 6 ctx1fabcde

56 raise AttributeError(name) 6 ctx1fabcde

57 

58 try: 6 ctx1fabcde

59 # Try to access model_fields without triggering __getattr__ for it 

60 # In Pydantic v2, fields are stored in __pydantic_fields__ 

61 fields = cls.__dict__.get("__pydantic_fields__") 6 ctx1fabcde

62 

63 if fields and name in fields: 6 ctx1fabcde

64 if name == "id": 6 ctx1fabcde

65 return "_id" 6 ctx1fabcde

66 return name 6 ctx1fabcde

67 except Exception: 

68 pass 

69 

70 raise AttributeError(f"type object '{cls.__name__}' has no attribute '{name}'") 6 ctx1fabcde

71 

72 

73class _BaseDocument(BaseModel, metaclass=ODMMeta): 6 ctx1fabcde

74 """ 

75 This class will handle all database and model-related configuration. 

76 """ 

77 

78 class ODMConfig: 6 ctx1fabcde

79 # Those fields will work as the default value of any child class. 

80 collection_name: Optional[str] = None 6 ctx1fabcde

81 allow_inheritance: bool = False 6 ctx1fabcde

82 index_inheritance_field: bool = True 6 ctx1fabcde

83 indexes: list[IndexModel] = [] 6 ctx1fabcde

84 database: Optional[str] = None 6 ctx1fabcde

85 

86 """ 6 ctx1fabcde

87 Definition of ODMConfig fields: 

88 

89 collection_name: This field will overwrite collection name. 

90 

91 allow_inheritance: It enables a model to have a child. 

92 We will manage a single field for parent and child collection. 

93 The collection name will be according to the parent model configuration. 

94 Child collection data will be separated by field name '_cls'. 

95 

96 index_inheritance_field: By default, we create an index for '_cls' 

97 if the model has a child. Users may control that by using this field. 

98 

99 indexes: Define list of IndexModel to manage indexes for a model 

100 

101 database: Handle multiple database configurations using this field. 

102 The default database will be None. 

103 """ 

104 

105 def __setattr__(self, key: str, value: Any) -> None: 6 ctx1fabcde

106 """ 

107 Add '# type: ignore' as a comment if get type error while getting this value 

108 """ 

109 if key == "_id" or key == "id": 6 ctx1fabcde

110 self.__dict__["id"] = value 6 ctx1fabcde

111 self.__dict__["_id"] = value 6 ctx1fabcde

112 else: 

113 self.__dict__[key] = value 6 ctx1fabcde

114 

115 def dict(self, **kwargs: Any) -> DICT_TYPE: 6 ctx1fabcde

116 return self.model_dump(**kwargs) 6 ctx1fabcde

117 

118 @classmethod 6 ctx1fabcde

119 def __get_collection_class(cls) -> tuple[str, Optional[str]]: 6 ctx1fabcde

120 """ 

121 Get model class. 

122 if called from a child class: 

123 return (parent-class, child-class) 

124 else: 

125 return (class, None) 

126 """ 

127 model: Any = cls 6 ctx1fabcde

128 if model.__base__ != Document: 6 ctx1fabcde

129 base_model = model.__base__ 6 ctx1fabcde

130 if ( 3 ctx1fab

131 not hasattr(base_model.ODMConfig, "allow_inheritance") 

132 or base_model.ODMConfig.allow_inheritance is not True 

133 ): 

134 raise InvalidConfiguration( 6 ctx1fabcde

135 f"Invalid model inheritance. {base_model} does not allow model inheritance." 

136 ) 

137 if ( 1 ctx1f

138 base_model.ODMConfig.allow_inheritance is True 

139 and model.ODMConfig.allow_inheritance is True 

140 ): 

141 raise InvalidConfiguration( 6 ctx1fabcde

142 f"Child Model{model.__name__} should declare a separate ODMConfig class." 

143 ) 

144 return base_model, model 6 ctx1fabcde

145 else: 

146 return model, None 6 ctx1fabcde

147 

148 @classmethod 6 ctx1fabcde

149 def __get_collection_config(cls) -> CollectionConfig: 6 ctx1fabcde

150 """ 

151 Get collection configuration for a model. 

152 Get data from the cache if it is already calculated. 

153 """ 

154 global _cashed_collection 

155 if cls in _cashed_collection: 6 ctx1fabcde

156 return _cashed_collection[cls] 6 ctx1fabcde

157 

158 model, child_model = cls.__get_collection_class() 6 ctx1fabcde

159 

160 has_children = False 6 ctx1fabcde

161 if ( 3 ctx1fab

162 hasattr(cls.ODMConfig, "allow_inheritance") 

163 and cls.ODMConfig.allow_inheritance is True 

164 ): 

165 """Check if this is a model that allows inheritance and has a child model.""" 5 ctx1abcde

166 has_children = len(cls.__subclasses__()) > 0 6 ctx1fabcde

167 

168 child_collection_name = ( 6 ctx1fabcde

169 convert_model_to_collection(child_model) if child_model else None 

170 ) 

171 _cashed_collection[cls] = CollectionConfig( 6 ctx1fabcde

172 collection_name=convert_model_to_collection(model), 

173 child_collection_name=child_collection_name, 

174 database_name=get_database_name(model), 

175 has_children=has_children, 

176 ) 

177 return _cashed_collection[cls] 6 ctx1fabcde

178 

179 @classmethod 6 ctx1fabcde

180 def _database_name(cls) -> Optional[str]: 6 ctx1fabcde

181 config = cls.__get_collection_config() 6 ctx1fabcde

182 return config.database_name 6 ctx1fabcde

183 

184 @classmethod 6 ctx1fabcde

185 def _get_collection_name(cls) -> str: 6 ctx1fabcde

186 return cls.__get_collection_config().collection_name 6 ctx1fabcde

187 

188 @classmethod 6 ctx1fabcde

189 def _get_child(cls) -> Optional[str]: 6 ctx1fabcde

190 """ 

191 Get the child collection name if it has a parent class. 

192 """ 

193 return cls.__get_collection_config().child_collection_name 6 ctx1fabcde

194 

195 @classmethod 6 ctx1fabcde

196 def _has_children(cls) -> bool: 6 ctx1fabcde

197 """ 

198 Check if a model has child class 

199 """ 

200 return cls.__get_collection_config().has_children 6 ctx1fabcde

201 

202 @classmethod 6 ctx1fabcde

203 def _get_collection(cls) -> Collection[Any]: 6 ctx1fabcde

204 db_connection = db(cls._database_name()) 6 ctx1fabcde

205 collection = db_connection[cls._get_collection_name()] 6 ctx1fabcde

206 

207 return cast(Collection[Any], collection) 6 ctx1fabcde

208 

209 @classmethod 6 ctx1fabcde

210 def _async_get_collection(cls) -> AsyncCollection[Any]: 6 ctx1fabcde

211 db_connection = db(cls._database_name(), is_async_action=True) 6 ctx1fabcde

212 collection = db_connection[cls._get_collection_name()] 6 ctx1fabcde

213 

214 return cast(AsyncCollection[Any], collection) 6 ctx1fabcde

215 

216 @classmethod 6 ctx1fabcde

217 def _db(cls) -> str: 6 ctx1fabcde

218 return cls._get_collection_name() 6 ctx1fabcde

219 

220 @classmethod 6 ctx1fabcde

221 def get_inheritance_key(cls) -> DICT_TYPE: 6 ctx1fabcde

222 """ 

223 Get child filter keys 

224 """ 

225 return {INHERITANCE_FIELD_NAME: cls._get_child()} 6 ctx1fabcde

226 

227 @classmethod 6 ctx1fabcde

228 def get_parent_child_fields(cls) -> DICT_TYPE: 6 ctx1fabcde

229 fields = get_model_fields(cls) 6 ctx1fabcde

230 if cls._has_children(): 6 ctx1fabcde

231 for model in cls.__subclasses__(): 6 ctx1fabcde

232 child_fields = get_model_fields(model) 6 ctx1fabcde

233 fields.update(child_fields) 6 ctx1fabcde

234 return fields 6 ctx1fabcde

235 

236 @classmethod 6 ctx1fabcde

237 def get_relational_field_info(cls) -> RELATION_TYPE: 6 ctx1fabcde

238 """ 

239 Get all relational field information. 

240 Get data from from cache if it's already calculated. 

241 """ 

242 global _cashed_field_info 

243 field_info_key = f"{hash(cls)}-field_info" # unique key for each model 6 ctx1fabcde

244 

245 cached_field_info = _cashed_field_info.get(field_info_key) 6 ctx1fabcde

246 if not cached_field_info: 6 ctx1fabcde

247 cached_field_info = get_relationship_fields_info(cls) 6 ctx1fabcde

248 _cashed_field_info[field_info_key] = cached_field_info 6 ctx1fabcde

249 return cached_field_info 6 ctx1fabcde

250 

251 @classmethod 6 ctx1fabcde

252 def get_exclude_fields(cls) -> set[str]: 6 ctx1fabcde

253 """ 

254 Get all fields that should not pass while creating or updating an object. 

255 """ 

256 relational_fields = cls.get_relational_field_info().keys() 6 ctx1fabcde

257 return {"_id", "id", *relational_fields} 6 ctx1fabcde

258 

259 def to_mongo(self) -> DICT_TYPE: 6 ctx1fabcde

260 """ 

261 relational-field: The relational field is only for representational use only. 

262 Do not store those fields in the database. 

263 

264 id: For '_id' and 'id' fields will be excluded on the object creation. 

265 So that 'id' creation happens on the database only. 

266 

267 None: Exclude null fields to improve database storage efficiency. 

268 """ 

269 return self.dict(exclude=self.get_exclude_fields()) 6 ctx1fabcde

270 

271 @classmethod 6 ctx1fabcde

272 def start_session(cls, **kwargs: Any) -> ClientSession: 6 ctx1fabcde

273 client = get_client() 6 ctx1fabcde

274 if not isinstance(client, MongoClient): 6 ctx1fabcde

275 raise InvalidConfiguration( 6 ctx1fabcde

276 "Client is not configured for sync operations use 'astart_session' instead." 

277 ) 

278 

279 return client.start_session(**kwargs) 6 ctx1fabcde

280 

281 @classmethod 6 ctx1fabcde

282 def astart_session(cls, **kwargs: Any) -> AsyncClientSession: 6 ctx1fabcde

283 client = get_client() 6 ctx1fabcde

284 if not isinstance(client, AsyncMongoClient): 6 ctx1fabcde

285 raise InvalidConfiguration( 6 ctx1fabcde

286 "Client is not configured for async operations use 'start_session' instead." 

287 ) 

288 

289 return client.start_session(**kwargs) 6 ctx1fabcde

290 

291 def __str__(self) -> str: 6 ctx1fabcde

292 return super().__repr__() 6 ctx1fabcde

293 

294 

295class Document(_BaseDocument): 6 ctx1fabcde

296 """ 

297 id: For '_id' and 'id' fields will be excluded on the object creation. 

298 So that 'id' creation happens on the database only. 

299 """ 

300 

301 _id: ODMObjectId = PrivateAttr(default_factory=ODMObjectId) 6 ctx1fabcde

302 id: ODMObjectId = Field(default_factory=ODMObjectId) 6 ctx1fabcde

303 

304 def __init__(self, *args: list[Any], **kwargs: Any) -> None: 6 ctx1fabcde

305 super().__init__(*args, **kwargs) 6 ctx1fabcde

306 

307 if "_id" in kwargs: 6 ctx1fabcde

308 id = ODMObjectId(kwargs["_id"]) 6 ctx1fabcde

309 elif "id" in kwargs: 6 ctx1fabcde

310 id = ODMObjectId(kwargs["id"]) 6 ctx1fabcde

311 else: 

312 id = ODMObjectId() 6 ctx1fabcde

313 

314 object.__setattr__(self, "id", id) 6 ctx1fabcde

315 object.__setattr__(self, "_id", id) 6 ctx1fabcde

316 

317 def _prepare_crate_data(self, **kwargs: Any) -> DICT_TYPE: 6 ctx1fabcde

318 data = self.to_mongo() 6 ctx1fabcde

319 if self._get_child() is not None: 6 ctx1fabcde

320 # Assign the '_cls' field if the model is a child. 

321 data = {**self.get_inheritance_key(), **data} 6 ctx1fabcde

322 

323 return data 6 ctx1fabcde

324 

325 def _update_new_id(self, new_id: ODMObjectId) -> None: 6 ctx1fabcde

326 self.__dict__.update({"_id": new_id, "id": new_id}) 6 ctx1fabcde

327 

328 def create(self, **kwargs: Any) -> Self: 6 ctx1fabcde

329 data = self._prepare_crate_data(**kwargs) 6 ctx1fabcde

330 

331 _collection = self._get_collection() 6 ctx1fabcde

332 result = _collection.insert_one(data, **kwargs) 6 ctx1fabcde

333 inserted_id = result.inserted_id 6 ctx1fabcde

334 self._update_new_id(inserted_id) 6 ctx1fabcde

335 

336 return self 6 ctx1fabcde

337 

338 async def acreate(self, **kwargs: Any) -> Self: 6 ctx1fabcde

339 data = self._prepare_crate_data(**kwargs) 6 ctx1fabcde

340 

341 _collection = self._async_get_collection() 6 ctx1fabcde

342 inserted_id = (await _collection.insert_one(data, **kwargs)).inserted_id 6 ctx1fabcde

343 self._update_new_id(inserted_id) 6 ctx1fabcde

344 

345 return self 6 ctx1fabcde

346 

347 @classmethod 6 ctx1fabcde

348 def _validate_and_prepare_filter(cls, filter: Optional[DICT_TYPE]) -> DICT_TYPE: 6 ctx1fabcde

349 if filter is None: 6 ctx1fabcde

350 filter = {} 6 ctx1fabcde

351 

352 validate_filter_dict(cls, filter) 6 ctx1fabcde

353 

354 if cls._get_child() is not None: 6 ctx1fabcde

355 filter = {**cls.get_inheritance_key(), **filter} 6 ctx1fabcde

356 

357 return filter 6 ctx1fabcde

358 

359 @classmethod 6 ctx1fabcde

360 def find_raw( 6 ctx1fabcde

361 cls, 

362 filter: Optional[DICT_TYPE] = None, 

363 projection: Optional[DICT_TYPE] = None, 

364 **kwargs: Any, 

365 ) -> Cursor[Any]: 

366 if projection is None: 6 ctx1fabcde

367 projection = {} 6 ctx1fabcde

368 

369 filter = cls._validate_and_prepare_filter(filter) 6 ctx1fabcde

370 

371 _collection = cls._get_collection() 6 ctx1fabcde

372 

373 if projection: 6 ctx1fabcde

374 return _collection.find(filter, projection, **kwargs) 6 ctx1fabcde

375 

376 return _collection.find(filter, **kwargs) 6 ctx1fabcde

377 

378 @classmethod 6 ctx1fabcde

379 def afind_raw( 6 ctx1fabcde

380 cls, 

381 filter: Optional[DICT_TYPE] = None, 

382 projection: Optional[DICT_TYPE] = None, 

383 **kwargs: Any, 

384 ) -> AsyncCursor[Any]: 

385 if projection is None: 6 ctx1fabcde

386 projection = {} 6 ctx1fabcde

387 

388 filter = cls._validate_and_prepare_filter(filter) 6 ctx1fabcde

389 _collection = cls._async_get_collection() 6 ctx1fabcde

390 

391 if projection: 6 ctx1fabcde

392 return _collection.find(filter, projection, **kwargs) 6 ctx1fabcde

393 

394 return _collection.find(filter, **kwargs) 6 ctx1fabcde

395 

396 @classmethod 6 ctx1fabcde

397 def _prepare_query( 6 ctx1fabcde

398 cls, 

399 query_set: Any, 

400 sort: Optional[SORT_TYPE] = None, 

401 skip: Optional[int] = None, 

402 limit: Optional[int] = None, 

403 ) -> Any: 

404 """Helper method to prepare query with sort, skip and limit parameters.""" 

405 if sort: 6 ctx1fabcde

406 query_set = query_set.sort(sort) 6 ctx1fabcde

407 if skip: 6 ctx1fabcde

408 query_set = query_set.skip(skip) 6 ctx1fabcde

409 if limit: 6 ctx1fabcde

410 query_set = query_set.limit(limit) 6 ctx1fabcde

411 

412 return query_set 6 ctx1fabcde

413 

414 @classmethod 6 ctx1fabcde

415 def _get_child_models(cls) -> dict[str, Self]: 6 ctx1fabcde

416 """Helper method to get child models mapping.""" 

417 model_children: dict[str, Self] = {} 6 ctx1fabcde

418 

419 for model in cls.__subclasses__(): 6 ctx1fabcde

420 child_model_name = model._get_child() 6 ctx1fabcde

421 if child_model_name is None: 6 ctx1fabcde

422 continue 

423 

424 model_children[child_model_name] = model # type: ignore 6 ctx1fabcde

425 

426 return model_children 6 ctx1fabcde

427 

428 @classmethod 6 ctx1fabcde

429 def _prepare_class_instance( 6 ctx1fabcde

430 cls, 

431 model_children: dict[str, Self], 

432 data: DICT_TYPE, 

433 ) -> Self: 

434 if data.get(INHERITANCE_FIELD_NAME) in model_children: 6 ctx1fabcde

435 """If this is a child model then convert it to that child model.""" 5 ctx1abcde

436 kls = model_children[data[INHERITANCE_FIELD_NAME]] 6 ctx1fabcde

437 return kls(**data) # type: ignore 6 ctx1fabcde

438 

439 return cls(**data) 6 ctx1fabcde

440 

441 @classmethod 6 ctx1fabcde

442 def find( 6 ctx1fabcde

443 cls, 

444 filter: Optional[DICT_TYPE] = None, 

445 projection: Optional[DICT_TYPE] = None, 

446 sort: Optional[SORT_TYPE] = None, 

447 skip: Optional[int] = None, 

448 limit: Optional[int] = None, 

449 **kwargs: Any, 

450 ) -> Iterator[Self]: 

451 qs = cls.find_raw(filter, projection, **kwargs) 6 ctx1fabcde

452 qs = cls._prepare_query(qs, sort, skip, limit) 6 ctx1fabcde

453 

454 if cls._has_children(): 6 ctx1fabcde

455 model_children = cls._get_child_models() 6 ctx1fabcde

456 for data in qs: 6 ctx1fabcde

457 yield cls._prepare_class_instance(model_children, data) 6 ctx1fabcde

458 else: 

459 for data in qs: 6 ctx1fabcde

460 yield cls(**data) 6 ctx1fabcde

461 

462 @classmethod 6 ctx1fabcde

463 async def afind( 6 ctx1fabcde

464 cls, 

465 filter: Optional[DICT_TYPE] = None, 

466 projection: Optional[DICT_TYPE] = None, 

467 sort: Optional[SORT_TYPE] = None, 

468 skip: Optional[int] = None, 

469 limit: Optional[int] = None, 

470 **kwargs: Any, 

471 ) -> AsyncIterator[Self]: 

472 qs = cls.afind_raw(filter, projection, **kwargs) 6 ctx1fabcde

473 qs = cls._prepare_query(qs, sort, skip, limit) 6 ctx1fabcde

474 

475 if cls._has_children(): 6 ctx1fabcde

476 model_children = cls._get_child_models() 6 ctx1fabcde

477 async for data in qs: 6 ctx1fabcde

478 yield cls._prepare_class_instance(model_children, data) 6 ctx1fabcde

479 else: 

480 async for data in qs: 6 ctx1fabcde

481 yield cls(**data) 6 ctx1fabcde

482 

483 @classmethod 6 ctx1fabcde

484 def find_one( 6 ctx1fabcde

485 cls, 

486 filter: Optional[DICT_TYPE] = None, 

487 projection: Optional[DICT_TYPE] = None, 

488 sort: Optional[SORT_TYPE] = None, 

489 **kwargs: Any, 

490 ) -> Optional[Self]: 

491 qs = cls.find_raw(filter, projection=projection, **kwargs) 6 ctx1fabcde

492 if sort: 6 ctx1fabcde

493 qs = qs.sort(sort) 6 ctx1fabcde

494 

495 obj = None 6 ctx1fabcde

496 for data in qs.limit(1): 6 ctx1fabcde

497 """limit 1 is equivalent to find_one and that is implemented in pymongo find_one""" 5 ctx1abcde

498 obj = data 6 ctx1fabcde

499 

500 if not obj: 6 ctx1fabcde

501 return None 6 ctx1fabcde

502 

503 if cls._has_children(): 6 ctx1fabcde

504 model_children = cls._get_child_models() 6 ctx1fabcde

505 return cls._prepare_class_instance(model_children, data) 6 ctx1fabcde

506 

507 return cls(**data) 6 ctx1fabcde

508 

509 @classmethod 6 ctx1fabcde

510 async def afind_one( 6 ctx1fabcde

511 cls, 

512 filter: Optional[DICT_TYPE] = None, 

513 projection: Optional[DICT_TYPE] = None, 

514 sort: Optional[SORT_TYPE] = None, 

515 **kwargs: Any, 

516 ) -> Optional[Self]: 

517 qs = cls.afind_raw(filter, projection=projection, **kwargs) 6 ctx1fabcde

518 if sort: 6 ctx1fabcde

519 qs = qs.sort(sort) 6 ctx1fabcde

520 

521 obj = None 6 ctx1fabcde

522 async for data in qs.limit(1): 6 ctx1fabcde

523 """limit 1 is equivalent to find_one and that is implemented in pymongo find_one""" 5 ctx1abcde

524 obj = data 6 ctx1fabcde

525 

526 if not obj: 6 ctx1fabcde

527 return None 6 ctx1fabcde

528 

529 if cls._has_children(): 6 ctx1fabcde

530 model_children = cls._get_child_models() 6 ctx1fabcde

531 return cls._prepare_class_instance(model_children, data) 6 ctx1fabcde

532 

533 return cls(**data) 6 ctx1fabcde

534 

535 @classmethod 6 ctx1fabcde

536 def get( 6 ctx1fabcde

537 cls, 

538 filter: DICT_TYPE, 

539 sort: Optional[SORT_TYPE] = None, 

540 **kwargs: Any, 

541 ) -> Self: 

542 obj = cls.find_one(filter, sort=sort, **kwargs) 6 ctx1fabcde

543 if obj: 6 ctx1fabcde

544 return obj 6 ctx1fabcde

545 

546 raise ObjectDoesNotExist("Object not found.") 6 ctx1fabcde

547 

548 @classmethod 6 ctx1fabcde

549 async def aget( 6 ctx1fabcde

550 cls, 

551 filter: DICT_TYPE, 

552 sort: Optional[SORT_TYPE] = None, 

553 **kwargs: Any, 

554 ) -> Self: 

555 obj = await cls.afind_one(filter, sort=sort, **kwargs) 6 ctx1fabcde

556 if obj: 6 ctx1fabcde

557 return obj 6 ctx1fabcde

558 

559 raise ObjectDoesNotExist("Object not found.") 6 ctx1fabcde

560 

561 @classmethod 6 ctx1fabcde

562 def get_or_create( 6 ctx1fabcde

563 cls, 

564 filter: DICT_TYPE, 

565 sort: Optional[SORT_TYPE] = None, 

566 **kwargs: Any, 

567 ) -> tuple[Self, bool]: 

568 obj = cls.find_one(filter, sort=sort, **kwargs) 6 ctx1fabcde

569 if obj: 6 ctx1fabcde

570 return obj, False 6 ctx1fabcde

571 

572 return cls(**filter).create(), True 6 ctx1fabcde

573 

574 @classmethod 6 ctx1fabcde

575 async def aget_or_create( 6 ctx1fabcde

576 cls, 

577 filter: DICT_TYPE, 

578 sort: Optional[SORT_TYPE] = None, 

579 **kwargs: Any, 

580 ) -> tuple[Self, bool]: 

581 obj = await cls.afind_one(filter, sort=sort, **kwargs) 6 ctx1fabcde

582 if obj: 6 ctx1fabcde

583 return obj, False 6 ctx1fabcde

584 

585 new_instance = await cls(**filter).acreate() 6 ctx1fabcde

586 

587 return new_instance, True 6 ctx1fabcde

588 

589 @classmethod 6 ctx1fabcde

590 def count_documents(cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any) -> int: 6 ctx1fabcde

591 filter = cls._validate_and_prepare_filter(filter) 6 ctx1fabcde

592 

593 _collection = cls._get_collection() 6 ctx1fabcde

594 

595 return _collection.count_documents(filter, **kwargs) 6 ctx1fabcde

596 

597 @classmethod 6 ctx1fabcde

598 async def acount_documents( 6 ctx1fabcde

599 cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any 

600 ) -> int: 

601 filter = cls._validate_and_prepare_filter(filter) 6 ctx1fabcde

602 

603 _collection = cls._async_get_collection() 6 ctx1fabcde

604 

605 return await _collection.count_documents(filter, **kwargs) 6 ctx1fabcde

606 

607 @classmethod 6 ctx1fabcde

608 def exists(cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any) -> bool: 6 ctx1fabcde

609 """ 

610 It does not need to count all documents and avoid unnecessary db overhead. 

611 """ 

612 for _ in cls.find_raw(filter, projection={"_id": 1}, **kwargs).limit(1): 6 ctx1fabcde

613 return True 6 ctx1fabcde

614 

615 return False 6 ctx1fabcde

616 

617 @classmethod 6 ctx1fabcde

618 async def aexists(cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any) -> bool: 6 ctx1fabcde

619 """ 

620 It does not need to count all documents and avoid unnecessary db overhead. 

621 """ 

622 async for _ in cls.afind_raw(filter, projection={"_id": 1}, **kwargs).limit(1): 6 ctx1fabcde

623 return True 6 ctx1fabcde

624 

625 return False 6 ctx1fabcde

626 

627 @classmethod 6 ctx1fabcde

628 def _prepare_aggregation_pipeline( 6 ctx1fabcde

629 cls, 

630 pipeline: list[Any], 

631 inheritance_filter: bool = True, 

632 ) -> list[Any]: 

633 if inheritance_filter and cls._get_child() is not None: 6 ctx1fabcde

634 """ 5 ctx1abcde

635 If aggregate was called from the child model then add the "$match" stage 

636 to separate the document from other child and parent. 

637 """ 

638 if len(pipeline) > 0 and "$match" in pipeline[0]: 6 ctx1fabcde

639 """ 5 ctx1abcde

640 If the first stage of the pipeline is "$match" 

641 update it with the '_cls' field. 

642 """ 

643 pipeline[0]["$match"] = { 6 ctx1fabcde

644 f"{INHERITANCE_FIELD_NAME}": cls._get_child(), 

645 **pipeline[0]["$match"], 

646 } 

647 else: 

648 pipeline = [ 6 ctx1fabcde

649 {"$match": {f"{INHERITANCE_FIELD_NAME}": cls._get_child()}} 

650 ] + pipeline 

651 

652 return pipeline 6 ctx1fabcde

653 

654 @classmethod 6 ctx1fabcde

655 def _get_aggregation_obj(cls, obj: DICT_TYPE, get_raw: bool = False) -> Any: 6 ctx1fabcde

656 if get_raw is True: 6 ctx1fabcde

657 return obj 6 ctx1fabcde

658 

659 return dict2obj(obj) 6 ctx1fabcde

660 

661 @classmethod 6 ctx1fabcde

662 def aggregate( 6 ctx1fabcde

663 cls, 

664 pipeline: list[Any], 

665 get_raw: bool = False, 

666 inheritance_filter: bool = True, 

667 **kwargs: Any, 

668 ) -> Iterator[Any]: 

669 pipeline = cls._prepare_aggregation_pipeline( 6 ctx1fabcde

670 pipeline, inheritance_filter=inheritance_filter 

671 ) 

672 _collection = cls._get_collection() 6 ctx1fabcde

673 

674 query = _collection.aggregate(pipeline, **kwargs) 6 ctx1fabcde

675 

676 for obj in query: 6 ctx1fabcde

677 yield cls._get_aggregation_obj(obj, get_raw=get_raw) 6 ctx1fabcde

678 

679 @classmethod 6 ctx1fabcde

680 async def aaggregate( 6 ctx1fabcde

681 cls, 

682 pipeline: list[Any], 

683 get_raw: bool = False, 

684 inheritance_filter: bool = True, 

685 **kwargs: Any, 

686 ) -> AsyncIterator[Any]: 

687 """ 

688 Return an async iterator for MongoDB aggregation results. 

689 """ 

690 pipeline = cls._prepare_aggregation_pipeline( 6 ctx1fabcde

691 pipeline, inheritance_filter=inheritance_filter 

692 ) 

693 _collection = cls._async_get_collection() 6 ctx1fabcde

694 query = await _collection.aggregate(pipeline, **kwargs) 6 ctx1fabcde

695 

696 async for obj in query: 6 ctx1fabcde

697 yield cls._get_aggregation_obj(obj, get_raw=get_raw) 6 ctx1fabcde

698 

699 @classmethod 6 ctx1fabcde

700 def _get_pipeline_for_random_one(cls, filter: DICT_TYPE) -> list[DICT_TYPE]: 6 ctx1fabcde

701 return [{"$match": filter}, {"$sample": {"size": 1}}] 6 ctx1fabcde

702 

703 @classmethod 6 ctx1fabcde

704 def get_random_one(cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any) -> Self: 6 ctx1fabcde

705 filter = cls._validate_and_prepare_filter(filter) 6 ctx1fabcde

706 pipeline = cls._get_pipeline_for_random_one(filter) 6 ctx1fabcde

707 

708 for data in cls.aggregate(pipeline, get_raw=True, **kwargs): 6 ctx1fabcde

709 return cls(**data) 6 ctx1fabcde

710 

711 raise ObjectDoesNotExist("Object not found.") 6 ctx1fabcde

712 

713 @classmethod 6 ctx1fabcde

714 async def aget_random_one( 6 ctx1fabcde

715 cls, filter: Optional[DICT_TYPE] = None, **kwargs: Any