Fix belongs_to polymorphic with custom primary key on target.

Closes #3104.

5 files changed, 30 insertions, 18 deletions

diff --git a/activerecord/lib/active_record/associations/association_scope.rb b/activerecord/lib/active_record/associations/association_scope.rb
index 9e6d9e73c5..a0efdbc2a4 100644
--- a/activerecord/lib/active_record/associations/association_scope.rb
+++ b/activerecord/lib/active_record/associations/association_scope.rb
@@ -68,7 +68,12 @@ module ActiveRecord
           end
 
           if reflection.source_macro == :belongs_to
-            key         = reflection.association_primary_key
+            if reflection.options[:polymorphic]
+              key = reflection.association_primary_key(klass)
+            else
+              key = reflection.association_primary_key
+            end
+
             foreign_key = reflection.foreign_key
           else
             key         = reflection.foreign_key
diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 58c9648ce8..97f531d064 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -45,7 +45,11 @@ module ActiveRecord
         end
 
         def replace_keys(record)
-          owner[reflection.foreign_key] = record && record[reflection.association_primary_key]
+          if record
+            owner[reflection.foreign_key] = record[reflection.association_primary_key(record.class)]
+          else
+            owner[reflection.foreign_key] = nil
+          end
         end
 
         def foreign_key_present?
diff --git a/activerecord/lib/active_record/reflection.rb b/activerecord/lib/active_record/reflection.rb
index 89179779e3..1929a808ed 100644
--- a/activerecord/lib/active_record/reflection.rb
+++ b/activerecord/lib/active_record/reflection.rb
@@ -211,11 +211,9 @@ module ActiveRecord
         @association_foreign_key ||= options[:association_foreign_key] || class_name.foreign_key
       end
 
-      def association_primary_key
-        @association_primary_key ||=
-          options[:primary_key] ||
-          !options[:polymorphic] && klass.primary_key ||
-          'id'
+      # klass option is necessary to support loading polymorphic associations
+      def association_primary_key(klass = self.klass)
+        options[:primary_key] || klass.primary_key
       end
 
       def active_record_primary_key
@@ -463,17 +461,15 @@ module ActiveRecord
       # We want to use the klass from this reflection, rather than just delegate straight to
       # the source_reflection, because the source_reflection may be polymorphic. We still
       # need to respect the source_reflection's :primary_key option, though.
-      def association_primary_key
-        @association_primary_key ||= begin
-          # Get the "actual" source reflection if the immediate source reflection has a
-          # source reflection itself
-          source_reflection = self.source_reflection
-          while source_reflection.source_reflection
-            source_reflection = source_reflection.source_reflection
-          end
-
-          source_reflection.options[:primary_key] || klass.primary_key
+      def association_primary_key(klass = self.klass)
+        # Get the "actual" source reflection if the immediate source reflection has a
+        # source reflection itself
+        source_reflection = self.source_reflection
+        while source_reflection.source_reflection
+          source_reflection = source_reflection.source_reflection
         end
+
+        source_reflection.options[:primary_key] || klass.primary_key
       end
 
       # Gets an array of possible <tt>:through</tt> source reflection names:
diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index 866a3cca10..1160d236c9 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -13,6 +13,7 @@ require 'models/comment'
 require 'models/sponsor'
 require 'models/member'
 require 'models/essay'
+require 'models/toy'
 
 class BelongsToAssociationsTest < ActiveRecord::TestCase
   fixtures :accounts, :companies, :developers, :projects, :topics,
@@ -696,4 +697,11 @@ class BelongsToAssociationsTest < ActiveRecord::TestCase
     assert_equal nil, comment.reload.parent
     assert_equal 0, comments(:greetings).reload.children_count
   end
+
+  def test_polymorphic_with_custom_primary_key
+    toy = Toy.create!
+    sponsor = Sponsor.create!(:sponsorable => toy)
+
+    assert_equal toy, sponsor.reload.sponsorable
+  end
 end
diff --git a/activerecord/test/cases/reflection_test.rb b/activerecord/test/cases/reflection_test.rb
index 41312e8661..0a48f418b1 100644
--- a/activerecord/test/cases/reflection_test.rb
+++ b/activerecord/test/cases/reflection_test.rb
@@ -244,7 +244,6 @@ class ReflectionTest < ActiveRecord::TestCase
     # Normal association
     assert_equal "id",   Author.reflect_on_association(:posts).association_primary_key.to_s
     assert_equal "name", Author.reflect_on_association(:essay).association_primary_key.to_s
-    assert_equal "id",   Tagging.reflect_on_association(:taggable).association_primary_key.to_s
 
     # Through association (uses the :primary_key option from the source reflection)
     assert_equal "nick", Author.reflect_on_association(:subscribers).association_primary_key.to_s