Coverage for src/kwai/modules/teams/repositories/team_member_db_query.py: 92%

1"""Module that defines a database query for team members.""" 

2 

3from collections import defaultdict 

4from dataclasses import dataclass 

5from typing import Self 

6 

7from sql_smith.functions import on 

8 

9from kwai.core.db.database_query import DatabaseQuery 

10from kwai.core.db.table_row import JoinedTableRow 

11from kwai.core.domain.value_objects.date import Date 

12from kwai.core.domain.value_objects.name import Name 

13from kwai.core.domain.value_objects.timestamp import Timestamp 

14from kwai.core.domain.value_objects.traceable_time import TraceableTime 

15from kwai.core.domain.value_objects.unique_id import UniqueId 

16from kwai.modules.club.domain.value_objects import Birthdate, Gender, License 

17from kwai.modules.teams.domain.team import TeamIdentifier 

18from kwai.modules.teams.domain.team_member import ( 

19 MemberEntity, 

20 MemberIdentifier, 

21 TeamMember, 

22) 

23from kwai.modules.teams.repositories._tables import ( 

24 CountryRow, 

25 MemberPersonRow, 

26 MemberRow, 

27 TeamMemberRow, 

28) 

29 

30 

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

32class TeamMemberQueryRow(JoinedTableRow): 

33 """A data transfer object for the team member query.""" 

34 

35 team_member: TeamMemberRow 

36 member: MemberRow 

37 person: MemberPersonRow 

38 country: CountryRow 

39 

40 def create_team_member(self) -> TeamMember: 

41 """Create a team member from a row.""" 

42 return TeamMember( 

43 active=self.team_member.active == 1, 

44 member=MemberEntity( 

45 id_=MemberIdentifier(self.member.id), 

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

47 name=Name( 

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

49 ), 

50 license=License( 

51 number=self.member.license, 

52 end_date=Date.create_from_date(self.member.license_end_date), 

53 ), 

54 birthdate=Birthdate(Date.create_from_date(self.person.birthdate)), 

55 gender=Gender(self.person.gender), 

56 nationality=self.country.create_country(), 

57 active_in_club=self.member.active == 1, 

58 ), 

59 traceable_time=TraceableTime( 

60 created_at=Timestamp(self.team_member.created_at), 

61 updated_at=Timestamp(self.team_member.updated_at), 

62 ), 

63 ) 

64 

65 

66class TeamMemberDbQuery(DatabaseQuery): 

67 """A database query for getting team members.""" 

68 

69 def init(self): 

70 self._query.from_(TeamMemberRow.__table_name__).left_join( 

71 MemberRow.__table_name__, 

72 on(TeamMemberRow.column("member_id"), MemberRow.column("id")), 

73 ).join( 

74 MemberPersonRow.__table_name__, 

75 on(MemberRow.column("person_id"), MemberPersonRow.column("id")), 

76 ).inner_join( 

77 CountryRow.__table_name__, 

78 on(CountryRow.column("id"), MemberPersonRow.column("nationality_id")), 

79 ) 

80 

81 @property 

82 def columns(self): 

83 return TeamMemberQueryRow.get_aliases() 

84 

85 def filter_by_teams(self, *ids: TeamIdentifier) -> Self: 

86 """Filter by teams. 

87 

88 Only the rows that belong to the teams with the given ids, will be returned. 

89 """ 

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

91 self._query.and_where(TeamMemberRow.field("team_id").in_(*unpacked_ids)) 

92 return self 

93 

94 async def fetch_team_members(self) -> dict[TeamIdentifier, list[TeamMember]]: 

95 """Fetch team members. 

96 

97 A specialized fetch method that already transforms the rows into TeamMember 

98 objects. 

99 

100 Returns: 

101 A dictionary that contains the list of team members for each team. 

102 The key is the identifier of the team. 

103 """ 

104 result: dict[TeamIdentifier, list[TeamMember]] = defaultdict(list) 

105 

106 async for row in self.fetch(): 

107 team_member_row = TeamMemberQueryRow.map(row) 

108 result[TeamIdentifier(team_member_row.team_member.team_id)].append( 

109 team_member_row.create_team_member() 

110 ) 

111 

112 return result