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

1"""Module that implements a refresh token query for a database.""" 

2from sql_smith.functions import on 

3 

4from kwai.core.db.database_query import DatabaseQuery 

5from kwai.modules.identity.tokens.refresh_token_query import RefreshTokenQuery 

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

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

8 AccessTokensTable, 

9 RefreshTokensTable, 

10) 

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

12 

13 

14class RefreshTokenDbQuery(RefreshTokenQuery, DatabaseQuery): 

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

16 

17 def init(self): 

18 self._query.from_(RefreshTokensTable.table_name).join( 

19 AccessTokensTable.table_name, 

20 on( 

21 RefreshTokensTable.column("access_token_id"), 

22 AccessTokensTable.column("id"), 

23 ), 

24 ).join( 

25 UserAccountsTable.table_name, 

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

27 ) 

28 

29 @property 

30 def columns(self): 

31 return ( 

32 RefreshTokensTable.aliases() 

33 + AccessTokensTable.aliases() 

34 + UserAccountsTable.aliases() 

35 ) 

36 

37 def filter_by_id(self, id_: int) -> "RefreshTokenQuery": 

38 self._query.and_where(RefreshTokensTable.field("id").eq(id_)) 

39 return self 

40 

41 def filter_by_token_identifier( 

42 self, identifier: TokenIdentifier 

43 ) -> "RefreshTokenQuery": 

44 self._query.and_where( 

45 RefreshTokensTable.field("identifier").eq(str(identifier)) 

46 ) 

47 return self