Coverage for api/utils/moderation_lock.py: 36%

1import functools 

2import time 

3 

4import django_redis 

5import structlog 

6from redis.exceptions import ConnectionError 

7 

8from api.models.moderation import get_moderators 

9 

10 

11LOCK_PREFIX = "moderation_lock" 

12TTL = 10 # seconds 

13 

14logger = structlog.get_logger(__name__) 

15 

16 

17def handle_redis_exception(func): 

18 @functools.wraps(func) 

19 def wrapper(*args, **kwargs): 

20 try: 

21 return func(*args, **kwargs) 

22 except ConnectionError: 

23 return None 

24 

25 return wrapper 

26 

27 

28class LockManager: 

29 """ 

30 Kudos to this Google Group discussion for the solution using a 

31 ranked-set: 

32 https://web.archive.org/web/20211205091916/https://groups.google.com/g/redis-db/c/rXXMCLNkNSs 

33 """ 

34 

35 def __init__(self, media_type): 

36 self.media_type = media_type 

37 

38 @handle_redis_exception 

39 def prune(self) -> dict[str, set[str]]: 

40 """ 

41 Delete all expired locks and get a mapping of usernames to 

42 media items that have active locks. 

43 

44 :return: a mapping of moderators to media items they are viewing 

45 """ 

46 

47 redis = django_redis.get_redis_connection("default") 

48 valid_locks = {} 

49 

50 now = int(time.time()) 

51 pipe = redis.pipeline() 

52 for username in get_moderators().values_list("username", flat=True): 

53 key = f"{LOCK_PREFIX}:{username}" 

54 for value, score in redis.zrange(key, 0, -1, withscores=True): 

55 if score <= now: 

56 logger.info("Deleting expired lock", key=key, value=value) 

57 pipe.zrem(key, value) 

58 else: 

59 logger.info("Keeping valid lock", key=key, value=value) 

60 valid_locks.setdefault(username, set()).add(value.decode()) 

61 pipe.execute() 

62 

63 return valid_locks 

64 

65 @handle_redis_exception 

66 def add_locks(self, username, object_id) -> int: 

67 """ 

68 Add a soft-lock for a given media item to the given moderator. 

69 

70 :param username: the username of the moderator viewing a media item 

71 :param object_id: the ID of the media item being viewed 

72 """ 

73 

74 redis = django_redis.get_redis_connection("default") 

75 

76 object = f"{self.media_type}:{object_id}" 

77 

78 expiration = int(time.time()) + TTL 

79 logger.info("Adding lock", object=object, user=username, expiration=expiration) 

80 redis.zadd(f"{LOCK_PREFIX}:{username}", {object: expiration}) 

81 return expiration 

82 

83 @handle_redis_exception 

84 def remove_locks(self, username, object_id): 

85 """ 

86 Remove the soft-lock for a given media item from the given moderator. 

87 

88 :param username: the username of the moderator not viewing a media item 

89 :param object_id: the ID of the media item not being viewed 

90 """ 

91 

92 redis = django_redis.get_redis_connection("default") 

93 

94 object = f"{self.media_type}:{object_id}" 

95 

96 logger.info("Removing lock", object=object, user=username) 

97 redis.zrem(f"{LOCK_PREFIX}:{username}", object) 

98 

99 def moderator_set(self, object_id) -> set[str]: 

100 """ 

101 Get the list of moderators on a particular item. 

102 

103 :param object_id: the ID of the media item being viewed 

104 :return: the list of moderators on a particular item 

105 """ 

106 

107 valid_locks = self.prune() or {} 

108 

109 object = f"{self.media_type}:{object_id}" 

110 mods = {mod for mod, objects in valid_locks.items() if object in objects} 

111 logger.info("Retrieved moderators", object=object, mods=mods) 

112 return mods 

113 

114 @handle_redis_exception 

115 def score(self, username, object_id) -> int: 

116 """ 

117 Get the score of a particular moderator on a particular item. 

118 

119 :param username: the username of the moderator viewing a media item 

120 :param object_id: the ID of the media item being viewed 

121 :return: the score of a particular moderator on a particular item 

122 """ 

123 

124 redis = django_redis.get_redis_connection("default") 

125 

126 object = f"{self.media_type}:{object_id}" 

127 score = redis.zscore(f"{LOCK_PREFIX}:{username}", object) 

128 logger.info("Retrieved score", object=object, user=username, score=score) 

129 return score