Add helper method for automatic paging

lib/google/apis/core/base_service.rb

@@ -28,6 +28,43 @@
 module Google
   module Apis
     module Core
+      # Helper class for enumerating over a result set requiring multiple fetches
+      class PagedResults
+        include Enumerable
+
+        attr_reader :last_result
+
+        # @param [BaseService] service
+        #   Current service instance
+        # @param [Fixnum] max
+        #   Maximum number of items to iterate over. Nil if no limit
+        # @param [Symbol] items
+        #   Name of the field in the result containing the items. Defaults to :items
+        def initialize(service, max: nil, items: :items, &block)
+          @service = service
+          @block = block
+          @max = max
+          @items_field = items
+        end
+
+        # Iterates over result set, fetching additional pages as needed
+        def each
+          page_token = nil
+          item_count = 0
+          loop do
+            @last_result = @block.call(page_token, @service)
+            for item in @last_result.send(@items_field)
+              item_count = item_count + 1
+              break if @max && item_count > @max
+              yield item
+            end
+            break if @max && item_count >= @max
+            break if @last_result.next_page_token.nil? || @last_result.next_page_token == page_token
+            page_token = @last_result.next_page_token
+          end
+        end
+      end
+
       # Base service for all APIs. Not to be used directly.
       #
       class BaseService
@@ -90,12 +127,12 @@ def authorization
         # request to the server.
         #
         # @example
-        #   service.batch do |service|
-        #     service.get_item(id1) do |res, err|
+        #   service.batch do |s|
+        #     s.get_item(id1) do |res, err|
         #       # process response for 1st call
         #     end
         #     # ...
-        #     service.get_item(idN) do |res, err|
+        #     s.get_item(idN) do |res, err|
         #       # process response for Nth call
         #     end
         #   end
@@ -122,12 +159,12 @@ def batch(options = nil)
         # files, use single requests which use a resumable upload protocol.
         #
         # @example
-        #   service.batch do |service|
-        #     service.insert_item(upload_source: 'file1.txt') do |res, err|
+        #   service.batch do |s|
+        #     s.insert_item(upload_source: 'file1.txt') do |res, err|
         #       # process response for 1st call
         #     end
         #     # ...
-        #     service.insert_item(upload_source: 'fileN.txt') do |res, err|
+        #     s.insert_item(upload_source: 'fileN.txt') do |res, err|
         #       # process response for Nth call
         #     end
         #   end
@@ -191,6 +228,34 @@ def http(method, url, params: nil, body: nil, download_dest: nil, options: nil,
           execute_or_queue_command(command, &block)
         end
 
+        # Executes a given query with paging, automatically retrieving
+        # additional pages as necessary. Requires a block that returns the
+        # result set of a page. The current page token is supplied as an argument
+        # to the block.
+        #
+        # Note: The returned enumerable also contains a `last_result` field
+        # containing the full result of the last query executed.
+        #
+        # @param [Fixnum] max
+        #   Maximum number of items to iterate over. Defaults to nil -- no upper bound.
+        # @param [Symbol] items
+        #   Name of the field in the result containing the items. Defaults to :items
+        # @return [Enumerble]
+        # @yield [token, service]
+        #   Current page token & service instance
+        # @yieldparam [String] token
+        #   Current page token to be used in the query
+        # @yieldparm [service]
+        #   Current service instance
+        #
+        # @example Retrieve files,
+        #   file_list = service.fetch_all { |token, s| s.list_files(page_token: token) }
+        #   file_list.each { |f| ... }
+        def fetch_all(max: nil, items: :items, &block)
+          fail "fetch_all may no be used inside a batch" if batch?
+          return PagedResults.new(self, max: max, items: items, &block)
+        end
+
         protected
 
         # Create a new upload command.

spec/google/apis/core/service_spec.rb

@@ -17,6 +17,7 @@
 require 'google/apis/core/base_service'
 require 'google/apis/core/json_representation'
 require 'hurley/test'
+require 'ostruct'
 
 RSpec.describe Google::Apis::Core::BaseService do
   include TestHelpers
@@ -243,5 +244,43 @@
         end
       end.to raise_error(Google::Apis::ClientError)
     end
+
+    context 'with fetch_all' do
+      let(:responses) do
+        data = {}
+        data[nil] = OpenStruct.new(next_page_token: 'p1', items: ['a', 'b', 'c'], alt_items: [1, 2 , 3])
+        data['p1'] = OpenStruct.new(next_page_token: 'p2', items: ['d', 'e', 'f'], alt_items: [4, 5, 6])
+        data['p2'] = OpenStruct.new(next_page_token: nil, items: ['g', 'h', 'i'], alt_items: [7,8, 9])
+        data
+      end
+
+      let(:items) { service.fetch_all { |token| responses[token] } }
+
+      it 'should fetch pages until next page token is nil' do
+        expect(items).to contain_exactly('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i')
+      end
+
+      it 'should stop on repeated page token' do
+        responses['p2'].next_page_token = 'p2'
+        expect(items).to contain_exactly('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i')
+      end
+
+      it 'should allow selecting another field for items' do
+        expect(service.fetch_all(items: :alt_items) { |token| responses[token] } ).to contain_exactly(1, 2, 3, 4, 5, 6, 7, 8, 9)
+      end
+
+      it 'should allow limiting the number of items to fetch' do
+        expect(service.fetch_all(max: 5) { |token| responses[token] } ).to contain_exactly('a', 'b', 'c', 'd', 'e')
+      end
+
+      it 'should yield the next token' do
+        expect do |b|
+          service.fetch_all do |token|
+            b.to_proc.call(token)
+            responses[token]
+          end.count
+        end.to yield_successive_args(nil, 'p1', 'p2')
+      end
+    end
   end
 end