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

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

2from collections import defaultdict 

3 

4from sql_smith.functions import on 

5 

6from kwai.core.db.database_query import DatabaseQuery 

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

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

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

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

11 

12 

13class TrainingTeamDbQuery(DatabaseQuery): 

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

15 

16 def init(self): 

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

18 TeamsTable.table_name, 

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

20 ) 

21 

22 @property 

23 def columns(self): 

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

25 

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

27 """Filter by trainings. 

28 

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

30 """ 

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

32 self._query.and_where( 

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

34 ) 

35 return self 

36 

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

38 """Fetch teams. 

39 

40 A specialized fetch method that already transforms the records into 

41 Team objects. 

42 

43 Returns: 

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

45 is the identifier of a training. 

46 """ 

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

48 

49 async for team_record in self.fetch(): 

50 training_team = TrainingTeamsTable(team_record) 

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

52 TeamsTable(team_record).create_entity() 

53 ) 

54 

55 return result