Merge pull request #462 from dhermes/fix-451-part6

tseaver · tseaver · commit e86ba09dfe06 · 2014-12-31T17:02:31.000-05:00
Address sixth part of 451: Implement Key.compare_to_proto to check pb keys against existing.
diff --git 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
@@ -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
@@ -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
@@ -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
