Coverage for bc/kwai-bc-training/src/kwai_bc_training/trainings/training_coach_db_query.py: 100%

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

2 

3from collections import defaultdict 

4from dataclasses import dataclass 

5 

6from kwai_core.db.database_query import DatabaseQuery 

7from kwai_core.db.rows import OwnerRow 

8from kwai_core.db.table_row import JoinedTableRow 

9from kwai_core.domain.value_objects.name import Name 

10from kwai_core.domain.value_objects.unique_id import UniqueId 

11from sql_smith.functions import on 

12 

13from kwai_bc_training.coaches._tables import ( # noqa 

14 CoachRow, 

15 MemberRow, 

16 PersonRow, 

17) 

18from kwai_bc_training.coaches.coach import CoachEntity, CoachIdentifier 

19from kwai_bc_training.trainings._tables import ( 

20 TrainingCoachRow, 

21) 

22from kwai_bc_training.trainings.training import ( 

23 TrainingCoachEntity, 

24 TrainingCoachIdentifier, 

25 TrainingIdentifier, 

26) 

27 

28 

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

30class TrainingCoachQueryRow(JoinedTableRow): 

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

32 

33 training_coach: TrainingCoachRow 

34 member: MemberRow 

35 person: PersonRow 

36 coach: CoachRow 

37 owner: OwnerRow 

38 

39 def create_coach(self) -> TrainingCoachEntity: 

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

41 return TrainingCoachEntity( 

42 id=TrainingCoachIdentifier(self.training_coach.id), 

43 coach=CoachEntity( 

44 id=CoachIdentifier(self.coach.id), 

45 name=Name( 

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

47 ), 

48 active=self.coach.active == 1, 

49 uuid=UniqueId.create_from_string(self.member.uuid), 

50 ), 

51 owner=self.owner.create_owner(), 

52 present=self.training_coach.present == 1, 

53 type=self.training_coach.coach_type, 

54 payed=self.training_coach.payed == 1, 

55 remark=( 

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

57 ), 

58 ) 

59 

60 

61class TrainingCoachDbQuery(DatabaseQuery): 

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

63 

64 def init(self): 

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

66 CoachRow.__table_name__, 

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

68 ).join( 

69 MemberRow.__table_name__, 

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

71 ).join( 

72 PersonRow.__table_name__, 

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

74 ).join( 

75 OwnerRow.__table_name__, 

76 on(TrainingCoachRow.column("user_id"), OwnerRow.column("id")), 

77 ) 

78 

79 @property 

80 def columns(self): 

81 return TrainingCoachQueryRow.get_aliases() 

82 

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

84 """Filter by trainings. 

85 

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

87 """ 

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

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

90 return self 

91 

92 async def fetch_coaches( 

93 self, 

94 ) -> dict[TrainingIdentifier, list[TrainingCoachEntity]]: 

95 """Fetch coaches. 

96 

97 A specialized fetch method that already transforms the records into 

98 TrainingCoach objects. 

99 

100 Returns: 

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

102 is the identifier of a training. 

103 """ 

104 result: dict[TrainingIdentifier, list[TrainingCoachEntity]] = defaultdict(list) 

105 

106 async for row in self.fetch(): 

107 training_coach_row = TrainingCoachQueryRow.map(row) 

108 result[ 

109 TrainingIdentifier(training_coach_row.training_coach.training_id) 

110 ].append(training_coach_row.create_coach()) 

111 return result