From 38fae1f2505179123ce59cf19a3f5571559ead35 Mon Sep 17 00:00:00 2001
From: Ryuta Kamizono <kamipo@gmail.com>
Date: Wed, 26 Sep 2018 22:47:49 +0900
Subject: [PATCH] Update counter cache in memory if parent target is existed

Fixes #19550.
---
 .../associations/builder/belongs_to.rb        | 13 ++------
 .../has_many_associations_test.rb             | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index da4cc343eb..374247ffec 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -31,24 +31,17 @@ module ActiveRecord::Associations::Builder # :nodoc:
 
       mixin.class_eval do
         def belongs_to_counter_cache_after_update(reflection)
-          foreign_key  = reflection.foreign_key
-          cache_column = reflection.counter_cache_column
-
           if association(reflection.name).target_changed?
             if reflection.polymorphic?
-              model     = attribute_in_database(reflection.foreign_type).try(:constantize)
               model_was = attribute_before_last_save(reflection.foreign_type).try(:constantize)
             else
-              model     = reflection.klass
               model_was = reflection.klass
             end
 
-            foreign_key_was = attribute_before_last_save foreign_key
-            foreign_key     = attribute_in_database foreign_key
+            foreign_key_was = attribute_before_last_save(reflection.foreign_key)
+            cache_column = reflection.counter_cache_column
 
-            if foreign_key && model < ActiveRecord::Base
-              counter_cache_target(reflection, model, foreign_key).update_counters(cache_column => 1)
-            end
+            association(reflection.name).increment_counters
 
             if foreign_key_was && model_was < ActiveRecord::Base
               counter_cache_target(reflection, model_was, foreign_key_was).update_counters(cache_column => -1)
diff --git a/activerecord/test/cases/associations/has_many_associations_test.rb b/activerecord/test/cases/associations/has_many_associations_test.rb
index a2f6174dc1..f8d27f1c56 100644
--- a/activerecord/test/cases/associations/has_many_associations_test.rb
+++ b/activerecord/test/cases/associations/has_many_associations_test.rb
@@ -1197,6 +1197,38 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
     assert_equal 2, topic.reload.replies.size
   end
 
+  def test_counter_cache_updates_in_memory_after_update_with_inverse_of_disabled
+    topic = Topic.create!(title: "Zoom-zoom-zoom")
+
+    assert_equal 0, topic.replies_count
+
+    reply1 = Reply.create!(title: "re: zoom", content: "speedy quick!")
+    reply2 = Reply.create!(title: "re: zoom 2", content: "OMG lol!")
+
+    assert_queries(6) do
+      topic.replies << [reply1, reply2]
+    end
+
+    assert_equal 2, topic.replies_count
+    assert_equal 2, topic.reload.replies_count
+  end
+
+  def test_counter_cache_updates_in_memory_after_update_with_inverse_of_enabled
+    category = Category.create!(name: "Counter Cache")
+
+    assert_nil category.categorizations_count
+
+    categorization1 = Categorization.create!
+    categorization2 = Categorization.create!
+
+    assert_queries(4) do
+      category.categorizations << [categorization1, categorization2]
+    end
+
+    assert_equal 2, category.categorizations_count
+    assert_equal 2, category.reload.categorizations_count
+  end
+
   def test_pushing_association_updates_counter_cache
     topic = Topic.order("id ASC").first
     reply = Reply.create!