Using Iterator as response for Pub/Sub list_subscriptions().

dhermes · dhermes · commit 0859cc38b215 · 2016-10-27T11:29:17.000-07:00
Also remove GAX confusion in Pub/Sub between max_results and
page_size.
diff --git a/pubsub/google/cloud/pubsub/_gax.py b/pubsub/google/cloud/pubsub/_gax.py
@@ -14,6 +14,8 @@
 
 """GAX wrapper for Pubsub API requests."""
 
+import functools
+
 from google.cloud.gapic.pubsub.v1.publisher_api import PublisherApi
 from google.cloud.gapic.pubsub.v1.subscriber_api import SubscriberApi
 from google.gax import CallOptions
@@ -78,12 +80,7 @@ def list_topics(self, project, page_size=0, page_token=None):
         path = 'projects/%s' % (project,)
         page_iter = self._gax_api.list_topics(
             path, page_size=page_size, options=options)
-
-        iter_kwargs = {}
-        if page_size:  # page_size can be 0 or explicit None.
-            iter_kwargs['max_results'] = page_size
-        return GAXIterator(self._client, page_iter, _item_to_topic,
-                           **iter_kwargs)
+        return GAXIterator(self._client, page_iter, _item_to_topic)
 
     def topic_create(self, topic_path):
         """API call:  create a topic
@@ -215,12 +212,8 @@ def topic_list_subscriptions(self, topic, page_size=0, page_token=None):
                 raise NotFound(topic_path)
             raise
 
-        iter_kwargs = {}
-        if page_size:  # page_size can be 0 or explicit None.
-            iter_kwargs['max_results'] = page_size
-        iterator = GAXIterator(
-            self._client, page_iter,
-            _item_to_subscription_for_topic, **iter_kwargs)
+        iterator = GAXIterator(self._client, page_iter,
+                               _item_to_subscription_for_topic)
         iterator.topic = topic
         return iterator
 
@@ -256,22 +249,25 @@ def list_subscriptions(self, project, page_size=0, page_token=None):
                            If not passed, the API will return the first page
                            of subscriptions.
 
-        :rtype: tuple, (list, str)
-        :returns: list of ``Subscription`` resource dicts, plus a
-                  "next page token" string:  if not None, indicates that
-                  more topics can be retrieved with another call (pass that
-                  value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of
+                  :class:`~google.cloud.pubsub.subscription.Subscription`
+                  accessible to the current API.
         """
         if page_token is None:
             page_token = INITIAL_PAGE
         options = CallOptions(page_token=page_token)
         path = 'projects/%s' % (project,)
         page_iter = self._gax_api.list_subscriptions(
             path, page_size=page_size, options=options)
-        subscriptions = [MessageToDict(sub_pb)
-                         for sub_pb in page_iter.next()]
-        token = page_iter.page_token or None
-        return subscriptions, token
+
+        # We attach a mutable topics dictionary so that as topic
+        # objects are created by Subscription.from_api_repr, they
+        # can be re-used by other subscriptions from the same topic.
+        topics = {}
+        item_to_value = functools.partial(
+            _item_to_subscription_for_client, topics=topics)
+        return GAXIterator(self._client, page_iter, item_to_value)
 
     def subscription_create(self, subscription_path, topic_path,
                             ack_deadline=None, push_endpoint=None):
@@ -579,3 +575,33 @@ def _item_to_subscription_for_topic(iterator, subscription_path):
     subscription_name = subscription_name_from_path(
         subscription_path, iterator.client.project)
     return Subscription(subscription_name, iterator.topic)
+
+
+def _item_to_subscription_for_client(iterator, sub_pb, topics):
+    """Convert a subscription protobuf to the native object.
+
+    .. note::
+
+       This method does not have the correct signature to be used as
+       the ``item_to_value`` argument to
+       :class:`~google.cloud.iterator.Iterator`. It is intended to be
+       patched with a mutable topics argument that can be updated
+       on subsequent calls. For an example, see how the method is
+       used above in :meth:`_SubscriberAPI.list_subscriptions`.
+
+    :type iterator: :class:`~google.cloud.iterator.Iterator`
+    :param iterator: The iterator that is currently in use.
+
+    :type sub_pb: :class:`~google.pubsub.v1.pubsub_pb2.Subscription`
+    :param sub_pb: A subscription returned from the API.
+
+    :type topics: dict
+    :param topics: A dictionary of topics to be used (and modified)
+                   as new subscriptions are created bound to topics.
+
+    :rtype: :class:`~google.cloud.pubsub.subscription.Subscription`
+    :returns: The next subscription in the page.
+    """
+    resource = MessageToDict(sub_pb)
+    return Subscription.from_api_repr(
+        resource, iterator.client, topics=topics)
diff --git a/pubsub/google/cloud/pubsub/client.py b/pubsub/google/cloud/pubsub/client.py
@@ -160,20 +160,14 @@ def list_subscriptions(self, page_size=None, page_token=None):
                            passed, the API will return the first page of
                            topics.
 
-        :rtype: tuple, (list, str)
-        :returns: list of :class:`~.pubsub.subscription.Subscription`,
-                  plus a "next page token" string:  if not None, indicates that
-                  more topics can be retrieved with another call (pass that
-                  value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of
+                  :class:`~google.cloud.pubsub.subscription.Subscription`
+                  accessible to the current client.
         """
         api = self.subscriber_api
-        resources, next_token = api.list_subscriptions(
+        return api.list_subscriptions(
             self.project, page_size, page_token)
-        topics = {}
-        subscriptions = [Subscription.from_api_repr(resource, self,
-                                                    topics=topics)
-                         for resource in resources]
-        return subscriptions, next_token
 
     def topic(self, name, timestamp_messages=False):
         """Creates a topic bound to the current client.
diff --git a/pubsub/google/cloud/pubsub/connection.py b/pubsub/google/cloud/pubsub/connection.py
@@ -15,6 +15,7 @@
 """Create / interact with Google Cloud Pub/Sub connections."""
 
 import base64
+import functools
 import os
 
 from google.cloud import connection as base_connection
@@ -274,24 +275,26 @@ def list_subscriptions(self, project, page_size=None, page_token=None):
                            If not passed, the API will return the first page
                            of subscriptions.
 
-        :rtype: tuple, (list, str)
-        :returns: list of ``Subscription`` resource dicts, plus a
-                  "next page token" string:  if not None, indicates that
-                  more subscriptions can be retrieved with another call (pass
-                  that value as ``page_token``).
+        :rtype: :class:`~google.cloud.iterator.Iterator`
+        :returns: Iterator of
+                  :class:`~google.cloud.pubsub.subscription.Subscription`
+                  accessible to the current API.
         """
-        conn = self._connection
-        params = {}
-
+        extra_params = {}
         if page_size is not None:
-            params['pageSize'] = page_size
-
-        if page_token is not None:
-            params['pageToken'] = page_token
-
+            extra_params['pageSize'] = page_size
         path = '/projects/%s/subscriptions' % (project,)
-        resp = conn.api_request(method='GET', path=path, query_params=params)
-        return resp.get('subscriptions', ()), resp.get('nextPageToken')
+
+        # We attach a mutable topics dictionary so that as topic
+        # objects are created by Subscription.from_api_repr, they
+        # can be re-used by other subscriptions from the same topic.
+        topics = {}
+        item_to_value = functools.partial(
+            _item_to_subscription_for_client, topics=topics)
+        return HTTPIterator(
+            client=self._client, path=path, item_to_value=item_to_value,
+            items_key='subscriptions', page_token=page_token,
+            extra_params=extra_params)
 
     def subscription_create(self, subscription_path, topic_path,
                             ack_deadline=None, push_endpoint=None):
@@ -607,3 +610,32 @@ def _item_to_subscription_for_topic(iterator, subscription_path):
     subscription_name = subscription_name_from_path(
         subscription_path, iterator.client.project)
     return Subscription(subscription_name, iterator.topic)
+
+
+def _item_to_subscription_for_client(iterator, resource, topics):
+    """Convert a subscription to the native object.
+
+    .. note::
+
+       This method does not have the correct signature to be used as
+       the ``item_to_value`` argument to
+       :class:`~google.cloud.iterator.Iterator`. It is intended to be
+       patched with a mutable topics argument that can be updated
+       on subsequent calls. For an example, see how the method is
+       used above in :meth:`_SubscriberAPI.list_subscriptions`.
+
+    :type iterator: :class:`~google.cloud.iterator.Iterator`
+    :param iterator: The iterator that is currently in use.
+
+    :type resource: dict
+    :param resource: A subscription returned from the API.
+
+    :type topics: dict
+    :param topics: A dictionary of topics to be used (and modified)
+                   as new subscriptions are created bound to topics.
+
+    :rtype: :class:`~google.cloud.pubsub.subscription.Subscription`
+    :returns: The next subscription in the page.
+    """
+    return Subscription.from_api_repr(
+        resource, iterator.client, topics=topics)
