Refactor BelongsToAssociation to allow BelongsToPolymorphicAssociation to inherit from it

5 files changed, 91 insertions, 114 deletions

diff --git a/activerecord/lib/active_record/associations/association_collection.rb b/activerecord/lib/active_record/associations/association_collection.rb
index 64e10f56e2..0959ea2c5a 100644
--- a/activerecord/lib/active_record/associations/association_collection.rb
+++ b/activerecord/lib/active_record/associations/association_collection.rb
@@ -31,10 +31,6 @@ module ActiveRecord
         end
       end
 
-      def scoped
-        with_scope(@scope) { @reflection.klass.scoped }
-      end
-
       def find(*args)
         options = args.extract_options!
 
diff --git a/activerecord/lib/active_record/associations/association_proxy.rb b/activerecord/lib/active_record/associations/association_proxy.rb
index 97eab8a793..7e68241a2c 100644
--- a/activerecord/lib/active_record/associations/association_proxy.rb
+++ b/activerecord/lib/active_record/associations/association_proxy.rb
@@ -8,7 +8,7 @@ module ActiveRecord
     #
     #   AssociationProxy
     #     BelongsToAssociation
-    #     BelongsToPolymorphicAssociation
+    #       BelongsToPolymorphicAssociation
     #     AssociationCollection + HasAssociation
     #       HasAndBelongsToManyAssociation
     #       HasManyAssociation
@@ -116,6 +116,7 @@ module ActiveRecord
       # Reloads the \target and returns +self+ on success.
       def reload
         reset
+        construct_scope
         load_target
         self unless @target.nil?
       end
@@ -166,6 +167,10 @@ module ActiveRecord
         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 @@ module ActiveRecord
 
         # 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 @@ module ActiveRecord
         end
 
         def aliased_table
-          @reflection.klass.arel_table
+          target_klass.arel_table
         end
 
         # Set the inverse association, if possible
@@ -224,6 +233,12 @@ module ActiveRecord
           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 @@ module ActiveRecord
         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
 
diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 98c1c13524..6b29e3ef92 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -11,29 +11,16 @@ module ActiveRecord
       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 @@ module ActiveRecord
 
       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 @@ module ActiveRecord
       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 @@ module ActiveRecord
           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
diff --git a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
index 90eff7399c..6293c4ca42 100644
--- a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
@@ -1,28 +1,7 @@
 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 @@ module ActiveRecord
 
       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
diff --git a/activerecord/lib/active_record/associations/through_association.rb b/activerecord/lib/active_record/associations/through_association.rb
index 57718285f8..65ca4b1282 100644
--- a/activerecord/lib/active_record/associations/through_association.rb
+++ b/activerecord/lib/active_record/associations/through_association.rb
@@ -158,8 +158,6 @@ module ActiveRecord
       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