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

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

2from sql_smith.functions import on 

3 

4from kwai.core.db.database_query import DatabaseQuery 

5from kwai.core.db.rows import OwnersTable 

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

7 TrainingDefinitionIdentifier, 

8) 

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

10 TrainingDefinitionQuery, 

11) 

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

13 TrainingDefinitionsTable, 

14) 

15 

16 

17class TrainingDefinitionDbQuery(DatabaseQuery, TrainingDefinitionQuery): 

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

19 

20 def init(self): 

21 return self._query.from_(TrainingDefinitionsTable.table_name).left_join( 

22 OwnersTable.table_name, 

23 on(OwnersTable.column("id"), TrainingDefinitionsTable.column("user_id")), 

24 ) 

25 

26 @property 

27 def columns(self): 

28 return TrainingDefinitionsTable.aliases() + OwnersTable.aliases() 

29 

30 def filter_by_id( 

31 self, id_: TrainingDefinitionIdentifier 

32 ) -> TrainingDefinitionQuery: 

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

34 return self 

35 

36 def filter_by_ids( 

37 self, *ids: TrainingDefinitionIdentifier 

38 ) -> "TrainingDefinitionQuery": 

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

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

41 return self