Merge pull request #429 from tseaver/306-connection_lookup_missing_deferred

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

Author: tseaver

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)

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(