Coverage for src/kwai/modules/club/repositories/member_db_query.py: 100%

1"""Module that implements a MemberQuery for a database.""" 

2 

3from dataclasses import dataclass 

4from typing import Self 

5 

6from sql_smith.functions import alias, criteria, func, group, literal, on 

7 

8from kwai.core.db.database import Database 

9from kwai.core.db.database_query import DatabaseQuery 

10from kwai.core.db.table_row import JoinedTableRow 

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

12from kwai.modules.club.domain.member import MemberEntity, MemberIdentifier 

13from kwai.modules.club.repositories._tables import ( 

14 ContactRow, 

15 CountryRow, 

16 MemberRow, 

17 PersonRow, 

18) 

19from kwai.modules.club.repositories.member_query import MemberQuery 

20 

21 

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

23class MemberQueryRow(JoinedTableRow): 

24 """A data transfer object for the Member query.""" 

25 

26 member: MemberRow 

27 person: PersonRow 

28 nationality: CountryRow 

29 contact: ContactRow 

30 country: CountryRow 

31 

32 def create_entity(self) -> MemberEntity: 

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

34 return self.member.create_entity( 

35 person=self.person.create_entity( 

36 nationality=self.nationality.create_country(), 

37 contact=self.contact.create_entity( 

38 country=self.country.create_country() 

39 ), 

40 ) 

41 ) 

42 

43 

44class MemberDbQuery(MemberQuery, DatabaseQuery): 

45 """A database query for members.""" 

46 

47 def __init__(self, database: Database): 

48 super().__init__(database) 

49 

50 @property 

51 def count_column(self): 

52 return MemberRow.column("id") 

53 

54 def init(self): 

55 self._query.from_(MemberRow.__table_name__).inner_join( 

56 PersonRow.__table_name__, 

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

58 ).inner_join( 

59 alias(CountryRow.__table_name__, "nationality"), 

60 on( 

61 "nationality.id", 

62 PersonRow.column("nationality_id"), 

63 ), 

64 ).inner_join( 

65 ContactRow.__table_name__, 

66 on(ContactRow.column("id"), PersonRow.column("contact_id")), 

67 ).inner_join( 

68 CountryRow.__table_name__, 

69 on(CountryRow.column("id"), ContactRow.column("country_id")), 

70 ) 

71 

72 @property 

73 def columns(self): 

74 return MemberQueryRow.get_aliases() 

75 

76 def filter_by_id(self, id_: MemberIdentifier) -> Self: 

77 self._query.and_where(MemberRow.field("id").eq(id_.value)) 

78 return self 

79 

80 def filter_by_license(self, license: str) -> Self: 

81 self._query.and_where(MemberRow.field("license").eq(license)) 

82 return self 

83 

84 def filter_by_license_date( 

85 self, license_end_month: int, license_end_year: int 

86 ) -> Self: 

87 condition = criteria( 

88 "{} = {}", 

89 func("YEAR", MemberRow.column("license_end_date")), 

90 literal(license_end_year), 

91 ).and_( 

92 criteria( 

93 "{} = {}", 

94 func("MONTH", MemberRow.column("license_end_date")), 

95 literal(license_end_month), 

96 ), 

97 ) 

98 self._query.and_where(group(condition)) 

99 return self 

100 

101 def filter_by_active(self) -> Self: 

102 self._query.and_where(MemberRow.field("active").eq(1)) 

103 return self 

104 

105 def filter_by_uuid(self, uuid: UniqueId) -> Self: 

106 self._query.and_where(MemberRow.field("uuid").eq(str(uuid))) 

107 return self