Skip to content

Commit 7865b29

Browse files
dhermesdhermes
committed
Making Key fully immutable and requiring a dataset ID.
1 parent 0aec96c commit 7865b29

19 files changed

Lines changed: 851 additions & 617 deletions

‎gcloud/datastore/connection.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -241,8 +241,7 @@ def lookup(self, dataset_id, key_pbs,
241241
if single_key:
242242
key_pbs = [key_pbs]
243243

244-
for key_pb in key_pbs:
245-
lookup_request.key.add().CopyFrom(key_pb)
244+
helpers._add_keys_to_request(lookup_request.key, key_pbs)
246245

247246
results, missing_found, deferred_found = self._lookup(
248247
lookup_request, dataset_id, deferred is not None)
@@ -417,8 +416,7 @@ def allocate_ids(self, dataset_id, key_pbs):
417416
:returns: An equal number of keys, with IDs filled in by the backend.
418417
"""
419418
request = datastore_pb.AllocateIdsRequest()
420-
for key_pb in key_pbs:
421-
request.key.add().CopyFrom(key_pb)
419+
helpers._add_keys_to_request(request.key, key_pbs)
422420
# Nothing to do with this response, so just execute the method.
423421
response = self._rpc(dataset_id, 'allocateIds', request,
424422
datastore_pb.AllocateIdsResponse)
@@ -444,8 +442,14 @@ def save_entity(self, dataset_id, key_pb, properties,
444442
445443
:type exclude_from_indexes: sequence of str
446444
:param exclude_from_indexes: Names of properties *not* to be indexed.
445+
446+
:rtype: bool or :class:`gcloud.datastore.datastore_v1_pb2.Key`
447+
:returns: True if the save succeeds, unless a new ID has been
448+
automatically allocated. In the auto ID case, the newly
449+
created key protobuf is returned.
447450
"""
448451
mutation = self.mutation()
452+
key_pb = helpers._prepare_key_for_request(key_pb)
449453

450454
# If the Key is complete, we should upsert
451455
# instead of using insert_auto_id.
@@ -506,10 +510,7 @@ def delete_entities(self, dataset_id, key_pbs):
506510
:returns: True
507511
"""
508512
mutation = self.mutation()
509-
510-
for key_pb in key_pbs:
511-
delete = mutation.delete.add()
512-
delete.CopyFrom(key_pb)
513+
helpers._add_keys_to_request(mutation.delete, key_pbs)
513514

514515
if not self.transaction():
515516
self.commit(dataset_id, mutation)

‎gcloud/datastore/dataset.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ def get_entity(self, key_or_path):
136136
if isinstance(key_or_path, Key):
137137
entities = self.get_entities([key_or_path])
138138
else:
139-
key = Key.from_path(*key_or_path)
139+
key = Key(*key_or_path)
140140
entities = self.get_entities([key])
141141

142142
if entities:
@@ -196,7 +196,7 @@ def allocate_ids(self, incomplete_key, num_ids):
196196
:return: The (complete) keys allocated with `incomplete_key` as root.
197197
:raises: `ValueError` if `incomplete_key` is not a partial key.
198198
"""
199-
if not incomplete_key.is_partial():
199+
if not incomplete_key.is_partial:
200200
raise ValueError(('Key is not partial.', incomplete_key))
201201

202202
incomplete_key_pb = incomplete_key.to_protobuf()
@@ -206,5 +206,5 @@ def allocate_ids(self, incomplete_key, num_ids):
206206
self.id(), incomplete_key_pbs)
207207
allocated_ids = [allocated_key_pb.path_element[-1].id
208208
for allocated_key_pb in allocated_key_pbs]
209-
return [incomplete_key.id(allocated_id)
209+
return [incomplete_key.complete_key(allocated_id)
210210
for allocated_id in allocated_ids]

‎gcloud/datastore/entity.py‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,10 @@ def __init__(self, dataset=None, kind=None, exclude_from_indexes=()):
9797
super(Entity, self).__init__(dataset=dataset)
9898
self._data = {}
9999
if kind:
100-
self._key = Key().kind(kind)
100+
# This is temporary since the dataset will eventually be 100%
101+
# removed from the Entity and the Dataset class may be
102+
# destroyed.
103+
self._key = Key(kind, dataset_id=self.dataset().id())
101104
else:
102105
self._key = None
103106
self._exclude_from_indexes = set(exclude_from_indexes)
@@ -193,7 +196,7 @@ def kind(self):
193196
"""
194197

195198
if self._key:
196-
return self._key.kind()
199+
return self._key.kind
197200

198201
def exclude_from_indexes(self):
199202
"""Names of fields which are *not* to be indexed for this entity.
@@ -284,29 +287,18 @@ def save(self):
284287
key_pb = connection.save_entity(
285288
dataset_id=dataset.id(),
286289
key_pb=key.to_protobuf(),
287-
properties=self._data,
290+
properties=self.to_dict(),
288291
exclude_from_indexes=self.exclude_from_indexes())
289292

290293
# If we are in a transaction and the current entity needs an
291294
# automatically assigned ID, tell the transaction where to put that.
292295
transaction = connection.transaction()
293-
if transaction and key.is_partial():
296+
if transaction and key.is_partial:
294297
transaction.add_auto_id_entity(self)
295298

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

311303
return self
312304

@@ -327,6 +319,6 @@ def delete(self):
327319

328320
def __repr__(self):
329321
if self._key:
330-
return '<Entity%s %r>' % (self._key.path(), self._data)
322+
return '<Entity%s %r>' % (self._key.path, self._data)
331323
else:
332324
return '<Entity %r>' % (self._data,)

‎gcloud/datastore/helpers.py‎

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import pytz
2626
import six
2727

28+
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
2829
from gcloud.datastore.entity import Entity
2930
from gcloud.datastore.key import Key
3031

@@ -65,19 +66,15 @@ def key_from_protobuf(pb):
6566
:rtype: :class:`gcloud.datastore.key.Key`
6667
:returns: a new `Key` instance
6768
"""
68-
path = []
69+
path_args = []
6970
for element in pb.path_element:
70-
element_dict = {'kind': element.kind}
71-
71+
path_args.append(element.kind)
7272
if element.HasField('id'):
73-
element_dict['id'] = element.id
74-
73+
path_args.append(element.id)
7574
# This is safe: we expect proto objects returned will only have
7675
# one of `name` or `id` set.
7776
if element.HasField('name'):
78-
element_dict['name'] = element.name
79-
80-
path.append(element_dict)
77+
path_args.append(element.name)
8178

8279
dataset_id = None
8380
if pb.partition_id.HasField('dataset_id'):
@@ -86,7 +83,7 @@ def key_from_protobuf(pb):
8683
if pb.partition_id.HasField('namespace'):
8784
namespace = pb.partition_id.namespace
8885

89-
return Key(path, namespace, dataset_id)
86+
return Key(*path_args, namespace=namespace, dataset_id=dataset_id)
9087

9188

9289
def _pb_attr_value(val):
@@ -263,3 +260,44 @@ def _set_protobuf_value(value_pb, val):
263260
_set_protobuf_value(i_pb, item)
264261
else: # scalar, just assign
265262
setattr(value_pb, attr, val)
263+
264+
265+
def _prepare_key_for_request(key_pb):
266+
"""Add protobuf keys to a request object.
267+
268+
:type key_pb: :class:`gcloud.datastore.datastore_v1_pb2.Key`
269+
:param key_pb: A key to be added to a request.
270+
271+
:rtype: :class:`gcloud.datastore.datastore_v1_pb2.Key`
272+
:returns: A key which will be added to a request. It will be the
273+
original if nothing needs to be changed.
274+
"""
275+
if key_pb.partition_id.HasField('dataset_id'):
276+
# We remove the dataset_id from the protobuf. This is because
277+
# the backend fails a request if the key contains un-prefixed
278+
# dataset ID. The backend fails because requests to
279+
# /datastore/.../datasets/foo/...
280+
# and
281+
# /datastore/.../datasets/s~foo/...
282+
# both go to the datastore given by 's~foo'. So if the key
283+
# protobuf in the request body has dataset_id='foo', the
284+
# backend will reject since 'foo' != 's~foo'.
285+
new_key_pb = datastore_pb.Key()
286+
new_key_pb.CopyFrom(key_pb)
287+
new_key_pb.partition_id.ClearField('dataset_id')
288+
key_pb = new_key_pb
289+
return key_pb
290+
291+
292+
def _add_keys_to_request(request_field_pb, key_pbs):
293+
"""Add protobuf keys to a request object.
294+
295+
:type request_field_pb: `RepeatedCompositeFieldContainer`
296+
:param request_field_pb: A repeated proto field that contains keys.
297+
298+
:type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
299+
:param key_pbs: The keys to add to a request.
300+
"""
301+
for key_pb in key_pbs:
302+
key_pb = _prepare_key_for_request(key_pb)
303+
request_field_pb.add().CopyFrom(key_pb)

0 commit comments

Comments
 (0)