Coverage for /home/runner/work/viur-core/viur-core/viur/src/viur/core/db/utils.py: 17%

1import datetime 

2import fnmatch 

3import sys 

4import typing as t 

5 

6from deprecated.sphinx import deprecated 

7from google.cloud.datastore.transaction import Transaction 

8 

9from viur.core import current 

10from viur.core.config import conf 

11from .transport import __client__, get, put, run_in_transaction 

12from .types import Entity, Key, current_db_access_log 

13 

14 

15def fix_unindexable_properties(entry: Entity, *, keep_exclusions: bool = True) -> Entity: 

16 """ 

17 Recursively walk the given Entity and add all properties to the list of unindexed properties if they contain 

18 a string longer than 1500 bytes (which is maximum size of a string that can be indexed). The datastore would 

19 return an error otherwise. 

20 https://cloud.google.com/datastore/docs/concepts/limits?hl=en#limits 

21 

22 :param entry: The entity to fix (inplace) 

23 :param keep_exclusions: If true, keep the properties already included in ``exclude_from_indexes``. 

24 Otherwise, ignore them and exclude only non-indexable properties. 

25 :return: The fixed entity 

26 """ 

27 

28 def has_unindexable_property(prop): 

29 if isinstance(prop, dict): 

30 return any(has_unindexable_property(x) for x in prop.values()) 

31 elif isinstance(prop, list): 

32 return any(has_unindexable_property(x) for x in prop) 

33 elif isinstance(prop, (str, bytes)): 

34 return sys.getsizeof(prop) >= 1500 

35 else: 

36 return False 

37 

38 unindexable_properties = set() 

39 for key, value in entry.items(): 

40 if not has_unindexable_property(value): 

41 continue 

42 if isinstance(value, dict): 

43 inner_entity = Entity() 

44 inner_entity.update(value) 

45 entry[key] = fix_unindexable_properties(inner_entity) 

46 if isinstance(value, Entity): 

47 inner_entity.key = value.key 

48 else: 

49 unindexable_properties.add(key) 

50 if keep_exclusions: 

51 entry.exclude_from_indexes.update(unindexable_properties) # type:ignore 

52 else: 

53 entry.exclude_from_indexes = unindexable_properties 

54 return entry 

55 

56 

57def normalize_key(key: t.Union[None, Key, str]) -> t.Union[None, Key]: 

58 """ 

59 Normalizes a datastore key (replacing the key's project with conf.instance.project_id) 

60 

61 The key's project is only allowed to be normalized when it matches one of the patterns 

62 configured in `conf.valid_application_ids`; otherwise a ValueError is raised. 

63 

64 :param key: Key to be normalized. 

65 :return: Normalized key in string representation. 

66 """ 

67 if key is None: 

68 return None 

69 

70 if isinstance(key, str): 

71 key = Key.from_legacy_urlsafe(key) 

72 

73 if key.project != conf.instance.project_id and not any( 

74 fnmatch.fnmatch(key.project, application_id) for application_id in conf.valid_application_ids 

75 ): 

76 raise ValueError(f"{key=} cannot be normalized; Only keys from conf.valid_application_ids can be provided.") 

77 

78 if key.parent: 

79 parent = normalize_key(key.parent) 

80 else: 

81 parent = None 

82 

83 return Key(key.kind, key.id_or_name, parent=parent) 

84 

85 

86@deprecated(version="3.8.0", reason="Use 'db.normalize_key' instead") 

87def normalizeKey(key: t.Union[None, Key]) -> t.Union[None, Key]: 

88 return normalize_key(key) 

89 

90 

91def key_helper( 

92 in_key: t.Union[Key, str, int], 

93 target_kind: str, 

94 additional_allowed_kinds: t.Union[t.List[str], t.Tuple[str]] = (), 

95 adjust_kind: bool = False, 

96) -> Key: 

97 if isinstance(in_key, Key): 

98 if in_key.kind != target_kind and in_key.kind not in additional_allowed_kinds: 

99 if not adjust_kind: 

100 raise ValueError( 

101 f"Kind mismatch: {in_key.kind!r} != {target_kind!r} (or in {additional_allowed_kinds!r})") 

102 in_key = Key(target_kind, in_key.id_or_name, parent=in_key.parent) 

103 return in_key 

104 elif isinstance(in_key, str): 

105 # Try to parse key from str 

106 try: 

107 decoded_key = normalize_key(in_key) 

108 except Exception: 

109 decoded_key = None 

110 

111 # If it did decode, recall keyHelper with Key object 

112 if decoded_key: 

113 return key_helper( 

114 decoded_key, 

115 target_kind=target_kind, 

116 additional_allowed_kinds=additional_allowed_kinds, 

117 adjust_kind=adjust_kind 

118 ) 

119 

120 # otherwise, construct key from str or int 

121 if in_key.isdigit(): 

122 in_key = int(in_key) 

123 

124 return Key(target_kind, in_key) 

125 elif isinstance(in_key, int): 

126 return Key(target_kind, in_key) 

127 

128 raise NotImplementedError(f"Unsupported key type {type(in_key)}") 

129 

130 

131def keyHelper( 

132 inKey: t.Union[Key, str, int], 

133 targetKind: str, 

134 additionalAllowedKinds: t.Union[t.List[str], t.Tuple[str]] = (), 

135 adjust_kind: bool = False, 

136) -> Key: 

137 return key_helper( 

138 in_key=inKey, 

139 target_kind=targetKind, 

140 additional_allowed_kinds=additionalAllowedKinds, 

141 adjust_kind=adjust_kind 

142 ) 

143 

144 

145def is_in_transaction() -> bool: 

146 return __client__.current_transaction is not None 

147 

148 

149@deprecated(version="3.8.0", reason="Use 'db.utils.is_in_transaction' instead") 

150def IsInTransaction() -> bool: 

151 return is_in_transaction() 

152 

153 

154def get_or_insert(key: Key, **kwargs) -> Entity: 

155 """ 

156 Either creates a new entity with the given key, or returns the existing one. 

157 

158 Its guaranteed that there is no race-condition here; it will never overwrite a 

159 previously created entity. Extra keyword arguments passed to this function will be 

160 used to populate the entity if it has to be created; otherwise they are ignored. 

161 

162 :param key: The key which will be fetched or created. 

163 :returns: Returns the fetched or newly created Entity. 

164 """ 

165 

166 def txn(key, kwargs): 

167 obj = get(key) 

168 if not obj: 

169 obj = Entity(key) 

170 for k, v in kwargs.items(): 

171 obj[k] = v 

172 put(obj) 

173 return obj 

174 

175 if is_in_transaction(): 

176 return txn(key, kwargs) 

177 return run_in_transaction(txn, key, kwargs) 

178 

179 

180@deprecated(version="3.8.0", reason="Use 'db.get_or_insert' instead") 

181def GetOrInsert(key: Key, **kwargs: t.Any) -> Entity: 

182 return get_or_insert(key, **kwargs) 

183 

184 

185@deprecated(version="3.8.0", reason="Use 'str(key)' instead") 

186def encodeKey(key: Key) -> str: 

187 """ 

188 Return the given key encoded as string (mimicking the old str() behaviour of keys) 

189 """ 

190 return str(key) 

191 

192 

193def acquire_transaction_success_marker() -> str: 

194 """ 

195 Generates a token that will be written to the datastore (under "viur-transactionmarker") if the transaction 

196 completes successfully. Currently only used by deferredTasks to check if the task should actually execute 

197 or if the transaction it was created in failed. 

198 :return: Name of the entry in viur-transactionmarker 

199 """ 

200 txn: Transaction | None = __client__.current_transaction 

201 assert txn, "acquire_transaction_success_marker cannot be called outside an transaction" 

202 marker = str(txn.id) 

203 request_data = current.request_data.get() 

204 if not request_data.get("__viur-transactionmarker__"): 

205 db_obj = Entity(Key("viur-transactionmarker", marker)) 

206 db_obj["creationdate"] = datetime.datetime.now(datetime.timezone.utc) 

207 put(db_obj) 

208 request_data["__viur-transactionmarker__"] = True 

209 return marker 

210 

211 

212def start_data_access_log() -> t.Set[t.Union[Key, str]]: 

213 """ 

214 Clears our internal access log (which keeps track of which entries have been accessed in the current 

215 request). The old set of accessed entries is returned so that it can be restored with 

216 :func:`server.db.popAccessData` in case of nested caching. You must call popAccessData afterwards, otherwise 

217 we'll continue to log all entries accessed in subsequent request on the same thread! 

218 :return: t.Set of old accessed entries 

219 """ 

220 old = current_db_access_log.get(set()) 

221 current_db_access_log.set(set()) 

222 return old 

223 

224 

225def startDataAccessLog() -> t.Set[t.Union[Key, str]]: 

226 return start_data_access_log() 

227 

228 

229def end_data_access_log( 

230 outer_access_log: t.Optional[t.Set[t.Union[Key, str]]] = None, 

231) -> t.Optional[t.Set[t.Union[Key, str]]]: 

232 """ 

233 Retrieves the set of entries accessed so far. 

234 

235 To clean up and restart the log, call :func:`viur.datastore.startAccessDataLog`. 

236 

237 If you called :func:`server.db.startAccessDataLog` before, you can re-apply the old log using 

238 the outerAccessLog param. Otherwise, it will disable the access log. 

239 

240 :param outerAccessLog: State of your log returned by :func:`server.db.startAccessDataLog` 

241 :return: t.Set of entries accessed 

242 """ 

243 res = current_db_access_log.get() 

244 if isinstance(outer_access_log, set): 

245 current_db_access_log.set((outer_access_log or set()).union(res)) 

246 else: 

247 current_db_access_log.set(None) 

248 return res 

249 

250 

251def endDataAccessLog( 

252 outerAccessLog: t.Optional[t.Set[t.Union[Key, str]]] = None, 

253) -> t.Optional[t.Set[t.Union[Key, str]]]: 

254 return end_data_access_log(outer_access_log=outerAccessLog)