Coverage for kwai/modules/training/trainings/training_coach_db_query.py: 81%

1"""Module that defines a database query to get coaches 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.core.db.rows import OwnersTable 

8from kwai.modules.training.coaches.coach_tables import CoachesTable, PersonsTable 

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

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

11from kwai.modules.training.trainings.value_objects import TrainingCoach 

12 

13 

14class TrainingCoachDbQuery(DatabaseQuery): 

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

16 

17 def init(self): 

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

19 CoachesTable.table_name, 

20 on(TrainingCoachesTable.column("coach_id"), CoachesTable.column("id")), 

21 ).join( 

22 PersonsTable.table_name, 

23 on(CoachesTable.column("person_id"), PersonsTable.column("id")), 

24 ).join( 

25 OwnersTable.table_name, 

26 on(CoachesTable.column("user_id"), OwnersTable.column("id")), 

27 ) 

28 

29 @property 

30 def columns(self): 

31 return ( 

32 TrainingCoachesTable.aliases() 

33 + CoachesTable.aliases() 

34 + PersonsTable.aliases() 

35 + OwnersTable.aliases() 

36 ) 

37 

38 def filter_by_trainings(self, *ids: TrainingIdentifier) -> "TrainingCoachDbQuery": 

39 """Filter by trainings. 

40 

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

42 """ 

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

44 self._query.and_where( 

45 TrainingCoachesTable.field("training_id").in_(*unpacked_ids) 

46 ) 

47 return self 

48 

49 async def fetch_coaches(self) -> dict[TrainingIdentifier, list[TrainingCoach]]: 

50 """Fetch coaches. 

51 

52 A specialized fetch method that already transforms the records into 

53 TrainingCoach objects. 

54 

55 Returns: 

56 A dictionary that contains the list of coaches for trainings. The key 

57 is the identifier of a training. 

58 """ 

59 result: dict[TrainingIdentifier, list[TrainingCoach]] = defaultdict(list) 

60 

61 async for record in self.fetch(): 

62 training_coach = TrainingCoachesTable(record) 

63 owner_row = OwnersTable(record).create_owner() 

64 coach_row = CoachesTable(record) 

65 person_row = PersonsTable(record) 

66 result[TrainingIdentifier(training_coach.training_id)].append( 

67 training_coach.create_coach( 

68 coach_row.create_entity(person_row), 

69 owner_row, 

70 ) 

71 ) 

72 return result