Get rid of unnecessary conversion between Entity and EntityPb in commit

function and make sure we handle partial keys correctly for
hierarchical / nested keys.

Also add test cases which cover hierarchical keys.

Author: Kami

datastore/google/cloud/datastore/batch.py

@@ -80,6 +80,8 @@ class Batch(object):
     def __init__(self, client):
         self._client = client
         self._mutations = []
+        # Stores a list of entity.Entity and .entity_pb2.Entity objects with
+        # partial keys
         self._partial_key_entities = []
         self._status = self._INITIAL
 
@@ -210,10 +212,10 @@ def put_entity_pb(self, entity_pb):
         if self.project != entity_pb.key.partition_id.project_id:
             raise ValueError("Key must be from same project as batch")
 
-        # Key is considered partial if neither "id" nor "name" values of the key
-        # "path" field are set
+        # Key is considered partial if neither "id" nor "name" values of the last
+        # key path element are set
         is_key_partial = not entity_pb.key.path or (
-            not entity_pb.key.path[0].name and not entity_pb.key.path[0].id
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
         )
 
         if is_key_partial:
@@ -280,17 +282,11 @@ def _commit(self):
         # order) directly ``_partial_key_entities``.
         for new_key_pb, entity in zip(updated_keys, self._partial_key_entities):
             if isinstance(entity, entity_pb2.Entity):
-                # Entity Protobuf instance
+                # entity_pb2.Entity Protobuf instance
                 new_id = new_key_pb.path[-1].id
-
-                # TODO: Make this efficient, don't rely on conversion back and forth
-                entity_ = helpers.entity_from_protobuf(entity)
-                entity_.key = entity_.key.completed_key(new_id)
-
-                entity.key.CopyFrom(entity_.key.to_protobuf())
-                pass
+                entity.key.path[-1].id = new_id
             else:
-                # Entity class
+                # entity.Entity class
                 new_id = new_key_pb.path[-1].id
                 entity.key = entity.key.completed_key(new_id)
 

datastore/tests/unit/test_batch.py

@@ -17,6 +17,7 @@
 import mock
 
 from google.cloud.datastore_v1.proto import entity_pb2
+from google.cloud.datastore.key import Key
 
 
 class TestBatch(unittest.TestCase):
@@ -463,6 +464,90 @@ def test_put_entity_w_completed_key(self):
         mutated_entity = _mutated_pb(self, batch.mutations, "upsert")
         self.assertEqual(mutated_entity.key, key._key)
 
+    def test_commit_w_completed_key_hierarchical_key_integer_id(self):
+        from google.cloud.datastore_v1.proto import datastore_pb2
+
+        project = "PROJECT"
+        new_id = 1234
+        ds_api = _make_datastore_api(new_id)
+        client = _Client(project, datastore_api=ds_api)
+        batch = self._make_one(client)
+
+        # Complete key where ID of last key path is set
+        key1 = Key("Parent", "foo", "Child", 1234, project=project)
+        self.assertFalse(key1.is_partial)
+
+        entity_pb = entity_pb2.Entity()
+        entity_pb.key.CopyFrom(key1.to_protobuf())
+
+        self.assertEqual(len(entity_pb.key.path), 2)
+        self.assertEqual(entity_pb.key.path[0].kind, "Parent")
+        self.assertEqual(entity_pb.key.path[0].name, "foo")
+        self.assertEqual(entity_pb.key.path[1].kind, "Child")
+        self.assertEqual(entity_pb.key.path[1].id, 1234)
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertFalse(is_key_partial)
+
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+        batch.commit()
+        self.assertEqual(batch._status, batch._FINISHED)
+
+        mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        ds_api.commit.assert_called_once_with(project, mode, [], transaction=None)
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertFalse(is_key_partial)
+        self.assertEqual(entity_pb.key.path[-1].id, 1234)
+
+    def test_commit_w_completed_key_hierarchical_key_string_name(self):
+        from google.cloud.datastore_v1.proto import datastore_pb2
+
+        project = "PROJECT"
+        new_id = 1234
+        ds_api = _make_datastore_api(new_id)
+        client = _Client(project, datastore_api=ds_api)
+        batch = self._make_one(client)
+
+        # Complete key where ID of last key path is set
+        key1 = Key("Parent", "foo", "Child", "name1", project=project)
+        self.assertFalse(key1.is_partial)
+
+        entity_pb = entity_pb2.Entity()
+        entity_pb.key.CopyFrom(key1.to_protobuf())
+
+        self.assertEqual(len(entity_pb.key.path), 2)
+        self.assertEqual(entity_pb.key.path[0].kind, "Parent")
+        self.assertEqual(entity_pb.key.path[0].name, "foo")
+        self.assertEqual(entity_pb.key.path[1].kind, "Child")
+        self.assertEqual(entity_pb.key.path[1].name, "name1")
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertFalse(is_key_partial)
+
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+        batch.commit()
+        self.assertEqual(batch._status, batch._FINISHED)
+
+        mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        ds_api.commit.assert_called_once_with(project, mode, [], transaction=None)
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertFalse(is_key_partial)
+        self.assertEqual(entity_pb.key.path[-1].name, "name1")
+
     def test_commit_w_partial_key_entities(self):
         from google.cloud.datastore_v1.proto import datastore_pb2
 
@@ -493,11 +578,66 @@ def test_commit_w_partial_key_entities(self):
         ds_api.commit.assert_called_once_with(project, mode, [], transaction=None)
 
         is_key_partial = not entity_pb.key.path or (
-            not entity_pb.key.path[0].name and not entity_pb.key.path[0].id
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
         )
         self.assertFalse(is_key_partial)
         self.assertEqual(entity_pb.key.path[0].id, new_id)
 
+    def test_commit_w_partial_key_entities_hierarchical_key(self):
+        from google.cloud.datastore_v1.proto import datastore_pb2
+
+        project = "PROJECT"
+        new_id = 1234
+        ds_api = _make_datastore_api(new_id)
+        client = _Client(project, datastore_api=ds_api)
+        batch = self._make_one(client)
+
+        key = Key("Parent", "foo", "Child", project=project)
+        self.assertTrue(key.is_partial)
+
+        key_completed_1 = Key("Parent", "foo", "Child", 1234, project=project)
+        self.assertFalse(key_completed_1.is_partial)
+
+        key_completed_2 = Key("Parent", "foo", "Child", "child1", project=project)
+        self.assertFalse(key_completed_2.is_partial)
+
+        key._id = None
+        entity_pb = entity_pb2.Entity()
+        entity_pb.key.CopyFrom(key.to_protobuf())
+        batch._partial_key_entities.append(entity_pb)
+
+        self.assertEqual(len(entity_pb.key.path), 2)
+        self.assertEqual(entity_pb.key.path[0].kind, "Parent")
+        self.assertEqual(entity_pb.key.path[0].name, "foo")
+        self.assertEqual(entity_pb.key.path[1].kind, "Child")
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertTrue(is_key_partial)
+
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+        batch.commit()
+        self.assertEqual(batch._status, batch._FINISHED)
+
+        mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        ds_api.commit.assert_called_once_with(project, mode, [], transaction=None)
+
+        is_key_partial = not entity_pb.key.path or (
+            not entity_pb.key.path[-1].name and not entity_pb.key.path[-1].id
+        )
+        self.assertFalse(is_key_partial)
+        self.assertEqual(entity_pb.key.path[-1].id, new_id)
+
+        self.assertEqual(len(entity_pb.key.path), 2)
+        self.assertEqual(entity_pb.key.path[0].kind, "Parent")
+        self.assertEqual(entity_pb.key.path[0].name, "foo")
+        self.assertEqual(entity_pb.key.path[1].kind, "Child")
+        self.assertEqual(entity_pb.key.path[-1].id, new_id)
+        self.assertEqual(entity_pb.key.path[1].id, new_id)
+
 
 class Test__parse_commit_response(unittest.TestCase):
     def _call_fut(self, commit_response_pb):