More gax paging fixes

gcloud/pubsub/_gax.py

@@ -46,7 +46,7 @@ class _PublisherAPI(object):
     def __init__(self, gax_api):
         self._gax_api = gax_api
 
-    def list_topics(self, project, page_token=None):
+    def list_topics(self, project, page_size=0, page_token=None):
         """List topics for the project associated with this API.
 
         See:
@@ -55,6 +55,10 @@ def list_topics(self, project, page_token=None):
         :type project: string
         :param project: project ID
 
+        :type page_size: int
+        :param page_size: maximum number of topics to return, If not passed,
+                          defaults to a value set by the API.
+
         :type page_token: string
         :param page_token: opaque marker for the next "page" of topics. If not
                            passed, the API will return the first page of
@@ -68,9 +72,11 @@ def list_topics(self, project, page_token=None):
         """
         options = _build_paging_options(page_token)
         path = 'projects/%s' % (project,)
-        response = self._gax_api.list_topics(path, options)
-        topics = [{'name': topic_pb.name} for topic_pb in response.topics]
-        return topics, response.next_page_token
+        page_iter = self._gax_api.list_topics(
+            path, page_size=page_size, options=options)
+        topics = [{'name': topic_pb.name} for topic_pb in page_iter.next()]
+        token = page_iter.page_token or None
+        return topics, token
 
     def topic_create(self, topic_path):
         """API call:  create a topic
@@ -159,14 +165,17 @@ def topic_publish(self, topic_path, messages):
         message_pbs = [_message_pb_from_dict(message)
                        for message in messages]
         try:
-            response = self._gax_api.publish(topic_path, message_pbs)
+            event = self._gax_api.publish(topic_path, message_pbs)
+            if not event.is_set():
+                event.wait()
         except GaxError as exc:
             if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
                 raise NotFound(topic_path)
             raise
-        return response.message_ids
+        return event.result.message_ids
 
-    def topic_list_subscriptions(self, topic_path, page_token=None):
+    def topic_list_subscriptions(self, topic_path, page_size=0,
+                                 page_token=None):
         """API call:  list subscriptions bound to a topic
 
         See:
@@ -176,6 +185,10 @@ def topic_list_subscriptions(self, topic_path, page_token=None):
         :param topic_path: fully-qualified path of the topic, in format
                             ``projects/<PROJECT>/topics/<TOPIC_NAME>``.
 
+        :type page_size: int
+        :param page_size: maximum number of subscriptions to return, If not
+                          passed, defaults to a value set by the API.
+
         :type page_token: string
         :param page_token: opaque marker for the next "page" of subscriptions.
                            If not passed, the API will return the first page
@@ -189,15 +202,15 @@ def topic_list_subscriptions(self, topic_path, page_token=None):
         """
         options = _build_paging_options(page_token)
         try:
-            response = self._gax_api.list_topic_subscriptions(
-                topic_path, options)
+            page_iter = self._gax_api.list_topic_subscriptions(
+                topic_path, page_size=page_size, options=options)
         except GaxError as exc:
             if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
                 raise NotFound(topic_path)
             raise
-        subs = [{'topic': topic_path, 'name': subscription}
-                for subscription in response.subscriptions]
-        return subs, response.next_page_token
+        subs = page_iter.next()
+        token = page_iter.page_token or None
+        return subs, token
 
 
 class _SubscriberAPI(object):
@@ -209,7 +222,7 @@ class _SubscriberAPI(object):
     def __init__(self, gax_api):
         self._gax_api = gax_api
 
-    def list_subscriptions(self, project, page_token=None):
+    def list_subscriptions(self, project, page_size=0, page_token=None):
         """List subscriptions for the project associated with this API.
 
         See:
@@ -218,6 +231,10 @@ def list_subscriptions(self, project, page_token=None):
         :type project: string
         :param project: project ID
 
+        :type page_size: int
+        :param page_size: maximum number of subscriptions to return, If not
+                          passed, defaults to a value set by the API.
+
         :type page_token: string
         :param page_token: opaque marker for the next "page" of subscriptions.
                            If not passed, the API will return the first page
@@ -231,10 +248,12 @@ def list_subscriptions(self, project, page_token=None):
         """
         options = _build_paging_options(page_token)
         path = 'projects/%s' % (project,)
-        response = self._gax_api.list_subscriptions(path, options)
+        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 response.subscriptions]
-        return subscriptions, response.next_page_token
+                         for sub_pb in page_iter.next()]
+        token = page_iter.page_token or None
+        return subscriptions, token
 
     def subscription_create(self, subscription_path, topic_path,
                             ack_deadline=None, push_endpoint=None):

gcloud/pubsub/connection.py

@@ -217,8 +217,8 @@ def topic_list_subscriptions(self, topic_path, page_size=None,
                            ``projects/<PROJECT>/topics/<TOPIC_NAME>``.
 
         :type page_size: int
-        :param page_size: maximum number of topics to return, If not passed,
-                          defaults to a value set by the API.
+        :param page_size: maximum number of subscriptions to return, If not
+                          passed, defaults to a value set by the API.
 
         :type page_token: string
         :param page_token: opaque marker for the next "page" of topics. If not