-
Notifications
You must be signed in to change notification settings - Fork 44
/
key.py
658 lines (508 loc) · 21 KB
/
key.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
# Copyright 2014 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Create / interact with Google Cloud Datastore keys."""
import base64
import copy
from google.cloud.datastore_v1.types import entity as _entity_pb2
from google.cloud._helpers import _to_bytes
from google.cloud.datastore import _app_engine_key_pb2
_DATABASE_ID_TEMPLATE = (
"Received non-empty database ID: {!r}.\n"
"urlsafe strings are not expected to encode a Reference that "
"contains a database ID."
)
_BAD_ELEMENT_TEMPLATE = (
"At most one of ID and name can be set on an element. Received "
"id = {!r} and name = {!r}."
)
_EMPTY_ELEMENT = (
"Exactly one of ID and name must be set on an element. "
"Encountered an element with neither set that was not the last "
"element of a path."
)
class Key(object):
"""An immutable representation of a datastore Key.
.. testsetup:: key-ctor
from google.cloud import datastore
project = 'my-special-pony'
client = datastore.Client(project=project)
Key = datastore.Key
parent_key = client.key('Parent', 'foo')
To create a basic key directly:
.. doctest:: key-ctor
>>> Key('EntityKind', 1234, project=project)
<Key('EntityKind', 1234), project=...>
>>> Key('EntityKind', 'foo', project=project)
<Key('EntityKind', 'foo'), project=...>
Though typical usage comes via the
:meth:`~google.cloud.datastore.client.Client.key` factory:
.. doctest:: key-ctor
>>> client.key('EntityKind', 1234)
<Key('EntityKind', 1234), project=...>
>>> client.key('EntityKind', 'foo')
<Key('EntityKind', 'foo'), project=...>
To create a key with a parent:
.. doctest:: key-ctor
>>> client.key('Parent', 'foo', 'Child', 1234)
<Key('Parent', 'foo', 'Child', 1234), project=...>
>>> client.key('Child', 1234, parent=parent_key)
<Key('Parent', 'foo', 'Child', 1234), project=...>
To create a partial key:
.. doctest:: key-ctor
>>> client.key('Parent', 'foo', 'Child')
<Key('Parent', 'foo', 'Child'), project=...>
:type path_args: tuple of string and integer
:param path_args: May represent a partial (odd length) or full (even
length) key path.
:param kwargs: Keyword arguments to be passed in.
Accepted keyword arguments are
* namespace (string): A namespace identifier for the key.
* project (string): The project associated with the key.
* parent (:class:`~google.cloud.datastore.key.Key`): The parent of the key.
The project argument is required unless it has been set implicitly.
"""
def __init__(self, *path_args, **kwargs):
self._flat_path = path_args
parent = self._parent = kwargs.get("parent")
self._namespace = kwargs.get("namespace")
project = kwargs.get("project")
self._project = _validate_project(project, parent)
# _flat_path, _parent, _namespace and _project must be set before
# _combine_args() is called.
self._path = self._combine_args()
def __eq__(self, other):
"""Compare two keys for equality.
Incomplete keys never compare equal to any other key.
Completed keys compare equal if they have the same path, project,
and namespace.
:rtype: bool
:returns: True if the keys compare equal, else False.
"""
if not isinstance(other, Key):
return NotImplemented
if self.is_partial or other.is_partial:
return False
return (
self.flat_path == other.flat_path
and self.project == other.project
and self.namespace == other.namespace
)
def __ne__(self, other):
"""Compare two keys for inequality.
Incomplete keys never compare equal to any other key.
Completed keys compare equal if they have the same path, project,
and namespace.
:rtype: bool
:returns: False if the keys compare equal, else True.
"""
return not self == other
def __hash__(self):
"""Hash a keys for use in a dictionary lookp.
:rtype: int
:returns: a hash of the key's state.
"""
return hash(self.flat_path) + hash(self.project) + hash(self.namespace)
@staticmethod
def _parse_path(path_args):
"""Parses positional arguments into key path with kinds and IDs.
:type path_args: tuple
:param path_args: A tuple from positional arguments. Should be
alternating list of kinds (string) and ID/name
parts (int or string).
:rtype: :class:`list` of :class:`dict`
:returns: A list of key parts with kind and ID or name set.
:raises: :class:`ValueError` if there are no ``path_args``, if one of
the kinds is not a string or if one of the IDs/names is not
a string or an integer.
"""
if len(path_args) == 0:
raise ValueError("Key path must not be empty.")
kind_list = path_args[::2]
id_or_name_list = path_args[1::2]
# Dummy sentinel value to pad incomplete key to even length path.
partial_ending = object()
if len(path_args) % 2 == 1:
id_or_name_list += (partial_ending,)
result = []
for kind, id_or_name in zip(kind_list, id_or_name_list):
curr_key_part = {}
if isinstance(kind, str):
curr_key_part["kind"] = kind
else:
raise ValueError(kind, "Kind was not a string.")
if isinstance(id_or_name, str):
curr_key_part["name"] = id_or_name
elif isinstance(id_or_name, int):
curr_key_part["id"] = id_or_name
elif id_or_name is not partial_ending:
raise ValueError(id_or_name, "ID/name was not a string or integer.")
result.append(curr_key_part)
return result
def _combine_args(self):
"""Sets protected data by combining raw data set from the constructor.
If a ``_parent`` is set, updates the ``_flat_path`` and sets the
``_namespace`` and ``_project`` if not already set.
:rtype: :class:`list` of :class:`dict`
:returns: A list of key parts with kind and ID or name set.
:raises: :class:`ValueError` if the parent key is not complete.
"""
child_path = self._parse_path(self._flat_path)
if self._parent is not None:
if self._parent.is_partial:
raise ValueError("Parent key must be complete.")
# We know that _parent.path() will return a copy.
child_path = self._parent.path + child_path
self._flat_path = self._parent.flat_path + self._flat_path
if (
self._namespace is not None
and self._namespace != self._parent.namespace
):
raise ValueError("Child namespace must agree with parent's.")
self._namespace = self._parent.namespace
if self._project is not None and self._project != self._parent.project:
raise ValueError("Child project must agree with parent's.")
self._project = self._parent.project
return child_path
def _clone(self):
"""Duplicates the Key.
Most attributes are simple types, so don't require copying. Other
attributes like ``parent`` are long-lived and so we re-use them.
:rtype: :class:`google.cloud.datastore.key.Key`
:returns: A new ``Key`` instance with the same data as the current one.
"""
cloned_self = self.__class__(
*self.flat_path, project=self.project, namespace=self.namespace
)
# If the current parent has already been set, we re-use
# the same instance
cloned_self._parent = self._parent
return cloned_self
def completed_key(self, id_or_name):
"""Creates new key from existing partial key by adding final ID/name.
:type id_or_name: str or integer
:param id_or_name: ID or name to be added to the key.
:rtype: :class:`google.cloud.datastore.key.Key`
:returns: A new ``Key`` instance with the same data as the current one
and an extra ID or name added.
:raises: :class:`ValueError` if the current key is not partial or if
``id_or_name`` is not a string or integer.
"""
if not self.is_partial:
raise ValueError("Only a partial key can be completed.")
if isinstance(id_or_name, str):
id_or_name_key = "name"
elif isinstance(id_or_name, int):
id_or_name_key = "id"
else:
raise ValueError(id_or_name, "ID/name was not a string or integer.")
new_key = self._clone()
new_key._path[-1][id_or_name_key] = id_or_name
new_key._flat_path += (id_or_name,)
return new_key
def to_protobuf(self):
"""Return a protobuf corresponding to the key.
:rtype: :class:`.entity_pb2.Key`
:returns: The protobuf representing the key.
"""
key = _entity_pb2.Key()
key.partition_id.project_id = self.project
if self.namespace:
key.partition_id.namespace_id = self.namespace
for item in self.path:
element = key.PathElement()
if "kind" in item:
element.kind = item["kind"]
if "id" in item:
element.id = item["id"]
if "name" in item:
element.name = item["name"]
key.path.append(element)
return key
def to_legacy_urlsafe(self, location_prefix=None):
"""Convert to a base64 encode urlsafe string for App Engine.
This is intended to work with the "legacy" representation of a
datastore "Key" used within Google App Engine (a so-called
"Reference"). The returned string can be used as the ``urlsafe``
argument to ``ndb.Key(urlsafe=...)``. The base64 encoded values
will have padding removed.
.. note::
The string returned by ``to_legacy_urlsafe`` is equivalent, but
not identical, to the string returned by ``ndb``. The location
prefix may need to be specified to obtain identical urlsafe
keys.
:type location_prefix: str
:param location_prefix: The location prefix of an App Engine project
ID. Often this value is 's~', but may also be
'e~', or other location prefixes currently
unknown.
:rtype: bytes
:returns: A bytestring containing the key encoded as URL-safe base64.
"""
if location_prefix is None:
project_id = self.project
else:
project_id = location_prefix + self.project
reference = _app_engine_key_pb2.Reference(
app=project_id,
path=_to_legacy_path(self._path), # Avoid the copy.
name_space=self.namespace,
)
raw_bytes = reference.SerializeToString()
return base64.urlsafe_b64encode(raw_bytes).strip(b"=")
@classmethod
def from_legacy_urlsafe(cls, urlsafe):
"""Convert urlsafe string to :class:`~google.cloud.datastore.key.Key`.
This is intended to work with the "legacy" representation of a
datastore "Key" used within Google App Engine (a so-called
"Reference"). This assumes that ``urlsafe`` was created within an App
Engine app via something like ``ndb.Key(...).urlsafe()``.
:type urlsafe: bytes or unicode
:param urlsafe: The base64 encoded (ASCII) string corresponding to a
datastore "Key" / "Reference".
:rtype: :class:`~google.cloud.datastore.key.Key`.
:returns: The key corresponding to ``urlsafe``.
"""
urlsafe = _to_bytes(urlsafe, encoding="ascii")
padding = b"=" * (-len(urlsafe) % 4)
urlsafe += padding
raw_bytes = base64.urlsafe_b64decode(urlsafe)
reference = _app_engine_key_pb2.Reference()
reference.ParseFromString(raw_bytes)
project = _clean_app(reference.app)
namespace = _get_empty(reference.name_space, "")
_check_database_id(reference.database_id)
flat_path = _get_flat_path(reference.path)
return cls(*flat_path, project=project, namespace=namespace)
@property
def is_partial(self):
"""Boolean indicating if the key has an ID (or name).
:rtype: bool
:returns: ``True`` if the last element of the key's path does not have
an ``id`` or a ``name``.
"""
return self.id_or_name is None
@property
def namespace(self):
"""Namespace getter.
:rtype: str
:returns: The namespace of the current key.
"""
return self._namespace
@property
def path(self):
"""Path getter.
Returns a copy so that the key remains immutable.
:rtype: :class:`list` of :class:`dict`
:returns: The (key) path of the current key.
"""
return copy.deepcopy(self._path)
@property
def flat_path(self):
"""Getter for the key path as a tuple.
:rtype: tuple of string and integer
:returns: The tuple of elements in the path.
"""
return self._flat_path
@property
def kind(self):
"""Kind getter. Based on the last element of path.
:rtype: str
:returns: The kind of the current key.
"""
return self.path[-1]["kind"]
@property
def id(self):
"""ID getter. Based on the last element of path.
:rtype: int
:returns: The (integer) ID of the key.
"""
return self.path[-1].get("id")
@property
def name(self):
"""Name getter. Based on the last element of path.
:rtype: str
:returns: The (string) name of the key.
"""
return self.path[-1].get("name")
@property
def id_or_name(self):
"""Getter. Based on the last element of path.
:rtype: int (if ``id``) or string (if ``name``)
:returns: The last element of the key's path if it is either an ``id``
or a ``name``.
"""
if self.id is None:
return self.name
return self.id
@property
def project(self):
"""Project getter.
:rtype: str
:returns: The key's project.
"""
return self._project
def _make_parent(self):
"""Creates a parent key for the current path.
Extracts all but the last element in the key path and creates a new
key, while still matching the namespace and the project.
:rtype: :class:`google.cloud.datastore.key.Key` or :class:`NoneType`
:returns: A new ``Key`` instance, whose path consists of all but the
last element of current path. If the current key has only
one path element, returns ``None``.
"""
if self.is_partial:
parent_args = self.flat_path[:-1]
else:
parent_args = self.flat_path[:-2]
if parent_args:
return self.__class__(
*parent_args, project=self.project, namespace=self.namespace
)
@property
def parent(self):
"""The parent of the current key.
:rtype: :class:`google.cloud.datastore.key.Key` or :class:`NoneType`
:returns: A new ``Key`` instance, whose path consists of all but the
last element of current path. If the current key has only
one path element, returns ``None``.
"""
if self._parent is None:
self._parent = self._make_parent()
return self._parent
def __repr__(self):
return "<Key%s, project=%s>" % (self._flat_path, self.project)
def _validate_project(project, parent):
"""Ensure the project is set appropriately.
If ``parent`` is passed, skip the test (it will be checked / fixed up
later).
If ``project`` is unset, attempt to infer the project from the environment.
:type project: str
:param project: A project.
:type parent: :class:`google.cloud.datastore.key.Key`
:param parent: (Optional) The parent of the key or ``None``.
:rtype: str
:returns: The ``project`` passed in, or implied from the environment.
:raises: :class:`ValueError` if ``project`` is ``None`` and no project
can be inferred from the parent.
"""
if parent is None:
if project is None:
raise ValueError("A Key must have a project set.")
return project
def _clean_app(app_str):
"""Clean a legacy (i.e. from App Engine) app string.
:type app_str: str
:param app_str: The ``app`` value stored in a "Reference" pb.
:rtype: str
:returns: The cleaned value.
"""
parts = app_str.split("~", 1)
return parts[-1]
def _get_empty(value, empty_value):
"""Check if a protobuf field is "empty".
:type value: object
:param value: A basic field from a protobuf.
:type empty_value: object
:param empty_value: The "empty" value for the same type as
``value``.
"""
if value == empty_value:
return None
else:
return value
def _check_database_id(database_id):
"""Make sure a "Reference" database ID is empty.
:type database_id: unicode
:param database_id: The ``database_id`` field from a "Reference" protobuf.
:raises: :exc:`ValueError` if the ``database_id`` is not empty.
"""
if database_id != "":
msg = _DATABASE_ID_TEMPLATE.format(database_id)
raise ValueError(msg)
def _add_id_or_name(flat_path, element_pb, empty_allowed):
"""Add the ID or name from an element to a list.
:type flat_path: list
:param flat_path: List of accumulated path parts.
:type element_pb: :class:`._app_engine_key_pb2.Path.Element`
:param element_pb: The element containing ID or name.
:type empty_allowed: bool
:param empty_allowed: Indicates if neither ID or name need be set. If
:data:`False`, then **exactly** one of them must be.
:raises: :exc:`ValueError` if 0 or 2 of ID/name are set (unless
``empty_allowed=True`` and 0 are set).
"""
id_ = element_pb.id
name = element_pb.name
# NOTE: Below 0 and the empty string are the "null" values for their
# respective types, indicating that the value is unset.
if id_ == 0:
if name == "":
if not empty_allowed:
raise ValueError(_EMPTY_ELEMENT)
else:
flat_path.append(name)
else:
if name == "":
flat_path.append(id_)
else:
msg = _BAD_ELEMENT_TEMPLATE.format(id_, name)
raise ValueError(msg)
def _get_flat_path(path_pb):
"""Convert a legacy "Path" protobuf to a flat path.
For example
.. code:: python
Element {
type: "parent"
id: 59
}
Element {
type: "child"
name: "naem"
}
would convert to ``('parent', 59, 'child', 'naem')``.
:type path_pb: :class:`._app_engine_key_pb2.Path`
:param path_pb: Legacy protobuf "Path" object (from a "Reference").
:rtype: tuple
:returns: The path parts from ``path_pb``.
"""
num_elts = len(path_pb.element)
last_index = num_elts - 1
result = []
for index, element in enumerate(path_pb.element):
result.append(element.type)
_add_id_or_name(result, element, index == last_index)
return tuple(result)
def _to_legacy_path(dict_path):
"""Convert a tuple of ints and strings in a legacy "Path".
.. note:
This assumes, but does not verify, that each entry in
``dict_path`` is valid (i.e. doesn't have more than one
key out of "name" / "id").
:type dict_path: lsit
:param dict_path: The "structured" path for a key, i.e. it
is a list of dictionaries, each of which has
"kind" and one of "name" / "id" as keys.
:rtype: :class:`._app_engine_key_pb2.Path`
:returns: The legacy path corresponding to ``dict_path``.
"""
elements = []
for part in dict_path:
element_kwargs = {"type": part["kind"]}
if "id" in part:
element_kwargs["id"] = part["id"]
elif "name" in part:
element_kwargs["name"] = part["name"]
element = _app_engine_key_pb2.Path.Element(**element_kwargs)
elements.append(element)
return _app_engine_key_pb2.Path(element=elements)