Coverage for /home/runner/work/viur-core/viur-core/viur/src/viur/core/db/cache.py: 76%

1import datetime 

2import logging 

3import sys 

4import typing as t 

5 

6from viur.core.config import conf 

7from viur.core import utils 

8from .types import Entity, Key 

9 

10MEMCACHE_MAX_BATCH_SIZE = 30 

11MEMCACHE_NAMESPACE = "viur-datastore" 

12MEMCACHE_TIMEOUT: int | datetime.timedelta = datetime.timedelta(days=1) 

13MEMCACHE_MAX_SIZE: t.Final[int] = 1_000_000 

14TESTBED = None 

15""" 

16 

17 This Module controls the Interaction with the Memcache from Google 

18 To activate the cache copy this code in your main.py 

19 .. code-block:: python 

20 # Example 

21 from viur.core import conf 

22 from google.appengine.api.memcache import Client 

23 conf.db.memcache_client = Client() 

24""" 

25 

26__all__ = [ 

27 "MEMCACHE_MAX_BATCH_SIZE", 

28 "MEMCACHE_NAMESPACE", 

29 "MEMCACHE_TIMEOUT", 

30 "MEMCACHE_MAX_SIZE", 

31 "get", 

32 "put", 

33 "delete", 

34 "flush", 

35] 

36 

37 

38def get(keys: t.Union[Key, t.Iterable[Key]], namespace: t.Optional[str] = None) -> list[Entity]: 

39 """ 

40 Reads data form the memcache. 

41 :param keys: Unique identifier(s) for one or more entry(s). 

42 :param namespace: Optional namespace to use. 

43 :return: The entities that were found, in arbitrary order. Always a list, even for a single key, 

44 because an Entity is dict-like and callers must not have to tell a hit from an iterable. 

45 """ 

46 # Inside a transaction reads must go straight to the datastore, so the cache 

47 # cannot serve a (potentially stale) value. Lazy import to avoid a cycle. 

48 from .utils import is_in_transaction 

49 if is_in_transaction(): 

50 return [] 

51 

52 if not check_for_memcache(): 

53 return [] 

54 

55 namespace = namespace or MEMCACHE_NAMESPACE 

56 keys = utils.ensure_iterable(keys) 

57 keys = [str(key) for key in keys] # Enforce that all keys are strings 

58 cached_data_result = {} 

59 result = [] 

60 try: 

61 while keys: 

62 if cached_data := conf.db.memcache_client.get_multi(keys[:MEMCACHE_MAX_BATCH_SIZE], namespace=namespace): 

63 cached_data_result |= cached_data 

64 keys = keys[MEMCACHE_MAX_BATCH_SIZE:] 

65 except Exception as e: 

66 logging.error(f"""Failed to get keys form the memcache with {e=}""") 

67 for key, value in cached_data_result.items(): 

68 entity = Entity(Key.from_legacy_urlsafe(key)) 

69 entity |= value 

70 result.append(entity) 

71 

72 return result 

73 

74 

75def put( 

76 data: t.Union[Entity, t.Dict[Key, Entity], t.Iterable[Entity]], 

77 namespace: t.Optional[str] = None, 

78 timeout: t.Optional[t.Union[int, datetime.timedelta]] = None 

79) -> bool: 

80 """ 

81 Writes Data to the memcache. 

82 :param data: Data to write 

83 :param namespace: Optional namespace to use. 

84 :param timeout: Optional timeout in seconds or a timedelta object. 

85 :return: A boolean indicating success. 

86 """ 

87 # Inside a transaction the write is not committed yet; caching it now would 

88 # serve values that may be rolled back. Lazy import to avoid a cycle. 

89 from .utils import is_in_transaction 

90 if is_in_transaction(): 

91 return False 

92 

93 if not check_for_memcache(): 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true

94 return False 

95 if not data: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true

96 return False 

97 namespace = namespace or MEMCACHE_NAMESPACE 

98 timeout = timeout or MEMCACHE_TIMEOUT 

99 if isinstance(timeout, datetime.timedelta): 99 ↛ 101line 99 didn't jump to line 101 because the condition on line 99 was always true

100 timeout = timeout.total_seconds() 

101 if isinstance(data, (list, tuple, set)): 

102 data = {item.key: item for item in data} 

103 elif isinstance(data, Entity): 103 ↛ 105line 103 didn't jump to line 105 because the condition on line 103 was always true

104 data = {data.key: data} 

105 elif not isinstance(data, dict): 

106 raise TypeError(f"Invalid type {type(data)}. Expected a db.Entity, list or dict.") 

107 

108 # Add only values to cache <= MEMMAX_SIZE (1.000.000) 

109 data = {str(key): value for key, value in data.items() if get_size(value) <= MEMCACHE_MAX_SIZE} 

110 

111 keys = list(data.keys()) 

112 try: 

113 while keys: 

114 data_batch = {key: data[key] for key in keys[:MEMCACHE_MAX_BATCH_SIZE]} 

115 conf.db.memcache_client.set_multi(data_batch, namespace=namespace, time=timeout) 

116 keys = keys[MEMCACHE_MAX_BATCH_SIZE:] 

117 return True 

118 except Exception as e: 

119 logging.error(f"""Failed to put data to the memcache with {e=}""") 

120 return False 

121 

122 

123def delete(keys: t.Union[Key, t.Iterable[Key]], namespace: t.Optional[str] = None) -> None: 

124 """ 

125 Deletes an Entry form memcache. 

126 :param keys: Unique identifier(s) for one or more entry(s). 

127 :param namespace: Optional namespace to use. 

128 """ 

129 if not check_for_memcache(): 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true

130 return None 

131 if not keys: 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true

132 return None 

133 namespace = namespace or MEMCACHE_NAMESPACE 

134 keys = utils.ensure_iterable(keys) 

135 keys = [str(key) for key in keys] # Enforce that all keys are strings 

136 try: 

137 while keys: 

138 conf.db.memcache_client.delete_multi(keys[:MEMCACHE_MAX_BATCH_SIZE], namespace=namespace) 

139 keys = keys[MEMCACHE_MAX_BATCH_SIZE:] 

140 except Exception as e: 

141 logging.error(f"""Failed to delete keys form the memcache with {e=}""") 

142 

143 

144def flush() -> bool: 

145 """ 

146 Deletes everything in memcache. 

147 :return: A boolean indicating success. 

148 """ 

149 if not check_for_memcache(): 

150 return False 

151 try: 

152 conf.db.memcache_client.flush_all() 

153 except Exception as e: 

154 logging.error(f"""Failed to flush the memcache with {e=}""") 

155 return False 

156 return True 

157 

158 

159def get_size(obj: t.Any) -> int: 

160 """ 

161 Utility function that counts the size of an object in bytes. 

162 """ 

163 if isinstance(obj, dict): 

164 return sum(get_size([k, v]) for k, v in obj.items()) 

165 elif isinstance(obj, list): 

166 return sum(get_size(x) for x in obj) 

167 

168 return sys.getsizeof(obj) 

169 

170 

171def check_for_memcache() -> bool: 

172 if conf.db.memcache_client is None: 

173 # logging.warning(f"""conf.db.memcache_client is 'None'. It can not be used.""") 

174 return False 

175 

176 init_testbed() 

177 return True 

178 

179 

180def init_testbed() -> None: 

181 global TESTBED 

182 if TESTBED is None and conf.instance.is_dev_server and conf.db.memcache_client: 182 ↛ 183line 182 didn't jump to line 183 because the condition on line 182 was never true

183 from google.appengine.ext.testbed import Testbed 

184 TESTBED = Testbed() 

185 TESTBED.activate() 

186 TESTBED.init_memcache_stub()