Coverage for src/couchers/rate_limits/definitions.py: 96%

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 dataclasses import dataclass 

7from datetime import timedelta 

8from typing import TYPE_CHECKING, Callable 

9 

10from sqlalchemy import func, select 

11 

12from couchers.models import ( 

13 Conversation, 

14 FriendRelationship, 

15 GroupChat, 

16 GroupChatSubscription, 

17 HostRequest, 

18 RateLimitAction, 

19 User, 

20) 

21from couchers.utils import now 

22 

23if TYPE_CHECKING: 

24 from sqlalchemy.orm import Session 

25 

26 

27@dataclass 

28class RateLimitDefinition: 

29 warning_limit: int 

30 hard_limit: int 

31 count_actions_query: Callable[["Session", int], int] 

32 mod_email_information_query: Callable[["Session", int], list[dict]] 

33 

34 

35RATE_LIMIT_INTERVAL = timedelta(hours=24) 

36RATE_LIMIT_INTERVAL_STRING = "24 hours" 

37 

38 

39def _get_user_host_requests_in_past_time_interval(session, user_id) -> list[dict]: 

40 return ( 

41 session.execute( 

42 select( 

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

44 HostRequest.host_user_id.label("host ID"), 

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

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

47 ) 

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

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

50 .where(HostRequest.surfer_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, user_id) -> list[dict]: 

59 return ( 

60 session.execute( 

61 select( 

62 FriendRelationship.time_sent, 

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

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

65 FriendRelationship.status, 

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

67 ) 

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

69 .where(FriendRelationship.from_user_id == user_id) 

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

71 ) 

72 .mappings() 

73 .all() 

74 ) 

75 

76 

77def _get_user_initiated_chats_in_past_time_interval(session, user_id) -> list[dict]: 

78 return ( 

79 session.execute( 

80 select( 

81 Conversation.id, 

82 Conversation.created, 

83 GroupChat.title, 

84 GroupChat.is_dm, 

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

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

87 ) 

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

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

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

91 .where(GroupChat.creator_id == user_id) 

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

93 .where(GroupChatSubscription.left == None) 

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

95 .where(User.id != user_id) 

96 ) 

97 .mappings() 

98 .all() 

99 ) 

100 

101 

102RATE_LIMIT_DEFINITIONS = { 

103 RateLimitAction.host_request: RateLimitDefinition( 

104 warning_limit=20, 

105 hard_limit=80, 

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

107 select(func.count()) 

108 .select_from(HostRequest) 

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

110 .where(HostRequest.surfer_user_id == user_id) 

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

112 ).scalar_one(), 

113 mod_email_information_query=_get_user_host_requests_in_past_time_interval, 

114 ), 

115 RateLimitAction.friend_request: RateLimitDefinition( 

116 warning_limit=10, 

117 hard_limit=40, 

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

119 select(func.count()) 

120 .select_from(FriendRelationship) 

121 .where(FriendRelationship.from_user_id == user_id) 

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

123 ).scalar_one(), 

124 mod_email_information_query=_get_user_friend_requests_in_past_time_interval, 

125 ), 

126 RateLimitAction.chat_initiation: RateLimitDefinition( 

127 warning_limit=15, 

128 hard_limit=150, 

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

130 select(func.count()) 

131 .select_from(GroupChat) 

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

133 .where(GroupChat.creator_id == user_id) 

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

135 ).scalar_one(), 

136 mod_email_information_query=_get_user_initiated_chats_in_past_time_interval, 

137 ), 

138}