[MODEL] Support :scope, :query and :preprocess importing options for Mongoid

elasticsearch-model/lib/elasticsearch/model/adapters/mongoid.rb

@@ -63,10 +63,17 @@ module Importing
           # @see https://github.com/karmi/retire/pull/724
           #
           def __find_in_batches(options={}, &block)
-            options[:batch_size] ||= 1_000
+            batch_size = options[:batch_size] || 1_000
+            query = options[:query]
+            named_scope = options[:scope]
+            preprocess = options[:preprocess]
+
+            scope = all
+            scope = scope.send(named_scope) if named_scope
+            scope = query.is_a?(Proc) ? scope.class_exec(&query) : scope.where(query) if query
 
-            all.no_timeout.each_slice(options[:batch_size]) do |items|
-              yield items
+            scope.no_timeout.each_slice(batch_size) do |items|
+              yield (preprocess ? self.__send__(preprocess, items) : items)
             end
           end
 

elasticsearch-model/test/integration/mongoid_basic_test.rb

@@ -168,6 +168,70 @@ def as_indexed_json(options={})
               assert response.results.any?, "Search has not returned results: #{response.to_a}"
             end
           end
+
+          context "importing when the model has a default scope" do
+            class ::MongoidArticleWithDefaultScope
+              include Mongoid::Document
+              include Elasticsearch::Model
+
+              default_scope -> { where(status: 'active') }
+
+              field :id, type: String
+              field :title, type: String
+              field :status, type: String, default: 'active'
+
+              attr_accessible :title if respond_to? :attr_accessible
+              attr_accessible :status if respond_to? :attr_accessible
+
+              settings index: { number_of_shards: 1, number_of_replicas: 0 } do
+                mapping do
+                  indexes :title,      type: 'text', analyzer: 'snowball'
+                  indexes :status,     type: 'text'
+                  indexes :created_at, type: 'date'
+                end
+              end
+
+              def as_indexed_json(options={})
+                as_json(except: [:id, :_id])
+              end
+            end
+
+            setup do
+              Elasticsearch::Model::Adapter.register \
+                Elasticsearch::Model::Adapter::Mongoid,
+                lambda { |klass| !!defined?(::Mongoid::Document) && klass.respond_to?(:ancestors) && klass.ancestors.include?(::Mongoid::Document) }
+
+              MongoidArticleWithDefaultScope.__elasticsearch__.create_index! force: true
+
+              MongoidArticleWithDefaultScope.delete_all
+
+              MongoidArticleWithDefaultScope.create! title: 'Test'
+              MongoidArticleWithDefaultScope.create! title: 'Testing Coding'
+              MongoidArticleWithDefaultScope.create! title: 'Coding'
+              MongoidArticleWithDefaultScope.create! title: 'Test legacy code', status: 'removed'
+
+              MongoidArticleWithDefaultScope.__elasticsearch__.refresh_index!
+              MongoidArticleWithDefaultScope.__elasticsearch__.client.cluster.health wait_for_status: 'yellow'
+            end
+
+            should "import only documents from the default scope" do
+              assert_equal 3, MongoidArticleWithDefaultScope.count
+
+              assert_equal 0, MongoidArticleWithDefaultScope.import
+
+              MongoidArticleWithDefaultScope.__elasticsearch__.refresh_index!
+              assert_equal 3, MongoidArticleWithDefaultScope.search('*').results.total
+            end
+
+            should "import only documents from a specific query combined with the default scope" do
+              assert_equal 3, MongoidArticleWithDefaultScope.count
+
+              assert_equal 0, MongoidArticleWithDefaultScope.import(query: -> { where(title: /^Test/) })
+
+              MongoidArticleWithDefaultScope.__elasticsearch__.refresh_index!
+              assert_equal 2, MongoidArticleWithDefaultScope.search('*').results.total
+            end
+          end
         end
 
       end

elasticsearch-model/test/unit/adapter_mongoid_test.rb

@@ -98,7 +98,64 @@ def ids
           assert_equal @transform.call(model), { index: { _id: "1", data: {} } }
         end
       end
-    end
 
+      should "limit the relation to a specific scope" do
+        relation = mock()
+        relation.stubs(:no_timeout).returns(relation)
+        relation.expects(:published).returns(relation)
+        relation.expects(:each_slice).returns([])
+        DummyClassForMongoid.expects(:all).returns(relation)
+
+        DummyClassForMongoid.__send__ :extend, Elasticsearch::Model::Adapter::Mongoid::Importing
+        DummyClassForMongoid.__find_in_batches(scope: :published) do; end
+      end
+
+      context "when limit the relation with proc" do
+        setup do
+          @query = Proc.new { where(color: "red") }
+        end
+        should "query with a specific criteria" do
+          relation = mock()
+          relation.stubs(:no_timeout).returns(relation)
+          relation.expects(:class_exec).returns(relation)
+          relation.expects(:each_slice).returns([])
+          DummyClassForMongoid.expects(:all).returns(relation)
+
+          DummyClassForMongoid.__find_in_batches(query: @query) do; end
+        end
+      end
+
+      context "when limit the relation with hash" do
+        setup do
+          @query = { color: "red" }
+        end
+        should "query with a specific criteria" do
+          relation = mock()
+          relation.stubs(:no_timeout).returns(relation)
+          relation.expects(:where).with(@query).returns(relation)
+          relation.expects(:each_slice).returns([])
+          DummyClassForMongoid.expects(:all).returns(relation)
+
+          DummyClassForMongoid.__find_in_batches(query: @query) do; end
+        end
+      end
+
+      should "preprocess the batch if option provided" do
+        class << DummyClassForMongoid
+          # Updates/transforms the batch while fetching it from the database
+          # (eg. with information from an external system)
+          #
+          def update_batch(batch)
+            batch.collect { |b| b.to_s + '!' }
+          end
+        end
+
+        DummyClassForMongoid.expects(:__find_in_batches).returns( [:a, :b] )
+
+        DummyClassForMongoid.__find_in_batches(preprocess: :update_batch) do |batch|
+          assert_same_elements ["a!", "b!"], batch
+        end
+      end
+    end
   end
 end