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

39 head=self.coach.head == 1, 

40 ) 

41 

42 

43class TrainingScheduleCoachDbQuery(DatabaseQuery): 

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

45 

46 def init(self): 

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

48 CoachRow.__table_name__, 

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

50 ).join( 

51 MemberRow.__table_name__, 

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

53 ).join( 

54 PersonRow.__table_name__, 

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

56 ) 

57 

58 @property 

59 def columns(self): 

60 return TrainingScheduleCoachQueryRow.get_aliases() 

61 

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

63 """Filter by training schedules.""" 

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

65 self._query.and_where( 

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

67 ) 

68 return self 

69 

70 async def fetch_coaches( 

71 self, 

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

73 """Fetch coaches. 

74 

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

76 

77 Returns: 

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

79 is the identifier of the schedule. 

80 """ 

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

82 

83 async for row in self.fetch(): 

84 training_schedule_coach_row = TrainingScheduleCoachQueryRow.map(row) 

85 result[ 

86 TrainingScheduleIdentifier( 

87 training_schedule_coach_row.training_schedule.training_schedule_id 

88 ) 

89 ].append(training_schedule_coach_row.create_coach()) 

90 

91 return result