Coverage for /home/runner/work/viur-core/viur-core/viur/src/viur/core/db/types.py: 85%

1""" 

2The constants, global variables and container classes used in the datastore api 

3""" 

4from __future__ import annotations 

5 

6import copy 

7import datetime 

8import enum 

9import itertools 

10import typing as t 

11from contextvars import ContextVar 

12from dataclasses import dataclass, field 

13from ..config import conf 

14 

15from google.cloud.datastore import Entity as Datastore_entity, Key as Datastore_key 

16 

17KEY_SPECIAL_PROPERTY = "__key__" 

18"""The property name pointing to an entities key in a query""" 

19 

20DATASTORE_BASE_TYPES = t.Union[None, str, int, float, bool, datetime.datetime, datetime.date, datetime.time, "Key"] 

21"""Types that can be used in a datastore query""" 

22 

23current_db_access_log: ContextVar[t.Optional[set[t.Union[Key, str]]]] = ContextVar("Database-Accesslog", default=None) 

24"""If set to a set for the current thread/request, we'll log all entities / kinds accessed""" 

25 

26"""The current projectID, which can't be imported from transport.py""" 

27 

28 

29class SortOrder(enum.Enum): 

30 """ 

31 Defines possible types of sort orders for queries. 

32 """ 

33 

34 Ascending = 1 

35 """Sort A->Z""" 

36 Descending = 2 

37 """Sort Z->A""" 

38 InvertedAscending = 3 

39 """Fetch Z->A, then flip the results (useful in pagination to go from a start cursor backwards)""" 

40 InvertedDescending = 4 

41 """Fetch A->Z, then flip the results (useful in pagination)""" 

42 

43 @classmethod 

44 def from_str(cls, ident: str | int) -> SortOrder: 

45 """ 

46 Parses a string defining a sort order into a db.SortOrder. 

47 """ 

48 match str(ident or "").lower(): 

49 case "desc" | "descending" | "1": 

50 return SortOrder.Descending 

51 case "inverted_asc" | "inverted_ascending" | "2": 

52 return SortOrder.InvertedAscending 

53 case "inverted_desc" | "inverted_descending" | "3": 

54 return SortOrder.InvertedDescending 

55 case _: # everything else 

56 return SortOrder.Ascending 

57 

58 

59class QueryOrder(t.NamedTuple): 

60 """A named tuple describing a single sort order for a datastore query.""" 

61 name: str 

62 order: SortOrder = SortOrder.Ascending 

63 

64 

65class Key(Datastore_key): 

66 """ 

67 The python representation of one datastore key. Unlike the original implementation, we don't store a 

68 reference to the project the key lives in. This is always expected to be the current project as ViUR 

69 does not support accessing data in multiple projects. 

70 """ 

71 

72 def __init__(self, *path_args, project: str | None = None, **kwargs): 

73 # Convert digit-only id_or_name attributes to int 

74 # See https://github.com/viur-framework/viur-core/issues/1636 

75 new_path_args = [] 

76 for pair in itertools.batched(path_args, 2): 

77 try: 

78 kind, id_or_name = pair 

79 except ValueError: # it's a incomplete key 

80 new_path_args.append(pair[0]) 

81 continue 

82 if isinstance(id_or_name, str) and id_or_name.isdigit(): 

83 id_or_name = int(id_or_name) 

84 new_path_args.extend((kind, id_or_name)) 

85 

86 from .transport import __client__ # noqa: E402 # import works only here because circular imports 

87 

88 if project is None: 

89 project = __client__.project 

90 

91 # Keys must match the client's db/namespace or Datastore rejects the 

92 # request as cross-database. Default from the client; caller wins. 

93 if __client__.database: 

94 kwargs.setdefault("database", __client__.database) 

95 if __client__.namespace: 

96 kwargs.setdefault("namespace", __client__.namespace) 

97 

98 super().__init__(*new_path_args, project=project, **kwargs) 

99 

100 def __str__(self): 

101 return self.to_legacy_urlsafe().decode("ASCII") 

102 

103 def to_legacy_urlsafe(self, location_prefix=None): 

104 # Upstream to_legacy_urlsafe() rejects keys carrying a database, but 

105 # str(key)/session paths hit it constantly. Encode a database-less copy — 

106 # unambiguous to restore since the process talks to a single database. 

107 # A copy keeps this thread-safe: mutating self._database in place would 

108 # let concurrent encodes of the same Key clobber each other's state. 

109 clone = copy.copy(self) 

110 clone._database = None 

111 return super(Key, clone).to_legacy_urlsafe(location_prefix=location_prefix) 

112 

113 ''' 

114 def __repr__(self): 

115 return "<viur.datastore.Key %s/%s, parent=%s>" % (self.kind, self.id_or_name, self.parent) 

116 

117 def __hash__(self): 

118 return hash("%s.%s.%s" % (self.kind, self.id, self.name)) 

119 

120 def __eq__(self, other): 

121 return isinstance(other, Key) and self.kind == other.kind and self.id == other.id and self.name == other.name \ 

122 and self.parent == other.parent 

123 

124 @staticmethod 

125 def _parse_path(path_args): 

126 """Parses positional arguments into key path with kinds and IDs. 

127 

128 :type path_args: tuple 

129 :param path_args: A tuple from positional arguments. Should be 

130 alternating list of kinds (string) and ID/name 

131 parts (int or string). 

132 

133 :rtype: :class:`list` of :class:`dict` 

134 :returns: A list of key parts with kind and ID or name set. 

135 :raises: :class:`ValueError` if there are no ``path_args``, if one of 

136 the kinds is not a string or if one of the IDs/names is not 

137 a string or an integer. 

138 """ 

139 if len(path_args) == 0: 

140 raise ValueError("Key path must not be empty.") 

141 

142 kind_list = path_args[::2] 

143 id_or_name_list = path_args[1::2] 

144 # Dummy sentinel value to pad incomplete key to even length path. 

145 partial_ending = object() 

146 if len(path_args) % 2 == 1: 

147 id_or_name_list += (partial_ending,) 

148 

149 result = [] 

150 for kind, id_or_name in zip(kind_list, id_or_name_list): 

151 curr_key_part = {} 

152 if isinstance(kind, str): 

153 curr_key_part["kind"] = kind 

154 else: 

155 raise ValueError(kind, "Kind was not a string.") 

156 

157 if isinstance(id_or_name, str): 

158 if (id_or_name.isdigit()): # !!! VIUR 

159 curr_key_part["id"] = int(id_or_name) 

160 else: 

161 curr_key_part["name"] = id_or_name 

162 

163 elif isinstance(id_or_name, int): 

164 curr_key_part["id"] = id_or_name 

165 elif id_or_name is not partial_ending: 

166 raise ValueError(id_or_name, "ID/name was not a string or integer.") 

167 

168 result.append(curr_key_part) 

169 return result 

170 

171 @classmethod 

172 def from_legacy_urlsafe(cls, strKey: str) -> Key: 

173 """ 

174 Parses the string representation generated by :meth:to_legacy_urlsafe into a new Key object 

175 :param strKey: The string key to parse 

176 :return: The new Key object constructed from the string key 

177 """ 

178 urlsafe = strKey.encode("ASCII") 

179 padding = b"=" * (-len(urlsafe) % 4) 

180 urlsafe += padding 

181 raw_bytes = base64.urlsafe_b64decode(urlsafe) 

182 reference = _app_engine_key_pb2.Reference() 

183 reference.ParseFromString(raw_bytes) 

184 resultKey = None 

185 for elem in reference.path.element: 

186 resultKey = Key(elem.type, elem.id or elem.name, parent=resultKey) 

187 return resultKey 

188 ''' 

189 

190 

191class Entity(Datastore_entity): 

192 """ 

193 The python representation of one datastore entity. The values of this entity are stored inside this dictionary, 

194 while the meta-data (it's key, the list of properties excluded from indexing and our version) as property values. 

195 """ 

196 

197 def __init__( 

198 self, 

199 key: t.Optional[Key] = None, 

200 exclude_from_indexes: t.Optional[list[str]] = None, 

201 ) -> None: 

202 super().__init__(key, exclude_from_indexes or []) 

203 if not (key is None or isinstance(key, Key)): 203 ↛ 204line 203 didn't jump to line 204 because the condition on line 203 was never true

204 raise ValueError(f"key must be a Key-Object (or None for an embedded entity). Got {key!r} ({type(key)})") 

205 

206 

207KeyType: t.TypeAlias = Key | str | int 

208""" 

209Alias that describes a key-type. 

210""" 

211 

212TOrders: t.TypeAlias = list[QueryOrder] 

213TFilters: t.TypeAlias = dict[str, DATASTORE_BASE_TYPES | list[DATASTORE_BASE_TYPES]] 

214TOrFilters: t.TypeAlias = list[list[tuple[str, DATASTORE_BASE_TYPES | list[DATASTORE_BASE_TYPES]]]] 

215 

216 

217@dataclass 

218class QueryDefinition: 

219 """ 

220 A single Query that will be run against the datastore. 

221 """ 

222 

223 kind: t.Optional[str] 

224 """The datastore kind to run the query on. Can be None for kindles queries.""" 

225 

226 filters: TFilters 

227 """A dictionary of constrains to apply to the query.""" 

228 

229 orders: t.Optional[TOrders] 

230 """The list of fields to sort the results by.""" 

231 

232 distinct: t.Optional[list[str]] = None 

233 """If set, a list of fields that we should return distinct values of""" 

234 

235 or_filters: "TOrFilters" = field(default_factory=list) 

236 """Each entry is a list of (filterStr, value) pairs that are OR-ed together. 

237 Multiple entries are AND-ed with each other and with the AND filters.""" 

238 

239 limit: int = field(init=False) 

240 """The maximum amount of entities that should be returned""" 

241 

242 startCursor: t.Optional[str] = None 

243 """If set, we'll only return entities that appear after this cursor in the index.""" 

244 

245 endCursor: t.Optional[str] = None 

246 """If set, we'll only return entities up to this cursor in the index.""" 

247 

248 currentCursor: t.Optional[str] = None 

249 """Will be set after this query has been run, pointing after the last entity returned""" 

250 

251 def __post_init__(self): 

252 self.limit = conf.db.query_default_limit