Merge pull request #593 from estolfo/RUBY-882-id

RUBY 882 Add id to document before inserting if it's missing one

Author: durran

lib/mongo/operation/result.rb

@@ -119,7 +119,7 @@ def each(&block)
         documents.each(&block)
       end
 
-      # Initialize a new result result.
+      # Initialize a new result.
       #
       # @example Instantiate the result.
       #   Result.new(replies)

lib/mongo/operation/write.rb

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/operation/write/idable'
 require 'mongo/operation/write/bulk'
 require 'mongo/operation/write/delete'
 require 'mongo/operation/write/insert'

lib/mongo/operation/write/bulk/bulk_insert.rb

@@ -47,6 +47,7 @@ module Write
       # @since 2.0.0
       class BulkInsert
         include Specifiable
+        include Idable
 
         # Execute the bulk insert operation.
         #
@@ -69,21 +70,22 @@ def execute(context)
         private
 
         def execute_write_command(context)
-          Result.new(Command::Insert.new(spec).execute(context))
+          command_spec = spec.merge(:documents => ensure_ids(documents))
+          Result.new(Command::Insert.new(command_spec).execute(context), @ids)
         end
 
         def execute_message(context)
           replies = []
           messages.map do |m|
             context.with_connection do |connection|
-              result = LegacyResult.new(connection.dispatch([ m, gle ].compact))
+              result = LegacyResult.new(connection.dispatch([ m, gle ].compact), @ids)
               replies << result.reply
               if stop_sending?(result)
-                return LegacyResult.new(replies)
+                return LegacyResult.new(replies, @ids)
               end
             end
           end
-          LegacyResult.new(replies.compact.empty? ? nil : replies)
+          LegacyResult.new(replies.compact.empty? ? nil : replies, @ids)
         end
 
         def stop_sending?(result)
@@ -117,10 +119,10 @@ def gle
         def messages
           if ordered? || gle
             documents.collect do |doc|
-              Protocol::Insert.new(db_name, coll_name, [ doc ], options)
+              Protocol::Insert.new(db_name, coll_name, ensure_ids([ doc ]), options)
             end
           else
-            [ Protocol::Insert.new(db_name, coll_name, documents, { :flags => [:continue_on_error] }) ]
+            [ Protocol::Insert.new(db_name, coll_name, ensure_ids(documents), { :flags => [:continue_on_error] }) ]
           end
         end
       end

lib/mongo/operation/write/bulk/bulk_insert/result.rb

@@ -26,6 +26,25 @@ class BulkInsert
         class Result < Operation::Result
           include BulkMergable
 
+          # Get the ids of the inserted documents.
+          #
+          # @since 2.0.0
+          attr_reader :inserted_ids
+
+          # Initialize a new result.
+          #
+          # @example Instantiate the result.
+          #   Result.new(replies, inserted_ids)
+          #
+          # @param [ Protocol::Reply ] replies The wire protocol replies.
+          # @params [ Array<Object> ] ids The ids of the inserted documents.
+          #
+          # @since 2.0.0
+          def initialize(replies, ids)
+            @replies = replies.is_a?(Protocol::Reply) ? [ replies ] : replies
+            @inserted_ids = ids
+          end
+
           # Gets the number of documents inserted.
           #
           # @example Get the number of documents inserted.
@@ -37,6 +56,18 @@ class Result < Operation::Result
           def n_inserted
             written_count
           end
+
+          # Gets the id of the document inserted.
+          #
+          # @example Get id of the document inserted.
+          #   result.inserted_id
+          #
+          # @return [ Object ] The id of the document inserted.
+          #
+          # @since 2.0.0
+          def inserted_id
+            inserted_ids.first
+          end
         end
 
         # Defines custom behaviour of results when inserting.
@@ -46,6 +77,25 @@ def n_inserted
         class LegacyResult < Operation::Result
           include LegacyBulkMergable
 
+          # Get the ids of the inserted documents.
+          #
+          # @since 2.0.0
+          attr_reader :inserted_ids
+
+          # Initialize a new result.
+          #
+          # @example Instantiate the result.
+          #   Result.new(replies, inserted_ids)
+          #
+          # @param [ Protocol::Reply ] replies The wire protocol replies.
+          # @params [ Array<Object> ] ids The ids of the inserted documents.
+          #
+          # @since 2.0.0
+          def initialize(replies, ids)
+            @replies = replies.is_a?(Protocol::Reply) ? [ replies ] : replies
+            @inserted_ids = ids
+          end
+
           # Gets the number of documents inserted.
           #
           # @example Get the number of documents inserted.
@@ -61,6 +111,18 @@ def n_inserted
               n
             end
           end
+
+          # Gets the id of the document inserted.
+          #
+          # @example Get id of the document inserted.
+          #   result.inserted_id
+          #
+          # @return [ Object ] The id of the document inserted.
+          #
+          # @since 2.0.0
+          def inserted_id
+            inserted_ids.first
+          end
         end
       end
     end

lib/mongo/operation/write/idable.rb

@@ -0,0 +1,41 @@
+# Copyright (C) 2014-2015 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
+    module Write
+      module Idable
+
+        private
+
+        def id(doc)
+          doc.respond_to?(:id) ? doc.id : (doc['_id'] || doc[:_id])
+        end
+
+        def has_id?(doc)
+          id(doc)
+        end
+
+        def ensure_ids(documents)
+          @ids ||= []
+          documents.collect do |doc|
+            doc_with_id = has_id?(doc) ? doc : doc.merge(_id: BSON::ObjectId.new)
+            @ids << id(doc_with_id)
+            doc_with_id
+          end
+        end
+      end
+    end
+  end
+end

lib/mongo/operation/write/insert.rb

@@ -44,6 +44,7 @@ module Write
       class Insert
         include Executable
         include Specifiable
+        include Idable
 
         # Execute the insert operation.
         #
@@ -66,12 +67,13 @@ def execute(context)
         private
 
         def execute_write_command(context)
-          Result.new(Command::Insert.new(spec).execute(context)).validate!
+          command_spec = spec.merge(:documents => ensure_ids(documents))
+          Result.new(Command::Insert.new(command_spec).execute(context), @ids).validate!
         end
 
         def execute_message(context)
           context.with_connection do |connection|
-            Result.new(connection.dispatch([ message, gle ].compact)).validate!
+            Result.new(connection.dispatch([ message, gle ].compact), @ids).validate!
           end
         end
 
@@ -82,7 +84,7 @@ def initialize_copy(original)
 
         def message
           opts = !!options[:continue_on_error] ? { :flags => [:continue_on_error] } : {}
-          Protocol::Insert.new(db_name, coll_name, documents, opts)
+          Protocol::Insert.new(db_name, coll_name, ensure_ids(documents), opts)
         end
       end
     end

lib/mongo/operation/write/insert/result.rb

@@ -25,6 +25,36 @@ class Insert
         # @since 2.0.0
         class Result < Operation::Result
 
+          # Get the ids of the inserted documents.
+          #
+          # @since 2.0.0
+          attr_reader :inserted_ids
+
+          # Initialize a new result.
+          #
+          # @example Instantiate the result.
+          #   Result.new(replies, inserted_ids)
+          #
+          # @param [ Protocol::Reply ] replies The wire protocol replies.
+          # @params [ Array<Object> ] ids The ids of the inserted documents.
+          #
+          # @since 2.0.0
+          def initialize(replies, ids)
+            @replies = replies.is_a?(Protocol::Reply) ? [ replies ] : replies
+            @inserted_ids = ids
+          end
+
+          # Gets the id of the document inserted.
+          #
+          # @example Get id of the document inserted.
+          #   result.inserted_id
+          #
+          # @return [ Object ] The id of the document inserted.
+          #
+          # @since 2.0.0
+          def inserted_id
+            inserted_ids.first
+          end
         end
       end
     end

spec/mongo/operation/write/bulk_delete_spec.rb renamed to spec/mongo/operation/write/bulk/bulk_delete_spec.rb

@@ -19,6 +19,10 @@
     described_class.new(spec)
   end
 
+  after do
+    authorized_collection.find.delete_many
+  end
+
   describe '#initialize' do
 
     context 'spec' do
@@ -81,6 +85,29 @@
     end
   end
 
+  describe 'document ids' do
+
+    context 'when documents do not contain an id' do
+
+      let(:documents) do
+        [{ 'field' => 'test' },
+         { 'field' => 'test' }]
+      end
+
+      let(:inserted_ids) do
+        op.execute(authorized_primary.context).inserted_ids
+      end
+
+      let(:collection_ids) do
+        authorized_collection.find(field: 'test').collect { |d| d['_id'] }
+      end
+
+      it 'adds an id to the documents' do
+        expect(inserted_ids).to eq(collection_ids)
+      end
+    end
+  end
+
   describe '#execute' do
 
     before do