Coverage for database\connection.py: 95%

44 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2023-02-05 19:00 +0800

1from typing import Optional, Any, List 

2 

3from beanie import init_beanie, PydanticObjectId 

4from models.events import Event 

5from models.users import User 

6from motor.motor_asyncio import AsyncIOMotorClient 

7from pydantic import BaseSettings, BaseModel 

8 

9 

10class Settings(BaseSettings): 

11 DATABASE_URL: Optional[str] = None 

12 SECRET_KEY: Optional[str] = "default" 

13 

14 async def initialize_database(self): 

15 client = AsyncIOMotorClient(self.DATABASE_URL) 

16 await init_beanie(database=client.get_default_database(), 

17 document_models=[Event, User]) 

18 

19 class Config: 

20 env_file = ".env" 

21 

22 

23class Database: 

24 def __init__(self, model): 

25 self.model = model 

26 

27 async def save(self, document): 

28 await document.create() 

29 return 

30 

31 async def get(self, id: PydanticObjectId) -> bool: 

32 doc = await self.model.get(id) 

33 if doc: 

34 return doc 

35 return False 

36 

37 async def get_all(self) -> List[Any]: 

38 docs = await self.model.find_all().to_list() 

39 return docs 

40 

41 async def update(self, id: PydanticObjectId, body: BaseModel) -> Any: 

42 doc_id = id 

43 des_body = body.dict() 

44 

45 des_body = {k: v for k, v in des_body.items() if v is not None} 

46 update_query = {"$set": { 

47 field: value for field, value in des_body.items() 

48 }} 

49 

50 doc = await self.get(doc_id) 

51 if not doc: 

52 return False 

53 await doc.update(update_query) 

54 return doc 

55 

56 async def delete(self, id: PydanticObjectId) -> bool: 

57 doc = await self.get(id) 

58 if not doc: 

59 return False 

60 await doc.delete() 

61 return True