Skip to content

Commit 031b864

Browse files
dhermesdhermes
committed
Implement Key.compare_to_proto to check pb keys against existing.
Addresses sixth part of #451.
1 parent 64ab5ee commit 031b864

File tree

4 files changed

+173
-27
lines changed

4 files changed

+173
-27
lines changed

‎gcloud/datastore/entity.py‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ class NoDataset(RuntimeError):
2727
"""Exception raised by Entity methods which require a dataset."""
2828

2929

30-
class Entity(object):
30+
class Entity(_implicit_environ._DatastoreBase):
3131
"""Entities are akin to rows in a relational database
3232
3333
An entity storing the actual instance of data.
@@ -94,9 +94,7 @@ class Entity(object):
9494
"""
9595

9696
def __init__(self, dataset=None, kind=None, exclude_from_indexes=()):
97-
# Does not inherit from object, so we don't use
98-
# _implicit_environ._DatastoreBase to avoid split MRO.
99-
self._dataset = dataset or _implicit_environ.DATASET
97+
super(Entity, self).__init__(dataset=dataset)
10098
self._data = {}
10199
if kind:
102100
self._key = Key(kind)
@@ -286,7 +284,7 @@ def save(self):
286284
key_pb = connection.save_entity(
287285
dataset_id=dataset.id(),
288286
key_pb=key.to_protobuf(),
289-
properties=self._data,
287+
properties=self.to_dict(),
290288
exclude_from_indexes=self.exclude_from_indexes())
291289

292290
# If we are in a transaction and the current entity needs an
@@ -296,22 +294,8 @@ def save(self):
296294
transaction.add_auto_id_entity(self)
297295

298296
if isinstance(key_pb, datastore_pb.Key):
299-
# Update the path (which may have been altered).
300-
# NOTE: The underlying namespace can't have changed in a save().
301-
# The value of the dataset ID may have changed from implicit
302-
# (i.e. None, with the ID implied from the dataset.Dataset
303-
# object associated with the Entity/Key), but if it was
304-
# implicit before the save() we leave it as implicit.
305-
path = []
306-
for element in key_pb.path_element:
307-
key_part = {}
308-
for descriptor, value in element._fields.items():
309-
key_part[descriptor.name] = value
310-
path.append(key_part)
311-
# This is temporary. Will be addressed throughout #451.
312-
clone = key._clone()
313-
clone._path = path
314-
self._key = clone
297+
# Update the key (which may have been altered).
298+
self.key(self.key().compare_to_proto(key_pb))
315299

316300
return self
317301

‎gcloud/datastore/key.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,92 @@ def complete_key(self, id_or_name):
152152
new_key._flat_path += (id_or_name,)
153153
return new_key
154154

155+
def _validate_protobuf_dataset_id(self, protobuf):
156+
"""Checks that dataset ID on protobuf matches current one.
157+
158+
The value of the dataset ID may have changed from unprefixed
159+
(e.g. 'foo') to prefixed (e.g. 's~foo' or 'e~foo').
160+
161+
:type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key`
162+
:param protobuf: A protobuf representation of the key. Expected to be
163+
returned after a datastore operation.
164+
165+
:rtype: :class:`str`
166+
"""
167+
proto_dataset_id = protobuf.partition_id.dataset_id
168+
if proto_dataset_id == self.dataset_id:
169+
return
170+
171+
# Since they don't match, we check to see if `proto_dataset_id` has a
172+
# prefix.
173+
unprefixed = None
174+
prefix = proto_dataset_id[:2]
175+
if prefix in ('s~', 'e~'):
176+
unprefixed = proto_dataset_id[2:]
177+
178+
if unprefixed != self.dataset_id:
179+
raise ValueError('Dataset ID on protobuf does not match.',
180+
proto_dataset_id, self.dataset_id)
181+
182+
def compare_to_proto(self, protobuf):
183+
"""Checks current key against a protobuf; updates if partial.
184+
185+
If the current key is partial, returns a new key that has been
186+
completed otherwise returns the current key.
187+
188+
The value of the dataset ID may have changed from implicit (i.e. None,
189+
with the ID implied from the dataset.Dataset object associated with the
190+
Entity/Key), but if it was implicit before, we leave it as implicit.
191+
192+
:type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key`
193+
:param protobuf: A protobuf representation of the key. Expected to be
194+
returned after a datastore operation.
195+
196+
:rtype: :class:`gcloud.datastore.key.Key`
197+
:returns: The current key if not partial.
198+
:raises: `ValueError` if the namespace or dataset ID of `protobuf`
199+
don't match the current values or if the path from `protobuf`
200+
doesn't match.
201+
"""
202+
if self.namespace is None:
203+
if protobuf.partition_id.HasField('namespace'):
204+
raise ValueError('Namespace unset on key but set on protobuf.')
205+
elif protobuf.partition_id.namespace != self.namespace:
206+
raise ValueError('Namespace on protobuf does not match.',
207+
protobuf.partition_id.namespace, self.namespace)
208+
209+
# Check that dataset IDs match if not implicit.
210+
if self.dataset_id is not None:
211+
self._validate_protobuf_dataset_id(protobuf)
212+
213+
path = []
214+
for element in protobuf.path_element:
215+
key_part = {}
216+
for descriptor, value in element._fields.items():
217+
key_part[descriptor.name] = value
218+
path.append(key_part)
219+
220+
if path == self.path:
221+
return self
222+
223+
if not self.is_partial:
224+
raise ValueError('Proto path does not match completed key.',
225+
path, self.path)
226+
227+
last_part = path[-1]
228+
id_or_name = None
229+
if 'id' in last_part:
230+
id_or_name = last_part.pop('id')
231+
elif 'name' in last_part:
232+
id_or_name = last_part.pop('name')
233+
234+
# We have edited path by popping from the last part, so check again.
235+
if path != self.path:
236+
raise ValueError('Proto path does not match partial key.',
237+
path, self.path)
238+
239+
return self.complete_key(id_or_name)
240+
155241
def to_protobuf(self):
156242
"""Return a protobuf corresponding to the key.
157243

‎gcloud/datastore/test_entity.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -345,12 +345,7 @@ def get_entities(self, keys):
345345
return [self.get(key) for key in keys]
346346

347347
def allocate_ids(self, incomplete_key, num_ids):
348-
def clone_with_new_id(key, new_id):
349-
clone = key._clone()
350-
clone._path[-1]['id'] = new_id
351-
return clone
352-
return [clone_with_new_id(incomplete_key, i + 1)
353-
for i in range(num_ids)]
348+
return [incomplete_key.complete_key(i + 1) for i in range(num_ids)]
354349

355350

356351
class _Connection(object):

‎gcloud/datastore/test_key.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,87 @@ def test_complete_key_on_complete(self):
8484
key = self._makeOne('KIND', 1234)
8585
self.assertRaises(ValueError, key.complete_key, 5678)
8686

87+
def test_compare_to_proto_incomplete_w_id(self):
88+
_ID = 1234
89+
key = self._makeOne('KIND')
90+
pb = key.to_protobuf()
91+
pb.path_element[0].id = _ID
92+
new_key = key.compare_to_proto(pb)
93+
self.assertFalse(new_key is key)
94+
self.assertEqual(new_key.id, _ID)
95+
self.assertEqual(new_key.name, None)
96+
97+
def test_compare_to_proto_incomplete_w_name(self):
98+
_NAME = 'NAME'
99+
key = self._makeOne('KIND')
100+
pb = key.to_protobuf()
101+
pb.path_element[0].name = _NAME
102+
new_key = key.compare_to_proto(pb)
103+
self.assertFalse(new_key is key)
104+
self.assertEqual(new_key.id, None)
105+
self.assertEqual(new_key.name, _NAME)
106+
107+
def test_compare_to_proto_incomplete_w_incomplete(self):
108+
key = self._makeOne('KIND')
109+
pb = key.to_protobuf()
110+
new_key = key.compare_to_proto(pb)
111+
self.assertTrue(new_key is key)
112+
113+
def test_compare_to_proto_incomplete_w_bad_path(self):
114+
key = self._makeOne('KIND1', 1234, 'KIND2')
115+
pb = key.to_protobuf()
116+
pb.path_element[0].kind = 'NO_KIND'
117+
self.assertRaises(ValueError, key.compare_to_proto, pb)
118+
119+
def test_compare_to_proto_complete_w_id(self):
120+
key = self._makeOne('KIND', 1234)
121+
pb = key.to_protobuf()
122+
pb.path_element[0].id = 5678
123+
self.assertRaises(ValueError, key.compare_to_proto, pb)
124+
125+
def test_compare_to_proto_complete_w_name(self):
126+
key = self._makeOne('KIND', 1234)
127+
pb = key.to_protobuf()
128+
pb.path_element[0].name = 'NAME'
129+
self.assertRaises(ValueError, key.compare_to_proto, pb)
130+
131+
def test_compare_to_proto_complete_w_incomplete(self):
132+
key = self._makeOne('KIND', 1234)
133+
pb = key.to_protobuf()
134+
pb.path_element[0].ClearField('id')
135+
self.assertRaises(ValueError, key.compare_to_proto, pb)
136+
137+
def test_compare_to_proto_complete_diff_dataset(self):
138+
key = self._makeOne('KIND', 1234, dataset_id='DATASET')
139+
pb = key.to_protobuf()
140+
pb.partition_id.dataset_id = 's~' + key.dataset_id
141+
new_key = key.compare_to_proto(pb)
142+
self.assertTrue(new_key is key)
143+
144+
def test_compare_to_proto_complete_bad_dataset(self):
145+
key = self._makeOne('KIND', 1234, dataset_id='DATASET')
146+
pb = key.to_protobuf()
147+
pb.partition_id.dataset_id = 'BAD_PRE~' + key.dataset_id
148+
self.assertRaises(ValueError, key.compare_to_proto, pb)
149+
150+
def test_compare_to_proto_complete_valid_namespace(self):
151+
key = self._makeOne('KIND', 1234, namespace='NAMESPACE')
152+
pb = key.to_protobuf()
153+
new_key = key.compare_to_proto(pb)
154+
self.assertTrue(new_key is key)
155+
156+
def test_compare_to_proto_complete_namespace_unset_on_pb(self):
157+
key = self._makeOne('KIND', 1234, namespace='NAMESPACE')
158+
pb = key.to_protobuf()
159+
pb.partition_id.ClearField('namespace')
160+
self.assertRaises(ValueError, key.compare_to_proto, pb)
161+
162+
def test_compare_to_proto_complete_namespace_unset_on_key(self):
163+
key = self._makeOne('KIND', 1234)
164+
pb = key.to_protobuf()
165+
pb.partition_id.namespace = 'NAMESPACE'
166+
self.assertRaises(ValueError, key.compare_to_proto, pb)
167+
87168
def test_to_protobuf_defaults(self):
88169
from gcloud.datastore.datastore_v1_pb2 import Key as KeyPB
89170
_KIND = 'KIND'

0 commit comments

Comments
 (0)