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

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

2 

3from typing import Self 

4 

5from sql_smith.functions import on 

6 

7from kwai.core.db.database_query import DatabaseQuery 

8from kwai.modules.identity.users.user_account import UserAccountEntity 

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

10 

11from .access_token import AccessTokenIdentifier 

12from .access_token_query import AccessTokenQuery 

13from .token_identifier import TokenIdentifier 

14from .token_tables import AccessTokenRow 

15 

16 

17class AccessTokenDbQuery(AccessTokenQuery, DatabaseQuery): 

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

19 

20 def init(self): 

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

22 UserAccountRow.__table_name__, 

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

24 ) 

25 

26 @property 

27 def columns(self): 

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

29 

30 def filter_by_id(self, id_: AccessTokenIdentifier) -> "AccessTokenQuery": 

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

32 return self 

33 

34 def filter_by_token_identifier( 

35 self, identifier: TokenIdentifier 

36 ) -> "AccessTokenQuery": 

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

38 return self 

39 

40 def filter_by_user_account(self, user_account: UserAccountEntity) -> Self: 

41 self._query.and_where(AccessTokenRow.field("user_id").eq(user_account.id.value)) 

42 return self