[STORE] Added the find_in_batches method for models

    Person.find_in_batches { |batch| puts batch.map(&:name) }

See: http://api.rubyonrails.org/classes/ActiveRecord/Batches.html#method-i-find_in_batches

Author: karmi

elasticsearch-persistence/lib/elasticsearch/persistence/model/find.rb

@@ -5,7 +5,7 @@ module Model
       module Find
         module ClassMethods
 
-          # Return all documents (up to 10,000) for a model
+          # Returns all models (up to 10,000)
           #
           # @example Retrieve all people
           #
@@ -20,6 +20,81 @@ module ClassMethods
           def all(options={})
             gateway.search( { query: { match_all: {} }, size: 10_000 }.merge(options) )
           end
+
+          # Returns all models efficiently via the Elasticsearch's scan/scroll API
+          #
+          # You can restrict the models being returned with a query.
+          #
+          # The full {Persistence::Repository::Response::Result} instance is yielded to the passed
+          # block in each batch, so you can access any of its properties. Calling `to_a` will
+          # convert the object to an Array of model instances.
+          #
+          # @example Return all models in batches of 20 x number of primary shards
+          #
+          #     Person.find_in_batches { |batch| puts batch.map(&:name) }
+          #
+          # @example Return all models in batches of 100 x number of primary shards
+          #
+          #     Person.find_in_batches(size: 100) { |batch| puts batch.map(&:name) }
+          #
+          # @example Return all models matching a specific query
+          #
+          #      Person.find_in_batches(query: { match: { name: 'test' } }) { |batch| puts batch.map(&:name) }
+          #
+          # @example Return all models, fetching only the `name` attribute from Elasticsearch
+          #
+          #      Person.find_in_batches( _source_include: 'name') { |_| puts _.response.hits.hits.map(&:to_hash) }
+          #
+          # @return [String] The `scroll_id` for the request
+          #
+          def find_in_batches(options={}, &block)
+            search_params = options.extract!(
+              :index,
+              :type,
+              :scroll,
+              :size,
+              :explain,
+              :ignore_indices,
+              :ignore_unavailable,
+              :allow_no_indices,
+              :expand_wildcards,
+              :preference,
+              :q,
+              :routing,
+              :source,
+              :_source,
+              :_source_include,
+              :_source_exclude,
+              :stats,
+              :timeout)
+
+            scroll = search_params.delete(:scroll) || '5m'
+
+            body = options
+
+            # Get the initial scroll_id
+            #
+            response = gateway.client.search( { index: gateway.index_name,
+                                         type:  gateway.document_type,
+                                         search_type: 'scan',
+                                         scroll:      scroll,
+                                         size:        20,
+                                         body:        body }.merge(search_params) )
+
+            # Get the initial batch of documents
+            #
+            response = gateway.client.scroll( { scroll_id: response['_scroll_id'], scroll: scroll } )
+
+            # Break when receiving an empty array of hits
+            #
+            while response['hits']['hits'].any? do
+              yield Repository::Response::Results.new(gateway, response)
+
+              response = gateway.client.scroll( { scroll_id: response['_scroll_id'], scroll: scroll } )
+            end
+
+            return response['_scroll_id']
+          end
         end
       end
 

elasticsearch-persistence/test/integration/model/model_basic_test.rb

@@ -26,6 +26,10 @@ class ::Person
       end
 
       context "A basic persistence model" do
+        setup do
+          Person.gateway.create_index! force: true
+        end
+
         should "save and find the object" do
           person = Person.new name: 'John Smith', birthday: Date.parse('1970-01-01')
           person.save
@@ -94,6 +98,20 @@ class ::Person
 
           assert people.map_with_hit { |o,h| h._score }.all? { |s| s > 0 }
         end
+
+        should "find instances in batches" do
+          100.times { |i| Person.create name: "John #{i+1}" }
+          Person.gateway.refresh_index!
+
+          @results = []
+
+          Person.find_in_batches(_source_include: 'name') do |batch|
+            @results += batch.map(&:name)
+          end
+
+          assert_equal 100, @results.size
+          assert_contains @results, 'John 1'
+        end
       end
 
     end

elasticsearch-persistence/test/unit/model_find_test.rb

@@ -34,7 +34,7 @@ def set_id(id); self.id = id; end
     end
 
     setup do
-      @gateway         = stub
+      @gateway         = stub(client: stub(), index_name: 'foo', document_type: 'bar')
       DummyFindModel.stubs(:gateway).returns(@gateway)
 
       @response = MultiJson.load <<-JSON
@@ -83,5 +83,35 @@ def set_id(id); self.id = id; end
       DummyFindModel.all( { query: { match: { title: 'test' } }, routing: 'abc123' } )
     end
 
+    should "find all records in batches" do
+      @gateway
+        .expects(:deserialize)
+        .with('_source' => {'foo' => 'bar'})
+        .returns('_source' => {'foo' => 'bar'})
+
+      @gateway.client
+        .expects(:search)
+        .with do |arguments|
+          assert_equal 'scan', arguments[:search_type]
+          assert_equal 'foo',  arguments[:index]
+          assert_equal 'bar',  arguments[:type]
+        end
+        .returns(MultiJson.load('{"_scroll_id":"abc123==", "hits":{"hits":[]}}'))
+
+      @gateway.client
+        .expects(:scroll)
+        .twice
+        .returns(MultiJson.load('{"_scroll_id":"abc456==", "hits":{"hits":[{"_source":{"foo":"bar"}}]}}'))
+        .then
+        .returns(MultiJson.load('{"_scroll_id":"abc789==", "hits":{"hits":[]}}'))
+
+      @doc = nil
+
+      result = DummyFindModel.find_in_batches { |batch| @doc = batch.first['_source']['foo'] }
+
+      assert_equal 'abc789==', result
+      assert_equal 'bar',      @doc
+    end
+
   end
 end