Coverage for /home/runner/work/viur-core/viur-core/viur/src/viur/core/db/transport.py: 69%
206 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 15:02 +0000
1"""
2Datastore transport layer: the process-wide client and the CRUD helpers.
4**Named database and namespace**
6The datastore client (:data:`__client__`) is built once at import time and
7kept for the whole process lifetime. Its target database and namespace come
8from :attr:`conf.db.name <viur.core.config.Database.name>` and
9:attr:`conf.db.namespace <viur.core.config.Database.namespace>`, which are sourced
10from the ``VIUR_DB_NAME`` / ``VIUR_DB_NAMESPACE`` environment variables. Both
11default to ``None`` — the standard ``(default)`` database and empty namespace —
12so existing deployments are unaffected.
14Because the client is created from the environment at import time, the target
15cannot be retargeted at runtime: a single process always talks to exactly one
16database. :class:`~viur.core.db.types.Key` objects inherit that database and
17namespace from the client, keeping every request on the configured target.
19The legacy urlsafe key encoding (App Engine "Reference") predates named
20databases and only supports the default one. Therefore
21:meth:`Key.to_legacy_urlsafe <viur.core.db.types.Key.to_legacy_urlsafe>`
22encodes a database-less copy of the key, and the client's database is restored
23on decoding — unambiguous precisely because the process is bound to a single
24database.
25"""
26from __future__ import annotations
28import contextvars
29import itertools
30import logging
31import time
32import typing as t
34from deprecated.sphinx import deprecated
35from google.cloud import datastore, exceptions
37from .overrides import entity_from_protobuf, key_from_protobuf
38from .types import Entity, Key, QueryDefinition, SortOrder, current_db_access_log
39from . import cache
40from viur.core.config import conf
42# patching our key and entity classes
43datastore.helpers.key_from_protobuf = key_from_protobuf
44datastore.helpers.entity_from_protobuf = entity_from_protobuf
46# Built once at import, kept for the process lifetime — so db/namespace have to
47# come from env (via conf.db); nothing can retarget the client afterwards.
48# Both default to None, which is the same as datastore.Client(): no change for
49# default deployments.
50__client__ = datastore.Client(database=conf.db.name, namespace=conf.db.namespace)
52_transaction_outdated: contextvars.ContextVar[list[t.Union[Entity, Key]] | None] = contextvars.ContextVar(
53 "Transaction-outdated-entities", default=None
54)
55"""Entities and keys the transaction running in this context has outdated.
57The cache cannot be updated while a transaction is open -- the write is not
58committed yet and may still roll back -- so :func:`put` and :func:`delete`
59record what they touched here and :func:`run_in_transaction` invalidates those
60entries once the transaction is over.
61"""
64def _mark_as_outdated(data: t.Union[Entity, Key, t.Iterable[t.Union[Entity, Key]]]) -> None:
65 """Record entities/keys whose cache entry a running transaction outdates."""
66 if (outdated := _transaction_outdated.get()) is None:
67 return # no transaction in this context, the caller updates the cache itself
69 outdated.extend(data if isinstance(data, (list, set, tuple)) else [data])
72def _invalidate_outdated(outdated: list[t.Union[Entity, Key]]) -> None:
73 """Drop everything a finished transaction outdated from the cache.
75 The keys are read *after* the transaction, so entities written with a
76 partial key carry their final, datastore-assigned key by now.
77 """
78 keys = [entry.key if isinstance(entry, Entity) else entry for entry in outdated]
79 if keys := [key for key in keys if key is not None and not key.is_partial]: 79 ↛ exitline 79 didn't return from function '_invalidate_outdated' because the condition on line 79 was always true
80 cache.delete(keys)
83MAX_LOOKUP_KEYS: t.Final[int] = 1000
84"""Maximum number of keys the datastore accepts for a single Lookup operation.
86Unlike a Lookup, a Commit has no comparable cap on the number of mutations - it is bounded by
87the 10 MiB request size instead. :func:`put` and :func:`delete` therefore stay a single commit
88of whatever they are handed, which keeps them atomic.
89"""
92def allocate_ids(kind_name: str, num_ids: int = 1, retry=None, timeout=None) -> list[Key]:
93 if type(kind_name) is not str:
94 raise TypeError("kind_name must be a string")
95 return __client__.allocate_ids(Key(kind_name), num_ids, retry, timeout)
98@deprecated(version="3.8.0", reason="Use 'db.allocate_ids' instead")
99def AllocateIDs(kind_name):
100 """
101 Allocates a new, free unique id for a given kind_name.
102 """
103 if isinstance(kind_name, Key): # so ein Murks...
104 kind_name = kind_name.kind
106 return allocate_ids(kind_name)[0]
109def get(keys: t.Union[Key, t.Iterable[Key]]) -> t.Union[list[Entity], Entity, None]:
110 """
111 Retrieves an entity (or a list thereof) from datastore.
112 If only a single key has been given we'll return the entity or none in case the key has not been found,
113 otherwise a list of all entities that have been looked up (which may be empty)
114 :param keys: A datastore key (or a list thereof) to lookup
115 :return: The entity (or None if it has not been found), or a list of entities.
116 """
117 _write_to_access_log(keys)
119 is_multiple = isinstance(keys, (list, set, tuple))
120 key_list = list(keys) if is_multiple else [keys]
122 # Serve whatever we can from the cache, indexed by its stringified key.
123 entities_by_key = {str(entity.key): entity for entity in cache.get(key_list)}
125 # Fetch the keys that were not cached and write them back into the cache
126 missing = [key for key in key_list if str(key) not in entities_by_key]
127 if missing:
128 # A Lookup accepts at most MAX_LOOKUP_KEYS keys, so ask in chunks and merge the answers
129 fetched = []
130 for chunk in itertools.batched(missing, MAX_LOOKUP_KEYS):
131 fetched.extend(__client__.get_multi(list(chunk)))
132 if fetched: 132 ↛ 134line 132 didn't jump to line 134 because the condition on line 132 was always true
133 cache.put(fetched)
134 for entity in fetched:
135 entities_by_key[str(entity.key)] = entity
137 # Reassemble in the original key order, dropping keys that were not found
138 result = [entities_by_key[str(key)] for key in key_list if str(key) in entities_by_key]
140 if conf.debug.trace_queries: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true
141 logging.info(f"db.get: {len(result)}/{len(key_list)} entities found")
143 if is_multiple:
144 return result
145 return result[0] if result else None
149@deprecated(version="3.8.0", reason="Use 'db.get' instead")
150def Get(keys: t.Union[Key, t.List[Key]]) -> t.Union[t.List[Entity], Entity, None]:
151 return get(keys)
154def put(entities: t.Union[Entity, t.List[Entity]]):
155 """
156 Save an entity in the Cloud Datastore.
157 Also ensures that no string-key with a digit-only name can be used.
159 A list of entities is written in one commit and therefore atomically, however long it is.
161 :param entities: The entities to be saved to the datastore.
162 """
163 _write_to_access_log(entities)
165 # Cache only after the datastore accepted the write: a failed write must not
166 # leave a value in the cache that was never persisted. The datastore also
167 # completes partial keys during the write, so caching afterwards stores the
168 # entity under its final key.
169 if isinstance(entities, Entity): 169 ↛ 174line 169 didn't jump to line 174 because the condition on line 169 was always true
170 res = __client__.put(entities)
171 if conf.debug.trace_queries: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 logging.info(f"db.put: saved {entities.key}")
173 else:
174 res = __client__.put_multi(entities=entities)
175 if conf.debug.trace_queries:
176 logging.info(f"db.put: saved {len(entities)} entities")
178 # Inside a transaction cache.put() is a no-op (nothing is committed yet), so
179 # remember the entities and let run_in_transaction() invalidate them afterwards.
180 _mark_as_outdated(entities)
181 cache.put(entities)
182 return res
185@deprecated(version="3.8.0", reason="Use 'db.put' instead")
186def Put(entities: t.Union[Entity, t.List[Entity]]) -> t.Union[Entity, None]:
187 return put(entities)
190def delete(keys: t.Union[Entity, t.Iterable[Entity], Key, t.Iterable[Key]]):
191 """
192 Deletes the entities with the given key(s) from the datastore.
194 A list of keys is deleted in one commit and therefore atomically, however long it is.
196 :param keys: A Key (or a t.List of Keys) to delete
197 """
199 _write_to_access_log(keys)
200 # Invalidate right away, and again once a surrounding transaction has finished:
201 # until then a concurrent read could pull the still-current value back in.
202 _mark_as_outdated(keys)
203 cache.delete(keys)
204 if not isinstance(keys, (set, list, tuple)): 204 ↛ 211line 204 didn't jump to line 211 because the condition on line 204 was always true
205 res = __client__.delete(keys)
206 if conf.debug.trace_queries: 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 logging.info(f"db.delete: deleted {keys}")
208 return res
211 res = __client__.delete_multi(keys)
212 if conf.debug.trace_queries:
213 logging.info(f"db.delete: deleted {len(keys)} keys")
214 return res
217@deprecated(version="3.8.0", reason="Use 'db.delete' instead")
218def Delete(keys: t.Union[Entity, t.List[Entity], Key, t.List[Key]]):
219 return delete(keys)
222def run_in_transaction(func: t.Callable, *args, **kwargs) -> t.Any:
223 """
224 Runs the function given in :param:callee inside a transaction.
225 Inside a transaction it's guaranteed that
226 - either all or no changes are written to the datastore
227 - no other transaction is currently reading/writing the entities accessed
229 See (transactions)[https://cloud.google.com/datastore/docs/concepts/cloud-datastore-transactions] for more
230 information.
232 ..Warning: The datastore may produce unexpected results if an entity that have been written inside a transaction
233 is read (or returned in a query) again. In this case you will the the *old* state of that entity. Keep that
234 in mind if wrapping functions to run in a transaction that may have not been designed to handle this case.
235 :param func: The function that will be run inside a transaction
236 :param args: All args will be passed into the callee
237 :param kwargs: All kwargs will be passed into the callee
238 :return: Whatever the callee function returned
239 :raises RuntimeError: If the maximum transaction retries exceeded
240 """
241 if __client__.current_transaction: 241 ↛ 243line 241 didn't jump to line 243 because the condition on line 241 was never true
242 # Nested call: the outermost one owns the cache invalidation.
243 return func(*args, **kwargs)
245 token = _transaction_outdated.set([])
246 try:
247 for i in range(conf.db.transaction_attempts): 247 ↛ 263line 247 didn't jump to line 263 because the loop on line 247 didn't complete
248 try:
249 with __client__.transaction():
250 res = func(*args, **kwargs)
251 break
253 except exceptions.Conflict:
254 if i + 1 >= conf.db.transaction_attempts:
255 # Last attempt failed: raise right away instead of sleeping first.
256 raise RuntimeError("Maximum transaction retries exceeded")
258 logging.error(f"Transaction failed with a conflict, trying again in {2 ** i} seconds")
259 time.sleep(2 ** i)
260 continue
262 else:
263 raise RuntimeError("Maximum transaction retries exceeded")
265 finally:
266 # Also invalidate when the transaction failed: a retry may have written
267 # before the conflict, and one invalidation too many only costs a lookup.
268 outdated = _transaction_outdated.get()
269 _transaction_outdated.reset(token)
270 if outdated:
271 _invalidate_outdated(outdated)
273 return res
276@deprecated(version="3.8.0", reason="Use 'db.run_in_transaction' instead")
277def RunInTransaction(callee: t.Callable, *args, **kwargs) -> t.Any:
278 return run_in_transaction(callee, *args, **kwargs)
281NATIVE_FILTER_OPERATORS = frozenset({"IN", "!=", "NOT_IN"})
282"""Operators the Datastore evaluates natively; their value goes into a single PropertyFilter, unsplit."""
285def _normalize_filter_value(op: str, value: t.Any) -> t.Any:
286 """
287 google-cloud-datastore only encodes a ``list`` as array value, any other collection raises
288 ``ValueError: Unknown protobuf attr type``. Convert tuples and sets for the array operators.
289 The order of a set is kept as it is; IN and NOT_IN don't depend on it.
290 """
291 if op in ("IN", "NOT_IN") and isinstance(value, (tuple, set, frozenset)):
292 return list(value)
294 return value
297def _build_property_filters(key: str, op: str, value: t.Any) -> list[datastore.query.PropertyFilter]:
298 """
299 Build the PropertyFilters for one entry of :attr:`QueryDefinition.filters`.
301 Native operators result in exactly one filter; any other operator with a list value
302 results in one filter per element (multi equal filters).
303 """
304 if op in NATIVE_FILTER_OPERATORS: 304 ↛ 307line 304 didn't jump to line 307 because the condition on line 304 was always true
305 return [datastore.query.PropertyFilter(key, op, _normalize_filter_value(op, value))]
307 if not isinstance(value, list):
308 value = [value]
310 return [datastore.query.PropertyFilter(key, op, val) for val in value]
313def _build_or_filter(or_group: list[tuple[str, t.Any]]) -> datastore.query.Or:
314 """
315 Build the Or composite filter for one entry of :attr:`QueryDefinition.or_filters`.
316 """
317 or_conditions = []
318 for filter_str, value in or_group:
319 key, op = filter_str.split(" ", 1)
320 or_conditions.append(datastore.query.PropertyFilter(key, op, _normalize_filter_value(op, value)))
322 return datastore.query.Or(or_conditions)
325def count(kind: str = None, up_to=2 ** 31 - 1, queryDefinition: QueryDefinition = None) -> int:
326 if not kind: 326 ↛ 329line 326 didn't jump to line 329 because the condition on line 326 was always true
327 kind = queryDefinition.kind
329 query = __client__.query(kind=kind)
330 if queryDefinition and queryDefinition.filters:
331 for k, v in queryDefinition.filters.items():
332 key, op = k.split(" ")
333 for f in _build_property_filters(key, op, v):
334 query.add_filter(filter=f)
336 if queryDefinition and queryDefinition.or_filters:
337 for or_group in queryDefinition.or_filters:
338 query.add_filter(filter=_build_or_filter(or_group))
340 aggregation_query = __client__.aggregation_query(query)
342 result = aggregation_query.count(alias="total").fetch(limit=up_to)
343 return list(result)[0][0].value
346@deprecated(version="3.8.0", reason="Use 'db.count' instead")
347def Count(kind: str = None, up_to=2 ** 31 - 1, queryDefinition: QueryDefinition = None) -> int:
348 return count(kind, up_to, queryDefinition)
351def run_single_filter(query: QueryDefinition, limit: int, keys_only: bool) -> t.List[Entity | Key]:
352 """
353 Internal helper function that runs a single query definition on the datastore and returns a list of
354 entities found.
355 :param query: The querydefinition (filters, orders, distinct etc.) to run against the datastore
356 :param limit: How many results should at most be returned
357 :return: The first *limit* entities that matches this query
358 """
360 qry = __client__.query(kind=query.kind)
361 startCursor = None
362 endCursor = None
363 hasInvertedOrderings = None
364 if conf.debug.trace_queries: 364 ↛ 365line 364 didn't jump to line 365 because the condition on line 364 was never true
365 logging.info(f"Running query: {query}")
367 if query: 367 ↛ 393line 367 didn't jump to line 393 because the condition on line 367 was always true
368 if query.filters:
369 for k, v in query.filters.items():
370 key, op = k.split(" ")
371 for f in _build_property_filters(key, op, v):
372 qry.add_filter(filter=f)
374 if query.or_filters:
375 for or_group in query.or_filters:
376 qry.add_filter(filter=_build_or_filter(or_group))
378 if query.orders: 378 ↛ 379line 378 didn't jump to line 379 because the condition on line 378 was never true
379 hasInvertedOrderings = any(
380 order.order in (SortOrder.InvertedAscending, SortOrder.InvertedDescending)
381 for order in query.orders
382 )
383 qry.order = [
384 order.name if order.order in (SortOrder.Ascending, SortOrder.InvertedDescending) else f"-{order.name}"
385 for order in query.orders
386 ]
388 if query.distinct: 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true
389 qry.distinct_on = query.distinct
391 startCursor = query.startCursor
392 endCursor = query.endCursor
393 if keys_only: 393 ↛ 394line 393 didn't jump to line 394 because the condition on line 393 was never true
394 qry.keys_only()
395 qryRes = qry.fetch(limit=limit, start_cursor=startCursor, end_cursor=endCursor)
396 res = list(qryRes)
397 query.currentCursor = qryRes.next_page_token
398 if hasInvertedOrderings: 398 ↛ 399line 398 didn't jump to line 399 because the condition on line 398 was never true
399 res.reverse()
401 if conf.debug.trace_queries: 401 ↛ 402line 401 didn't jump to line 402 because the condition on line 401 was never true
402 distinct_on = f" distinct on {query.distinct}" if query.distinct else ""
403 logging.debug(
404 f"Queried {query.kind} with filter {query.filters} and orders {query.orders}{distinct_on}."
405 f" Returned {len(res)} results"
406 )
408 return res
411@deprecated(version="3.8.0", reason="Use 'run_single_filter' instead")
412def runSingleFilter(query: QueryDefinition, limit: int) -> t.List[Entity]:
413 run_single_filter(query, limit)
416# helper function for access log
417def _write_to_access_log(data: t.Union[Key, list[Key], Entity, list[Entity]]) -> None:
418 if not conf.db.create_access_log: 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true
419 return
420 access_log = current_db_access_log.get()
421 if not isinstance(access_log, set): 421 ↛ 423line 421 didn't jump to line 423 because the condition on line 421 was always true
422 return # access log not exist
423 if not data:
424 return
425 if isinstance(data, Entity):
426 access_log.add(data.key)
427 elif isinstance(data, Key):
428 access_log.add(data)
429 else:
430 for entry in data:
431 if isinstance(entry, Entity):
432 access_log.add(entry.key)
433 elif isinstance(entry, Key):
434 access_log.add(entry)
437__all__ = [allocate_ids, delete, get, put, run_in_transaction, count]