Coverage for app / backend / src / couchers / rate_limits / definitions.py: 100%

1"""Rate limit definitions: 

2In order to add a new rate limit definition, extend RateLimitAction and RATE_LIMIT_DEFINITIONS and call 

3rate_limits.check.process_rate_limits_and_check_abort in the relevant endpoint. 

4""" 

5 

6from collections.abc import Callable, Sequence 

7from dataclasses import dataclass 

8from datetime import timedelta 

9 

10from sqlalchemy import RowMapping, func, select 

11from sqlalchemy.orm import Session 

12 

13from couchers.models import ( 

14 Conversation, 

15 FriendRelationship, 

16 GroupChat, 

17 GroupChatSubscription, 

18 HostRequest, 

19 RateLimitAction, 

20 User, 

21) 

22from couchers.utils import now 

23 

24 

25@dataclass 

26class RateLimitDefinition: 

27 warning_limit: int 

28 hard_limit: int 

29 count_actions_query: Callable[[Session, int], int] 

30 mod_email_information_query: Callable[[Session, int], Sequence[RowMapping]] 

31 

32 

33RATE_LIMIT_HOURS = 24 

34RATE_LIMIT_INTERVAL = timedelta(hours=RATE_LIMIT_HOURS) 

35 

36 

37def _get_user_host_requests_in_past_time_interval(session: Session, user_id: int) -> Sequence[RowMapping]: 

38 return ( 

39 session.execute( 

40 select( 

41 Conversation.created.label("created"), 

42 HostRequest.recipient_user_id.label("host ID"), 

43 User.gender.label("host gender"), 

44 User.username.label("host username"), 

45 HostRequest.status, 

46 User.city.label("host city"), 

47 ) 

48 .join(Conversation, HostRequest.conversation_id == Conversation.id) 

49 .join(User, HostRequest.recipient_user_id == User.id) 

50 .where(HostRequest.initiator_user_id == user_id) 

51 .where(Conversation.created >= now() - RATE_LIMIT_INTERVAL) 

52 ) 

53 .mappings() 

54 .all() 

55 ) 

56 

57 

58def _get_user_friend_requests_in_past_time_interval(session: Session, user_id: int) -> Sequence[RowMapping]: 

59 return ( 

60 session.execute( 

61 select( 

62 FriendRelationship.time_sent, 

63 User.id.label("recipient ID"), 

64 User.gender.label("recipient gender"), 

65 User.username.label("recipient username"), 

66 FriendRelationship.status, 

67 User.city.label("recipient city"), 

68 ) 

69 .join(User, FriendRelationship.to_user_id == User.id) 

70 .where(FriendRelationship.from_user_id == user_id) 

71 .where(FriendRelationship.time_sent >= now() - RATE_LIMIT_INTERVAL) 

72 ) 

73 .mappings() 

74 .all() 

75 ) 

76 

77 

78def _get_user_initiated_chats_in_past_time_interval(session: Session, user_id: int) -> Sequence[RowMapping]: 

79 return ( 

80 session.execute( 

81 select( 

82 Conversation.id, 

83 Conversation.created, 

84 GroupChat.title, 

85 GroupChat.is_dm, 

86 func.array_agg(User.username).label("participants"), 

87 func.array_agg(User.gender).label("participants genders"), 

88 func.array_agg(User.city).label("participants cities"), 

89 ) 

90 .join(Conversation, GroupChat.conversation_id == Conversation.id) 

91 .join(GroupChatSubscription, Conversation.id == GroupChatSubscription.group_chat_id) 

92 .join(User, GroupChatSubscription.user_id == User.id) 

93 .where(GroupChat.creator_id == user_id) 

94 .where(Conversation.created >= now() - RATE_LIMIT_INTERVAL) 

95 .where(GroupChatSubscription.left == None) 

96 .group_by(Conversation.id, Conversation.created, GroupChat.title, GroupChat.is_dm) 

97 .where(User.id != user_id) 

98 ) 

99 .mappings() 

100 .all() 

101 ) 

102 

103 

104RATE_LIMIT_DEFINITIONS = { 

105 RateLimitAction.host_request: RateLimitDefinition( 

106 warning_limit=20, 

107 hard_limit=80, 

108 count_actions_query=lambda session, user_id: session.execute( 

109 select(func.count()) 

110 .select_from(HostRequest) 

111 .join(Conversation, HostRequest.conversation_id == Conversation.id) 

112 .where(HostRequest.initiator_user_id == user_id) 

113 .where(Conversation.created >= now() - RATE_LIMIT_INTERVAL) 

114 ).scalar_one(), 

115 mod_email_information_query=_get_user_host_requests_in_past_time_interval, 

116 ), 

117 RateLimitAction.friend_request: RateLimitDefinition( 

118 warning_limit=10, 

119 hard_limit=40, 

120 count_actions_query=lambda session, user_id: session.execute( 

121 select(func.count()) 

122 .select_from(FriendRelationship) 

123 .where(FriendRelationship.from_user_id == user_id) 

124 .where(FriendRelationship.time_sent >= now() - RATE_LIMIT_INTERVAL) 

125 ).scalar_one(), 

126 mod_email_information_query=_get_user_friend_requests_in_past_time_interval, 

127 ), 

128 RateLimitAction.chat_initiation: RateLimitDefinition( 

129 warning_limit=15, 

130 hard_limit=150, 

131 count_actions_query=lambda session, user_id: session.execute( 

132 select(func.count()) 

133 .select_from(GroupChat) 

134 .join(Conversation, GroupChat.conversation_id == Conversation.id) 

135 .where(GroupChat.creator_id == user_id) 

136 .where(Conversation.created >= now() - RATE_LIMIT_INTERVAL) 

137 ).scalar_one(), 

138 mod_email_information_query=_get_user_initiated_chats_in_past_time_interval, 

139 ), 

140}