diff --git a/lib/google/apis/core/base_service.rb b/lib/google/apis/core/base_service.rb index 79c631c7a..f692cd17e 100644 --- a/lib/google/apis/core/base_service.rb +++ b/lib/google/apis/core/base_service.rb @@ -28,6 +28,43 @@ require 'hurley/addressable' 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 @@ module Google # 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 @@ module Google # 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 @@ module Google 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. diff --git a/spec/google/apis/core/service_spec.rb b/spec/google/apis/core/service_spec.rb index bd79b8c1f..093fb5e1a 100644 --- a/spec/google/apis/core/service_spec.rb +++ b/spec/google/apis/core/service_spec.rb @@ -17,6 +17,7 @@ require 'google/apis/options' 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 @@ EOF 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