Coverage for src/kwai/modules/identity/users/user_db_repository.py: 91%

1"""Module for implementing a user repository with a database.""" 

2 

3from kwai.core.db.database import Database 

4from kwai.core.domain.value_objects.email_address import EmailAddress 

5from kwai.core.domain.value_objects.unique_id import UniqueId 

6from kwai.modules.identity.users.user import UserEntity, UserIdentifier 

7from kwai.modules.identity.users.user_db_query import UserDbQuery 

8from kwai.modules.identity.users.user_repository import ( 

9 UserNotFoundException, 

10 UserRepository, 

11) 

12from kwai.modules.identity.users.user_tables import UserRow 

13 

14 

15class UserDbRepository(UserRepository): 

16 """Database repository for the user entity.""" 

17 

18 def __init__(self, database: Database): 

19 self._database = database 

20 

21 async def update(self, user: UserEntity) -> None: 

22 await self._database.update( 

23 user.id.value, UserRow.__table_name__, UserRow.persist(user) 

24 ) 

25 

26 def create_query(self) -> UserDbQuery: 

27 """Create a user database query.""" 

28 return UserDbQuery(self._database) 

29 

30 async def get_user_by_id(self, id_: UserIdentifier) -> UserEntity: 

31 """Get the user with the given id. 

32 

33 UserNotFoundException is raised when the user does not exist. 

34 """ 

35 query = self.create_query() 

36 query.filter_by_id(id_) 

37 

38 row = await query.fetch_one() 

39 if row: 

40 return UserRow.map(row).create_entity() 

41 

42 raise UserNotFoundException() 

43 

44 async def get_user_by_uuid(self, uuid: UniqueId) -> UserEntity: 

45 """Get the user with the given uuid. 

46 

47 UserNotFoundException is raised when the user does not exist. 

48 """ 

49 query = self.create_query() 

50 query.filter_by_uuid(uuid) 

51 

52 row = await query.fetch_one() 

53 if row: 

54 return UserRow.map(row).create_entity() 

55 

56 raise UserNotFoundException() 

57 

58 async def get_user_by_email(self, email: EmailAddress) -> UserEntity: 

59 """Get the user with the given email. 

60 

61 UserNotFoundException is raised when the user does not exist. 

62 """ 

63 query = self.create_query() 

64 query.filter_by_email(email) 

65 

66 row = await query.fetch_one() 

67 if row: 

68 return UserRow.map(row).create_entity() 

69 

70 raise UserNotFoundException()