Coverage for bc/kwai-bc-identity/src/kwai_bc_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 

6 

7from kwai_bc_identity.users.user import UserEntity, UserIdentifier 

8from kwai_bc_identity.users.user_db_query import UserDbQuery 

9from kwai_bc_identity.users.user_repository import ( 

10 UserNotFoundException, 

11 UserRepository, 

12) 

13from kwai_bc_identity.users.user_tables import UserRow 

14 

15 

16class UserDbRepository(UserRepository): 

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

18 

19 def __init__(self, database: Database): 

20 self._database = database 

21 

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

23 await self._database.update( 

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

25 ) 

26 

27 def create_query(self) -> UserDbQuery: 

28 """Create a user database query.""" 

29 return UserDbQuery(self._database) 

30 

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

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

33 

34 UserNotFoundException is raised when the user does not exist. 

35 """ 

36 query = self.create_query() 

37 query.filter_by_id(id_) 

38 

39 row = await query.fetch_one() 

40 if row: 

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

42 

43 raise UserNotFoundException() 

44 

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

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

47 

48 UserNotFoundException is raised when the user does not exist. 

49 """ 

50 query = self.create_query() 

51 query.filter_by_uuid(uuid) 

52 

53 row = await query.fetch_one() 

54 if row: 

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

56 

57 raise UserNotFoundException() 

58 

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

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

61 

62 UserNotFoundException is raised when the user does not exist. 

63 """ 

64 query = self.create_query() 

65 query.filter_by_email(email) 

66 

67 row = await query.fetch_one() 

68 if row: 

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

70 

71 raise UserNotFoundException()