Merge pull request #619 from estolfo/RUBY-909-parallel-scan

RUBY-909 Add #parallel_scan back onto Collection::View

Author: estolfo

lib/mongo/collection/view/readable.rb

@@ -53,6 +53,35 @@ def aggregate(pipeline, options = {})
           Aggregation.new(self, pipeline, options)
         end
 
+        # Execute a parallel scan on the collection view.
+        # Returns a list of up to cursor_count cursors that can be iterated concurrently.
+        # As long as the collection is not modified during scanning, each document appears once
+        # in one of the cursors' result sets.
+        #
+        # @example Execute a parallel collection scan.
+        #   view.parallel_scan(2)
+        #
+        # @param [ Integer ] cursor_count The max number of cursors to return.
+        #
+        # @return [ Array<Cursor> ] An array of cursors.
+        #
+        # @since 2.1
+        def parallel_scan(cursor_count)
+          server = read.select_server(cluster)
+          Operation::ParallelScan.new(
+            :coll_name => collection.name,
+            :db_name => database.name,
+            :cursor_count => cursor_count
+          ).execute(server.context).cursor_ids.map do |cursor_id|
+            result = Operation::Read::GetMore.new({ :to_return => 0,
+                                                    :cursor_id => cursor_id,
+                                                    :db_name   => database.name,
+                                                    :coll_name => collection.name
+              }).execute(server.context)
+            Cursor.new(self, result, server)
+          end
+        end
+
         # Allows the query to get partial results if some shards are down.
         #
         # @example Allow partial results.

lib/mongo/operation.rb

@@ -23,3 +23,4 @@
 require 'mongo/operation/command'
 require 'mongo/operation/kill_cursors'
 require 'mongo/operation/map_reduce'
+require 'mongo/operation/parallel_scan'

lib/mongo/operation/parallel_scan.rb

@@ -0,0 +1,75 @@
+# Copyright (C) 2009-2014 MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require 'mongo/operation/parallel_scan/result'
+
+module Mongo
+  module Operation
+
+    # A MongoDB parallel scan operation.
+    #
+    # @example Create the parallel scan operation.
+    #   ParallelScan.new({
+    #     :db_name  => 'test_db',
+    #     :coll_name = > 'test_collection',
+    #     :cursor_count => 5
+    #   })
+    #
+    # @param [ Hash ] spec The specifications for the operation.
+    #
+    # @option spec :db_name [ String ] The name of the database on which
+    #   the operation should be executed.
+    # @option spec :coll_name [ String ] The collection to scan.
+    # @option spec :cursor_count [ Integer ] The number of cursors to use.
+    # @option spec :options [ Hash ] Options for the command.
+    #
+    # @since 2.0.0
+    class ParallelScan
+      include Executable
+      include Specifiable
+      include Limited
+      include ReadPreferrable
+
+      # Execute the parallel scan operation.
+      #
+      # @example Execute the operation.
+      #   operation.execute(context)
+      #
+      # @params [ Mongo::Server::Context ] The context for this operation.
+      #
+      # @return [ Result ] The operation response, if there is one.
+      #
+      # @since 2.0.0
+      def execute(context)
+        execute_message(context)
+      end
+
+      private
+
+      def execute_message(context)
+        context.with_connection do |connection|
+          Result.new(connection.dispatch([ message(context) ])).validate!
+        end
+      end
+
+      def selector
+        { :parallelCollectionScan => coll_name, :numCursors => cursor_count }
+      end
+
+      def query_coll
+        Database::COMMAND
+      end
+    end
+  end
+end

lib/mongo/operation/parallel_scan/result.rb

@@ -0,0 +1,72 @@
+
+# Copyright (C) 2009-2014 MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  module Operation
+    class ParallelScan
+
+      # Defines custom behaviour of results in a parallel scan.
+      #
+      # @since 2.0.0
+      class Result < Operation::Result
+
+        # Get the name of the cursor field.
+        #
+        # @since 2.0.0
+        CURSOR = 'cursor'.freeze
+
+        # The name of the cursors field in the result.
+        #
+        # @since 2.0.0
+        CURSORS = 'cursors'.freeze
+
+        # Get the name of the id field.
+        #
+        # @since 2.0.0
+        ID = 'id'.freeze
+
+        # Get all the cursor ids from the result.
+        #
+        # @example Get the cursor ids.
+        #   result.cursor_ids
+        #
+        # @return [ Array<Integer> ] The cursor ids.
+        #
+        # @since 2.0.0
+        def cursor_ids
+          documents.map{ |doc| doc[CURSOR][ID] }
+        end
+
+        # Get the documents from parallel scan.
+        #
+        # @example Get the documents.
+        #   result.documents
+        #
+        # @return [ Array<BSON::Document> ] The documents.
+        #
+        # @since 2.0.0
+        def documents
+          reply.documents[0][CURSORS]
+        end
+
+        private
+
+        def first
+          @first ||= reply.documents[0] || {}
+        end
+      end
+    end
+  end
+end

spec/mongo/collection/view/readable_spec.rb

@@ -434,6 +434,41 @@
     end
   end
 
+  describe '#parallel_scan' do
+
+    let(:documents) do
+      (1..200).map do |i|
+        { name: "testing-scan-#{i}" }
+      end
+    end
+
+    before do
+      authorized_collection.insert_many(documents)
+    end
+
+    let(:cursors) do
+      view.parallel_scan(2)
+    end
+
+    it 'returns an array of cursors', if: write_command_enabled? do
+      cursors.each do |cursor|
+        expect(cursor.class).to be(Mongo::Cursor)
+      end
+    end
+
+    it 'returns the correct number of documents', if: write_command_enabled? do
+      expect(
+        cursors.reduce(0) { |total, cursor| total + cursor.to_a.size }
+      ).to eq(200)
+    end
+
+    it 'raises an error', unless: write_command_enabled? do
+      expect {
+        cursors
+      }.to raise_error(Mongo::Error::OperationFailure)
+    end
+  end
+
   describe '#projection' do
 
     let(:options) do