diff --git a/gcloud/datastore/entity.py b/gcloud/datastore/entity.py index 97efc729dd8ef..0298a2dedd5f8 100644 --- a/gcloud/datastore/entity.py +++ b/gcloud/datastore/entity.py @@ -27,7 +27,7 @@ class NoDataset(RuntimeError): """Exception raised by Entity methods which require a dataset.""" -class Entity(object): +class Entity(_implicit_environ._DatastoreBase): """Entities are akin to rows in a relational database An entity storing the actual instance of data. @@ -94,9 +94,7 @@ class Entity(object): """ def __init__(self, dataset=None, kind=None, exclude_from_indexes=()): - # Does not inherit from object, so we don't use - # _implicit_environ._DatastoreBase to avoid split MRO. - self._dataset = dataset or _implicit_environ.DATASET + super(Entity, self).__init__(dataset=dataset) self._data = {} if kind: self._key = Key(kind) @@ -286,7 +284,7 @@ def save(self): key_pb = connection.save_entity( dataset_id=dataset.id(), key_pb=key.to_protobuf(), - properties=self._data, + properties=self.to_dict(), exclude_from_indexes=self.exclude_from_indexes()) # If we are in a transaction and the current entity needs an @@ -296,22 +294,8 @@ def save(self): transaction.add_auto_id_entity(self) if isinstance(key_pb, datastore_pb.Key): - # Update the path (which may have been altered). - # NOTE: The underlying namespace can't have changed in a save(). - # The value of the dataset ID may have changed from implicit - # (i.e. None, with the ID implied from the dataset.Dataset - # object associated with the Entity/Key), but if it was - # implicit before the save() we leave it as implicit. - path = [] - for element in key_pb.path_element: - key_part = {} - for descriptor, value in element._fields.items(): - key_part[descriptor.name] = value - path.append(key_part) - # This is temporary. Will be addressed throughout #451. - clone = key._clone() - clone._path = path - self._key = clone + # Update the key (which may have been altered). + self.key(self.key().compare_to_proto(key_pb)) return self diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py index de31d7a2aebf1..1e04e910e1a69 100644 --- a/gcloud/datastore/key.py +++ b/gcloud/datastore/key.py @@ -152,6 +152,92 @@ def complete_key(self, id_or_name): new_key._flat_path += (id_or_name,) return new_key + def _validate_protobuf_dataset_id(self, protobuf): + """Checks that dataset ID on protobuf matches current one. + + The value of the dataset ID may have changed from unprefixed + (e.g. 'foo') to prefixed (e.g. 's~foo' or 'e~foo'). + + :type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key` + :param protobuf: A protobuf representation of the key. Expected to be + returned after a datastore operation. + + :rtype: :class:`str` + """ + proto_dataset_id = protobuf.partition_id.dataset_id + if proto_dataset_id == self.dataset_id: + return + + # Since they don't match, we check to see if `proto_dataset_id` has a + # prefix. + unprefixed = None + prefix = proto_dataset_id[:2] + if prefix in ('s~', 'e~'): + unprefixed = proto_dataset_id[2:] + + if unprefixed != self.dataset_id: + raise ValueError('Dataset ID on protobuf does not match.', + proto_dataset_id, self.dataset_id) + + def compare_to_proto(self, protobuf): + """Checks current key against a protobuf; updates if partial. + + If the current key is partial, returns a new key that has been + completed otherwise returns the current key. + + The value of the dataset ID may have changed from implicit (i.e. None, + with the ID implied from the dataset.Dataset object associated with the + Entity/Key), but if it was implicit before, we leave it as implicit. + + :type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key` + :param protobuf: A protobuf representation of the key. Expected to be + returned after a datastore operation. + + :rtype: :class:`gcloud.datastore.key.Key` + :returns: The current key if not partial. + :raises: `ValueError` if the namespace or dataset ID of `protobuf` + don't match the current values or if the path from `protobuf` + doesn't match. + """ + if self.namespace is None: + if protobuf.partition_id.HasField('namespace'): + raise ValueError('Namespace unset on key but set on protobuf.') + elif protobuf.partition_id.namespace != self.namespace: + raise ValueError('Namespace on protobuf does not match.', + protobuf.partition_id.namespace, self.namespace) + + # Check that dataset IDs match if not implicit. + if self.dataset_id is not None: + self._validate_protobuf_dataset_id(protobuf) + + path = [] + for element in protobuf.path_element: + key_part = {} + for descriptor, value in element._fields.items(): + key_part[descriptor.name] = value + path.append(key_part) + + if path == self.path: + return self + + if not self.is_partial: + raise ValueError('Proto path does not match completed key.', + path, self.path) + + last_part = path[-1] + id_or_name = None + if 'id' in last_part: + id_or_name = last_part.pop('id') + elif 'name' in last_part: + id_or_name = last_part.pop('name') + + # We have edited path by popping from the last part, so check again. + if path != self.path: + raise ValueError('Proto path does not match partial key.', + path, self.path) + + return self.complete_key(id_or_name) + def to_protobuf(self): """Return a protobuf corresponding to the key. diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py index f785954ecac8a..80de95e713d98 100644 --- a/gcloud/datastore/test_entity.py +++ b/gcloud/datastore/test_entity.py @@ -345,12 +345,7 @@ def get_entities(self, keys): return [self.get(key) for key in keys] def allocate_ids(self, incomplete_key, num_ids): - def clone_with_new_id(key, new_id): - clone = key._clone() - clone._path[-1]['id'] = new_id - return clone - return [clone_with_new_id(incomplete_key, i + 1) - for i in range(num_ids)] + return [incomplete_key.complete_key(i + 1) for i in range(num_ids)] class _Connection(object): diff --git a/gcloud/datastore/test_key.py b/gcloud/datastore/test_key.py index 2973b90ab765c..d301edd5075ce 100644 --- a/gcloud/datastore/test_key.py +++ b/gcloud/datastore/test_key.py @@ -84,6 +84,87 @@ def test_complete_key_on_complete(self): key = self._makeOne('KIND', 1234) self.assertRaises(ValueError, key.complete_key, 5678) + def test_compare_to_proto_incomplete_w_id(self): + _ID = 1234 + key = self._makeOne('KIND') + pb = key.to_protobuf() + pb.path_element[0].id = _ID + new_key = key.compare_to_proto(pb) + self.assertFalse(new_key is key) + self.assertEqual(new_key.id, _ID) + self.assertEqual(new_key.name, None) + + def test_compare_to_proto_incomplete_w_name(self): + _NAME = 'NAME' + key = self._makeOne('KIND') + pb = key.to_protobuf() + pb.path_element[0].name = _NAME + new_key = key.compare_to_proto(pb) + self.assertFalse(new_key is key) + self.assertEqual(new_key.id, None) + self.assertEqual(new_key.name, _NAME) + + def test_compare_to_proto_incomplete_w_incomplete(self): + key = self._makeOne('KIND') + pb = key.to_protobuf() + new_key = key.compare_to_proto(pb) + self.assertTrue(new_key is key) + + def test_compare_to_proto_incomplete_w_bad_path(self): + key = self._makeOne('KIND1', 1234, 'KIND2') + pb = key.to_protobuf() + pb.path_element[0].kind = 'NO_KIND' + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_w_id(self): + key = self._makeOne('KIND', 1234) + pb = key.to_protobuf() + pb.path_element[0].id = 5678 + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_w_name(self): + key = self._makeOne('KIND', 1234) + pb = key.to_protobuf() + pb.path_element[0].name = 'NAME' + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_w_incomplete(self): + key = self._makeOne('KIND', 1234) + pb = key.to_protobuf() + pb.path_element[0].ClearField('id') + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_diff_dataset(self): + key = self._makeOne('KIND', 1234, dataset_id='DATASET') + pb = key.to_protobuf() + pb.partition_id.dataset_id = 's~' + key.dataset_id + new_key = key.compare_to_proto(pb) + self.assertTrue(new_key is key) + + def test_compare_to_proto_complete_bad_dataset(self): + key = self._makeOne('KIND', 1234, dataset_id='DATASET') + pb = key.to_protobuf() + pb.partition_id.dataset_id = 'BAD_PRE~' + key.dataset_id + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_valid_namespace(self): + key = self._makeOne('KIND', 1234, namespace='NAMESPACE') + pb = key.to_protobuf() + new_key = key.compare_to_proto(pb) + self.assertTrue(new_key is key) + + def test_compare_to_proto_complete_namespace_unset_on_pb(self): + key = self._makeOne('KIND', 1234, namespace='NAMESPACE') + pb = key.to_protobuf() + pb.partition_id.ClearField('namespace') + self.assertRaises(ValueError, key.compare_to_proto, pb) + + def test_compare_to_proto_complete_namespace_unset_on_key(self): + key = self._makeOne('KIND', 1234) + pb = key.to_protobuf() + pb.partition_id.namespace = 'NAMESPACE' + self.assertRaises(ValueError, key.compare_to_proto, pb) + def test_to_protobuf_defaults(self): from gcloud.datastore.datastore_v1_pb2 import Key as KeyPB _KIND = 'KIND'