diff --git a/datastore/google/cloud/datastore/_http.py b/datastore/google/cloud/datastore/_http.py index 2650321068bf..08b8ac3d4964 100644 --- a/datastore/google/cloud/datastore/_http.py +++ b/datastore/google/cloud/datastore/_http.py @@ -326,23 +326,14 @@ def lookup(self, project, key_pbs, the given transaction. Incompatible with ``eventual==True``. - :rtype: tuple - :returns: A triple of (``results``, ``missing``, ``deferred``) where - both ``results`` and ``missing`` are lists of - :class:`.entity_pb2.Entity` - and ``deferred`` is a list of - :class:`.entity_pb2.Key`. + :rtype: :class:`.datastore_pb2.LookupResponse` + :returns: The returned protobuf for the lookup request. """ lookup_request = _datastore_pb2.LookupRequest() _set_read_options(lookup_request, eventual, transaction_id) _add_keys_to_request(lookup_request.keys, key_pbs) - lookup_response = self._datastore_api.lookup(project, lookup_request) - - results = [result.entity for result in lookup_response.found] - missing = [result.entity for result in lookup_response.missing] - - return results, missing, list(lookup_response.deferred) + return self._datastore_api.lookup(project, lookup_request) def run_query(self, project, query_pb, namespace=None, eventual=False, transaction_id=None): diff --git a/datastore/google/cloud/datastore/client.py b/datastore/google/cloud/datastore/client.py index 87ab8f6ee0c6..beb7c49dcb5f 100644 --- a/datastore/google/cloud/datastore/client.py +++ b/datastore/google/cloud/datastore/client.py @@ -116,28 +116,29 @@ def _extended_lookup(connection, project, key_pbs, while loop_num < _MAX_LOOPS: # loop against possible deferred. loop_num += 1 - results_found, missing_found, deferred_found = connection.lookup( + lookup_response = connection.lookup( project=project, key_pbs=key_pbs, eventual=eventual, transaction_id=transaction_id, ) - results.extend(results_found) + # Accumulate the new results. + results.extend(result.entity for result in lookup_response.found) if missing is not None: - missing.extend(missing_found) + missing.extend(result.entity for result in lookup_response.missing) if deferred is not None: - deferred.extend(deferred_found) + deferred.extend(lookup_response.deferred) break - if len(deferred_found) == 0: + if len(lookup_response.deferred) == 0: break # We have deferred keys, and the user didn't ask to know about # them, so retry (but only with the deferred ones). - key_pbs = deferred_found + key_pbs = lookup_response.deferred return results diff --git a/datastore/unit_tests/test__http.py b/datastore/unit_tests/test__http.py index 0103e4883479..4d1fdc1bc357 100644 --- a/datastore/unit_tests/test__http.py +++ b/datastore/unit_tests/test__http.py @@ -255,26 +255,32 @@ def test_build_api_url_w_explicit_base_version(self): def test_lookup_single_key_empty_response(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - key_pb = self._make_key_pb(PROJECT) + project = 'PROJECT' + key_pb = self._make_key_pb(project) rsp_pb = datastore_pb2.LookupResponse() + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb]) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - found, missing, deferred = conn.lookup(PROJECT, [key_pb]) - self.assertEqual(len(found), 0) - self.assertEqual(len(missing), 0) - self.assertEqual(len(deferred), 0) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.missing), 0) + self.assertEqual(len(response.deferred), 0) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 1) @@ -283,27 +289,32 @@ def test_lookup_single_key_empty_response(self): def test_lookup_single_key_empty_response_w_eventual(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - key_pb = self._make_key_pb(PROJECT) + project = 'PROJECT' + key_pb = self._make_key_pb(project) rsp_pb = datastore_pb2.LookupResponse() + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb], eventual=True) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - found, missing, deferred = conn.lookup(PROJECT, [key_pb], - eventual=True) - self.assertEqual(len(found), 0) - self.assertEqual(len(missing), 0) - self.assertEqual(len(deferred), 0) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.missing), 0) + self.assertEqual(len(response.deferred), 0) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 1) @@ -323,62 +334,75 @@ def test_lookup_single_key_empty_response_w_eventual_and_transaction(self): def test_lookup_single_key_empty_response_w_transaction(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - TRANSACTION = b'TRANSACTION' - key_pb = self._make_key_pb(PROJECT) + project = 'PROJECT' + transaction = b'TRANSACTION' + key_pb = self._make_key_pb(project) rsp_pb = datastore_pb2.LookupResponse() + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb], transaction_id=transaction) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - found, missing, deferred = conn.lookup(PROJECT, [key_pb], - transaction_id=TRANSACTION) - self.assertEqual(len(found), 0) - self.assertEqual(len(missing), 0) - self.assertEqual(len(deferred), 0) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.missing), 0) + self.assertEqual(len(response.deferred), 0) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 1) self.assertEqual(key_pb, keys[0]) - self.assertEqual(request.read_options.transaction, TRANSACTION) + self.assertEqual(request.read_options.transaction, transaction) def test_lookup_single_key_nonempty_response(self): from google.cloud.proto.datastore.v1 import datastore_pb2 from google.cloud.proto.datastore.v1 import entity_pb2 - PROJECT = 'PROJECT' - key_pb = self._make_key_pb(PROJECT) + project = 'PROJECT' + key_pb = self._make_key_pb(project) rsp_pb = datastore_pb2.LookupResponse() entity = entity_pb2.Entity() entity.key.CopyFrom(key_pb) rsp_pb.found.add(entity=entity) + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb]) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - (found,), missing, deferred = conn.lookup(PROJECT, [key_pb]) - self.assertEqual(len(missing), 0) - self.assertEqual(len(deferred), 0) + self.assertEqual(len(response.found), 1) + self.assertEqual(len(response.missing), 0) + self.assertEqual(len(response.deferred), 0) + found = response.found[0].entity self.assertEqual(found.key.path[0].kind, 'Kind') self.assertEqual(found.key.path[0].id, 1234) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 1) @@ -387,27 +411,33 @@ def test_lookup_single_key_nonempty_response(self): def test_lookup_multiple_keys_empty_response(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - key_pb1 = self._make_key_pb(PROJECT) - key_pb2 = self._make_key_pb(PROJECT, id_=2345) + project = 'PROJECT' + key_pb1 = self._make_key_pb(project) + key_pb2 = self._make_key_pb(project, id_=2345) rsp_pb = datastore_pb2.LookupResponse() + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb1, key_pb2]) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - found, missing, deferred = conn.lookup(PROJECT, [key_pb1, key_pb2]) - self.assertEqual(len(found), 0) - self.assertEqual(len(missing), 0) - self.assertEqual(len(deferred), 0) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.missing), 0) + self.assertEqual(len(response.deferred), 0) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 2) @@ -417,32 +447,38 @@ def test_lookup_multiple_keys_empty_response(self): def test_lookup_multiple_keys_w_missing(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - key_pb1 = self._make_key_pb(PROJECT) - key_pb2 = self._make_key_pb(PROJECT, id_=2345) + project = 'PROJECT' + key_pb1 = self._make_key_pb(project) + key_pb2 = self._make_key_pb(project, id_=2345) rsp_pb = datastore_pb2.LookupResponse() er_1 = rsp_pb.missing.add() er_1.entity.key.CopyFrom(key_pb1) er_2 = rsp_pb.missing.add() er_2.entity.key.CopyFrom(key_pb2) + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb1, key_pb2]) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - result, missing, deferred = conn.lookup(PROJECT, [key_pb1, key_pb2]) - self.assertEqual(result, []) - self.assertEqual(len(deferred), 0) - self.assertEqual([missed.key for missed in missing], - [key_pb1, key_pb2]) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.deferred), 0) + missing_keys = [result.entity.key for result in response.missing] + self.assertEqual(missing_keys, [key_pb1, key_pb2]) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 2) @@ -455,28 +491,35 @@ def test_lookup_multiple_keys_w_deferred(self): from google.cloud import _http as connection_module from google.cloud.datastore._http import _CLIENT_INFO - PROJECT = 'PROJECT' - key_pb1 = self._make_key_pb(PROJECT) - key_pb2 = self._make_key_pb(PROJECT, id_=2345) + project = 'PROJECT' + key_pb1 = self._make_key_pb(project) + key_pb2 = self._make_key_pb(project, id_=2345) rsp_pb = datastore_pb2.LookupResponse() rsp_pb.deferred.add().CopyFrom(key_pb1) rsp_pb.deferred.add().CopyFrom(key_pb2) + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.lookup(project, [key_pb1, key_pb2]) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':lookup', + project + ':lookup', ]) - result, missing, deferred = conn.lookup(PROJECT, [key_pb1, key_pb2]) - self.assertEqual(result, []) - self.assertEqual(len(missing), 0) - self.assertEqual([def_key for def_key in deferred], [key_pb1, key_pb2]) + self.assertEqual(len(response.found), 0) + self.assertEqual(len(response.missing), 0) + self.assertEqual(list(response.deferred), [key_pb1, key_pb2]) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - self.assertEqual(cw['uri'], URI) + self._verify_protobuf_call(cw, uri, conn) + self.assertEqual(cw['uri'], uri) self.assertEqual(cw['method'], 'POST') expected_headers = { 'Content-Type': 'application/x-protobuf', @@ -485,8 +528,7 @@ def test_lookup_multiple_keys_w_deferred(self): connection_module.CLIENT_INFO_HEADER: _CLIENT_INFO, } self.assertEqual(cw['headers'], expected_headers) - rq_class = datastore_pb2.LookupRequest - request = rq_class() + request = datastore_pb2.LookupRequest() request.ParseFromString(cw['body']) keys = list(request.keys) self.assertEqual(len(keys), 2) diff --git a/datastore/unit_tests/test_client.py b/datastore/unit_tests/test_client.py index 26a9c56dfcc6..9ffe2dfd4996 100644 --- a/datastore/unit_tests/test_client.py +++ b/datastore/unit_tests/test_client.py @@ -940,7 +940,17 @@ def lookup(self, project, key_pbs, eventual=False, transaction_id=None): self._lookup_cw.append((project, key_pbs, eventual, transaction_id)) triple, self._lookup = self._lookup[0], self._lookup[1:] results, missing, deferred = triple - return results, missing, deferred + + entity_results_found = [ + mock.Mock(entity=result, spec=['entity']) for result in results] + entity_results_missing = [ + mock.Mock(entity=missing_entity, spec=['entity']) + for missing_entity in missing] + return mock.Mock( + found=entity_results_found, + missing=entity_results_missing, + deferred=deferred, + spec=['found', 'missing', 'deferred']) def commit(self, project, commit_request, transaction_id): from google.cloud.proto.datastore.v1 import datastore_pb2