Merge pull request #2633 from dhermes/pubsub-iterators-3

Converting Pub/Sub client->list_subscriptions to iterator.

Author: dhermes

docs/pubsub_snippets.py

@@ -60,14 +60,8 @@ def do_something_with(sub):  # pylint: disable=unused-argument
         pass
 
     # [START client_list_subscriptions]
-    subscriptions, token = client.list_subscriptions()   # API request
-    while True:
-        for subscription in subscriptions:
-            do_something_with(subscription)
-        if token is None:
-            break
-        subscriptions, token = client.list_subscriptions(
-            page_token=token)                           # API request
+    for subscription in client.list_subscriptions():  # API request(s)
+        do_something_with(subscription)
     # [END client_list_subscriptions]
 
 

pubsub/google/cloud/pubsub/_gax.py

@@ -14,12 +14,15 @@
 
 """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
 from google.gax import INITIAL_PAGE
 from google.gax.errors import GaxError
 from google.gax.grpc import exc_to_code
+from google.protobuf.json_format import MessageToDict
 from google.pubsub.v1.pubsub_pb2 import PubsubMessage
 from google.pubsub.v1.pubsub_pb2 import PushConfig
 from grpc import insecure_channel
@@ -77,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
@@ -214,11 +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, **iter_kwargs)
+                               _item_to_subscription_for_topic)
         iterator.topic = topic
         return iterator
 
@@ -228,9 +223,13 @@ class _SubscriberAPI(object):
 
     :type gax_api: :class:`google.pubsub.v1.publisher_api.SubscriberApi`
     :param gax_api: API object used to make GAX requests.
+
+    :type client: :class:`~google.cloud.pubsub.client.Client`
+    :param client: The client that owns this API object.
     """
-    def __init__(self, gax_api):
+    def __init__(self, gax_api, client):
         self._gax_api = gax_api
+        self._client = client
 
     def list_subscriptions(self, project, page_size=0, page_token=None):
         """List subscriptions for the project associated with this API.
@@ -250,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 = [_subscription_pb_to_mapping(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_sub_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):
@@ -313,7 +315,7 @@ def subscription_create(self, subscription_path, topic_path,
             if exc_to_code(exc.cause) == StatusCode.FAILED_PRECONDITION:
                 raise Conflict(topic_path)
             raise
-        return _subscription_pb_to_mapping(sub_pb)
+        return MessageToDict(sub_pb)
 
     def subscription_get(self, subscription_path):
         """API call:  retrieve a subscription
@@ -335,7 +337,7 @@ def subscription_get(self, subscription_path):
             if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
                 raise NotFound(subscription_path)
             raise
-        return _subscription_pb_to_mapping(sub_pb)
+        return MessageToDict(sub_pb)
 
     def subscription_delete(self, subscription_path):
         """API call:  delete a subscription
@@ -474,24 +476,6 @@ def _message_pb_from_mapping(message):
                          attributes=message['attributes'])
 
 
-def _subscription_pb_to_mapping(sub_pb):
-    """Helper for :meth:`list_subscriptions`, et aliae
-
-    Performs "impedance matching" between the protobuf attrs and the keys
-    expected in the JSON API.
-    """
-    mapping = {
-        'name': sub_pb.name,
-        'topic': sub_pb.topic,
-        'ackDeadlineSeconds': sub_pb.ack_deadline_seconds,
-    }
-    if sub_pb.push_config.push_endpoint != '':
-        mapping['pushConfig'] = {
-            'pushEndpoint': sub_pb.push_config.push_endpoint,
-        }
-    return mapping
-
-
 def _message_pb_to_mapping(message_pb):
     """Helper for :meth:`pull`, et aliae
 
@@ -576,7 +560,7 @@ def _item_to_topic(iterator, resource):
         {'name': resource.name}, iterator.client)
 
 
-def _item_to_subscription(iterator, subscription_path):
+def _item_to_subscription_for_topic(iterator, subscription_path):
     """Convert a subscription name to the native object.
 
     :type iterator: :class:`~google.cloud.iterator.Iterator`
@@ -591,3 +575,33 @@ def _item_to_subscription(iterator, subscription_path):
     subscription_name = subscription_name_from_path(
         subscription_path, iterator.client.project)
     return Subscription(subscription_name, iterator.topic)
+
+
+def _item_to_sub_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)

pubsub/google/cloud/pubsub/client.py

@@ -22,7 +22,6 @@
 from google.cloud.pubsub.connection import _PublisherAPI as JSONPublisherAPI
 from google.cloud.pubsub.connection import _SubscriberAPI as JSONSubscriberAPI
 from google.cloud.pubsub.connection import _IAMPolicyAPI
-from google.cloud.pubsub.subscription import Subscription
 from google.cloud.pubsub.topic import Topic
 
 try:
@@ -98,9 +97,9 @@ def subscriber_api(self):
         if self._subscriber_api is None:
             if self._use_gax:
                 generated = make_gax_subscriber_api(self.connection)
-                self._subscriber_api = GAXSubscriberAPI(generated)
+                self._subscriber_api = GAXSubscriberAPI(generated, self)
             else:
-                self._subscriber_api = JSONSubscriberAPI(self.connection)
+                self._subscriber_api = JSONSubscriberAPI(self)
         return self._subscriber_api
 
     @property
@@ -160,20 +159,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.

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
@@ -238,7 +239,8 @@ def topic_list_subscriptions(self, topic, page_size=None, page_token=None):
 
         iterator = HTTPIterator(
             client=self._client, path=path,
-            item_to_value=_item_to_subscription, items_key='subscriptions',
+            item_to_value=_item_to_subscription_for_topic,
+            items_key='subscriptions',
             page_token=page_token, extra_params=extra_params)
         iterator.topic = topic
         return iterator
@@ -247,12 +249,13 @@ def topic_list_subscriptions(self, topic, page_size=None, page_token=None):
 class _SubscriberAPI(object):
     """Helper mapping subscriber-related APIs.
 
-    :type connection: :class:`Connection`
-    :param connection: the connection used to make API requests.
+    :type client: :class:`~google.cloud.pubsub.client.Client`
+    :param client: the client used to make API requests.
     """
 
-    def __init__(self, connection):
-        self._connection = connection
+    def __init__(self, client):
+        self._client = client
+        self._connection = client.connection
 
     def list_subscriptions(self, project, page_size=None, page_token=None):
         """API call:  list subscriptions for a given project
@@ -272,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_sub_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):
@@ -590,7 +595,7 @@ def _item_to_topic(iterator, resource):
     return Topic.from_api_repr(resource, iterator.client)
 
 
-def _item_to_subscription(iterator, subscription_path):
+def _item_to_subscription_for_topic(iterator, subscription_path):
     """Convert a subscription name to the native object.
 
     :type iterator: :class:`~google.cloud.iterator.Iterator`
@@ -605,3 +610,32 @@ def _item_to_subscription(iterator, subscription_path):
     subscription_name = subscription_name_from_path(
         subscription_path, iterator.client.project)
     return Subscription(subscription_name, iterator.topic)
+
+
+def _item_to_sub_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)