Merge pull request #5662 from arturopie/3-2_fixing_IM_when_using_find_select

Fixing Identity Map when using find select in rails 3.2 

3 files changed, 29 insertions, 21 deletions

diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 680d9ffea0..b1da547142 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -90,7 +90,7 @@ module ActiveRecord
       end
 
       def add(record)
-        repository[record.class.symbolized_sti_name][record.id] = record
+        repository[record.class.symbolized_sti_name][record.id] = record if contain_all_columns?(record)
       end
 
       def remove(record)
@@ -104,6 +104,12 @@ module ActiveRecord
       def clear
         repository.clear
       end
+
+      private
+
+        def contain_all_columns?(record)
+          (record.class.column_names - record.attribute_names).empty?
+        end
     end
 
     # Reinitialize an Identity Map model object from +coder+.
diff --git a/activerecord/lib/active_record/inheritance.rb b/activerecord/lib/active_record/inheritance.rb
index de9461982a..3e12a67c2a 100644
--- a/activerecord/lib/active_record/inheritance.rb
+++ b/activerecord/lib/active_record/inheritance.rb
@@ -63,15 +63,7 @@ module ActiveRecord
         record_id = sti_class.primary_key && record[sti_class.primary_key]
 
         if ActiveRecord::IdentityMap.enabled? && record_id
-          if (column = sti_class.columns_hash[sti_class.primary_key]) && column.number?
-            record_id = record_id.to_i
-          end
-          if instance = IdentityMap.get(sti_class, record_id)
-            instance.reinit_with('attributes' => record)
-          else
-            instance = sti_class.allocate.init_with('attributes' => record)
-            IdentityMap.add(instance)
-          end
+          instance = use_identity_map(sti_class, record_id, record)
         else
           instance = sti_class.allocate.init_with('attributes' => record)
         end
@@ -122,6 +114,21 @@ module ActiveRecord
 
       private
 
+      def use_identity_map(sti_class, record_id, record)
+        if (column = sti_class.columns_hash[sti_class.primary_key]) && column.number?
+          record_id = record_id.to_i
+        end
+
+        if instance = IdentityMap.get(sti_class, record_id)
+          instance.reinit_with('attributes' => record)
+        else
+          instance = sti_class.allocate.init_with('attributes' => record)
+          IdentityMap.add(instance)
+        end
+
+        instance
+      end
+
       def find_sti_class(type_name)
         if type_name.blank? || !columns_hash.include?(inheritance_column)
           self
diff --git a/activerecord/test/cases/identity_map_test.rb b/activerecord/test/cases/identity_map_test.rb
index 24132c7617..63d0bcf1fc 100644
--- a/activerecord/test/cases/identity_map_test.rb
+++ b/activerecord/test/cases/identity_map_test.rb
@@ -404,18 +404,13 @@ class IdentityMapTest < ActiveRecord::TestCase
     assert comment.save
   end
 
-  def test_find_using_select_and_identity_map
-    author_id, author = Author.select('id').order(:id).first, Author.order(:id).first
+  def test_do_not_add_to_repository_if_record_does_not_contain_all_columns
+    author = Author.select(:id).first
+    post = author.posts.first
 
-    assert_equal author_id, author
-    assert_same author_id, author
-    assert_not_nil author.name
-
-    post, post_id = Post.order(:id).first, Post.select('id').order(:id).first
-
-    assert_equal post_id, post
-    assert_same post_id, post
-    assert_not_nil post.title
+    assert_nothing_raised do
+      assert_not_nil post.author.name
+    end
   end
 
 # Currently AR is not allowing changing primary key (see Persistence#update)