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