Skip to content

Commit 4d2b91d

Browse files
dhermesdhermes
authored
Merge pull request googleapis#3064 from dhermes/fix-2746-commit
Making datastore Connection.commit() return low-level protobuf.
2 parents 5557e19 + 9488e3c commit 4d2b91d

6 files changed

Lines changed: 133 additions & 138 deletions

File tree

‎datastore/google/cloud/datastore/_http.py‎

Lines changed: 3 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -426,19 +426,16 @@ def commit(self, project, request, transaction_id):
426426
427427
This method will mutate ``request`` before using it.
428428
429-
:rtype: tuple
430-
:returns: The pair of the number of index updates and a list of
431-
:class:`.entity_pb2.Key` for each incomplete key
432-
that was completed in the commit.
429+
:rtype: :class:`.datastore_pb2.CommitResponse`
430+
:returns: The protobuf response from a commit request.
433431
"""
434432
if transaction_id:
435433
request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
436434
request.transaction = transaction_id
437435
else:
438436
request.mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
439437

440-
response = self._datastore_api.commit(project, request)
441-
return _parse_commit_response(response)
438+
return self._datastore_api.commit(project, request)
442439

443440
def rollback(self, project, transaction_id):
444441
"""Rollback the connection's existing transaction.
@@ -508,21 +505,3 @@ def _add_keys_to_request(request_field_pb, key_pbs):
508505
"""
509506
for key_pb in key_pbs:
510507
request_field_pb.add().CopyFrom(key_pb)
511-
512-
513-
def _parse_commit_response(commit_response_pb):
514-
"""Extract response data from a commit response.
515-
516-
:type commit_response_pb: :class:`.datastore_pb2.CommitResponse`
517-
:param commit_response_pb: The protobuf response from a commit request.
518-
519-
:rtype: tuple
520-
:returns: The pair of the number of index updates and a list of
521-
:class:`.entity_pb2.Key` for each incomplete key
522-
that was completed in the commit.
523-
"""
524-
mut_results = commit_response_pb.mutation_results
525-
index_updates = commit_response_pb.index_updates
526-
completed_keys = [mut_result.key for mut_result in mut_results
527-
if mut_result.HasField('key')] # Message field (Key)
528-
return index_updates, completed_keys

‎datastore/google/cloud/datastore/batch.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,8 +238,9 @@ def _commit(self):
238238
This is called by :meth:`commit`.
239239
"""
240240
# NOTE: ``self._commit_request`` will be modified.
241-
_, updated_keys = self._client._connection.commit(
241+
commit_response_pb = self._client._connection.commit(
242242
self.project, self._commit_request, self._id)
243+
_, updated_keys = _parse_commit_response(commit_response_pb)
243244
# If the back-end returns without error, we are guaranteed that
244245
# :meth:`Connection.commit` will return keys that match (length and
245246
# order) directly ``_partial_key_entities``.
@@ -311,3 +312,21 @@ def _assign_entity_to_pb(entity_pb, entity):
311312
bare_entity_pb = helpers.entity_to_protobuf(entity)
312313
bare_entity_pb.key.CopyFrom(bare_entity_pb.key)
313314
entity_pb.CopyFrom(bare_entity_pb)
315+
316+
317+
def _parse_commit_response(commit_response_pb):
318+
"""Extract response data from a commit response.
319+
320+
:type commit_response_pb: :class:`.datastore_pb2.CommitResponse`
321+
:param commit_response_pb: The protobuf response from a commit request.
322+
323+
:rtype: tuple
324+
:returns: The pair of the number of index updates and a list of
325+
:class:`.entity_pb2.Key` for each incomplete key
326+
that was completed in the commit.
327+
"""
328+
mut_results = commit_response_pb.mutation_results
329+
index_updates = commit_response_pb.index_updates
330+
completed_keys = [mut_result.key for mut_result in mut_results
331+
if mut_result.HasField('key')] # Message field (Key)
332+
return index_updates, completed_keys

‎datastore/unit_tests/test__http.py‎

Lines changed: 16 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -696,8 +696,8 @@ def test_commit_wo_transaction(self):
696696
from google.cloud.grpc.datastore.v1 import datastore_pb2
697697
from google.cloud.datastore.helpers import _new_value_pb
698698

699-
PROJECT = 'PROJECT'
700-
key_pb = self._make_key_pb(PROJECT)
699+
project = 'PROJECT'
700+
key_pb = self._make_key_pb(project)
701701
rsp_pb = datastore_pb2.CommitResponse()
702702
req_pb = datastore_pb2.CommitRequest()
703703
mutation = req_pb.mutations.add()
@@ -708,44 +708,32 @@ def test_commit_wo_transaction(self):
708708
http = Http({'status': '200'}, rsp_pb.SerializeToString())
709709
client = mock.Mock(_http=http, spec=['_http'])
710710
conn = self._make_one(client)
711-
URI = '/'.join([
711+
uri = '/'.join([
712712
conn.api_base_url,
713713
conn.API_VERSION,
714714
'projects',
715-
PROJECT + ':commit',
715+
project + ':commit',
716716
])
717717

718-
# Set up mock for parsing the response.
719-
expected_result = object()
720-
_parsed = []
721-
722-
def mock_parse(response):
723-
_parsed.append(response)
724-
return expected_result
725-
726-
patch = mock.patch(
727-
'google.cloud.datastore._http._parse_commit_response',
728-
new=mock_parse)
729-
with patch:
730-
result = conn.commit(PROJECT, req_pb, None)
718+
result = conn.commit(project, req_pb, None)
719+
self.assertEqual(result, rsp_pb)
731720

732-
self.assertIs(result, expected_result)
721+
# Verify the caller.
733722
cw = http._called_with
734-
self._verify_protobuf_call(cw, URI, conn)
723+
self._verify_protobuf_call(cw, uri, conn)
735724
rq_class = datastore_pb2.CommitRequest
736725
request = rq_class()
737726
request.ParseFromString(cw['body'])
738727
self.assertEqual(request.transaction, b'')
739728
self.assertEqual(list(request.mutations), [mutation])
740729
self.assertEqual(request.mode, rq_class.NON_TRANSACTIONAL)
741-
self.assertEqual(_parsed, [rsp_pb])
742730

743731
def test_commit_w_transaction(self):
744732
from google.cloud.grpc.datastore.v1 import datastore_pb2
745733
from google.cloud.datastore.helpers import _new_value_pb
746734

747-
PROJECT = 'PROJECT'
748-
key_pb = self._make_key_pb(PROJECT)
735+
project = 'PROJECT'
736+
key_pb = self._make_key_pb(project)
749737
rsp_pb = datastore_pb2.CommitResponse()
750738
req_pb = datastore_pb2.CommitRequest()
751739
mutation = req_pb.mutations.add()
@@ -756,37 +744,25 @@ def test_commit_w_transaction(self):
756744
http = Http({'status': '200'}, rsp_pb.SerializeToString())
757745
client = mock.Mock(_http=http, spec=['_http'])
758746
conn = self._make_one(client)
759-
URI = '/'.join([
747+
uri = '/'.join([
760748
conn.api_base_url,
761749
conn.API_VERSION,
762750
'projects',
763-
PROJECT + ':commit',
751+
project + ':commit',
764752
])
765753

766-
# Set up mock for parsing the response.
767-
expected_result = object()
768-
_parsed = []
754+
result = conn.commit(project, req_pb, b'xact')
755+
self.assertEqual(result, rsp_pb)
769756

770-
def mock_parse(response):
771-
_parsed.append(response)
772-
return expected_result
773-
774-
patch = mock.patch(
775-
'google.cloud.datastore._http._parse_commit_response',
776-
new=mock_parse)
777-
with patch:
778-
result = conn.commit(PROJECT, req_pb, b'xact')
779-
780-
self.assertIs(result, expected_result)
757+
# Verify the caller.
781758
cw = http._called_with
782-
self._verify_protobuf_call(cw, URI, conn)
759+
self._verify_protobuf_call(cw, uri, conn)
783760
rq_class = datastore_pb2.CommitRequest
784761
request = rq_class()
785762
request.ParseFromString(cw['body'])
786763
self.assertEqual(request.transaction, b'xact')
787764
self.assertEqual(list(request.mutations), [mutation])
788765
self.assertEqual(request.mode, rq_class.TRANSACTIONAL)
789-
self.assertEqual(_parsed, [rsp_pb])
790766

791767
def test_rollback_ok(self):
792768
from google.cloud.grpc.datastore.v1 import datastore_pb2
@@ -870,46 +846,6 @@ def test_allocate_ids_non_empty(self):
870846
self.assertEqual(key_before, key_after)
871847

872848

873-
class Test__parse_commit_response(unittest.TestCase):
874-
875-
def _call_fut(self, commit_response_pb):
876-
from google.cloud.datastore._http import _parse_commit_response
877-
878-
return _parse_commit_response(commit_response_pb)
879-
880-
def test_it(self):
881-
from google.cloud.grpc.datastore.v1 import datastore_pb2
882-
from google.cloud.grpc.datastore.v1 import entity_pb2
883-
884-
index_updates = 1337
885-
keys = [
886-
entity_pb2.Key(
887-
path=[
888-
entity_pb2.Key.PathElement(
889-
kind='Foo',
890-
id=1234,
891-
),
892-
],
893-
),
894-
entity_pb2.Key(
895-
path=[
896-
entity_pb2.Key.PathElement(
897-
kind='Bar',
898-
name='baz',
899-
),
900-
],
901-
),
902-
]
903-
response = datastore_pb2.CommitResponse(
904-
mutation_results=[
905-
datastore_pb2.MutationResult(key=key) for key in keys
906-
],
907-
index_updates=index_updates,
908-
)
909-
result = self._call_fut(response)
910-
self.assertEqual(result, (index_updates, keys))
911-
912-
913849
class Http(object):
914850

915851
_called_with = None

‎datastore/unit_tests/test_batch.py‎

Lines changed: 61 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -249,13 +249,13 @@ def test_commit_wrong_status(self):
249249
self.assertRaises(ValueError, batch.commit)
250250

251251
def test_commit_w_partial_key_entities(self):
252-
_PROJECT = 'PROJECT'
253-
_NEW_ID = 1234
254-
connection = _Connection(_NEW_ID)
255-
client = _Client(_PROJECT, connection)
252+
project = 'PROJECT'
253+
new_id = 1234
254+
connection = _Connection(new_id)
255+
client = _Client(project, connection)
256256
batch = self._make_one(client)
257257
entity = _Entity({})
258-
key = entity.key = _Key(_PROJECT)
258+
key = entity.key = _Key(project)
259259
key._id = None
260260
batch._partial_key_entities.append(entity)
261261

@@ -266,9 +266,9 @@ def test_commit_w_partial_key_entities(self):
266266
self.assertEqual(batch._status, batch._FINISHED)
267267

268268
self.assertEqual(connection._committed,
269-
[(_PROJECT, batch._commit_request, None)])
269+
[(project, batch._commit_request, None)])
270270
self.assertFalse(entity.key.is_partial)
271-
self.assertEqual(entity.key._id, _NEW_ID)
271+
self.assertEqual(entity.key._id, new_id)
272272

273273
def test_as_context_mgr_wo_error(self):
274274
_PROJECT = 'PROJECT'
@@ -369,30 +369,62 @@ def begin(self):
369369
self.assertEqual(client._batches, [])
370370

371371

372-
class _PathElementPB(object):
372+
class Test__parse_commit_response(unittest.TestCase):
373373

374-
def __init__(self, id_):
375-
self.id = id_
374+
def _call_fut(self, commit_response_pb):
375+
from google.cloud.datastore.batch import _parse_commit_response
376376

377+
return _parse_commit_response(commit_response_pb)
377378

378-
class _KeyPB(object):
379+
def test_it(self):
380+
from google.cloud.grpc.datastore.v1 import datastore_pb2
381+
from google.cloud.grpc.datastore.v1 import entity_pb2
379382

380-
def __init__(self, id_):
381-
self.path = [_PathElementPB(id_)]
383+
index_updates = 1337
384+
keys = [
385+
entity_pb2.Key(
386+
path=[
387+
entity_pb2.Key.PathElement(
388+
kind='Foo',
389+
id=1234,
390+
),
391+
],
392+
),
393+
entity_pb2.Key(
394+
path=[
395+
entity_pb2.Key.PathElement(
396+
kind='Bar',
397+
name='baz',
398+
),
399+
],
400+
),
401+
]
402+
response = datastore_pb2.CommitResponse(
403+
mutation_results=[
404+
datastore_pb2.MutationResult(key=key) for key in keys
405+
],
406+
index_updates=index_updates,
407+
)
408+
result = self._call_fut(response)
409+
self.assertEqual(result, (index_updates, keys))
382410

383411

384412
class _Connection(object):
385413
_marker = object()
386414
_save_result = (False, None)
387415

388-
def __init__(self, *new_keys):
389-
self._completed_keys = [_KeyPB(key) for key in new_keys]
416+
def __init__(self, *new_key_ids):
417+
from google.cloud.grpc.datastore.v1 import datastore_pb2
418+
390419
self._committed = []
391-
self._index_updates = 0
420+
mutation_results = [
421+
_make_mutation(key_id) for key_id in new_key_ids]
422+
self._commit_response_pb = datastore_pb2.CommitResponse(
423+
mutation_results=mutation_results)
392424

393425
def commit(self, project, commit_request, transaction_id):
394426
self._committed.append((project, commit_request, transaction_id))
395-
return self._index_updates, self._completed_keys
427+
return self._commit_response_pb
396428

397429

398430
class _Entity(dict):
@@ -472,3 +504,15 @@ def _mutated_pb(test_case, mutation_pb_list, mutation_type):
472504
mutation_type)
473505

474506
return getattr(mutated_pb, mutation_type)
507+
508+
509+
def _make_mutation(id_):
510+
from google.cloud.grpc.datastore.v1 import datastore_pb2
511+
from google.cloud.grpc.datastore.v1 import entity_pb2
512+
513+
key = entity_pb2.Key()
514+
key.partition_id.project_id = 'PROJECT'
515+
elem = key.path.add()
516+
elem.kind = 'Kind'
517+
elem.id = id_
518+
return datastore_pb2.MutationResult(key=key)

‎datastore/unit_tests/test_client.py‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -568,7 +568,8 @@ def test_put_multi_no_batch_w_partial_key(self):
568568

569569
creds = _make_credentials()
570570
client = self._make_one(credentials=creds)
571-
client._connection._commit.append([_KeyPB(key)])
571+
key_pb = _make_key(234)
572+
client._connection._commit.append([key_pb])
572573

573574
result = client.put_multi([entity])
574575
self.assertIsNone(result)
@@ -931,7 +932,6 @@ def __init__(self, credentials=None, http=None):
931932
self._commit = []
932933
self._alloc_cw = []
933934
self._alloc = []
934-
self._index_updates = 0
935935

936936
def _add_lookup_result(self, results=(), missing=(), deferred=()):
937937
self._lookup.append((list(results), list(missing), list(deferred)))
@@ -943,9 +943,13 @@ def lookup(self, project, key_pbs, eventual=False, transaction_id=None):
943943
return results, missing, deferred
944944

945945
def commit(self, project, commit_request, transaction_id):
946+
from google.cloud.grpc.datastore.v1 import datastore_pb2
947+
946948
self._commit_cw.append((project, commit_request, transaction_id))
947-
response, self._commit = self._commit[0], self._commit[1:]
948-
return self._index_updates, response
949+
keys, self._commit = self._commit[0], self._commit[1:]
950+
mutation_results = [
951+
datastore_pb2.MutationResult(key=key) for key in keys]
952+
return datastore_pb2.CommitResponse(mutation_results=mutation_results)
949953

950954
def allocate_ids(self, project, key_pbs):
951955
self._alloc_cw.append((project, key_pbs))
@@ -1058,3 +1062,12 @@ def _mutated_pb(test_case, mutation_pb_list, mutation_type):
10581062
mutation_type)
10591063

10601064
return getattr(mutated_pb, mutation_type)
1065+
1066+
1067+
def _make_key(id_):
1068+
from google.cloud.grpc.datastore.v1 import entity_pb2
1069+
1070+
key = entity_pb2.Key()
1071+
elem = key.path.add()
1072+
elem.id = id_
1073+
return key

0 commit comments

Comments
 (0)