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

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

2from sql_smith.functions import on 

3 

4from kwai.core.db.database_query import DatabaseQuery 

5from kwai.modules.identity.users.user_tables import UserAccountsTable 

6 

7from .access_token_query import AccessTokenQuery 

8from .token_identifier import TokenIdentifier 

9from .token_tables import AccessTokensTable 

10 

11 

12class AccessTokenDbQuery(AccessTokenQuery, DatabaseQuery): 

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

14 

15 def init(self): 

16 """Initialize the query.""" 

17 self._query.from_(AccessTokensTable.table_name).join( 

18 UserAccountsTable.table_name, 

19 on(AccessTokensTable.column("user_id"), UserAccountsTable.column("id")), 

20 ) 

21 

22 @property 

23 def columns(self): 

24 """Return the columns for the query.""" 

25 return AccessTokensTable.aliases() + UserAccountsTable.aliases() 

26 

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

28 """Add a filter for id.""" 

29 self._query.and_where(AccessTokensTable.field("id").eq(id_)) 

30 return self 

31 

32 def filter_by_token_identifier( 

33 self, identifier: TokenIdentifier 

34 ) -> "AccessTokenQuery": 

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

36 self._query.and_where(AccessTokensTable.field("identifier").eq(str(identifier))) 

37 return self