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

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 

131@deprecated(version="3.8.0", reason="Use 'db.key_helper' instead") 

132def keyHelper( 

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

134 targetKind: str, 

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

136 adjust_kind: bool = False, 

137) -> Key: 

138 return key_helper( 

139 in_key=inKey, 

140 target_kind=targetKind, 

141 additional_allowed_kinds=additionalAllowedKinds, 

142 adjust_kind=adjust_kind 

143 ) 

144 

145 

146def is_in_transaction() -> bool: 

147 return __client__.current_transaction is not None 

148 

149 

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

151def IsInTransaction() -> bool: 

152 return is_in_transaction() 

153 

154 

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

156 """ 

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

158 

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

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

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

162 

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

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

165 """ 

166 

167 def txn(key, kwargs): 

168 obj = get(key) 

169 if not obj: 

170 obj = Entity(key) 

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

172 obj[k] = v 

173 put(obj) 

174 return obj 

175 

176 if is_in_transaction(): 

177 return txn(key, kwargs) 

178 return run_in_transaction(txn, key, kwargs) 

179 

180 

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

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

183 return get_or_insert(key, **kwargs) 

184 

185 

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

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

188 """ 

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

190 """ 

191 return str(key) 

192 

193 

194def acquire_transaction_success_marker() -> str: 

195 """ 

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

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

198 or if the transaction it was created in failed. 

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

200 """ 

201 txn: Transaction | None = __client__.current_transaction 

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

203 marker = str(txn.id) 

204 request_data = current.request_data.get() 

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

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

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

208 put(db_obj) 

209 request_data["__viur-transactionmarker__"] = True 

210 return marker 

211 

212 

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

214 """ 

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

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

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

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

219 :return: t.Set of old accessed entries 

220 """ 

221 old = current_db_access_log.get(set()) 

222 current_db_access_log.set(set()) 

223 return old 

224 

225 

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

227 return start_data_access_log() 

228 

229 

230def end_data_access_log( 

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

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

233 """ 

234 Retrieves the set of entries accessed so far. 

235 

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

237 

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

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

240 

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

242 :return: t.Set of entries accessed 

243 """ 

244 res = current_db_access_log.get() 

245 if isinstance(outer_access_log, set): 

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

247 else: 

248 current_db_access_log.set(None) 

249 return res 

250 

251 

252def endDataAccessLog( 

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

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

255 return end_data_access_log(outer_access_log=outerAccessLog)