Fix counter_cache for polymorphic associations

Also removes a false positive test that depends on the fixed bug:

At this time, counter_cache does not work with polymorphic relationships
(which is a bug). The test was added to make sure that no
StaleObjectError is raised when the car is destroyed. No such error is
currently raised because the lock version is not incremented by
appending a wheel to the car.

Furthermore, `assert_difference` succeeds because `car.wheels.count`
does not check the counter cache, but the collection size. The test will
fail if it is replaced with `car.wheels_count || 0`.

6 files changed, 34 insertions, 16 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..839bbea5bb 100644
--- a/activerecord/test/cases/locking_test.rb
+++ b/activerecord/test/cases/locking_test.rb
@@ -263,18 +263,6 @@ class OptimisticLockingTest < ActiveRecord::TestCase
     end
   end
 
-  def test_polymorphic_destroy_with_dependencies_and_lock_version
-    car = Car.create!
-
-    assert_difference 'car.wheels.count'  do
-      car.wheels << Wheel.create!
-    end
-    assert_difference 'car.wheels.count', -1  do
-      car.destroy
-    end
-    assert car.destroyed?
-  end
-
   def test_removing_has_and_belongs_to_many_associations_upon_destroy
     p = RichPerson.create! first_name: 'Jon'
     p.treasures.create!
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|