Skip to content

Commit 56f12dd

Browse files
dhermesdhermes
committed
Making datastore Connection.commit() return low-level protobuf.
Towards #2746. This approach is to slowly transition from our current approach to use the GAPIC generated surface. It is unfortunately tangled quite a bit (partly because we may have too much mocked in the tests).
1 parent d4d0abc commit 56f12dd

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
@@ -434,19 +434,16 @@ def commit(self, project, request, transaction_id):
434434
435435
This method will mutate ``request`` before using it.
436436
437-
:rtype: tuple
438-
:returns: The pair of the number of index updates and a list of
439-
:class:`.entity_pb2.Key` for each incomplete key
440-
that was completed in the commit.
437+
:rtype: :class:`.datastore_pb2.CommitResponse`
438+
:returns: The protobuf response from a commit request.
441439
"""
442440
if transaction_id:
443441
request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
444442
request.transaction = transaction_id
445443
else:
446444
request.mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
447445

448-
response = self._datastore_api.commit(project, request)
449-
return _parse_commit_response(response)
446+
return self._datastore_api.commit(project, request)
450447

451448
def rollback(self, project, transaction_id):
452449
"""Rollback the connection's existing transaction.
@@ -516,21 +513,3 @@ def _add_keys_to_request(request_field_pb, key_pbs):
516513
"""
517514
for key_pb in key_pbs:
518515
request_field_pb.add().CopyFrom(key_pb)
519-
520-
521-
def _parse_commit_response(commit_response_pb):
522-
"""Extract response data from a commit response.
523-
524-
:type commit_response_pb: :class:`.datastore_pb2.CommitResponse`
525-
:param commit_response_pb: The protobuf response from a commit request.
526-
527-
:rtype: tuple
528-
:returns: The pair of the number of index updates and a list of
529-
:class:`.entity_pb2.Key` for each incomplete key
530-
that was completed in the commit.
531-
"""
532-
mut_results = commit_response_pb.mutation_results
533-
index_updates = commit_response_pb.index_updates
534-
completed_keys = [mut_result.key for mut_result in mut_results
535-
if mut_result.HasField('key')] # Message field (Key)
536-
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
@@ -684,8 +684,8 @@ def test_commit_wo_transaction(self):
684684
from google.cloud.grpc.datastore.v1 import datastore_pb2
685685
from google.cloud.datastore.helpers import _new_value_pb
686686

687-
PROJECT = 'PROJECT'
688-
key_pb = self._make_key_pb(PROJECT)
687+
project = 'PROJECT'
688+
key_pb = self._make_key_pb(project)
689689
rsp_pb = datastore_pb2.CommitResponse()
690690
req_pb = datastore_pb2.CommitRequest()
691691
mutation = req_pb.mutations.add()
@@ -696,44 +696,32 @@ def test_commit_wo_transaction(self):
696696
http = Http({'status': '200'}, rsp_pb.SerializeToString())
697697
client = mock.Mock(_http=http, spec=['_http'])
698698
conn = self._make_one(client)
699-
URI = '/'.join([
699+
uri = '/'.join([
700700
conn.api_base_url,
701701
conn.API_VERSION,
702702
'projects',
703-
PROJECT + ':commit',
703+
project + ':commit',
704704
])
705705

706-
# Set up mock for parsing the response.
707-
expected_result = object()
708-
_parsed = []
709-
710-
def mock_parse(response):
711-
_parsed.append(response)
712-
return expected_result
713-
714-
patch = mock.patch(
715-
'google.cloud.datastore._http._parse_commit_response',
716-
new=mock_parse)
717-
with patch:
718-
result = conn.commit(PROJECT, req_pb, None)
706+
result = conn.commit(project, req_pb, None)
707+
self.assertEqual(result, rsp_pb)
719708

720-
self.assertIs(result, expected_result)
709+
# Verify the caller.
721710
cw = http._called_with
722-
self._verifyProtobufCall(cw, URI, conn)
711+
self._verifyProtobufCall(cw, uri, conn)
723712
rq_class = datastore_pb2.CommitRequest
724713
request = rq_class()
725714
request.ParseFromString(cw['body'])
726715
self.assertEqual(request.transaction, b'')
727716
self.assertEqual(list(request.mutations), [mutation])
728717
self.assertEqual(request.mode, rq_class.NON_TRANSACTIONAL)
729-
self.assertEqual(_parsed, [rsp_pb])
730718

731719
def test_commit_w_transaction(self):
732720
from google.cloud.grpc.datastore.v1 import datastore_pb2
733721
from google.cloud.datastore.helpers import _new_value_pb
734722

735-
PROJECT = 'PROJECT'
736-
key_pb = self._make_key_pb(PROJECT)
723+
project = 'PROJECT'
724+
key_pb = self._make_key_pb(project)
737725
rsp_pb = datastore_pb2.CommitResponse()
738726
req_pb = datastore_pb2.CommitRequest()
739727
mutation = req_pb.mutations.add()
@@ -744,37 +732,25 @@ def test_commit_w_transaction(self):
744732
http = Http({'status': '200'}, rsp_pb.SerializeToString())
745733
client = mock.Mock(_http=http, spec=['_http'])
746734
conn = self._make_one(client)
747-
URI = '/'.join([
735+
uri = '/'.join([
748736
conn.api_base_url,
749737
conn.API_VERSION,
750738
'projects',
751-
PROJECT + ':commit',
739+
project + ':commit',
752740
])
753741

754-
# Set up mock for parsing the response.
755-
expected_result = object()
756-
_parsed = []
757-
758-
def mock_parse(response):
759-
_parsed.append(response)
760-
return expected_result
742+
result = conn.commit(project, req_pb, b'xact')
743+
self.assertEqual(result, rsp_pb)
761744

762-
patch = mock.patch(
763-
'google.cloud.datastore._http._parse_commit_response',
764-
new=mock_parse)
765-
with patch:
766-
result = conn.commit(PROJECT, req_pb, b'xact')
767-
768-
self.assertIs(result, expected_result)
745+
# Verify the caller.
769746
cw = http._called_with
770-
self._verifyProtobufCall(cw, URI, conn)
747+
self._verifyProtobufCall(cw, uri, conn)
771748
rq_class = datastore_pb2.CommitRequest
772749
request = rq_class()
773750
request.ParseFromString(cw['body'])
774751
self.assertEqual(request.transaction, b'xact')
775752
self.assertEqual(list(request.mutations), [mutation])
776753
self.assertEqual(request.mode, rq_class.TRANSACTIONAL)
777-
self.assertEqual(_parsed, [rsp_pb])
778754

779755
def test_rollback_ok(self):
780756
from google.cloud.grpc.datastore.v1 import datastore_pb2
@@ -858,46 +834,6 @@ def test_allocate_ids_non_empty(self):
858834
self.assertEqual(key_before, key_after)
859835

860836

861-
class Test__parse_commit_response(unittest.TestCase):
862-
863-
def _call_fut(self, commit_response_pb):
864-
from google.cloud.datastore._http import _parse_commit_response
865-
866-
return _parse_commit_response(commit_response_pb)
867-
868-
def test_it(self):
869-
from google.cloud.grpc.datastore.v1 import datastore_pb2
870-
from google.cloud.grpc.datastore.v1 import entity_pb2
871-
872-
index_updates = 1337
873-
keys = [
874-
entity_pb2.Key(
875-
path=[
876-
entity_pb2.Key.PathElement(
877-
kind='Foo',
878-
id=1234,
879-
),
880-
],
881-
),
882-
entity_pb2.Key(
883-
path=[
884-
entity_pb2.Key.PathElement(
885-
kind='Bar',
886-
name='baz',
887-
),
888-
],
889-
),
890-
]
891-
response = datastore_pb2.CommitResponse(
892-
mutation_results=[
893-
datastore_pb2.MutationResult(key=key) for key in keys
894-
],
895-
index_updates=index_updates,
896-
)
897-
result = self._call_fut(response)
898-
self.assertEqual(result, (index_updates, keys))
899-
900-
901837
class Http(object):
902838

903839
_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)