diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index 46fa55fc6363..4edc005be6c2 100644 --- a/gcloud/datastore/connection.py +++ b/gcloud/datastore/connection.py @@ -167,7 +167,7 @@ def dataset(self, *args, **kwargs): kwargs['connection'] = self return Dataset(*args, **kwargs) - def lookup(self, dataset_id, key_pbs): + def lookup(self, dataset_id, key_pbs, missing=None, deferred=None): """Lookup keys from a dataset in the Cloud Datastore. Maps the ``DatastoreService.Lookup`` protobuf RPC. @@ -201,6 +201,16 @@ def lookup(self, dataset_id, key_pbs): (or a single Key) :param key_pbs: The key (or keys) to retrieve from the datastore. + :type missing: an empty list or None. + :param missing: If a list is passed, the key-only entity protobufs + returned by the backend as "missing" will be copied + into it. Use only as a keyword param. + + :type deferred: an empty list or None. + :param deferred: If a list is passed, the key protobufs returned + by the backend as "deferred" will be copied into it. + Use only as a keyword param. + :rtype: list of :class:`gcloud.datastore.datastore_v1_pb2.Entity` (or a single Entity) :returns: The entities corresponding to the keys provided. @@ -209,6 +219,12 @@ def lookup(self, dataset_id, key_pbs): If multiple keys were provided and no results matched, this will return an empty list. """ + if missing is not None and missing != []: + raise ValueError('missing must be None or an empty list') + + if deferred is not None and deferred != []: + raise ValueError('deferred must be None or an empty list') + lookup_request = datastore_pb.LookupRequest() single_key = isinstance(key_pbs, datastore_pb.Key) @@ -219,10 +235,28 @@ def lookup(self, dataset_id, key_pbs): for key_pb in key_pbs: lookup_request.key.add().CopyFrom(key_pb) - lookup_response = self._rpc(dataset_id, 'lookup', lookup_request, - datastore_pb.LookupResponse) + results = [] + while True: # loop against possible deferred. + lookup_response = self._rpc(dataset_id, 'lookup', lookup_request, + datastore_pb.LookupResponse) + + results.extend( + [result.entity for result in lookup_response.found]) + + if missing is not None: + missing.extend( + [result.entity for result in lookup_response.missing]) - results = [result.entity for result in lookup_response.found] + if deferred is not None: + deferred.extend([key for key in lookup_response.deferred]) + break + + if not lookup_response.deferred: + break + + # We have deferred keys, and the user didn't ask to know about + # them, so retry (but only with the deferred ones). + _copy_deferred_keys(lookup_request, lookup_response) if single_key: if results: @@ -476,3 +510,14 @@ def delete_entities(self, dataset_id, key_pbs): self.commit(dataset_id, mutation) return True + + +def _copy_deferred_keys(lookup_request, lookup_response): + """Clear requested keys and copy deferred keys back in. + + Helper ``Connection.lookup()``. + """ + for old_key in list(lookup_request.key): + lookup_request.key.remove(old_key) + for def_key in lookup_response.deferred: + lookup_request.key.add().CopyFrom(def_key) diff --git a/gcloud/datastore/dataset.py b/gcloud/datastore/dataset.py index 98659e4e0aa2..c527c34aab1f 100644 --- a/gcloud/datastore/dataset.py +++ b/gcloud/datastore/dataset.py @@ -142,20 +142,41 @@ def get_entity(self, key_or_path): if entities: return entities[0] - def get_entities(self, keys): + def get_entities(self, keys, missing=None, deferred=None): """Retrieves entities from the dataset, along with their attributes. :type key: list of :class:`gcloud.datastore.key.Key` :param item_name: The name of the item to retrieve. + :type missing: an empty list or None. + :param missing: If a list is passed, the key-only entities returned + by the backend as "missing" will be copied into it. + Use only as a keyword param. + + :type deferred: an empty list or None. + :param deferred: If a list is passed, the keys returned + by the backend as "deferred" will be copied into it. + Use only as a keyword param. + :rtype: list of :class:`gcloud.datastore.entity.Entity` :return: The requested entities. """ entity_pbs = self.connection().lookup( dataset_id=self.id(), - key_pbs=[k.to_protobuf() for k in keys] + key_pbs=[k.to_protobuf() for k in keys], + missing=missing, deferred=deferred, ) + if missing is not None: + missing[:] = [ + helpers.entity_from_protobuf(missed_pb, dataset=self) + for missed_pb in missing] + + if deferred is not None: + deferred[:] = [ + helpers.key_from_protobuf(deferred_pb) + for deferred_pb in deferred] + entities = [] for entity_pb in entity_pbs: entities.append(helpers.entity_from_protobuf( diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py index a804c84c0461..ccd93f843dbd 100644 --- a/gcloud/datastore/test_connection.py +++ b/gcloud/datastore/test_connection.py @@ -25,6 +25,14 @@ def _getTargetClass(self): def _makeOne(self, *args, **kw): return self._getTargetClass()(*args, **kw) + def _verifyProtobufCall(self, called_with, URI, conn): + self.assertEqual(called_with['uri'], URI) + self.assertEqual(called_with['method'], 'POST') + self.assertEqual(called_with['headers']['Content-Type'], + 'application/x-protobuf') + self.assertEqual(called_with['headers']['User-Agent'], + conn.USER_AGENT) + def test_ctor_defaults(self): conn = self._makeOne() self.assertEqual(conn.credentials, None) @@ -75,12 +83,7 @@ def test__request_w_200(self): ]) http = conn._http = Http({'status': '200'}, 'CONTENT') self.assertEqual(conn._request(DATASET_ID, METHOD, DATA), 'CONTENT') - self.assertEqual(http._called_with['uri'], URI) - self.assertEqual(http._called_with['method'], 'POST') - self.assertEqual(http._called_with['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(http._called_with['headers']['User-Agent'], - conn.USER_AGENT) + self._verifyProtobufCall(http._called_with, URI, conn) self.assertEqual(http._called_with['body'], DATA) def test__request_not_200(self): @@ -126,12 +129,7 @@ def FromString(cls, pb): response = conn._rpc(DATASET_ID, METHOD, ReqPB(), RspPB) self.assertTrue(isinstance(response, RspPB)) self.assertEqual(response._pb, 'CONTENT') - self.assertEqual(http._called_with['uri'], URI) - self.assertEqual(http._called_with['method'], 'POST') - self.assertEqual(http._called_with['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(http._called_with['headers']['User-Agent'], - conn.USER_AGENT) + self._verifyProtobufCall(http._called_with, URI, conn) self.assertEqual(http._called_with['body'], REQPB) def test_build_api_url_w_default_base_version(self): @@ -225,11 +223,7 @@ def test_lookup_single_key_empty_response(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.lookup(DATASET_ID, key_pb), None) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -261,11 +255,7 @@ def test_lookup_single_key_nonempty_response(self): self.assertEqual(found.key.path_element[0].kind, 'Kind') self.assertEqual(found.key.path_element[0].id, 1234) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -293,6 +283,89 @@ def test_lookup_multiple_keys_empty_response(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.lookup(DATASET_ID, [key_pb1, key_pb2]), []) cw = http._called_with + self._verifyProtobufCall(cw, URI, conn) + rq_class = datastore_pb.LookupRequest + request = rq_class() + request.ParseFromString(cw['body']) + keys = list(request.key) + self.assertEqual(len(keys), 2) + self.assertEqual(keys[0], key_pb1) + self.assertEqual(keys[1], key_pb2) + + def test_lookup_multiple_keys_w_missing(self): + from gcloud.datastore.connection import datastore_pb + from gcloud.datastore.key import Key + + DATASET_ID = 'DATASET' + key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf() + key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf() + rsp_pb = datastore_pb.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) + conn = self._makeOne() + URI = '/'.join([ + conn.API_BASE_URL, + 'datastore', + conn.API_VERSION, + 'datasets', + DATASET_ID, + 'lookup', + ]) + http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) + missing = [] + result = conn.lookup(DATASET_ID, [key_pb1, key_pb2], missing=missing) + self.assertEqual(result, []) + self.assertEqual([missed.key for missed in missing], + [key_pb1, key_pb2]) + cw = http._called_with + self._verifyProtobufCall(cw, URI, conn) + rq_class = datastore_pb.LookupRequest + request = rq_class() + request.ParseFromString(cw['body']) + keys = list(request.key) + self.assertEqual(len(keys), 2) + self.assertEqual(keys[0], key_pb1) + self.assertEqual(keys[1], key_pb2) + + def test_lookup_multiple_keys_w_missing_non_empty(self): + from gcloud.datastore.key import Key + DATASET_ID = 'DATASET' + key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf() + key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf() + conn = self._makeOne() + missing = ['this', 'list', 'is', 'not', 'empty'] + self.assertRaises( + ValueError, + conn.lookup, DATASET_ID, [key_pb1, key_pb2], missing=missing) + + def test_lookup_multiple_keys_w_deferred(self): + from gcloud.datastore.connection import datastore_pb + from gcloud.datastore.key import Key + + DATASET_ID = 'DATASET' + key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf() + key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf() + rsp_pb = datastore_pb.LookupResponse() + rsp_pb.deferred.add().CopyFrom(key_pb1) + rsp_pb.deferred.add().CopyFrom(key_pb2) + conn = self._makeOne() + URI = '/'.join([ + conn.API_BASE_URL, + 'datastore', + conn.API_VERSION, + 'datasets', + DATASET_ID, + 'lookup', + ]) + http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) + deferred = [] + result = conn.lookup(DATASET_ID, [key_pb1, key_pb2], deferred=deferred) + self.assertEqual(result, []) + self.assertEqual([def_key for def_key in deferred], [key_pb1, key_pb2]) + cw = http._called_with + self._verifyProtobufCall(cw, URI, conn) self.assertEqual(cw['uri'], URI) self.assertEqual(cw['method'], 'POST') self.assertEqual(cw['headers']['Content-Type'], @@ -306,6 +379,70 @@ def test_lookup_multiple_keys_empty_response(self): self.assertEqual(keys[0], key_pb1) self.assertEqual(keys[1], key_pb2) + def test_lookup_multiple_keys_w_deferred_non_empty(self): + from gcloud.datastore.key import Key + DATASET_ID = 'DATASET' + key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf() + key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf() + conn = self._makeOne() + deferred = ['this', 'list', 'is', 'not', 'empty'] + self.assertRaises( + ValueError, + conn.lookup, DATASET_ID, [key_pb1, key_pb2], deferred=deferred) + + def test_lookup_multiple_keys_w_deferred_from_backend_but_not_passed(self): + from gcloud.datastore.connection import datastore_pb + from gcloud.datastore.key import Key + + DATASET_ID = 'DATASET' + key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf() + key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf() + rsp_pb1 = datastore_pb.LookupResponse() + entity1 = datastore_pb.Entity() + entity1.key.CopyFrom(key_pb1) + rsp_pb1.found.add(entity=entity1) + rsp_pb1.deferred.add().CopyFrom(key_pb2) + rsp_pb2 = datastore_pb.LookupResponse() + entity2 = datastore_pb.Entity() + entity2.key.CopyFrom(key_pb2) + rsp_pb2.found.add(entity=entity2) + conn = self._makeOne() + URI = '/'.join([ + conn.API_BASE_URL, + 'datastore', + conn.API_VERSION, + 'datasets', + DATASET_ID, + 'lookup', + ]) + http = conn._http = HttpMultiple( + ({'status': '200'}, rsp_pb1.SerializeToString()), + ({'status': '200'}, rsp_pb2.SerializeToString()), + ) + found = conn.lookup(DATASET_ID, [key_pb1, key_pb2]) + self.assertEqual(len(found), 2) + self.assertEqual(found[0].key.path_element[0].kind, 'Kind') + self.assertEqual(found[0].key.path_element[0].id, 1234) + self.assertEqual(found[1].key.path_element[0].kind, 'Kind') + self.assertEqual(found[1].key.path_element[0].id, 2345) + cw = http._called_with + rq_class = datastore_pb.LookupRequest + request = rq_class() + self.assertEqual(len(cw), 2) + + self._verifyProtobufCall(cw[0], URI, conn) + request.ParseFromString(cw[0]['body']) + keys = list(request.key) + self.assertEqual(len(keys), 2) + self.assertEqual(keys[0], key_pb1) + self.assertEqual(keys[1], key_pb2) + + self._verifyProtobufCall(cw[1], URI, conn) + request.ParseFromString(cw[1]['body']) + keys = list(request.key) + self.assertEqual(len(keys), 1) + self.assertEqual(keys[0], key_pb2) + def test_run_query_wo_namespace_empty_result(self): from gcloud.datastore.connection import datastore_pb from gcloud.datastore.query import Query @@ -330,11 +467,7 @@ def test_run_query_wo_namespace_empty_result(self): self.assertTrue(more) self.assertEqual(skipped, 0) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.RunQueryRequest request = rq_class() request.ParseFromString(cw['body']) @@ -366,11 +499,7 @@ def test_run_query_w_namespace_nonempty_result(self): pbs = conn.run_query(DATASET_ID, q_pb, 'NS')[0] self.assertEqual(len(pbs), 1) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.RunQueryRequest request = rq_class() request.ParseFromString(cw['body']) @@ -402,11 +531,7 @@ def test_begin_transaction_default_serialize(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.begin_transaction(DATASET_ID), TRANSACTION) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.BeginTransactionRequest request = rq_class() request.ParseFromString(cw['body']) @@ -431,11 +556,7 @@ def test_begin_transaction_explicit_serialize(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.begin_transaction(DATASET_ID, True), TRANSACTION) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.BeginTransactionRequest request = rq_class() request.ParseFromString(cw['body']) @@ -468,11 +589,7 @@ def test_commit_wo_transaction(self): self.assertEqual(result.index_updates, 0) self.assertEqual(list(result.insert_auto_id_key), []) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -511,11 +628,7 @@ def id(self): self.assertEqual(result.index_updates, 0) self.assertEqual(list(result.insert_auto_id_key), []) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -564,11 +677,7 @@ def id(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.rollback(DATASET_ID), None) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.RollbackRequest request = rq_class() request.ParseFromString(cw['body']) @@ -591,11 +700,7 @@ def test_allocate_ids_empty(self): http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString()) self.assertEqual(conn.allocate_ids(DATASET_ID, []), []) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.AllocateIdsRequest request = rq_class() request.ParseFromString(cw['body']) @@ -630,11 +735,7 @@ def test_allocate_ids_non_empty(self): self.assertEqual(conn.allocate_ids(DATASET_ID, before_key_pbs), after_key_pbs) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.AllocateIdsRequest request = rq_class() request.ParseFromString(cw['body']) @@ -660,11 +761,7 @@ def test_save_entity_wo_transaction_w_upsert(self): result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo'}) self.assertEqual(result, True) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -706,11 +803,7 @@ def test_save_entity_w_exclude_from_indexes(self): exclude_from_indexes=['foo', 'bar']) self.assertEqual(result, True) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -762,11 +855,7 @@ def test_save_entity_wo_transaction_w_auto_id(self): result = conn.save_entity(DATASET_ID, key_pb, {'foo': u'Foo'}) self.assertEqual(result, updated_key_pb) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -850,11 +939,7 @@ def test_delete_entities_wo_transaction(self): result = conn.delete_entities(DATASET_ID, [key_pb]) self.assertEqual(result, True) cw = http._called_with - self.assertEqual(cw['uri'], URI) - self.assertEqual(cw['method'], 'POST') - self.assertEqual(cw['headers']['Content-Type'], - 'application/x-protobuf') - self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT) + self._verifyProtobufCall(cw, URI, conn) rq_class = datastore_pb.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -901,3 +986,15 @@ def __init__(self, headers, content): def request(self, **kw): self._called_with = kw return self._headers, self._content + + +class HttpMultiple(object): + + def __init__(self, *responses): + self._called_with = [] + self._responses = list(responses) + + def request(self, **kw): + self._called_with.append(kw) + result, self._responses = self._responses[0], self._responses[1:] + return result diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py index e6924ebbb069..a59e4b307b17 100644 --- a/gcloud/datastore/test_dataset.py +++ b/gcloud/datastore/test_dataset.py @@ -76,15 +76,15 @@ def test_transaction_factory(self): self.assertIsInstance(transaction, Transaction) self.assertTrue(transaction.dataset() is dataset) - def test_get_entities_miss(self): + def test_get_entity_miss(self): from gcloud.datastore.key import Key DATASET_ID = 'DATASET' connection = _Connection() dataset = self._makeOne(DATASET_ID, connection) key = Key(path=[{'kind': 'Kind', 'id': 1234}]) - self.assertEqual(dataset.get_entities([key]), []) + self.assertEqual(dataset.get_entity(key), None) - def test_get_entities_hit(self): + def test_get_entity_hit(self): from gcloud.datastore.connection import datastore_pb from gcloud.datastore.key import Key DATASET_ID = 'DATASET' @@ -102,24 +102,15 @@ def test_get_entities_hit(self): connection = _Connection(entity_pb) dataset = self._makeOne(DATASET_ID, connection) key = Key(path=PATH) - result, = dataset.get_entities([key]) + result = dataset.get_entity(key) key = result.key() self.assertEqual(key._dataset_id, DATASET_ID) self.assertEqual(key.path(), PATH) self.assertEqual(list(result), ['foo']) self.assertEqual(result['foo'], 'Foo') - def test_get_entity_miss(self): - from gcloud.datastore.key import Key - DATASET_ID = 'DATASET' - connection = _Connection() - dataset = self._makeOne(DATASET_ID, connection) - key = Key(path=[{'kind': 'Kind', 'id': 1234}]) - self.assertEqual(dataset.get_entity(key), None) - - def test_get_entity_hit(self): + def test_get_entity_path(self): from gcloud.datastore.connection import datastore_pb - from gcloud.datastore.key import Key DATASET_ID = 'DATASET' KIND = 'Kind' ID = 1234 @@ -134,16 +125,72 @@ def test_get_entity_hit(self): prop.value.string_value = 'Foo' connection = _Connection(entity_pb) dataset = self._makeOne(DATASET_ID, connection) - key = Key(path=PATH) - result = dataset.get_entity(key) + result = dataset.get_entity([KIND, ID]) key = result.key() self.assertEqual(key._dataset_id, DATASET_ID) self.assertEqual(key.path(), PATH) self.assertEqual(list(result), ['foo']) self.assertEqual(result['foo'], 'Foo') - def test_get_entity_path(self): + def test_get_entity_odd_nonetype(self): + DATASET_ID = 'DATASET' + KIND = 'Kind' + connection = _Connection() + dataset = self._makeOne(DATASET_ID, connection) + with self.assertRaises(ValueError): + dataset.get_entity([KIND]) + with self.assertRaises(TypeError): + dataset.get_entity(None) + + def test_get_entities_miss(self): + from gcloud.datastore.key import Key + DATASET_ID = 'DATASET' + connection = _Connection() + dataset = self._makeOne(DATASET_ID, connection) + key = Key(path=[{'kind': 'Kind', 'id': 1234}]) + self.assertEqual(dataset.get_entities([key]), []) + + def test_get_entities_miss_w_missing(self): from gcloud.datastore.connection import datastore_pb + from gcloud.datastore.key import Key + DATASET_ID = 'DATASET' + KIND = 'Kind' + ID = 1234 + PATH = [{'kind': KIND, 'id': ID}] + missed = datastore_pb.Entity() + missed.key.partition_id.dataset_id = DATASET_ID + path_element = missed.key.path_element.add() + path_element.kind = KIND + path_element.id = ID + connection = _Connection() + connection._missing = [missed] + dataset = self._makeOne(DATASET_ID, connection) + key = Key(path=PATH, dataset_id=DATASET_ID) + missing = [] + entities = dataset.get_entities([key], missing=missing) + self.assertEqual(entities, []) + self.assertEqual([missed.key().to_protobuf() for missed in missing], + [key.to_protobuf()]) + + def test_get_entities_miss_w_deferred(self): + from gcloud.datastore.key import Key + DATASET_ID = 'DATASET' + KIND = 'Kind' + ID = 1234 + PATH = [{'kind': KIND, 'id': ID}] + connection = _Connection() + dataset = self._makeOne(DATASET_ID, connection) + key = Key(path=PATH, dataset_id=DATASET_ID) + connection._deferred = [key.to_protobuf()] + deferred = [] + entities = dataset.get_entities([key], deferred=deferred) + self.assertEqual(entities, []) + self.assertEqual([def_key.to_protobuf() for def_key in deferred], + [key.to_protobuf()]) + + def test_get_entities_hit(self): + from gcloud.datastore.connection import datastore_pb + from gcloud.datastore.key import Key DATASET_ID = 'DATASET' KIND = 'Kind' ID = 1234 @@ -158,30 +205,28 @@ def test_get_entity_path(self): prop.value.string_value = 'Foo' connection = _Connection(entity_pb) dataset = self._makeOne(DATASET_ID, connection) - result = dataset.get_entity([KIND, ID]) + key = Key(path=PATH) + result, = dataset.get_entities([key]) key = result.key() self.assertEqual(key._dataset_id, DATASET_ID) self.assertEqual(key.path(), PATH) self.assertEqual(list(result), ['foo']) self.assertEqual(result['foo'], 'Foo') - def test_get_entity_odd_nonetype(self): - DATASET_ID = 'DATASET' - KIND = 'Kind' - connection = _Connection() - dataset = self._makeOne(DATASET_ID, connection) - with self.assertRaises(ValueError): - dataset.get_entity([KIND]) - with self.assertRaises(TypeError): - dataset.get_entity(None) - class _Connection(object): _called_with = None + _missing = _deferred = () def __init__(self, *result): self._result = list(result) def lookup(self, **kw): self._called_with = kw + missing = kw.pop('missing', None) + if missing is not None: + missing.extend(self._missing) + deferred = kw.pop('deferred', None) + if deferred is not None: + deferred.extend(self._deferred) return self._result