Coverage for src/kwai/modules/identity/user_invitations/user_invitation_db_query.py: 100%

1"""Module that implements a UserInvitationQuery for a database.""" 

2 

3from sql_smith.functions import on 

4 

5from kwai.core.db.database_query import DatabaseQuery 

6from kwai.core.domain.value_objects.email_address import EmailAddress 

7from kwai.core.domain.value_objects.timestamp import Timestamp 

8from kwai.core.domain.value_objects.unique_id import UniqueId 

9from kwai.modules.identity.user_invitations.user_invitation import ( 

10 UserInvitationIdentifier, 

11) 

12from kwai.modules.identity.user_invitations.user_invitation_query import ( 

13 UserInvitationQuery, 

14) 

15from kwai.modules.identity.user_invitations.user_invitation_tables import ( 

16 UserInvitationRow, 

17) 

18from kwai.modules.identity.users.user_tables import UserRow 

19 

20 

21class UserInvitationDbQuery(UserInvitationQuery, DatabaseQuery): 

22 """A database query for a user invitation.""" 

23 

24 def init(self): 

25 return self._query.from_(UserInvitationRow.__table_name__).join( 

26 UserRow.__table_name__, 

27 on(UserInvitationRow.column("user_id"), UserRow.column("id")), 

28 ) 

29 

30 @property 

31 def columns(self): 

32 return UserInvitationRow.get_aliases() + UserRow.get_aliases() 

33 

34 @property 

35 def count_column(self) -> str: 

36 return UserInvitationRow.column("id") 

37 

38 def filter_by_id(self, id_: UserInvitationIdentifier) -> "UserInvitationQuery": 

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

40 return self 

41 

42 def filter_by_uuid(self, uuid: UniqueId) -> "UserInvitationQuery": 

43 self._query.and_where(UserInvitationRow.field("uuid").eq(str(uuid))) 

44 return self 

45 

46 def filter_by_email(self, email: EmailAddress) -> "UserInvitationQuery": 

47 self._query.and_where(UserInvitationRow.field("email").eq(str(email))) 

48 return self 

49 

50 def filter_active(self, timestamp: Timestamp) -> "UserInvitationQuery": 

51 self._query.and_where( 

52 UserInvitationRow.field("expired_at") 

53 .gt(str(timestamp)) 

54 .and_( 

55 UserInvitationRow.field("confirmed_at") 

56 .is_null() 

57 .and_(UserInvitationRow.field("revoked").not_eq(1)) 

58 ) 

59 ) 

60 return self