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

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

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

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

9 AccessTokenRow, 

10 RefreshTokenRow, 

11) 

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

13 

14 

15class RefreshTokenDbQuery(RefreshTokenQuery, DatabaseQuery): 

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

17 

18 def init(self): 

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

20 AccessTokenRow.__table_name__, 

21 on( 

22 RefreshTokenRow.column("access_token_id"), 

23 AccessTokenRow.column("id"), 

24 ), 

25 ).join( 

26 UserAccountRow.__table_name__, 

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

28 ) 

29 

30 @property 

31 def columns(self): 

32 return ( 

33 RefreshTokenRow.get_aliases() 

34 + AccessTokenRow.get_aliases() 

35 + UserAccountRow.get_aliases() 

36 ) 

37 

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

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

40 return self 

41 

42 def filter_by_token_identifier( 

43 self, identifier: TokenIdentifier 

44 ) -> "RefreshTokenQuery": 

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

46 return self