Coverage for bc/kwai-bc-club/src/kwai_bc_club/repositories/coach_db_repository.py: 93%

1"""Module for defining a coach repository using a database.""" 

2 

3from typing import AsyncGenerator 

4 

5from kwai_core.db.database import Database 

6 

7from kwai_bc_club.domain.club_coach import ClubCoachEntity, ClubCoachIdentifier 

8from kwai_bc_club.repositories._tables import CoachRow 

9from kwai_bc_club.repositories.coach_db_query import CoachDBQuery, CoachQueryRow 

10from kwai_bc_club.repositories.coach_query import CoachQuery 

11from kwai_bc_club.repositories.coach_repository import ( 

12 CoachNotFoundException, 

13 CoachRepository, 

14) 

15 

16 

17class CoachDbRepository(CoachRepository): 

18 """A coach repository using a database.""" 

19 

20 def __init__(self, database: Database) -> None: 

21 self._database = database 

22 

23 def create_query(self) -> CoachQuery: 

24 return CoachDBQuery(self._database) 

25 

26 async def get_by_id(self, id_: ClubCoachIdentifier) -> ClubCoachEntity: 

27 query = self.create_query().filter_by_id(id_) 

28 coach_iterator = self.get_all(query) 

29 try: 

30 return await anext(coach_iterator) 

31 except StopAsyncIteration: 

32 raise CoachNotFoundException(f"Coach not found with {id_}") from None 

33 

34 async def get_all( 

35 self, 

36 query: CoachQuery | None = None, 

37 limit: int | None = None, 

38 offset: int | None = None, 

39 ) -> AsyncGenerator[ClubCoachEntity, None]: 

40 query = query or self.create_query() 

41 async for row in query.fetch(limit, offset): 

42 yield CoachQueryRow.map(row).create_entity() 

43 

44 async def create(self, coach: ClubCoachEntity): 

45 new_coach_id = await self._database.insert( 

46 CoachRow.__table_name__, CoachRow.persist(coach) 

47 ) 

48 return coach.set_id(ClubCoachIdentifier(new_coach_id)) 

49 

50 async def delete(self, coach: ClubCoachEntity): 

51 await self._database.delete(coach.id.value, CoachRow.__table_name__) 

52 

53 async def update(self, coach: ClubCoachEntity) -> None: 

54 await self._database.update( 

55 coach.id.value, CoachRow.__table_name__, CoachRow.persist(coach) 

56 )