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

1"""Module that implements a refresh 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.tokens.refresh_token_query import RefreshTokenQuery 

9from kwai.modules.identity.tokens.token_identifier import TokenIdentifier 

10from kwai.modules.identity.tokens.token_tables import ( 

11 AccessTokenRow, 

12 RefreshTokenRow, 

13) 

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

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

16 

17 

18class RefreshTokenDbQuery(RefreshTokenQuery, DatabaseQuery): 

19 """A refresh token query for a database.""" 

20 

21 def init(self): 

22 self._query.from_(RefreshTokenRow.__table_name__).join( 

23 AccessTokenRow.__table_name__, 

24 on( 

25 RefreshTokenRow.column("access_token_id"), 

26 AccessTokenRow.column("id"), 

27 ), 

28 ).join( 

29 UserAccountRow.__table_name__, 

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

31 ) 

32 

33 @property 

34 def columns(self): 

35 return ( 

36 RefreshTokenRow.get_aliases() 

37 + AccessTokenRow.get_aliases() 

38 + UserAccountRow.get_aliases() 

39 ) 

40 

41 def filter_by_id(self, id_: int) -> Self: 

42 self._query.and_where(RefreshTokenRow.field("id").eq(id_)) 

43 return self 

44 

45 def filter_by_token_identifier(self, identifier: TokenIdentifier) -> Self: 

46 self._query.and_where(RefreshTokenRow.field("identifier").eq(str(identifier))) 

47 return self 

48 

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

50 self._query.and_where(UserAccountRow.field("id").eq(user_account.id.value)) 

51 return self