Refactor BelongsToAssociation to allow BelongsToPolymorphicAssociation to inherit from it

activerecord/lib/active_record/associations/association_collection.rb

@@ -31,10 +31,6 @@ def select(select = nil)
         end
       end
 
-      def scoped
-        with_scope(@scope) { @reflection.klass.scoped }
-      end
-
       def find(*args)
         options = args.extract_options!
 

activerecord/lib/active_record/associations/association_proxy.rb

@@ -8,7 +8,7 @@ module Associations
     #
     #   AssociationProxy
     #     BelongsToAssociation
-    #     BelongsToPolymorphicAssociation
+    #       BelongsToPolymorphicAssociation
     #     AssociationCollection + HasAssociation
     #       HasAndBelongsToManyAssociation
     #       HasManyAssociation
@@ -116,6 +116,7 @@ def reset
       # Reloads the \target and returns +self+ on success.
       def reload
         reset
+        construct_scope
         load_target
         self unless @target.nil?
       end
@@ -166,6 +167,10 @@ def send(method, *args)
         end
       end
 
+      def scoped
+        with_scope(@scope) { target_klass.scoped }
+      end
+
       protected
         def interpolate_sql(sql, record = nil)
           @owner.send(:interpolate_sql, sql, record)
@@ -192,15 +197,19 @@ def merge_options_from_reflection!(options)
 
         # Forwards +with_scope+ to the reflection.
         def with_scope(*args, &block)
-          @reflection.klass.send :with_scope, *args, &block
+          target_klass.send :with_scope, *args, &block
         end
 
         # Construct the scope used for find/create queries on the target
         def construct_scope
-          @scope = {
-            :find   => construct_find_scope,
-            :create => construct_create_scope
-          }
+          if target_klass
+            @scope = {
+              :find   => construct_find_scope,
+              :create => construct_create_scope
+            }
+          else
+            @scope = nil
+          end
         end
 
         # Implemented by subclasses
@@ -214,7 +223,7 @@ def construct_create_scope
         end
 
         def aliased_table
-          @reflection.klass.arel_table
+          target_klass.arel_table
         end
 
         # Set the inverse association, if possible
@@ -224,6 +233,12 @@ def set_inverse_instance(record)
           end
         end
 
+        # This class of the target. belongs_to polymorphic overrides this to look at the
+        # polymorphic_type field on the owner.
+        def target_klass
+          @reflection.klass
+        end
+
       private
         # Forwards any missing method call to the \target.
         def method_missing(method, *args)
@@ -254,7 +269,7 @@ def method_missing(method, *args)
         def load_target
           return nil unless defined?(@loaded)
 
-          if !loaded? && (!@owner.new_record? || foreign_key_present)
+          if !loaded? && (!@owner.new_record? || foreign_key_present) && @scope
             @target = find_target
           end
 

activerecord/lib/active_record/associations/belongs_to_association.rb

@@ -11,29 +11,16 @@ def build(attributes = {})
       end
 
       def replace(record)
-        counter_cache_name = @reflection.counter_cache_column
-
-        if record.nil?
-          if counter_cache_name && @owner.persisted?
-            @reflection.klass.decrement_counter(counter_cache_name, previous_record_id) if @owner[@reflection.primary_key_name]
-          end
-
-          @target = @owner[@reflection.primary_key_name] = nil
-        else
-          raise_on_type_mismatch(record)
-
-          if counter_cache_name && @owner.persisted? && record.id != @owner[@reflection.primary_key_name]
-            @reflection.klass.increment_counter(counter_cache_name, record.id)
-            @reflection.klass.decrement_counter(counter_cache_name, @owner[@reflection.primary_key_name]) if @owner[@reflection.primary_key_name]
-          end
-
-          @target = (AssociationProxy === record ? record.target : record)
-          @owner[@reflection.primary_key_name] = record_id(record) if record.persisted?
-          @updated = true
-        end
+        record = record.target if AssociationProxy === record
+        raise_on_type_mismatch(record) unless record.nil?
 
+        update_counters(record)
+        replace_keys(record)
         set_inverse_instance(record)
 
+        @target  = record
+        @updated = true if record
+
         loaded
         record
       end
@@ -44,8 +31,8 @@ def updated?
 
       def stale_target?
         if @target && @target.persisted?
-          target_id   = @target.send(@reflection.association_primary_key).to_s
-          foreign_key = @owner.send(@reflection.primary_key_name).to_s
+          target_id   = @target[@reflection.association_primary_key].to_s
+          foreign_key = @owner[@reflection.primary_key_name].to_s
 
           target_id != foreign_key
         else
@@ -54,27 +41,51 @@ def stale_target?
       end
 
       private
+        def update_counters(record)
+          counter_cache_name = @reflection.counter_cache_column
+
+          if counter_cache_name && @owner.persisted? && different_target?(record)
+            if record
+              target_klass.increment_counter(counter_cache_name, record.id)
+            end
+
+            if foreign_key_present
+              target_klass.decrement_counter(counter_cache_name, target_id)
+            end
+          end
+        end
+
+        # Checks whether record is different to the current target, without loading it
+        def different_target?(record)
+          record.nil? && @owner[@reflection.primary_key_name] ||
+          record.id   != @owner[@reflection.primary_key_name]
+        end
+
+        def replace_keys(record)
+          @owner[@reflection.primary_key_name] = record && record[@reflection.association_primary_key]
+        end
+
         def find_target
-          find_method = if @reflection.options[:primary_key]
-                          "find_by_#{@reflection.options[:primary_key]}"
-                        else
-                          "find"
-                        end
-
-          options = @reflection.options.dup.slice(:select, :include, :readonly)
-
-          the_target = with_scope(:find => @scope[:find]) do
-            @reflection.klass.send(find_method,
-              @owner[@reflection.primary_key_name],
-              options
-            ) if @owner[@reflection.primary_key_name]
+          if foreign_key_present
+            scoped.first.tap { |record| set_inverse_instance(record) }
           end
-          set_inverse_instance(the_target)
-          the_target
         end
 
         def construct_find_scope
-          { :conditions => conditions }
+          {
+            :conditions => construct_conditions,
+            :select     => @reflection.options[:select],
+            :include    => @reflection.options[:include],
+            :readonly   => @reflection.options[:readonly]
+          }
+        end
+
+        def construct_conditions
+          conditions = aliased_table[@reflection.association_primary_key].
+                       eq(@owner[@reflection.primary_key_name])
+
+          conditions = conditions.and(Arel.sql(sql_conditions)) if sql_conditions
+          conditions
         end
 
         def foreign_key_present
@@ -88,17 +99,12 @@ def invertible_for?(record)
           inverse && inverse.macro == :has_one
         end
 
-        def record_id(record)
-          record.send(@reflection.options[:primary_key] || :id)
-        end
-
-        def previous_record_id
-          @previous_record_id ||= if @reflection.options[:primary_key]
-                                    previous_record = @owner.send(@reflection.name)
-                                    previous_record.nil? ? nil : previous_record.id
-                                  else
-                                    @owner[@reflection.primary_key_name]
-                                  end
+        def target_id
+          if @reflection.options[:primary_key]
+            @owner.send(@reflection.name).try(:id)
+          else
+            @owner[@reflection.primary_key_name]
+          end
         end
     end
   end

activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb

 module ActiveRecord
   # = Active Record Belongs To Polymorphic Association
   module Associations
-    class BelongsToPolymorphicAssociation < AssociationProxy #:nodoc:
-      def replace(record)
-        if record.nil?
-          @target = @owner[@reflection.primary_key_name] = @owner[@reflection.options[:foreign_type]] = nil
-        else
-          @target = (AssociationProxy === record ? record.target : record)
-
-          @owner[@reflection.primary_key_name] = record_id(record)
-          @owner[@reflection.options[:foreign_type]] = record.class.base_class.name.to_s
-
-          @updated = true
-        end
-
-        set_inverse_instance(record)
-        loaded
-        record
-      end
-
-      def updated?
-        @updated
-      end
-
+    class BelongsToPolymorphicAssociation < BelongsToAssociation #:nodoc:
       def stale_target?
         if @target && @target.persisted?
           target_id    = @target.send(@reflection.association_primary_key).to_s
@@ -38,43 +17,26 @@ def stale_target?
 
       private
 
-        def inverse_reflection_for(record)
-          @reflection.polymorphic_inverse_of(record.class)
+        def replace_keys(record)
+          super
+          @owner[@reflection.options[:foreign_type]] = record && record.class.base_class.name
         end
 
-        def invertible_for?(record)
-          inverse = inverse_reflection_for(record)
-          inverse && inverse.macro == :has_one
+        def different_target?(record)
+          super || record.class != target_klass
         end
 
-        def construct_find_scope
-          { :conditions => conditions }
-        end
-
-        def find_target
-          return nil if association_class.nil?
-
-          target = association_class.send(:with_scope, :find => @scope[:find]) do
-            association_class.find(
-              @owner[@reflection.primary_key_name],
-              :select  => @reflection.options[:select],
-              :include => @reflection.options[:include]
-            )
-          end
-          set_inverse_instance(target)
-          target
-        end
-
-        def foreign_key_present
-          !@owner[@reflection.primary_key_name].nil?
+        def inverse_reflection_for(record)
+          @reflection.polymorphic_inverse_of(record.class)
         end
 
-        def record_id(record)
-          record.send(@reflection.options[:primary_key] || :id)
+        def target_klass
+          type = @owner[@reflection.options[:foreign_type]]
+          type && type.constantize
         end
 
-        def association_class
-          @owner[@reflection.options[:foreign_type]] ? @owner[@reflection.options[:foreign_type]].constantize : nil
+        def raise_on_type_mismatch(record)
+          # A polymorphic association cannot have a type mismatch, by definition
         end
     end
   end

activerecord/lib/active_record/associations/through_association.rb

@@ -158,8 +158,6 @@ def build_sti_condition
       alias_method :sql_conditions, :conditions
 
       def update_stale_state
-        construct_scope if stale_target?
-
         if @reflection.through_reflection.macro == :belongs_to
           @through_foreign_key = @owner.send(@reflection.through_reflection.primary_key_name)
         end

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -198,16 +198,23 @@ def test_belongs_to_counter_with_assigning_nil
     assert_equal 1, Post.find(p.id).comments.size
   end
 
-  def test_belongs_to_with_primary_key_counter_with_assigning_nil
-    debate = Topic.create("title" => "debate")
-    reply  = Reply.create("title" => "blah!", "content" => "world around!", "parent_title" => "debate")
+  def test_belongs_to_with_primary_key_counter
+    debate  = Topic.create("title" => "debate")
+    debate2 = Topic.create("title" => "debate2")
+    reply   = Reply.create("title" => "blah!", "content" => "world around!", "parent_title" => "debate")
+
+    assert_equal 1, debate.reload.replies_count
+    assert_equal 0, debate2.reload.replies_count
+
+    reply.topic_with_primary_key = debate2
 
-    assert_equal debate.title, reply.parent_title
-    assert_equal 1, Topic.find(debate.id).send(:read_attribute, "replies_count")
+    assert_equal 0, debate.reload.replies_count
+    assert_equal 1, debate2.reload.replies_count
 
     reply.topic_with_primary_key = nil
 
-    assert_equal 0, Topic.find(debate.id).send(:read_attribute, "replies_count")
+    assert_equal 0, debate.reload.replies_count
+    assert_equal 0, debate2.reload.replies_count
   end
 
   def test_belongs_to_counter_with_reassigning
@@ -419,6 +426,18 @@ def test_polymorphic_assignment_updates_foreign_id_field_for_new_and_saved_recor
     assert_nil sponsor.sponsorable_id
   end
 
+  def test_assignment_updates_foreign_id_field_for_new_and_saved_records
+    client = Client.new
+    saved_firm = Firm.create :name => "Saved"
+    new_firm = Firm.new
+
+    client.firm = saved_firm
+    assert_equal saved_firm.id, client.client_of
+
+    client.firm = new_firm
+    assert_nil client.client_of
+  end
+
   def test_polymorphic_assignment_with_primary_key_updates_foreign_id_field_for_new_and_saved_records
     essay = Essay.new
     saved_writer = Author.create(:name => "David")
@@ -537,4 +556,15 @@ def test_polymorphic_reassignment_of_associated_type_updates_the_object
     assert proxy.stale_target?
     assert_equal companies(:first_firm), sponsor.sponsorable
   end
+
+  def test_reloading_association_with_key_change
+    client = companies(:second_client)
+    firm = client.firm # note this is a proxy object
+
+    client.firm = companies(:another_firm)
+    assert_equal companies(:another_firm), firm.reload
+
+    client.client_of = companies(:first_firm).id
+    assert_equal companies(:first_firm), firm.reload
+  end
 end