Merge Pull Request #16445

6 files changed, 35 insertions, 5 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 997abcb48d..3fa24f3837 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,10 @@
+*   Fix a bug where counter_cache doesn't always work with  polymorphic
+    relations.
+
+    Fixes #16407.
+
+    *Stefan Kanev & Sean Griffin*
+
 *   Ensure that cyclic associations with autosave don't cause duplicate errors
     to be added to the parent record.
 
diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 97eb007f62..6e4a53f7fb 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -33,16 +33,24 @@ module ActiveRecord::Associations::Builder
 
           if (@_after_create_counter_called ||= false)
             @_after_create_counter_called = false
-          elsif attribute_changed?(foreign_key) && !new_record? && reflection.constructable?
-            model           = reflection.klass
+          elsif attribute_changed?(foreign_key) && !new_record?
+            if reflection.polymorphic?
+              model     = attribute(reflection.foreign_type).try(:constantize)
+              model_was = attribute_was(reflection.foreign_type).try(:constantize)
+            else
+              model     = reflection.klass
+              model_was = reflection.klass
+            end
+
             foreign_key_was = attribute_was foreign_key
             foreign_key     = attribute foreign_key
 
             if foreign_key && model.respond_to?(:increment_counter)
               model.increment_counter(cache_column, foreign_key)
             end
-            if foreign_key_was && model.respond_to?(:decrement_counter)
-              model.decrement_counter(cache_column, foreign_key_was)
+
+            if foreign_key_was && model_was.respond_to?(:decrement_counter)
+              model_was.decrement_counter(cache_column, foreign_key_was)
             end
           end
         end
diff --git a/activerecord/test/cases/counter_cache_test.rb b/activerecord/test/cases/counter_cache_test.rb
index 1f5055b2a2..922cb59280 100644
--- a/activerecord/test/cases/counter_cache_test.rb
+++ b/activerecord/test/cases/counter_cache_test.rb
@@ -1,6 +1,7 @@
 require 'cases/helper'
 require 'models/topic'
 require 'models/car'
+require 'models/aircraft'
 require 'models/wheel'
 require 'models/engine'
 require 'models/reply'
@@ -198,4 +199,16 @@ class CounterCacheTest < ActiveRecord::TestCase
     assert_equal 2, car.engines_count
     assert_equal 2, car.reload.engines_count
   end
+
+  test "update counters in a polymorphic relationship" do
+    aircraft = Aircraft.create!
+
+    assert_difference 'aircraft.reload.wheels_count' do
+      aircraft.wheels << Wheel.create!
+    end
+
+    assert_difference 'aircraft.reload.wheels_count', -1 do
+      aircraft.wheels.first.destroy
+    end
+  end
 end
diff --git a/activerecord/test/cases/locking_test.rb b/activerecord/test/cases/locking_test.rb
index dbdcc84b7d..2e1363334d 100644
--- a/activerecord/test/cases/locking_test.rb
+++ b/activerecord/test/cases/locking_test.rb
@@ -270,7 +270,7 @@ class OptimisticLockingTest < ActiveRecord::TestCase
       car.wheels << Wheel.create!
     end
     assert_difference 'car.wheels.count', -1  do
-      car.destroy
+      car.reload.destroy
     end
     assert car.destroyed?
   end
diff --git a/activerecord/test/models/aircraft.rb b/activerecord/test/models/aircraft.rb
index 1f35ef45da..c4404a8094 100644
--- a/activerecord/test/models/aircraft.rb
+++ b/activerecord/test/models/aircraft.rb
@@ -1,4 +1,5 @@
 class Aircraft < ActiveRecord::Base
   self.pluralize_table_names = false
   has_many :engines, :foreign_key => "car_id"
+  has_many :wheels, as: :wheelable
 end
diff --git a/activerecord/test/schema/schema.rb b/activerecord/test/schema/schema.rb
index 2fc806f181..7bab675b2a 100644
--- a/activerecord/test/schema/schema.rb
+++ b/activerecord/test/schema/schema.rb
@@ -37,6 +37,7 @@ ActiveRecord::Schema.define do
 
   create_table :aircraft, force: true do |t|
     t.string :name
+    t.integer :wheels_count, default: 0, null: false
   end
 
   create_table :articles, force: true do |t|