diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index 9642be52afc1..ba57de9269e0 100644 --- a/gcloud/datastore/connection.py +++ b/gcloud/datastore/connection.py @@ -444,6 +444,11 @@ 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: :class:`tuple` + :returns: The pair (`assigned`, `new_id`) where `assigned` is a boolean + indicating if a new ID has been assigned and `new_id` is + either `None` or an integer that has been assigned. """ mutation = self.mutation() @@ -477,14 +482,18 @@ def save_entity(self, dataset_id, key_pb, properties, # If this is in a transaction, we should just return True. The # transaction will handle assigning any keys as necessary. if self.transaction(): - return True + return False, None result = self.commit(dataset_id, mutation) - # If this was an auto-assigned ID, return the new Key. + # If this was an auto-assigned ID, return the new Key. We don't + # verify that this matches the original `key_pb` but trust the + # backend to uphold the values sent (e.g. dataset ID). if auto_id: - return result.insert_auto_id_key[0] + inserted_key_pb = result.insert_auto_id_key[0] + # Assumes the backend has set `id` without checking HasField('id'). + return True, inserted_key_pb.path_element[-1].id - return True + return False, None def delete_entities(self, dataset_id, key_pbs): """Delete keys from a dataset in the Cloud Datastore. diff --git a/gcloud/datastore/entity.py b/gcloud/datastore/entity.py index 31b41e1cb8eb..2e30821b9f65 100644 --- a/gcloud/datastore/entity.py +++ b/gcloud/datastore/entity.py @@ -15,7 +15,6 @@ """Class for representing a single entity in the Cloud Datastore.""" from gcloud.datastore import _implicit_environ -from gcloud.datastore import datastore_v1_pb2 as datastore_pb from gcloud.datastore.key import Key @@ -241,7 +240,7 @@ def save(self): key = self._must_key dataset = self._must_dataset connection = dataset.connection() - key_pb = connection.save_entity( + assigned, new_id = connection.save_entity( dataset_id=dataset.id(), key_pb=key.to_protobuf(), properties=dict(self), @@ -253,23 +252,9 @@ def save(self): 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) - # This is temporary. Will be addressed throughout #451. - clone = key._clone() - clone._path = path - self._key = clone + if assigned: + # Update the key (which may have been altered). + self.key(self.key().completed_key(new_id)) return self diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py index 857cc3132b3d..888a95392629 100644 --- a/gcloud/datastore/test_connection.py +++ b/gcloud/datastore/test_connection.py @@ -926,7 +926,7 @@ def test_save_entity_wo_transaction_w_upsert(self): ]) http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo'}) - self.assertEqual(result, True) + self.assertEqual(result, (False, None)) cw = http._called_with self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest @@ -967,7 +967,7 @@ def test_save_entity_w_exclude_from_indexes(self): result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo', 'bar': [u'bar1', u'bar2']}, exclude_from_indexes=['foo', 'bar']) - self.assertEqual(result, True) + self.assertEqual(result, (False, None)) cw = http._called_with self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest @@ -1018,7 +1018,7 @@ def test_save_entity_wo_transaction_w_auto_id(self): ]) http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo'}) - self.assertEqual(result, updated_key_pb) + self.assertEqual(result, (True, 1234)) cw = http._called_with self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest @@ -1054,7 +1054,7 @@ def mutation(self): conn.transaction(Xact()) http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo'}) - self.assertEqual(result, True) + self.assertEqual(result, (False, None)) self.assertEqual(http._called_with, None) mutation = conn.mutation() self.assertEqual(len(mutation.upsert), 1) @@ -1077,7 +1077,7 @@ def mutation(self): conn.transaction(Xact()) http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) result = conn.save_entity(DATASET_ID, key_pb, {'foo': nested}) - self.assertEqual(result, True) + self.assertEqual(result, (False, None)) self.assertEqual(http._called_with, None) mutation = conn.mutation() self.assertEqual(len(mutation.upsert), 1) diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py index 97971ca7ce4e..e4ae5575a587 100644 --- a/gcloud/datastore/test_entity.py +++ b/gcloud/datastore/test_entity.py @@ -194,7 +194,7 @@ def test_save_w_returned_key_exclude_from_indexes(self): key_pb.partition_id.dataset_id = _DATASET_ID key_pb.path_element.add(kind=_KIND, id=_ID) connection = _Connection() - connection._save_result = key_pb + connection._save_result = (True, _ID) dataset = _Dataset(connection) key = Key('KIND', dataset_id='DATASET') entity = self._makeOne(dataset, exclude_from_indexes=['foo']) @@ -287,17 +287,12 @@ 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.completed_key(i + 1) for i in range(num_ids)] class _Connection(object): _transaction = _saved = _deleted = None - _save_result = True + _save_result = (False, None) def transaction(self): return self._transaction