Coverage for src/kwai/modules/training/coaches/coach_db_query.py: 100%

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

2 

3from dataclasses import dataclass 

4 

5from sql_smith.functions import on 

6 

7from kwai.core.db.database_query import DatabaseQuery 

8from kwai.core.db.table_row import JoinedTableRow 

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

10from kwai.modules.training.coaches._tables import ( 

11 CoachRow, 

12 MemberRow, 

13 PersonRow, 

14) 

15from kwai.modules.training.coaches.coach import CoachEntity, CoachIdentifier 

16from kwai.modules.training.coaches.coach_query import CoachQuery 

17 

18 

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

20class CoachQueryRow(JoinedTableRow): 

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

22 

23 member: MemberRow 

24 person: PersonRow 

25 coach: CoachRow 

26 

27 def create_entity(self) -> CoachEntity: 

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

29 return CoachEntity( 

30 id_=CoachIdentifier(self.coach.id), 

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

32 active=self.coach.active == 1, 

33 ) 

34 

35 

36class CoachDbQuery(DatabaseQuery, CoachQuery): 

37 """A database query for coaches.""" 

38 

39 @property 

40 def count_column(self) -> str: 

41 return CoachRow.column("id") 

42 

43 def init(self): 

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

45 MemberRow.__table_name__, 

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

47 ).inner_join( 

48 PersonRow.__table_name__, 

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

50 ) 

51 

52 @property 

53 def columns(self): 

54 return CoachQueryRow.get_aliases() 

55 

56 def filter_by_ids(self, *ids: CoachIdentifier) -> "CoachQuery": 

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

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

59 return self 

60 

61 def filter_by_id(self, id_: CoachIdentifier) -> "CoachQuery": 

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

63 return self 

64 

65 def filter_by_active(self) -> "CoachQuery": 

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

67 return self