Skip to content

Commit cc59310

Browse files
committed
Implement Key.compare_to_proto to check pb keys against existing.
Addresses sixth part of #451.
1 parent f0e26b8 commit cc59310

File tree

4 files changed

+170
-22
lines changed

4 files changed

+170
-22
lines changed

gcloud/datastore/entity.py

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -254,22 +254,8 @@ def save(self):
254254
transaction.add_auto_id_entity(self)
255255

256256
if isinstance(key_pb, datastore_pb.Key):
257-
# Update the path (which may have been altered).
258-
# NOTE: The underlying namespace can't have changed in a save().
259-
# The value of the dataset ID may have changed from implicit
260-
# (i.e. None, with the ID implied from the dataset.Dataset
261-
# object associated with the Entity/Key), but if it was
262-
# implicit before the save() we leave it as implicit.
263-
path = []
264-
for element in key_pb.path_element:
265-
key_part = {}
266-
for descriptor, value in element._fields.items():
267-
key_part[descriptor.name] = value
268-
path.append(key_part)
269-
# This is temporary. Will be addressed throughout #451.
270-
clone = key._clone()
271-
clone._path = path
272-
self._key = clone
257+
# Update the key (which may have been altered).
258+
self.key(self.key().compare_to_proto(key_pb))
273259

274260
return self
275261

gcloud/datastore/key.py

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

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

gcloud/datastore/test_entity.py

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

289289
def allocate_ids(self, incomplete_key, num_ids):
290-
def clone_with_new_id(key, new_id):
291-
clone = key._clone()
292-
clone._path[-1]['id'] = new_id
293-
return clone
294-
return [clone_with_new_id(incomplete_key, i + 1)
295-
for i in range(num_ids)]
290+
return [incomplete_key.complete_key(i + 1) for i in range(num_ids)]
296291

297292

298293
class _Connection(object):

gcloud/datastore/test_key.py

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,87 @@ def test_completed_key_on_complete(self):
8686
key = self._makeOne('KIND', 1234)
8787
self.assertRaises(ValueError, key.completed_key, 5678)
8888

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

0 commit comments

Comments
 (0)