Coverage for src/kwai/modules/training/teams/team_db_query.py: 100%

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

2 

3from kwai.core.db.database_query import DatabaseQuery 

4from kwai.modules.training.teams.team import TeamIdentifier 

5from kwai.modules.training.teams.team_query import TeamQuery 

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

7 

8 

9class TeamDbQuery(DatabaseQuery, TeamQuery): 

10 """A database query for teams.""" 

11 

12 def init(self): 

13 self._query.from_(TeamsTable.table_name) 

14 

15 @property 

16 def columns(self): 

17 return TeamsTable.aliases() 

18 

19 def filter_by_ids(self, *ids: TeamIdentifier) -> "TeamQuery": 

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

21 self._query.and_where(TeamsTable.field("id").in_(*unpacked_ids)) 

22 return self 

23 

24 def filter_by_id(self, id_: TeamIdentifier) -> "TeamQuery": 

25 self._query.and_where(TeamsTable.field("id").eq(id_.value)) 

26 return self