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

1"""Module that defines a database query for getting the coaches of training schedules.""" 

2 

3from collections import defaultdict 

4from dataclasses import dataclass 

5from typing import Self 

6 

7from kwai_core.db.database_query import DatabaseQuery 

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 TrainingScheduleCoachRow 

20from kwai_bc_training.trainings.training_schedule import TrainingScheduleIdentifier 

21 

22 

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

24class TrainingScheduleCoachQueryRow(JoinedTableRow): 

25 """A data transfer object for the training schedules coach query.""" 

26 

27 person: PersonRow 

28 member: MemberRow 

29 coach: CoachRow 

30 training_schedule: TrainingScheduleCoachRow 

31 

32 def create_coach(self) -> CoachEntity: 

33 """Create a coach entity from a row.""" 

34 return CoachEntity( 

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

36 name=Name(first_name=self.person.firstname, last_name=self.person.lastname), 

37 active=self.coach.active == 1, 

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

39 ) 

40 

41 

42class TrainingScheduleCoachDbQuery(DatabaseQuery): 

43 """A database query for training schedule coaches.""" 

44 

45 def init(self): 

46 self._query.from_(TrainingScheduleCoachRow.__table_name__).left_join( 

47 CoachRow.__table_name__, 

48 on(TrainingScheduleCoachRow.column("coach_id"), CoachRow.column("id")), 

49 ).join( 

50 MemberRow.__table_name__, 

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

52 ).join( 

53 PersonRow.__table_name__, 

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

55 ) 

56 

57 @property 

58 def columns(self): 

59 return TrainingScheduleCoachQueryRow.get_aliases() 

60 

61 def filter_by_schedule(self, *ids: TrainingScheduleIdentifier) -> Self: 

62 """Filter by training schedules.""" 

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

64 self._query.and_where( 

65 TrainingScheduleCoachRow.field("training_schedule_id").in_(*unpacked_ids) 

66 ) 

67 return self 

68 

69 async def fetch_coaches( 

70 self, 

71 ) -> dict[TrainingScheduleIdentifier, list[CoachEntity]]: 

72 """Fetch coaches. 

73 

74 A specialized fetch method that already transforms the records into a CoachEntity. 

75 

76 Returns: 

77 A dictionary that contains the list of coaches for the training schedules. The key 

78 is the identifier of the schedule. 

79 """ 

80 result: dict[TrainingScheduleIdentifier, list[CoachEntity]] = defaultdict(list) 

81 

82 async for row in self.fetch(): 

83 training_schedule_coach_row = TrainingScheduleCoachQueryRow.map(row) 

84 result[ 

85 TrainingScheduleIdentifier( 

86 training_schedule_coach_row.training_schedule.training_schedule_id 

87 ) 

88 ].append(training_schedule_coach_row.create_coach()) 

89 

90 return result