Coverage for bc/kwai-bc-training/src/kwai_bc_training/teams/team_db_repository.py: 100%

1"""Module that defines a team repository for a database.""" 

2 

3from typing import AsyncIterator 

4 

5from kwai_core.db.database import Database 

6 

7from kwai_bc_training.teams.team import TeamEntity, TeamIdentifier 

8from kwai_bc_training.teams.team_db_query import TeamDbQuery 

9from kwai_bc_training.teams.team_query import TeamQuery 

10from kwai_bc_training.teams.team_repository import TeamRepository 

11from kwai_bc_training.teams.team_tables import TeamRow 

12 

13 

14class TeamDbRepository(TeamRepository): 

15 """A team repository for a database.""" 

16 

17 def __init__(self, database: Database): 

18 """Initialize the repository. 

19 

20 Args: 

21 database: The database for this repository. 

22 """ 

23 self._database = database 

24 

25 def create_query(self) -> TeamQuery: 

26 """Create the team query.""" 

27 return TeamDbQuery(self._database) 

28 

29 async def get_by_id(self, id: TeamIdentifier) -> TeamEntity: 

30 query = self.create_query() 

31 query.filter_by_id(id) 

32 

33 row = await query.fetch_one() 

34 

35 return TeamRow.map(row).create_entity() 

36 

37 async def get_all(self) -> AsyncIterator[TeamEntity]: 

38 query = self.create_query() 

39 async for row in query.fetch(): 

40 yield TeamRow.map(row).create_entity() 

41 

42 async def get_by_ids(self, *ids: TeamIdentifier) -> AsyncIterator[TeamEntity]: 

43 query = self.create_query() 

44 query.filter_by_ids(*ids) 

45 

46 async for row in query.fetch(): 

47 yield TeamRow.map(row).create_entity()