Reducing datastore commit reliance on structure of response.

gcloud/datastore/batch.py

@@ -185,13 +185,13 @@ def commit(self):
         however it can be called explicitly if you don't want to use a
         context manager.
         """
-        response = self.connection.commit(
+        _, updated_keys = self.connection.commit(
             self.dataset_id, self.mutation, self._id)
         # If the back-end returns without error, we are guaranteed that
         # the response's 'insert_auto_id_key' will match (length and order)
         # the request's 'insert_auto_id` entities, which are derived from
         # our '_partial_key_entities' (no partial success).
-        for new_key_pb, entity in zip(response.insert_auto_id_key,
+        for new_key_pb, entity in zip(updated_keys,
                                       self._partial_key_entities):
             new_id = new_key_pb.path_element[-1].id
             entity.key = entity.key.completed_key(new_id)

gcloud/datastore/connection.py

@@ -297,16 +297,18 @@ def commit(self, dataset_id, mutation_pb, transaction_id):
         :type dataset_id: string
         :param dataset_id: The ID dataset to which the transaction applies.
 
-        :type mutation_pb: :class:`._datastore_pb2.Mutation`.
+        :type mutation_pb: :class:`._datastore_pb2.Mutation`
         :param mutation_pb: The protobuf for the mutations being saved.
 
         :type transaction_id: string or None
         :param transaction_id: The transaction ID returned from
                                :meth:`begin_transaction`.  Non-transactional
                                batches must pass ``None``.
 
-        :rtype: :class:`gcloud.datastore._datastore_pb2.MutationResult`.
-        :returns': the result protobuf for the mutation.
+        :rtype: tuple
+        :returns': The pair of the number of index updates and a list of
+                   :class:`._entity_pb2.Key` for each incomplete key that was
+                   completed in the commit.
         """
         request = _datastore_pb2.CommitRequest()
 
@@ -319,7 +321,7 @@ def commit(self, dataset_id, mutation_pb, transaction_id):
         request.mutation.CopyFrom(mutation_pb)
         response = self._rpc(dataset_id, 'commit', request,
                              _datastore_pb2.CommitResponse)
-        return response.mutation_result
+        return _parse_commit_response(response)
 
     def rollback(self, dataset_id, transaction_id):
         """Rollback the connection's existing transaction.
@@ -415,3 +417,20 @@ def _add_keys_to_request(request_field_pb, key_pbs):
     for key_pb in key_pbs:
         key_pb = _prepare_key_for_request(key_pb)
         request_field_pb.add().CopyFrom(key_pb)
+
+
+def _parse_commit_response(commit_response_pb):
+    """Extract response data from a commit response.
+
+    :type commit_response_pb: :class:`._datastore_pb2.CommitResponse`
+    :param commit_response_pb: The protobuf response from a commit request.
+
+    :rtype: tuple
+    :returns': The pair of the number of index updates and a list of
+               :class:`._entity_pb2.Key` for each incomplete key that was
+               completed in the commit.
+    """
+    mut_result = commit_response_pb.mutation_result
+    index_updates = mut_result.index_updates
+    completed_keys = list(mut_result.insert_auto_id_key)
+    return index_updates, completed_keys

gcloud/datastore/test_batch.py

@@ -351,12 +351,6 @@ def test_as_context_mgr_w_error(self):
         self.assertEqual(connection._committed, [])
 
 
-class _CommitResult(object):
-
-    def __init__(self, *new_keys):
-        self.insert_auto_id_key = [_KeyPB(key) for key in new_keys]
-
-
 class _PathElementPB(object):
 
     def __init__(self, id):
@@ -374,12 +368,13 @@ class _Connection(object):
     _save_result = (False, None)
 
     def __init__(self, *new_keys):
-        self._commit_result = _CommitResult(*new_keys)
+        self._completed_keys = [_KeyPB(key) for key in new_keys]
         self._committed = []
+        self._index_updates = 0
 
     def commit(self, dataset_id, mutation, transaction_id):
         self._committed.append((dataset_id, mutation, transaction_id))
-        return self._commit_result
+        return self._index_updates, self._completed_keys
 
 
 class _Entity(dict):

gcloud/datastore/test_client.py

@@ -608,17 +608,17 @@ def test_put_multi_w_single_empty_entity(self):
         self.assertRaises(ValueError, client.put_multi, Entity())
 
     def test_put_multi_no_batch_w_partial_key(self):
-        from gcloud.datastore.test_batch import _CommitResult
         from gcloud.datastore.test_batch import _Entity
         from gcloud.datastore.test_batch import _Key
+        from gcloud.datastore.test_batch import _KeyPB
 
         entity = _Entity(foo=u'bar')
         key = entity.key = _Key(self.DATASET_ID)
         key._id = None
 
         creds = object()
         client = self._makeOne(credentials=creds)
-        client.connection._commit.append(_CommitResult(key))
+        client.connection._commit.append([_KeyPB(key)])
 
         result = client.put_multi([entity])
         self.assertTrue(result is None)
@@ -680,14 +680,13 @@ def test_delete_multi_no_keys(self):
         self.assertEqual(len(client.connection._commit_cw), 0)
 
     def test_delete_multi_no_batch(self):
-        from gcloud.datastore.test_batch import _CommitResult
         from gcloud.datastore.test_batch import _Key
 
         key = _Key(self.DATASET_ID)
 
         creds = object()
         client = self._makeOne(credentials=creds)
-        client.connection._commit.append(_CommitResult())
+        client.connection._commit.append([])
 
         result = client.delete_multi([key])
         self.assertEqual(result, None)
@@ -1000,6 +999,7 @@ def __init__(self, credentials=None, http=None):
         self._commit = []
         self._alloc_cw = []
         self._alloc = []
+        self._index_updates = 0
 
     def _add_lookup_result(self, results=(), missing=(), deferred=()):
         self._lookup.append((list(results), list(missing), list(deferred)))
@@ -1013,7 +1013,7 @@ def lookup(self, dataset_id, key_pbs, eventual=False, transaction_id=None):
     def commit(self, dataset_id, mutation, transaction_id):
         self._commit_cw.append((dataset_id, mutation, transaction_id))
         response, self._commit = self._commit[0], self._commit[1:]
-        return response
+        return self._index_updates, response
 
     def allocate_ids(self, dataset_id, key_pbs):
         from gcloud.datastore.test_connection import _KeyProto

gcloud/datastore/test_connection.py

@@ -667,7 +667,9 @@ def test_begin_transaction(self):
         self.assertEqual(request.isolation_level, rq_class.SERIALIZABLE)
 
     def test_commit_wo_transaction(self):
+        from gcloud._testing import _Monkey
         from gcloud.datastore import _datastore_pb2
+        from gcloud.datastore import connection as MUT
 
         DATASET_ID = 'DATASET'
         key_pb = self._make_key_pb(DATASET_ID)
@@ -688,9 +690,19 @@ def test_commit_wo_transaction(self):
             'commit',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        result = conn.commit(DATASET_ID, mutation, None)
-        self.assertEqual(result.index_updates, 0)
-        self.assertEqual(list(result.insert_auto_id_key), [])
+
+        # Set up mock for parsing the response.
+        expected_result = object()
+        _parsed = []
+
+        def mock_parse(response):
+            _parsed.append(response)
+            return expected_result
+
+        with _Monkey(MUT, _parse_commit_response=mock_parse):
+            result = conn.commit(DATASET_ID, mutation, None)
+
+        self.assertTrue(result is expected_result)
         cw = http._called_with
         self._verifyProtobufCall(cw, URI, conn)
         rq_class = _datastore_pb2.CommitRequest
@@ -699,9 +711,12 @@ def test_commit_wo_transaction(self):
         self.assertEqual(request.transaction, b'')
         self.assertEqual(request.mutation, mutation)
         self.assertEqual(request.mode, rq_class.NON_TRANSACTIONAL)
+        self.assertEqual(_parsed, [rsp_pb])
 
     def test_commit_w_transaction(self):
+        from gcloud._testing import _Monkey
         from gcloud.datastore import _datastore_pb2
+        from gcloud.datastore import connection as MUT
 
         DATASET_ID = 'DATASET'
         key_pb = self._make_key_pb(DATASET_ID)
@@ -722,9 +737,19 @@ def test_commit_w_transaction(self):
             'commit',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        result = conn.commit(DATASET_ID, mutation, b'xact')
-        self.assertEqual(result.index_updates, 0)
-        self.assertEqual(list(result.insert_auto_id_key), [])
+
+        # Set up mock for parsing the response.
+        expected_result = object()
+        _parsed = []
+
+        def mock_parse(response):
+            _parsed.append(response)
+            return expected_result
+
+        with _Monkey(MUT, _parse_commit_response=mock_parse):
+            result = conn.commit(DATASET_ID, mutation, b'xact')
+
+        self.assertTrue(result is expected_result)
         cw = http._called_with
         self._verifyProtobufCall(cw, URI, conn)
         rq_class = _datastore_pb2.CommitRequest
@@ -733,6 +758,7 @@ def test_commit_w_transaction(self):
         self.assertEqual(request.transaction, b'xact')
         self.assertEqual(request.mutation, mutation)
         self.assertEqual(request.mode, rq_class.TRANSACTIONAL)
+        self.assertEqual(_parsed, [rsp_pb])
 
     def test_rollback_ok(self):
         from gcloud.datastore import _datastore_pb2
@@ -818,6 +844,45 @@ def test_allocate_ids_non_empty(self):
             _compare_key_pb_after_request(self, key_before, key_after)
 
 
+class Test__parse_commit_response(unittest2.TestCase):
+
+    def _callFUT(self, commit_response_pb):
+        from gcloud.datastore.connection import _parse_commit_response
+        return _parse_commit_response(commit_response_pb)
+
+    def test_it(self):
+        from gcloud.datastore import _datastore_pb2
+        from gcloud.datastore import _entity_pb2
+
+        index_updates = 1337
+        keys = [
+            _entity_pb2.Key(
+                path_element=[
+                    _entity_pb2.Key.PathElement(
+                        kind='Foo',
+                        id=1234,
+                    ),
+                ],
+            ),
+            _entity_pb2.Key(
+                path_element=[
+                    _entity_pb2.Key.PathElement(
+                        kind='Bar',
+                        name='baz',
+                    ),
+                ],
+            ),
+        ]
+        response = _datastore_pb2.CommitResponse(
+            mutation_result=_datastore_pb2.MutationResult(
+                index_updates=index_updates,
+                insert_auto_id_key=keys,
+            ),
+        )
+        result = self._callFUT(response)
+        self.assertEqual(result, (index_updates, keys))
+
+
 class Http(object):
 
     _called_with = None

gcloud/datastore/test_transaction.py

@@ -113,8 +113,7 @@ def test_commit_w_partial_keys(self):
         _KIND = 'KIND'
         _ID = 123
         connection = _Connection(234)
-        connection._commit_result = _CommitResult(
-            _make_key(_KIND, _ID, _DATASET))
+        connection._completed_keys = [_make_key(_KIND, _ID, _DATASET)]
         client = _Client(_DATASET, connection)
         xact = self._makeOne(client)
         entity = _Entity()
@@ -177,7 +176,8 @@ class _Connection(object):
 
     def __init__(self, xact_id=123):
         self._xact_id = xact_id
-        self._commit_result = _CommitResult()
+        self._completed_keys = []
+        self._index_updates = 0
 
     def begin_transaction(self, dataset_id):
         self._begun = dataset_id
@@ -188,13 +188,7 @@ def rollback(self, dataset_id, transaction_id):
 
     def commit(self, dataset_id, mutation, transaction_id):
         self._committed = (dataset_id, mutation, transaction_id)
-        return self._commit_result
-
-
-class _CommitResult(object):
-
-    def __init__(self, *new_keys):
-        self.insert_auto_id_key = new_keys
+        return self._index_updates, self._completed_keys
 
 
 class _Entity(dict):