Coverage for bc/kwai-bc-training/src/kwai_bc_training/teams/team_db_query.py: 100%

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

2 

3from typing import Self 

4 

5from kwai_core.db.database_query import DatabaseQuery 

6 

7from kwai_bc_training.teams.team import TeamIdentifier 

8from kwai_bc_training.teams.team_query import TeamQuery 

9from kwai_bc_training.teams.team_tables import TeamRow 

10 

11 

12class TeamDbQuery(DatabaseQuery, TeamQuery): 

13 """A database query for teams.""" 

14 

15 def init(self): 

16 self._query.from_(TeamRow.__table_name__) 

17 

18 @property 

19 def columns(self): 

20 return TeamRow.get_aliases() 

21 

22 def filter_by_ids(self, *ids: TeamIdentifier) -> Self: 

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

24 self._query.and_where(TeamRow.field("id").in_(*unpacked_ids)) 

25 return self 

26 

27 def filter_by_id(self, id_: TeamIdentifier) -> Self: 

28 self._query.and_where(TeamRow.field("id").eq(id_.value)) 

29 return self