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

1"""Module that defines all tables for tokens.""" 

2from dataclasses import dataclass 

3from datetime import datetime 

4 

5from kwai.core.db.table import Table 

6from kwai.core.domain.value_objects.local_timestamp import LocalTimestamp 

7from kwai.core.domain.value_objects.traceable_time import TraceableTime 

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

9 AccessTokenEntity, 

10 AccessTokenIdentifier, 

11) 

12from kwai.modules.identity.tokens.refresh_token import ( 

13 RefreshTokenEntity, 

14 RefreshTokenIdentifier, 

15) 

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

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

18 

19 

20@dataclass(kw_only=True, frozen=True, slots=True) 

21class AccessTokenRow: 

22 """Represent a table row in the access tokens table.""" 

23 

24 id: int | None 

25 identifier: str 

26 expiration: datetime 

27 user_id: int 

28 revoked: bool 

29 created_at: datetime 

30 updated_at: datetime 

31 

32 def create_entity(self, user_account: UserAccountEntity) -> AccessTokenEntity: 

33 """Create an entity from the table row.""" 

34 return AccessTokenEntity( 

35 id_=AccessTokenIdentifier(self.id), 

36 identifier=TokenIdentifier(hex_string=self.identifier), 

37 expiration=self.expiration, 

38 user_account=user_account, 

39 revoked=self.revoked, 

40 traceable_time=TraceableTime( 

41 created_at=LocalTimestamp(timestamp=self.created_at), 

42 updated_at=LocalTimestamp(timestamp=self.updated_at), 

43 ), 

44 ) 

45 

46 @classmethod 

47 def persist(cls, access_token: AccessTokenEntity) -> "AccessTokenRow": 

48 """Persist an access token entity to a table record.""" 

49 return AccessTokenRow( 

50 id=access_token.id.value, 

51 identifier=str(access_token.identifier), 

52 expiration=access_token.expiration, 

53 user_id=access_token.user_account.id.value, 

54 revoked=access_token.revoked, 

55 created_at=access_token.traceable_time.created_at.timestamp, 

56 updated_at=access_token.traceable_time.updated_at.timestamp, 

57 ) 

58 

59 

60AccessTokensTable = Table("oauth_access_tokens", AccessTokenRow) 

61 

62 

63@dataclass(kw_only=True, frozen=True, slots=True) 

64class RefreshTokenRow: 

65 """Represent a table row in the refresh token table.""" 

66 

67 id: int | None 

68 identifier: str 

69 access_token_id: int 

70 expiration: datetime 

71 revoked: bool 

72 created_at: datetime 

73 updated_at: datetime 

74 

75 def create_entity(self, access_token: AccessTokenEntity) -> RefreshTokenEntity: 

76 """Create a refresh token entity from the table row.""" 

77 return RefreshTokenEntity( 

78 id_=RefreshTokenIdentifier(self.id), 

79 identifier=TokenIdentifier(hex_string=self.identifier), 

80 access_token=access_token, 

81 expiration=self.expiration, 

82 revoked=self.revoked, 

83 traceable_time=TraceableTime( 

84 created_at=LocalTimestamp(timestamp=self.created_at), 

85 updated_at=LocalTimestamp(timestamp=self.updated_at), 

86 ), 

87 ) 

88 

89 @classmethod 

90 def persist(cls, refresh_token: RefreshTokenEntity) -> "RefreshTokenRow": 

91 """Transform a refresh token entity into a table record.""" 

92 return RefreshTokenRow( 

93 id=refresh_token.id.value, 

94 identifier=str(refresh_token.identifier), 

95 access_token_id=refresh_token.access_token.id.value, 

96 expiration=refresh_token.expiration, 

97 revoked=refresh_token.revoked, 

98 created_at=refresh_token.traceable_time.created_at.timestamp, 

99 updated_at=refresh_token.traceable_time.updated_at.timestamp, 

100 ) 

101 

102 

103RefreshTokensTable = Table("oauth_refresh_tokens", RefreshTokenRow)