Coverage for mongo/utils.py: 100%

1import abc 

2import hashlib 

3import logging 

4from functools import wraps 

5from typing import Dict, Optional, Any, TYPE_CHECKING 

6from minio import Minio 

7import redis 

8from config import settings 

9from . import engine 

10 

11if TYPE_CHECKING: 

12 from .user import User # pragma: no cover 

13 from .problem import Problem # pragma: no cover 

14 

15__all__ = ( 

16 'hash_id', 

17 'is_testing', 

18 'perm', 

19 'RedisCache', 

20 'doc_required', 

21 'drop_none', 

22) 

23 

24 

25def is_testing() -> bool: 

26 '''Return True if the app is running in test mode.''' 

27 return settings.TESTING 

28 

29 

30def hash_id(salt, text): 

31 text = ((salt or '') + (text or '')).encode() 

32 sha = hashlib.sha3_512(text) 

33 return sha.hexdigest()[:24] 

34 

35 

36def perm(course, user): 

37 '''4: admin, 3: teacher, 2: TA, 1: student, 0: not found 

38 ''' 

39 return 4 - [ 

40 user.role == 0, user == course.teacher, user in course.tas, 

41 user.username in course.student_nicknames.keys(), True 

42 ].index(True) 

43 

44 

45class Cache(abc.ABC): 

46 

47 @abc.abstractmethod 

48 def exists(self, key: str) -> bool: 

49 ''' 

50 check whether a value exists 

51 ''' 

52 raise NotImplementedError # pragma: no cover 

53 

54 @abc.abstractmethod 

55 def get(self, key: str): 

56 ''' 

57 get value by key 

58 ''' 

59 raise NotImplementedError # pragma: no cover 

60 

61 @abc.abstractmethod 

62 def set(self, key: str, value, ex: Optional[int] = None): 

63 ''' 

64 set a value and set expire time in seconds 

65 ''' 

66 raise NotImplementedError # pragma: no cover 

67 

68 @abc.abstractmethod 

69 def delete(self, key: str): 

70 ''' 

71 delete a value by key 

72 ''' 

73 raise NotImplementedError # pragma: no cover 

74 

75 

76class RedisCache(Cache): 

77 POOL = None 

78 

79 def __new__(cls) -> Any: 

80 if cls.POOL is None: 

81 cls.HOST = settings.REDIS_HOST 

82 cls.PORT = settings.REDIS_PORT 

83 cls.POOL = redis.ConnectionPool( 

84 host=cls.HOST, 

85 port=cls.PORT, 

86 db=0, 

87 ) 

88 

89 return super().__new__(cls) 

90 

91 def __init__(self) -> None: 

92 self._client = None 

93 

94 @property 

95 def client(self): 

96 if self._client is None: 

97 if self.PORT is None: 

98 import fakeredis 

99 self._client = fakeredis.FakeStrictRedis() 

100 else: 

101 self._client = redis.Redis(connection_pool=self.POOL) 

102 return self._client 

103 

104 def exists(self, key: str) -> bool: 

105 return self.client.exists(key) 

106 

107 def get(self, key: str): 

108 return self.client.get(key) 

109 

110 def delete(self, key: str): 

111 return self.client.delete(key) 

112 

113 def set(self, key: str, value, ex: Optional[int] = None): 

114 return self.client.set(key, value, ex=ex) 

115 

116 

117def doc_required( 

118 src, 

119 des, 

120 cls=None, 

121 src_none_allowed=False, 

122): 

123 ''' 

124 query db to inject document into functions. 

125 if the document does not exist in db, raise `engine.DoesNotExist`. 

126 if `src` not in parameters, this funtcion will raise `TypeError` 

127 `doc_required` will check the existence of `des` in `func` parameters, 

128 if `des` is exist, this function will override it, so `src == des` 

129 are acceptable 

130 ''' 

131 # user the same name for `src` and `des` 

132 # e.g. `doc_required('user', User)` will replace parameter `user` 

133 if cls is None: 

134 cls = des 

135 des = src 

136 

137 def deco(func): 

138 

139 @wraps(func) 

140 def wrapper(*args, **ks): 

141 # try get source param 

142 if src not in ks: 

143 raise TypeError(f'{src} not found in function argument') 

144 src_param = ks.get(src) 

145 # convert it to document 

146 # TODO: add type checking, whether the cls is a subclass of `MongoBase` 

147 # or maybe it is not need 

148 if type(cls) != type: 

149 raise TypeError('cls must be a type') 

150 # process `None` 

151 if src_param is None: 

152 if not src_none_allowed: 

153 raise ValueError('src can not be None') 

154 doc = None 

155 elif not isinstance(src_param, cls): 

156 doc = cls(src_param) 

157 # or, it is already target class instance 

158 else: 

159 doc = src_param 

160 # not None and non-existent 

161 if doc is not None and not doc: 

162 raise engine.DoesNotExist(f'{doc} not found!') 

163 # replace original paramters 

164 del ks[src] 

165 if des in ks: 

166 logging.getLogger(__name__).warning( 

167 f'replace a existed argument in {func}') 

168 ks[des] = doc 

169 return func(*args, **ks) 

170 

171 return wrapper 

172 

173 return deco 

174 

175 

176def drop_none(d: Dict): 

177 return {k: v for k, v in d.items() if v is not None} 

178 

179 

180class MinioClient: 

181 

182 def __init__(self): 

183 self.client = Minio( 

184 settings.MINIO_HOST, 

185 access_key=settings.MINIO_ACCESS_KEY, 

186 secret_key=settings.MINIO_SECRET_KEY, 

187 secure=not settings.DEBUG, 

188 region=settings.MINIO_REGION, 

189 ) 

190 self.bucket = settings.MINIO_BUCKET