Coverage for bc/kwai-bc-training/src/kwai_bc_training/coaches/coach_db_query.py: 100%

1"""Module that defines a database query for coaches.""" 

2 

3from dataclasses import dataclass 

4from typing import Self 

5 

6from kwai_core.db.database_query import DatabaseQuery 

7from kwai_core.db.table_row import JoinedTableRow 

8from kwai_core.domain.value_objects.name import Name 

9from kwai_core.domain.value_objects.unique_id import UniqueId 

10from sql_smith.functions import on 

11 

12from kwai_bc_training.coaches._tables import ( 

13 CoachRow, 

14 MemberRow, 

15 PersonRow, 

16) 

17from kwai_bc_training.coaches.coach import CoachEntity, CoachIdentifier 

18from kwai_bc_training.coaches.coach_query import CoachQuery 

19 

20 

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

22class CoachQueryRow(JoinedTableRow): 

23 """A data transfer object for the coach query.""" 

24 

25 member: MemberRow 

26 person: PersonRow 

27 coach: CoachRow 

28 

29 def create_entity(self) -> CoachEntity: 

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

31 return CoachEntity( 

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

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

34 active=self.coach.active == 1, 

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

36 ) 

37 

38 

39class CoachDbQuery(DatabaseQuery, CoachQuery): 

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

41 

42 @property 

43 def count_column(self) -> str: 

44 return CoachRow.column("id") 

45 

46 def init(self): 

47 self._query.from_(CoachRow.__table_name__).join( 

48 MemberRow.__table_name__, 

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

50 ).inner_join( 

51 PersonRow.__table_name__, 

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

53 ) 

54 

55 @property 

56 def columns(self): 

57 return CoachQueryRow.get_aliases() 

58 

59 def filter_by_ids(self, *ids: CoachIdentifier) -> Self: 

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

61 self._query.and_where(CoachRow.field("id").in_(*unpacked_ids)) 

62 return self 

63 

64 def filter_by_id(self, id_: CoachIdentifier) -> Self: 

65 self._query.and_where(CoachRow.field("id").eq(id_.value)) 

66 return self 

67 

68 def filter_by_active(self) -> Self: 

69 self._query.and_where(CoachRow.field("active").eq(1)) 

70 return self 

71 

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

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

74 return self 

75 

76 def filter_by_uuids(self, *uuids: UniqueId) -> Self: 

77 unpacked_uuids = tuple(str(uuid) for uuid in uuids) 

78 self._query.and_where(MemberRow.field("uuid").in_(*unpacked_uuids)) 

79 return self