Coverage for src/kwai/modules/training/trainings/training_definition_db_query.py: 100%

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

2 

3from sql_smith.functions import on 

4 

5from kwai.core.db.database_query import DatabaseQuery 

6from kwai.core.db.rows import OwnersTable 

7from kwai.modules.training.teams.team_tables import TeamsTable 

8from kwai.modules.training.trainings.training_definition import ( 

9 TrainingDefinitionIdentifier, 

10) 

11from kwai.modules.training.trainings.training_definition_query import ( 

12 TrainingDefinitionQuery, 

13) 

14from kwai.modules.training.trainings.training_tables import ( 

15 TrainingDefinitionsTable, 

16) 

17 

18 

19class TrainingDefinitionDbQuery(DatabaseQuery, TrainingDefinitionQuery): 

20 """A database query for a training definition.""" 

21 

22 def init(self): 

23 return ( 

24 self._query.from_(TrainingDefinitionsTable.table_name) 

25 .join( 

26 OwnersTable.table_name, 

27 on( 

28 OwnersTable.column("id"), TrainingDefinitionsTable.column("user_id") 

29 ), 

30 ) 

31 .left_join( 

32 TeamsTable.table_name, 

33 on(TeamsTable.column("id"), TrainingDefinitionsTable.column("team_id")), 

34 ) 

35 ) 

36 

37 @property 

38 def columns(self): 

39 return ( 

40 TrainingDefinitionsTable.aliases() 

41 + TeamsTable.aliases() 

42 + OwnersTable.aliases() 

43 ) 

44 

45 @property 

46 def count_column(self) -> str: 

47 return TrainingDefinitionsTable.column("id") 

48 

49 def filter_by_id( 

50 self, id_: TrainingDefinitionIdentifier 

51 ) -> TrainingDefinitionQuery: 

52 self._query.and_where(TrainingDefinitionsTable.field("id").eq(id_.value)) 

53 return self 

54 

55 def filter_by_ids( 

56 self, *ids: TrainingDefinitionIdentifier 

57 ) -> "TrainingDefinitionQuery": 

58 unpacked_ids = (id_.value for id_ in ids) 

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

60 return self