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

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

2from sql_smith.functions import on 

3 

4from kwai.core.db.database_query import DatabaseQuery 

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

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

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

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

9 UserInvitationIdentifier, 

10) 

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

12 UserInvitationQuery, 

13) 

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

15 UserInvitationsTable, 

16) 

17from kwai.modules.identity.users.user_tables import UsersTable 

18 

19 

20class UserInvitationDbQuery(UserInvitationQuery, DatabaseQuery): 

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

22 

23 def init(self): 

24 return self._query.from_(UserInvitationsTable.table_name).join( 

25 UsersTable.table_name, 

26 on(UserInvitationsTable.column("user_id"), UsersTable.column("id")), 

27 ) 

28 

29 @property 

30 def columns(self): 

31 return UserInvitationsTable.aliases() + UsersTable.aliases() 

32 

33 @property 

34 def count_column(self) -> str: 

35 return UserInvitationsTable.column("id") 

36 

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

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

39 return self 

40 

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

42 self._query.and_where(UserInvitationsTable.field("uuid").eq(str(uuid))) 

43 return self 

44 

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

46 self._query.and_where(UserInvitationsTable.field("email").eq(str(email))) 

47 return self 

48 

49 def filter_active(self, timestamp: LocalTimestamp) -> "UserInvitationQuery": 

50 self._query.and_where( 

51 UserInvitationsTable.field("expired_at") 

52 .gt(str(timestamp)) 

53 .and_( 

54 UserInvitationsTable.field("confirmed_at") 

55 .is_null() 

56 .and_(UserInvitationsTable.field("revoked").not_eq(1)) 

57 ) 

58 ) 

59 return self