Making Key fully immutable and requiring a dataset ID.

Author: dhermes

gcloud/datastore/connection.py

@@ -241,8 +241,7 @@ def lookup(self, dataset_id, key_pbs,
  if single_key:
  key_pbs = [key_pbs]
 
- for key_pb in key_pbs:
- lookup_request.key.add().CopyFrom(key_pb)
+ helpers._add_keys_to_request(lookup_request.key, key_pbs)
 
  results, missing_found, deferred_found = self._lookup(
  lookup_request, dataset_id, deferred is not None)
@@ -417,8 +416,7 @@ def allocate_ids(self, dataset_id, key_pbs):
  :returns: An equal number of keys, with IDs filled in by the backend.
  """
  request = datastore_pb.AllocateIdsRequest()
- for key_pb in key_pbs:
- request.key.add().CopyFrom(key_pb)
+ helpers._add_keys_to_request(request.key, key_pbs)
  # Nothing to do with this response, so just execute the method.
  response = self._rpc(dataset_id, 'allocateIds', request,
  datastore_pb.AllocateIdsResponse)
@@ -444,8 +442,14 @@ def save_entity(self, dataset_id, key_pb, properties,
 
  :type exclude_from_indexes: sequence of str
  :param exclude_from_indexes: Names of properties *not* to be indexed.
+
+ :rtype: bool or :class:`gcloud.datastore.datastore_v1_pb2.Key`
+ :returns: True if the save succeeds, unless a new ID has been
+ automatically allocated. In the auto ID case, the newly
+ created key protobuf is returned.
  """
  mutation = self.mutation()
+ key_pb = helpers._prepare_key_for_request(key_pb)
 
  # If the Key is complete, we should upsert
  # instead of using insert_auto_id.
@@ -506,10 +510,7 @@ def delete_entities(self, dataset_id, key_pbs):
  :returns: True
  """
  mutation = self.mutation()
-
- for key_pb in key_pbs:
- delete = mutation.delete.add()
- delete.CopyFrom(key_pb)
+ helpers._add_keys_to_request(mutation.delete, key_pbs)
 
  if not self.transaction():
  self.commit(dataset_id, mutation)

gcloud/datastore/dataset.py

@@ -136,7 +136,7 @@ def get_entity(self, key_or_path):
  if isinstance(key_or_path, Key):
  entities = self.get_entities([key_or_path])
  else:
- key = Key.from_path(*key_or_path)
+ key = Key(*key_or_path)
  entities = self.get_entities([key])
 
  if entities:
@@ -196,7 +196,7 @@ def allocate_ids(self, incomplete_key, num_ids):
  :return: The (complete) keys allocated with `incomplete_key` as root.
  :raises: `ValueError` if `incomplete_key` is not a partial key.
  """
- if not incomplete_key.is_partial():
+ if not incomplete_key.is_partial:
  raise ValueError(('Key is not partial.', incomplete_key))
 
  incomplete_key_pb = incomplete_key.to_protobuf()
@@ -206,5 +206,5 @@ def allocate_ids(self, incomplete_key, num_ids):
  self.id(), incomplete_key_pbs)
  allocated_ids = [allocated_key_pb.path_element[-1].id
  for allocated_key_pb in allocated_key_pbs]
- return [incomplete_key.id(allocated_id)
+ return [incomplete_key.complete_key(allocated_id)
  for allocated_id in allocated_ids]

gcloud/datastore/entity.py

@@ -97,7 +97,10 @@ def __init__(self, dataset=None, kind=None, exclude_from_indexes=()):
  super(Entity, self).__init__(dataset=dataset)
  self._data = {}
  if kind:
- self._key = Key().kind(kind)
+ # This is temporary since the dataset will eventually be 100%
+ # removed from the Entity and the Dataset class may be
+ # destroyed.
+ self._key = Key(kind, dataset_id=self.dataset().id())
  else:
  self._key = None
  self._exclude_from_indexes = set(exclude_from_indexes)
@@ -193,7 +196,7 @@ def kind(self):
  """
 
  if self._key:
- return self._key.kind()
+ return self._key.kind
 
  def exclude_from_indexes(self):
  """Names of fields which are *not* to be indexed for this entity.
@@ -284,29 +287,18 @@ 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
  # automatically assigned ID, tell the transaction where to put that.
  transaction = connection.transaction()
- if transaction and key.is_partial():
+ if transaction and key.is_partial:
  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)
- self._key = key.path(path)
+ # Update the key (which may have been altered).
+ self._key = self.key().compare_to_proto(key_pb)
 
  return self
 
@@ -327,6 +319,6 @@ def delete(self):
 
  def __repr__(self):
  if self._key:
- return '<Entity%s %r>' % (self._key.path(), self._data)
+ return '<Entity%s %r>' % (self._key.path, self._data)
  else:
  return '<Entity %r>' % (self._data,)

gcloud/datastore/helpers.py

@@ -25,6 +25,7 @@
 import pytz
 import six
 
+from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 
@@ -65,19 +66,15 @@ def key_from_protobuf(pb):
  :rtype: :class:`gcloud.datastore.key.Key`
  :returns: a new `Key` instance
  """
- path = []
+ path_args = []
  for element in pb.path_element:
- element_dict = {'kind': element.kind}
-
+ path_args.append(element.kind)
  if element.HasField('id'):
- element_dict['id'] = element.id
-
+ path_args.append(element.id)
  # This is safe: we expect proto objects returned will only have
  # one of `name` or `id` set.
  if element.HasField('name'):
- element_dict['name'] = element.name
-
- path.append(element_dict)
+ path_args.append(element.name)
 
  dataset_id = None
  if pb.partition_id.HasField('dataset_id'):
@@ -86,7 +83,7 @@ def key_from_protobuf(pb):
  if pb.partition_id.HasField('namespace'):
  namespace = pb.partition_id.namespace
 
- return Key(path, namespace, dataset_id)
+ return Key(*path_args, namespace=namespace, dataset_id=dataset_id)
 
 
 def _pb_attr_value(val):
@@ -263,3 +260,44 @@ def _set_protobuf_value(value_pb, val):
  _set_protobuf_value(i_pb, item)
  else: # scalar, just assign
  setattr(value_pb, attr, val)
+
+
+def _prepare_key_for_request(key_pb):
+ """Add protobuf keys to a request object.
+
+ :type key_pb: :class:`gcloud.datastore.datastore_v1_pb2.Key`
+ :param key_pb: A key to be added to a request.
+
+ :rtype: :class:`gcloud.datastore.datastore_v1_pb2.Key`
+ :returns: A key which will be added to a request. It will be the
+ original if nothing needs to be changed.
+ """
+ if key_pb.partition_id.HasField('dataset_id'):
+ # We remove the dataset_id from the protobuf. This is because
+ # the backend fails a request if the key contains un-prefixed
+ # dataset ID. The backend fails because requests to
+ # /datastore/.../datasets/foo/...
+ # and
+ # /datastore/.../datasets/s~foo/...
+ # both go to the datastore given by 's~foo'. So if the key
+ # protobuf in the request body has dataset_id='foo', the
+ # backend will reject since 'foo' != 's~foo'.
+ new_key_pb = datastore_pb.Key()
+ new_key_pb.CopyFrom(key_pb)
+ new_key_pb.partition_id.ClearField('dataset_id')
+ key_pb = new_key_pb
+ return key_pb
+
+
+def _add_keys_to_request(request_field_pb, key_pbs):
+ """Add protobuf keys to a request object.
+
+ :type request_field_pb: `RepeatedCompositeFieldContainer`
+ :param request_field_pb: A repeated proto field that contains keys.
+
+ :type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+ :param key_pbs: The keys to add to a request.
+ """
+ for key_pb in key_pbs:
+ key_pb = _prepare_key_for_request(key_pb)
+ request_field_pb.add().CopyFrom(key_pb)