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 sql_smith.functions import on 

11 

12from kwai_bc_training.coaches._tables import ( # noqa 

13 CoachRow, 

14 MemberRow, 

15 PersonRow, 

16) 

17from kwai_bc_training.coaches.coach import CoachEntity, CoachIdentifier 

18from kwai_bc_training.trainings._tables import TrainingScheduleCoachRow 

19from kwai_bc_training.trainings.training_schedule import TrainingScheduleIdentifier 

20 

21 

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

23class TrainingScheduleCoachQueryRow(JoinedTableRow): 

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

25 

26 person: PersonRow 

27 coach: CoachRow 

28 training_schedule: TrainingScheduleCoachRow 

29 

30 def create_coach(self) -> CoachEntity: 

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

32 return CoachEntity( 

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

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

35 active=self.coach.active == 1, 

36 ) 

37 

38 

39class TrainingScheduleCoachDbQuery(DatabaseQuery): 

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

41 

42 def init(self): 

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

44 CoachRow.__table_name__, 

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

46 ).join( 

47 MemberRow.__table_name__, 

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

49 ).join( 

50 PersonRow.__table_name__, 

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

52 ) 

53 

54 @property 

55 def columns(self): 

56 return TrainingScheduleCoachQueryRow.get_aliases() 

57 

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

59 """Filter by training schedules.""" 

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

61 self._query.and_where( 

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

63 ) 

64 return self 

65 

66 async def fetch_coaches( 

67 self, 

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

69 """Fetch coaches. 

70 

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

72 

73 Returns: 

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

75 is the identifier of the schedule. 

76 """ 

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

78 

79 async for row in self.fetch(): 

80 training_schedule_coach_row = TrainingScheduleCoachQueryRow.map(row) 

81 result[ 

82 TrainingScheduleIdentifier( 

83 training_schedule_coach_row.training_schedule.training_schedule_id 

84 ) 

85 ].append(training_schedule_coach_row.create_coach()) 

86 

87 return result