Moving behavior in Connection.lookup to datastore.api.

This was so that Connection.lookup was only performing the
API call by the same name. The "lookup until all deferred return"
behavior has been moving into `datastore.get` (which is where
end users expect the behavior).

Author: dhermes

gcloud/datastore/api.py

@@ -24,6 +24,10 @@
 from gcloud.datastore import helpers
 
 
+_MAX_LOOPS = 128
+"""Maximum number of iterations to wait for deferred keys."""
+
+
 def _require_dataset_id(dataset_id=None, first_key=None):
     """Infer a dataset ID from the environment, if not passed explicitly.
 
@@ -80,6 +84,64 @@ def _require_connection(connection=None):
     return connection
 
 
+def _extended_lookup(connection, dataset_id, key_pbs,
+                     missing=None, deferred=None,
+                     eventual=False, transaction_id=None):
+    """Repeat lookup until all keys found (unless stop requested).
+
+    Helper method for ``lookup()``.
+
+    :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.
+
+    :raises: :class:`ValueError` if missing / deferred are not null or
+             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')
+
+    results = []
+
+    loop_num = 0
+    while loop_num < _MAX_LOOPS:  # loop against possible deferred.
+        loop_num += 1
+
+        results_found, missing_found, deferred_found = connection.lookup(
+            dataset_id=dataset_id,
+            key_pbs=key_pbs,
+            eventual=eventual,
+            transaction_id=transaction_id,
+        )
+
+        results.extend(results_found)
+
+        if missing is not None:
+            missing.extend(missing_found)
+
+        if deferred is not None:
+            deferred.extend(deferred_found)
+            break
+
+        if len(deferred_found) == 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
+
+    return results
+
+
 def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
     """Retrieves entities, along with their attributes.
 
@@ -122,7 +184,8 @@ def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
 
     transaction = Transaction.current()
 
-    entity_pbs = connection.lookup(
+    entity_pbs = _extended_lookup(
+        connection,
         dataset_id=dataset_id,
         key_pbs=[k.to_protobuf() for k in keys],
         missing=missing,

gcloud/datastore/connection.py

@@ -122,7 +122,6 @@ def build_api_url(cls, dataset_id, method, base_url=None,
             dataset_id=dataset_id, method=method)
 
     def lookup(self, dataset_id, key_pbs,
-               missing=None, deferred=None,
                eventual=False, transaction_id=None):
         """Lookup keys from a dataset in the Cloud Datastore.
 
@@ -150,16 +149,6 @@ def lookup(self, dataset_id, key_pbs,
         :type key_pbs: list of :class:`gcloud.datastore._datastore_v1_pb2.Key`
         :param key_pbs: The 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.
-
         :type eventual: boolean
         :param eventual: If False (the default), request ``STRONG`` read
                         consistency.  If True, request ``EVENTUAL`` read
@@ -170,35 +159,24 @@ def lookup(self, dataset_id, key_pbs,
                                the given transaction.  Incompatible with
                                ``eventual==True``.
 
-        :rtype: list of :class:`gcloud.datastore._datastore_v1_pb2.Entity`
-                (or a single Entity)
-        :returns: The entities corresponding to the keys provided.
-                  If a single key was provided and no results matched,
-                  this will return None.
-                  If multiple keys were provided and no results matched,
-                  this will return an empty list.
-        :raises: ValueError if ``eventual`` is True
+        :rtype: tuple
+        :returns: A triple of (``results``, ``missing``, ``deferred``) where
+                  both ``results`` and ``missing`` are lists of
+                  :class:`gcloud.datastore._datastore_v1_pb2.Entity` and
+                  ``deferred`` is a list of
+                  :class:`gcloud.datastore._datastore_v1_pb2.Key`.
         """
-        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()
         _set_read_options(lookup_request, eventual, transaction_id)
         helpers._add_keys_to_request(lookup_request.key, key_pbs)
 
-        results, missing_found, deferred_found = self._lookup(
-            lookup_request, dataset_id, deferred is not None)
-
-        if missing is not None:
-            missing.extend(missing_found)
+        lookup_response = self._rpc(dataset_id, 'lookup', lookup_request,
+                                    datastore_pb.LookupResponse)
 
-        if deferred is not None:
-            deferred.extend(deferred_found)
+        results = [result.entity for result in lookup_response.found]
+        missing = [result.entity for result in lookup_response.missing]
 
-        return results
+        return results, missing, lookup_response.deferred
 
     def run_query(self, dataset_id, query_pb, namespace=None,
                   eventual=False, transaction_id=None):
@@ -376,41 +354,14 @@ def allocate_ids(self, dataset_id, key_pbs):
                              datastore_pb.AllocateIdsResponse)
         return list(response.key)
 
-    def _lookup(self, lookup_request, dataset_id, stop_on_deferred):
-        """Repeat lookup until all keys found (unless stop requested).
-
-        Helper method for ``lookup()``.
-        """
-        results = []
-        missing = []
-        deferred = []
-        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])
-
-            missing.extend(
-                [result.entity for result in lookup_response.missing])
-
-            if stop_on_deferred:
-                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)
-        return results, missing, deferred
-
 
 def _set_read_options(request, eventual, transaction_id):
     """Validate rules for read options, and assign to the request.
 
     Helper method for ``lookup()`` and ``run_query``.
+
+    :raises: :class:`ValueError` if ``eventual`` is ``True`` and the
+             ``transaction_id`` is not ``None``.
     """
     if eventual and (transaction_id is not None):
         raise ValueError('eventual must be False when in a transaction')
@@ -420,14 +371,3 @@ def _set_read_options(request, eventual, transaction_id):
         opts.read_consistency = datastore_pb.ReadOptions.EVENTUAL
     elif transaction_id:
         opts.transaction = transaction_id
-
-
-def _copy_deferred_keys(lookup_request, lookup_response):
-    """Clear requested keys and copy deferred keys back in.
-
-    Helper for ``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)