MONGOID-5078 Fix shard_key_selector_in_db during post-persist callbacks - DRAFT

lib/mongoid/document.rb

@@ -29,7 +29,7 @@ module Document
     include Mongoid::Touchable::InstanceMethods
 
     attr_accessor :__selected_fields
-    attr_reader :new_record
+    attr_reader :new_record, :during_post_persist_callbacks
 
     included do
       Mongoid.register_model(self)
@@ -116,6 +116,7 @@ def identity
     def initialize(attrs = nil)
       @__parent = nil
       _building do
+        @during_post_persist_callbacks = false
         @new_record = true
         @attributes ||= {}
         apply_pre_processed_defaults
@@ -250,6 +251,14 @@ def becomes(klass)
       became
     end
 
+    def begin_post_persist_callbacks
+      @during_post_persist_callbacks = true
+    end
+
+    def end_post_persist_callbacks
+      @during_post_persist_callbacks = false
+    end
+
     private
 
     # Returns the logger

lib/mongoid/persistable/creatable.rb

@@ -120,10 +120,16 @@ def prepare_insert(options = {})
         result = run_callbacks(:save) do
           run_callbacks(:create) do
             yield(self)
+            begin_post_persist_callbacks
             post_process_insert
           end
         end
-        post_process_persist(result, options) and self
+        begin
+          post_process_persist(result, options) and result
+          self
+        ensure
+          end_post_persist_callbacks
+        end
       end
 
       module ClassMethods

lib/mongoid/persistable/updatable.rb

@@ -112,10 +112,15 @@ def prepare_update(options = {})
         result = run_callbacks(:save) do
           run_callbacks(:update) do
             yield(self)
+            begin_post_persist_callbacks
             true
           end
         end
-        post_process_persist(result, options) and result
+        begin
+          post_process_persist(result, options) and result
+        ensure
+          end_post_persist_callbacks
+        end
       end
 
       # Update the document in the database.

lib/mongoid/persistable/upsertable.rb

@@ -48,10 +48,15 @@ def prepare_upsert(options = {})
         return false if performing_validations?(options) && invalid?(:upsert)
         result = run_callbacks(:upsert) do
           yield(self)
+          begin_post_persist_callbacks
           true
         end
-        self.new_record = false
-        post_process_persist(result, options) and result
+        begin
+          self.new_record = false
+          post_process_persist(result, options) and result
+        ensure
+          end_post_persist_callbacks
+        end
       end
     end
   end

lib/mongoid/shardable.rb

@@ -79,7 +79,7 @@ def shard_key_selector
     def shard_key_selector_in_db
       selector = {}
       shard_key_fields.each do |field|
-        selector[field.to_s] = new_record? ? send(field) : attribute_was(field)
+        selector[field.to_s] = new_record? || during_post_persist_callbacks ? send(field) : attribute_was(field)
       end
       selector
     end

spec/mongoid/persistable/upsertable_spec.rb

@@ -106,4 +106,104 @@
       end
     end
   end
+
+  describe "#during_post_persist" do
+
+    context "when registered model lifecycle callbacks directly/indirectly use Mongoid::Document.during_post_persist_callbacks " do
+
+      context "when the document is sharded" do
+
+        class ShardedProfile
+          include Mongoid::Document
+
+          attr_reader(
+            :before_upsert_called,
+            :before_upsert_val,
+            :after_upsert_called,
+            :after_upsert_val,
+            :around_upsert_called,
+            :around_upsert_val_pre_yield,
+            :around_upsert_val_post_yield,
+          )
+
+          field :name, type: String
+
+          shard_key :name
+
+          before_upsert :beforeUpsertMethod
+
+          around_upsert :aroundUpsertMethod
+
+          after_upsert :afterUpsertMethod
+
+          def beforeUpsertMethod
+            @before_upsert_called = true
+
+            @before_upsert_val = self.during_post_persist_callbacks
+          end
+
+          def aroundUpsertMethod
+            @around_upsert_called = true
+            @around_upsert_val_pre_yield = self.during_post_persist_callbacks
+            yield
+            @around_upsert_val_post_yield = self.during_post_persist_callbacks
+          end
+
+          def afterUpsertMethod
+            @after_upsert_called = true
+
+            @after_upsert_val = self.during_post_persist_callbacks
+          end
+        end
+
+        let!(:profile) do
+          ShardedProfile.create(name: "Alice")
+        end
+
+        context "when before_upsert " do
+
+          it "returns true" do
+            expect(profile.before_upsert_called).to be nil
+            expect(profile.before_upsert_val).to be nil
+
+            profile.name = "Bob"
+            profile.upsert
+
+            expect(profile.before_upsert_called).to be true
+            expect(profile.before_upsert_val).to be false
+          end
+        end
+
+        context "when around_upsert" do
+
+          it "returns true" do
+            expect(profile.around_upsert_called).to be nil
+            expect(profile.around_upsert_val_pre_yield).to be nil
+            expect(profile.around_upsert_val_post_yield).to be nil
+
+            profile.name = "Bob"
+            profile.upsert
+
+            expect(profile.around_upsert_called).to be true
+            expect(profile.around_upsert_val_pre_yield).to be false
+            expect(profile.around_upsert_val_post_yield).to be true
+          end
+        end
+
+        context "when after_upsert" do
+
+          it "returns true" do
+            expect(profile.after_upsert_called).to be nil
+            expect(profile.after_upsert_val).to be nil
+
+            profile.name = "Bob"
+            profile.upsert
+
+            expect(profile.after_upsert_called).to be true
+            expect(profile.after_upsert_val).to be true
+          end
+        end
+      end
+    end
+  end
 end

spec/mongoid/reloadable_spec.rb

@@ -7,6 +7,131 @@
 
   describe "#reload" do
 
+    context "with new during_post_persist_callbacks functionality" do
+
+      context "when using non-sharded documents" do
+
+        class NonShardedProfile
+          include Mongoid::Document
+
+          attr_reader :after_save_count, :after_save_name_val
+
+          field :name, type: String
+
+          after_save :afterSaveMethod
+
+          def afterSaveMethod
+            @after_save_count = @after_save_count ? @after_save_count + 1 : 1
+
+            self.reload
+
+            @after_save_name_val = self.name
+          end
+        end
+
+        context "when using reload during a post-persist callback" do
+
+          context "when document is not yet persisted" do
+
+            context "when after_save" do
+              let(:profile) do
+                NonShardedProfile.new(name: "Alice")
+              end
+
+              it "reloads successfully" do
+                expect(profile.after_save_count).to be nil
+                expect(profile.after_save_name_val).to be nil
+                profile.name = "Bob"
+                profile.save
+                expect(profile.after_save_count).to eq(1)
+                expect(profile.after_save_name_val).to eq("Bob")
+              end
+            end
+          end
+
+          context "when document is already persisted" do
+
+            context "when after_save" do
+              let(:profile) do
+                NonShardedProfile.create(name: "Alice")
+              end
+
+              it "reloads successfully" do
+                expect(profile.after_save_count).to eq(1)
+                expect(profile.after_save_name_val).to eq("Alice")
+                profile.name = "Bob"
+                profile.save
+                expect(profile.after_save_count).to eq(2)
+                expect(profile.after_save_name_val).to eq("Bob")
+              end
+            end
+          end
+        end
+      end
+
+      context "when using sharded documents" do
+
+        class ShardedProfile
+          include Mongoid::Document
+
+          attr_reader :after_save_count, :after_save_name_val
+
+          field :name, type: String
+
+          shard_key :name
+
+          after_save :afterSaveMethod
+
+          def afterSaveMethod
+            @after_save_count = @after_save_count ? @after_save_count + 1 : 1
+
+            self.reload
+
+            @after_save_name_val = self.name
+          end
+        end
+
+        context "when using reload during a post-persist callback" do
+
+          context "when document is not yet persisted" do
+
+            context "when after_save" do
+              let(:profile) do
+                ShardedProfile.new(name: "Alice")
+              end
+
+              it "reloads successfully" do
+                expect(profile.after_save_count).to be nil
+                expect(profile.after_save_name_val).to be nil
+                profile.name = "Bob"
+                profile.save
+                expect(profile.after_save_count).to eq(1)
+                expect(profile.after_save_name_val).to eq("Bob")
+              end
+            end
+          end
+
+          context "when document is already persisted" do
+
+            context "when after_save" do
+              let(:profile) do
+                ShardedProfile.create(name: "Alice")
+              end
+
+              it "reloads successfully" do
+                expect(profile.after_save_count).to eq(1)
+                expect(profile.after_save_name_val).to eq("Alice")
+                profile.name = "Bob"
+                profile.save
+                expect(profile.after_save_count).to eq(2)
+                expect(profile.after_save_name_val).to eq("Bob")
+              end
+            end
+          end
+        end
+      end
+    end
+
     context 'when persistence options are set' do
 
       let(:person) do