habtm and possibly other association issues

lib/identity_map.rb

@@ -3,3 +3,4 @@
   require "identity_map/action_controller/dispatcher"
 end
 require "identity_map/active_record/base"
+require "identity_map/active_record/association_preload"

lib/identity_map/active_record/association_preload.rb

@@ -0,0 +1,64 @@
+module ActiveRecord
+  module AssociationPreload #:nodoc:
+    module ClassMethods
+      def preload_has_and_belongs_to_many_association_with_identity_map(records, reflection, preload_options={})
+        unless reflection.klass.respond_to?(:id_map)
+          return preload_has_and_belongs_to_many_association_without_identity_map(records, reflection, preload_options)
+        end
+        table_name = reflection.klass.quoted_table_name
+        records = records.find_all{|record| !record.send(reflection.name).loaded?}
+        return if records.empty?
+        id_to_record_map, ids = construct_id_map(records)
+        records.each {|record| record.send(reflection.name).loaded}
+        options = reflection.options
+
+        conditions = "t0.#{reflection.primary_key_name} #{in_or_equals_for_ids(ids)}"
+        conditions << append_conditions(reflection, preload_options)
+
+        joins = connection.select_all(sanitize_sql([
+            "select t0.#{reflection.primary_key_name} as prnt_id, t0.#{reflection.association_foreign_key} as chld_id
+             from #{connection.quote_table_name options[:join_table]} t0
+             where #{conditions}
+             ", ids]))
+        child_record_ids = joins.map{|j| j['chld_id']}.uniq
+
+        associated_records = reflection.klass.with_exclusive_scope do
+          reflection.klass.find(:all, :conditions => {reflection.klass.primary_key => child_record_ids},
+            :include => options[:include],
+            :select => options[:select].presence,
+            :order => options[:order])
+        end
+        associated_record_map = associated_records.inject({}){|h, r| h[r.id.to_s] = r; h}
+        joins.each do |j|
+          mapped_records = id_to_record_map[j['prnt_id'].to_s]
+          $stderr.puts(mapped_records.inspect + ' ' + associated_record_map[j['chld_id'].to_s].inspect)
+          add_preloaded_records_to_collection(mapped_records, reflection.name, associated_record_map[j['chld_id'].to_s])
+        end
+        $stderr.puts("exit")
+      end
+      alias_method_chain :preload_has_and_belongs_to_many_association, :identity_map
+
+      if Array.respond_to?(:wrap)
+        def add_preloaded_records_to_collection(parent_records, reflection_name, associated_record)
+          parent_records.each do |parent_record|
+            association_proxy = parent_record.send(reflection_name)
+            association_proxy.loaded
+            associated_records = Array.wrap(associated_record) - association_proxy.target
+            association_proxy.target.push(*associated_records)
+            association_proxy.__send__(:set_inverse_instance, associated_record, parent_record)
+          end
+        end
+      else
+        def add_preloaded_records_to_collection(parent_records, reflection_name, associated_record)
+          parent_records.each do |parent_record|
+            association_proxy = parent_record.send(reflection_name)
+            association_proxy.loaded
+            associated_records = [associated_record].flatten - association_proxy.target
+            association_proxy.target.push(*associated_records)
+            association_proxy.__send__(:set_inverse_instance, associated_record, parent_record)
+          end
+        end
+      end
+    end
+  end
+end

spec/identity_map_spec.rb

@@ -138,6 +138,23 @@
       c3.__id__.should == c2.__id__
     end
   end
+
+  context "has and belongs to many" do
+    before(:each) do
+      gotwo = Building.create(:name=>'GoTwo')
+      Address.find(:all).each do |address|
+        gotwo.addresses << address
+      end
+    end
+
+    it "should load habtm adequatly" do
+      buildings = Building.find(:all, :include=>:addresses)
+      buildings[0].addresses.loaded?.should be_true
+      buildings[0].addresses.to_a.size.should == 2
+      buildings[1].addresses.loaded?.should be_true
+      buildings[1].addresses.to_a.size.should == 2
+    end
+  end
 
 	after(:each) do
 		ActiveRecord::Base.drop_identity_map

spec/spec_helper.rb

@@ -31,7 +31,17 @@
   end
   create_table :phone_numbers, :force => true do |t|
     t.string :number
-    t.integer :customer_id
+    t.references :customer
+  end
+  create_table :buildings, :force => true do |t|
+    t.string :name
+  end
+  create_table :addresses, :force => true do |t|
+    t.string :name
+  end
+  create_table :addresses_buildings, :force => true, :id => false do |t|
+    t.references :building
+    t.references :address
   end
 end
 
@@ -49,3 +59,19 @@ class PhoneNumber < ActiveRecord::Base
 
 phone_number = customer.phone_numbers.create(:number => "8675309")
 
+class Building < ActiveRecord::Base
+  use_id_map
+  has_and_belongs_to_many :addresses
+end
+
+building = Building.create(:name => 'GoOne')
+
+class Address < ActiveRecord::Base
+  use_id_map
+  has_and_belongs_to_many :customers
+end
+
+address1 = Address.create(:name=>'volga')
+address2 = Address.create(:name=>'don')
+building.addresses << address1
+building.addresses << address2