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

36 head=self.coach.head == 1, 

37 ) 

38 

39 

40class CoachDbQuery(DatabaseQuery, CoachQuery): 

41 """A database query for coaches.""" 

42 

43 @property 

44 def count_column(self) -> str: 

45 return CoachRow.column("id") 

46 

47 def init(self): 

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

49 MemberRow.__table_name__, 

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

51 ).inner_join( 

52 PersonRow.__table_name__, 

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

54 ) 

55 

56 @property 

57 def columns(self): 

58 return CoachQueryRow.get_aliases() 

59 

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

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

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

63 return self 

64 

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

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

67 return self 

68 

69 def filter_by_active(self) -> Self: 

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

71 return self 

72 

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

74 self._query.and_where(CoachRow.field("uuid").eq(str(uuid))) 

75 return self 

76 

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

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

79 self._query.and_where(CoachRow.field("uuid").in_(*unpacked_uuids)) 

80 return self