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

10 

11from kwai_bc_training.coaches._tables import ( 

12 CoachRow, 

13 MemberRow, 

14 PersonRow, 

15) 

16from kwai_bc_training.coaches.coach import CoachEntity, CoachIdentifier 

17from kwai_bc_training.coaches.coach_query import CoachQuery 

18 

19 

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

21class CoachQueryRow(JoinedTableRow): 

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

23 

24 member: MemberRow 

25 person: PersonRow 

26 coach: CoachRow 

27 

28 def create_entity(self) -> CoachEntity: 

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

30 return CoachEntity( 

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

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

33 active=self.coach.active == 1, 

34 ) 

35 

36 

37class CoachDbQuery(DatabaseQuery, CoachQuery): 

38 """A database query for coaches.""" 

39 

40 @property 

41 def count_column(self) -> str: 

42 return CoachRow.column("id") 

43 

44 def init(self): 

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

46 MemberRow.__table_name__, 

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

48 ).inner_join( 

49 PersonRow.__table_name__, 

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

51 ) 

52 

53 @property 

54 def columns(self): 

55 return CoachQueryRow.get_aliases() 

56 

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

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

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

60 return self 

61 

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

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

64 return self 

65 

66 def filter_by_active(self) -> Self: 

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

68 return self