Coverage for src/kwai/modules/identity/tokens/access_token_db_query.py: 89%

1"""Module that implements an access token query for a database.""" 

2 

3from sql_smith.functions import on 

4 

5from kwai.core.db.database_query import DatabaseQuery 

6from kwai.modules.identity.users.user_tables import UserAccountRow 

7 

8from .access_token_query import AccessTokenQuery 

9from .token_identifier import TokenIdentifier 

10from .token_tables import AccessTokenRow 

11 

12 

13class AccessTokenDbQuery(AccessTokenQuery, DatabaseQuery): 

14 """An access token query for a database.""" 

15 

16 def init(self): 

17 """Initialize the query.""" 

18 self._query.from_(AccessTokenRow.__table_name__).join( 

19 UserAccountRow.__table_name__, 

20 on(AccessTokenRow.column("user_id"), UserAccountRow.column("id")), 

21 ) 

22 

23 @property 

24 def columns(self): 

25 """Return the columns for the query.""" 

26 return AccessTokenRow.get_aliases() + UserAccountRow.get_aliases() 

27 

28 def filter_by_id(self, id_: int) -> "AccessTokenQuery": 

29 """Add a filter for id.""" 

30 self._query.and_where(AccessTokenRow.field("id").eq(id_)) 

31 return self 

32 

33 def filter_by_token_identifier( 

34 self, identifier: TokenIdentifier 

35 ) -> "AccessTokenQuery": 

36 """Add a filter for a token identifier.""" 

37 self._query.and_where(AccessTokenRow.field("identifier").eq(str(identifier))) 

38 return self