Coverage for src/kwai/modules/training/trainings/training_team_db_query.py: 91%

1"""Module that defines a database query to get teams of training(s).""" 

2 

3from collections import defaultdict 

4 

5from sql_smith.functions import on 

6 

7from kwai.core.db.database_query import DatabaseQuery 

8from kwai.modules.training.teams.team import TeamEntity 

9from kwai.modules.training.teams.team_tables import TeamsTable 

10from kwai.modules.training.trainings.training import TrainingIdentifier 

11from kwai.modules.training.trainings.training_tables import TrainingTeamsTable 

12 

13 

14class TrainingTeamDbQuery(DatabaseQuery): 

15 """A database query for getting teams of training(s).""" 

16 

17 def init(self): 

18 self._query.from_(TrainingTeamsTable.table_name).left_join( 

19 TeamsTable.table_name, 

20 on(TrainingTeamsTable.column("team_id"), TeamsTable.column("id")), 

21 ) 

22 

23 @property 

24 def columns(self): 

25 return TrainingTeamsTable.aliases() + TeamsTable.aliases() 

26 

27 def filter_by_trainings(self, *ids: TrainingIdentifier) -> "TrainingTeamDbQuery": 

28 """Filter by trainings. 

29 

30 Only the rows of the trainings with the given ids, will be returned. 

31 """ 

32 unpacked_ids = tuple(i.value for i in ids) 

33 self._query.and_where( 

34 TrainingTeamsTable.field("training_id").in_(*unpacked_ids) 

35 ) 

36 return self 

37 

38 async def fetch_teams(self) -> dict[TrainingIdentifier, list[TeamEntity]]: 

39 """Fetch teams. 

40 

41 A specialized fetch method that already transforms the records into 

42 Team objects. 

43 

44 Returns: 

45 A dictionary that contains the list of teams for trainings. The key 

46 is the identifier of a training. 

47 """ 

48 result: dict[TrainingIdentifier, list[TeamEntity]] = defaultdict(list) 

49 

50 async for team_record in self.fetch(): 

51 training_team = TrainingTeamsTable(team_record) 

52 result[TrainingIdentifier(training_team.training_id)].append( 

53 TeamsTable(team_record).create_entity() 

54 ) 

55 

56 return result