Coverage for 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.host_user_id.label("host ID"), 

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

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

45 ) 

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

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

48 .where(HostRequest.surfer_user_id == user_id) 

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

50 ) 

51 .mappings() 

52 .all() 

53 ) 

54 

55 

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

57 return ( 

58 session.execute( 

59 select( 

60 FriendRelationship.time_sent, 

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

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

63 FriendRelationship.status, 

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

65 ) 

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

67 .where(FriendRelationship.from_user_id == user_id) 

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

69 ) 

70 .mappings() 

71 .all() 

72 ) 

73 

74 

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

76 return ( 

77 session.execute( 

78 select( 

79 Conversation.id, 

80 Conversation.created, 

81 GroupChat.title, 

82 GroupChat.is_dm, 

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

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

85 ) 

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

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

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

89 .where(GroupChat.creator_id == user_id) 

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

91 .where(GroupChatSubscription.left == None) 

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

93 .where(User.id != user_id) 

94 ) 

95 .mappings() 

96 .all() 

97 ) 

98 

99 

100RATE_LIMIT_DEFINITIONS = { 

101 RateLimitAction.host_request: RateLimitDefinition( 

102 warning_limit=20, 

103 hard_limit=80, 

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

105 select(func.count()) 

106 .select_from(HostRequest) 

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

108 .where(HostRequest.surfer_user_id == user_id) 

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

110 ).scalar_one(), 

111 mod_email_information_query=_get_user_host_requests_in_past_time_interval, 

112 ), 

113 RateLimitAction.friend_request: RateLimitDefinition( 

114 warning_limit=10, 

115 hard_limit=40, 

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

117 select(func.count()) 

118 .select_from(FriendRelationship) 

119 .where(FriendRelationship.from_user_id == user_id) 

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

121 ).scalar_one(), 

122 mod_email_information_query=_get_user_friend_requests_in_past_time_interval, 

123 ), 

124 RateLimitAction.chat_initiation: RateLimitDefinition( 

125 warning_limit=15, 

126 hard_limit=150, 

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

128 select(func.count()) 

129 .select_from(GroupChat) 

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

131 .where(GroupChat.creator_id == user_id) 

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

133 ).scalar_one(), 

134 mod_email_information_query=_get_user_initiated_chats_in_past_time_interval, 

135 ), 

136}