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

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

2 

3from collections import defaultdict 

4from dataclasses import dataclass 

5 

6from sql_smith.functions import on 

7 

8from kwai.core.db.database_query import DatabaseQuery 

9from kwai.core.db.rows import OwnersTable, OwnerTableRow 

10from kwai.core.db.table_row import JoinedTableRow 

11from kwai.core.domain.value_objects.name import Name 

12from kwai.modules.training.coaches._tables import ( # noqa 

13 CoachRow, 

14 MemberRow, 

15 PersonRow, 

16) 

17from kwai.modules.training.coaches.coach import CoachEntity, CoachIdentifier 

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

19from kwai.modules.training.trainings.training_tables import ( 

20 TrainingCoachRow, 

21) 

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

23 

24 

25@dataclass(kw_only=True, frozen=True, slots=True) 

26class TrainingCoachQueryRow(JoinedTableRow): 

27 """A data transfer object for the training coach query.""" 

28 

29 training_coach: TrainingCoachRow 

30 member: MemberRow 

31 person: PersonRow 

32 coach: CoachRow 

33 owner: OwnerTableRow 

34 

35 def create_coach(self) -> TrainingCoach: 

36 """Create a training coach from a row.""" 

37 return TrainingCoach( 

38 coach=CoachEntity( 

39 id_=CoachIdentifier(self.coach.id), 

40 name=Name( 

41 first_name=self.person.firstname, last_name=self.person.lastname 

42 ), 

43 active=self.coach.active == 1, 

44 ), 

45 owner=self.owner.create_owner(), 

46 present=self.training_coach.present == 1, 

47 type=self.training_coach.coach_type, 

48 payed=self.training_coach.payed == 1, 

49 remark=( 

50 "" if self.training_coach.remark is None else self.training_coach.remark 

51 ), 

52 ) 

53 

54 

55class TrainingCoachDbQuery(DatabaseQuery): 

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

57 

58 def init(self): 

59 self._query.from_(TrainingCoachRow.__table_name__).left_join( 

60 CoachRow.__table_name__, 

61 on(TrainingCoachRow.column("coach_id"), CoachRow.column("id")), 

62 ).join( 

63 MemberRow.__table_name__, 

64 on(CoachRow.column("member_id"), MemberRow.column("id")), 

65 ).join( 

66 PersonRow.__table_name__, 

67 on(MemberRow.column("person_id"), PersonRow.column("id")), 

68 ).join( 

69 OwnersTable.table_name, 

70 on(CoachRow.column("user_id"), OwnerTableRow.column("id")), 

71 ) 

72 

73 @property 

74 def columns(self): 

75 return TrainingCoachQueryRow.get_aliases() 

76 

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

78 """Filter by trainings. 

79 

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

81 """ 

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

83 self._query.and_where(TrainingCoachRow.field("training_id").in_(*unpacked_ids)) 

84 return self 

85 

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

87 """Fetch coaches. 

88 

89 A specialized fetch method that already transforms the records into 

90 TrainingCoach objects. 

91 

92 Returns: 

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

94 is the identifier of a training. 

95 """ 

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

97 

98 async for row in self.fetch(): 

99 training_coach_row = TrainingCoachQueryRow.map(row) 

100 result[ 

101 TrainingIdentifier(training_coach_row.training_coach.training_id) 

102 ].append(training_coach_row.create_coach()) 

103 return result