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.coach.uuid), 

50 head=False, 

51 ), 

52 owner=self.owner.create_owner(), 

53 present=self.training_coach.present == 1, 

54 type=self.training_coach.coach_type, 

55 payed=self.training_coach.payed == 1, 

56 remark=( 

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

58 ), 

59 ) 

60 

61 

62class TrainingCoachDbQuery(DatabaseQuery): 

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

64 

65 def init(self): 

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

67 CoachRow.__table_name__, 

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

69 ).join( 

70 MemberRow.__table_name__, 

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

72 ).join( 

73 PersonRow.__table_name__, 

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

75 ).join( 

76 OwnerRow.__table_name__, 

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

78 ) 

79 

80 @property 

81 def columns(self): 

82 return TrainingCoachQueryRow.get_aliases() 

83 

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

85 """Filter by trainings. 

86 

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

88 """ 

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

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

91 return self 

92 

93 async def fetch_coaches( 

94 self, 

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

96 """Fetch coaches. 

97 

98 A specialized fetch method that already transforms the records into 

99 TrainingCoach objects. 

100 

101 Returns: 

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

103 is the identifier of a training. 

104 """ 

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

106 

107 async for row in self.fetch(): 

108 training_coach_row = TrainingCoachQueryRow.map(row) 

109 result[ 

110 TrainingIdentifier(training_coach_row.training_coach.training_id) 

111 ].append(training_coach_row.create_coach()) 

112 return result