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("recipient ID"), 

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

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

45 HostRequest.status, 

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

47 # set when the user offered to host someone's public trip, so the recipient is the traveller 

48 HostRequest.public_trip_id.label("public trip ID"), 

49 ) 

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

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

52 .where(HostRequest.initiator_user_id == user_id) 

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

54 ) 

55 .mappings() 

56 .all() 

57 ) 

58 

59 

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

61 return ( 

62 session.execute( 

63 select( 

64 FriendRelationship.time_sent, 

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

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

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

68 FriendRelationship.status, 

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

70 ) 

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

72 .where(FriendRelationship.from_user_id == user_id) 

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

74 ) 

75 .mappings() 

76 .all() 

77 ) 

78 

79 

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

81 return ( 

82 session.execute( 

83 select( 

84 Conversation.id, 

85 Conversation.created, 

86 GroupChat.title, 

87 GroupChat.is_dm, 

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

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

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

91 ) 

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

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

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

95 .where(GroupChat.creator_id == user_id) 

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

97 .where(GroupChatSubscription.left == None) 

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

99 .where(User.id != user_id) 

100 ) 

101 .mappings() 

102 .all() 

103 ) 

104 

105 

106RATE_LIMIT_DEFINITIONS = { 

107 RateLimitAction.host_request: RateLimitDefinition( 

108 warning_limit=20, 

109 hard_limit=80, 

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

111 select(func.count()) 

112 .select_from(HostRequest) 

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

114 .where(HostRequest.initiator_user_id == user_id) 

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

116 ).scalar_one(), 

117 mod_email_information_query=_get_user_host_requests_in_past_time_interval, 

118 ), 

119 RateLimitAction.friend_request: RateLimitDefinition( 

120 warning_limit=10, 

121 hard_limit=40, 

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

123 select(func.count()) 

124 .select_from(FriendRelationship) 

125 .where(FriendRelationship.from_user_id == user_id) 

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

127 ).scalar_one(), 

128 mod_email_information_query=_get_user_friend_requests_in_past_time_interval, 

129 ), 

130 RateLimitAction.chat_initiation: RateLimitDefinition( 

131 warning_limit=15, 

132 hard_limit=150, 

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

134 select(func.count()) 

135 .select_from(GroupChat) 

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

137 .where(GroupChat.creator_id == user_id) 

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

139 ).scalar_one(), 

140 mod_email_information_query=_get_user_initiated_chats_in_past_time_interval, 

141 ), 

142}