Merge pull request #23 from tjdett/seamless-resumption

Seamless resumption

Author: tjdett

lib/oai/client.rb

@@ -18,6 +18,7 @@
 require 'oai/client/record'
 require 'oai/client/identify'
 require 'oai/client/get_record'
+require 'oai/client/resumable'
 require 'oai/client/list_identifiers'
 require 'oai/client/list_metadata_formats'
 require 'oai/client/list_records'
@@ -118,30 +119,30 @@ def initialize(base_url, options={})
     # parser then you will get an XML::Node object instead.
 
     def identify
-      return OAI::IdentifyResponse.new(do_request('Identify'))
+      OAI::IdentifyResponse.new(do_request('Identify'))
     end
 
     # Equivalent to a ListMetadataFormats request. A ListMetadataFormatsResponse
     # object is returned to you.
 
     def list_metadata_formats(opts={})
-      return OAI::ListMetadataFormatsResponse.new(do_request('ListMetadataFormats', opts))
+      OAI::ListMetadataFormatsResponse.new(do_request('ListMetadataFormats', opts))
     end
 
     # Equivalent to a ListIdentifiers request. Pass in :from, :until arguments
     # as Date or DateTime objects as appropriate depending on the granularity
     # supported by the server.
 
     def list_identifiers(opts={})
-      return OAI::ListIdentifiersResponse.new(do_request('ListIdentifiers', opts))
+      do_resumable(OAI::ListIdentifiersResponse, 'ListIdentifiers', opts)
     end
 
     # Equivalent to a GetRecord request. You must supply an identifier
     # argument. You should get back a OAI::GetRecordResponse object
     # which you can extract a OAI::Record object from.
 
     def get_record(opts={})
-      return OAI::GetRecordResponse.new(do_request('GetRecord', opts))
+      OAI::GetRecordResponse.new(do_request('GetRecord', opts))
     end
 
     # Equivalent to the ListRecords request. A ListRecordsResponse
@@ -152,7 +153,7 @@ def get_record(opts={})
     #   end
 
     def list_records(opts={})
-      return OAI::ListRecordsResponse.new(do_request('ListRecords', opts))
+      do_resumable(OAI::ListRecordsResponse, 'ListRecords', opts)
     end
 
     # Equivalent to the ListSets request. A ListSetsResponse object
@@ -164,7 +165,7 @@ def list_records(opts={})
     #   end
 
     def list_sets(opts={})
-      return OAI::ListSetsResponse.new(do_request('ListSets', opts))
+      do_resumable(OAI::ListSetsResponse, 'ListSets', opts)
     end
 
     private
@@ -183,6 +184,13 @@ def do_request(verb, opts = nil)
       return load_document(xml)
     end
 
+    def do_resumable(responseClass, verb, opts)
+      responseClass.new(do_request(verb, opts)) do |response|
+        responseClass.new \
+          do_request(verb, :resumption_token => response.resumption_token)
+      end
+    end
+
     def build_uri(verb, opts)
       opts = validate_options(verb, opts)
       uri = @base.clone

lib/oai/client/list_identifiers.rb

@@ -1,6 +1,7 @@
 module OAI
   class ListIdentifiersResponse < Response
     include Enumerable
+    include OAI::Resumable
     include OAI::XPath
 
     def each

lib/oai/client/list_records.rb

@@ -7,10 +7,11 @@ module OAI
   #   end
   #
   # you'll need to handle resumption tokens
-  
+
   class ListRecordsResponse < Response
-    include OAI::XPath
     include Enumerable
+    include OAI::Resumable
+    include OAI::XPath
 
     def each
       for record_element in xpath_all(@doc, './/ListRecords/record')

lib/oai/client/list_sets.rb

@@ -5,10 +5,11 @@ module OAI
   #   for set in client.list_sets
   #     puts set
   #   end
-  
+
   class ListSetsResponse < Response
-    include OAI::XPath
     include Enumerable
+    include OAI::Resumable
+    include OAI::XPath
 
     def each
       for set_element in xpath_all(@doc, './/set')

lib/oai/client/response.rb

@@ -1,6 +1,6 @@
 module OAI
 
-  # An OAI::Response contains entries and a resumption token. If a resumption token is present, 
+  # An OAI::Response contains entries and a resumption token. If a resumption token is present,
   # then you must use it to fetch the rest of the entries for your query. For example:
   #  # List all records in a given set
   #  client = OAI::Client.new 'http://my-oai-provider.example.com/oai'
@@ -16,11 +16,12 @@ module OAI
 
   class Response
     include OAI::XPath
-    attr_reader :doc, :resumption_token
+    attr_reader :doc, :resumption_token, :resumption_block
 
-    def initialize(doc)
+    def initialize(doc, &resumption_block)
       @doc = doc
       @resumption_token = xpath(doc, './/resumptionToken')
+      @resumption_block = resumption_block
 
       # throw an exception if there was an error
       error = xpath_first(doc, './/error')

lib/oai/client/resumable.rb

@@ -0,0 +1,42 @@
+module OAI
+  module Resumable
+
+    class ResumptionWrapper
+      include Enumerable
+
+      def initialize(response)
+        @response = response
+        @resumption_block = response.resumption_block
+      end
+
+      def each(&block)
+        yield_from_response &block
+        while resumable?
+          @response = @resumption_block.call @response
+          yield_from_response &block
+        end
+      end
+
+      private
+
+      def yield_from_response(&block)
+        @response.each do |obj|
+          block.call(obj)
+        end
+      end
+
+      def resumable?
+        @response.resumption_token and not @response.resumption_token.empty?
+      end
+
+    end
+
+    def full
+      if @resumption_block.nil?
+        raise NotImplementedError.new("Resumption block not provided")
+      end
+      ResumptionWrapper.new(self)
+    end
+
+  end
+end

test/client/tc_list_identifiers.rb

@@ -3,10 +3,10 @@
 class ListIdentifiersTest < Test::Unit::TestCase
 
   def test_list_with_resumption_token
-    client = OAI::Client.new 'http://localhost:3333/oai' 
+    client = OAI::Client.new 'http://localhost:3333/oai'
 
     # get a list of identifier headers
-    response = client.list_identifiers :metadata_prefix => 'oai_dc' 
+    response = client.list_identifiers :metadata_prefix => 'oai_dc'
     assert_kind_of OAI::ListIdentifiersResponse, response
     assert_kind_of OAI::Response, response
     assert response.entries.size > 0
@@ -27,6 +27,25 @@ def test_list_with_resumption_token
     assert_not_equal response.entries[0].identifier, first_identifier
   end
 
+  def test_list_full
+    client = OAI::Client.new 'http://localhost:3333/oai'
+
+    # get a list of identifier headers
+    response = client.list_identifiers :metadata_prefix => 'oai_dc'
+    assert_kind_of OAI::ListIdentifiersResponse, response
+    assert_kind_of OAI::Response, response
+    assert response.respond_to?(:full), "Should expose :full"
+
+    # Check that it runs through the pages
+    assert_equal 1150, response.full.count
+    response.full.each do |header|
+      assert_kind_of OAI::Header, header
+      assert header.identifier
+      assert header.datestamp
+      assert header.set_spec
+    end
+  end
+
   def test_list_with_date_range
     client = OAI::Client.new 'http://localhost:3333/oai'
     from_date = Date.new(1998,1,1)
@@ -48,5 +67,5 @@ def test_invalid_argument
     client = OAI::Client.new 'http://localhost:3333/oai'
     assert_raise(OAI::ArgumentException) {client.list_identifiers :foo => 'bar'}
   end
-  
+
 end

test/client/tc_list_records.rb

@@ -1,13 +1,26 @@
 require 'test_helper'
 
-class GetRecordsTest < Test::Unit::TestCase
+class ListRecordsTest < Test::Unit::TestCase
 
-  def test_get_records
+  def test_list
     client = OAI::Client.new 'http://localhost:3333/oai'
-    response = client.list_records 
+    response = client.list_records
     assert_kind_of OAI::ListRecordsResponse, response
     assert response.entries.size > 0
     assert_kind_of OAI::Record,  response.entries[0]
   end
-  
+
+  def test_list_full
+    client = OAI::Client.new 'http://localhost:3333/oai'
+
+    response = client.list_records
+    assert_kind_of OAI::ListRecordsResponse, response
+
+    # Check that it runs through the pages
+    assert_equal 1150, response.full.count
+    response.full.each do |record|
+      assert_kind_of OAI::Record, record
+    end
+  end
+
 end

test/client/tc_list_sets.rb

@@ -8,12 +8,27 @@ def test_list
     assert_kind_of OAI::ListSetsResponse, response
     assert response.entries.size > 0
     assert_kind_of OAI::Set,  response.entries[0]
-  
+
     # test iterator
     for set in response
       assert_kind_of OAI::Set, set
     end
   end
-  
+
+  def test_list_full
+    client = OAI::Client.new 'http://localhost:3333/oai'
+
+    response = client.list_sets
+    assert_kind_of OAI::ListSetsResponse, response
+    assert_kind_of OAI::Response, response
+    assert response.respond_to?(:full), "Should expose :full"
+
+    # This won't page, but it should work anyway
+    assert_equal 6, response.full.count
+    response.full.each do |set|
+      assert_kind_of OAI::Set, set
+    end
+  end
+
 end