Expose 'missing'/'deferred' in 'Connection.lookup'/'Dataset.get_entities'

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: 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: 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.datastore_v1_pb2.Entity`
                 (or a single Entity)
         :returns: The entities corresponding to the keys provided.
@@ -219,10 +229,32 @@ 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])
+
+            if deferred is not None:
+                deferred.extend([key for key in lookup_response.deferred])
+                break
+
+            if lookup_response.deferred:  # retry
+                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)
+            else:
+                break
 
-        results = [result.entity for result in lookup_response.found]
+            # Hmm, should we sleep here?  Asked in:
+            # https://github.com/GoogleCloudPlatform/gcloud-python/issues/306#issuecomment-67377587
 
         if single_key:
             if results:

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: 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: 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(

gcloud/datastore/test_connection.py

@@ -306,6 +306,145 @@ 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_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.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)
+        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_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.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)
+        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_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.assertEqual(cw[0]['uri'], URI)
+        self.assertEqual(cw[0]['method'], 'POST')
+        self.assertEqual(cw[0]['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw[0]['headers']['User-Agent'], conn.USER_AGENT)
+        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.assertEqual(cw[1]['uri'], URI)
+        self.assertEqual(cw[1]['method'], 'POST')
+        self.assertEqual(cw[1]['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw[1]['headers']['User-Agent'], conn.USER_AGENT)
+        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
@@ -901,3 +1040,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