Fix has_many_aggregate when association uses a custom primary key

lib/jit_preloader/active_record/base.rb

@@ -148,9 +148,10 @@ def has_many_aggregate(assoc, name, aggregate, field, table_alias_name: nil, def
               preloaded_data.merge!(data)
             end
 
+            aggregate_association_primary_key = aggregate_association.active_record_primary_key
             jit_preloader.records.each do |record|
               record.jit_preload_aggregates ||= {}
-              record.jit_preload_aggregates[key] = preloaded_data[record.id] || default
+              record.jit_preload_aggregates[key] = preloaded_data.fetch(record[aggregate_association_primary_key], default)
             end
           else
             self.jit_preload_aggregates[key] = send(assoc).where(conditions).send(aggregate, field) || default

spec/lib/jit_preloader/preloader_spec.rb

@@ -313,6 +313,16 @@
     end
   end
 
+  context "when preloading an aggregate where the association uses a custom primary key instad of 'id'" do 
+    let!(:new_jersey) { State.create(name: "New Jersey", region: usa.name) }
+    let!(:new_york) { State.create(name: "New York", region: usa.name) }
+
+    it "uses the custom primary key to join the data" do 
+      countries = Country.where(id: usa.id).jit_preload.to_a
+      expect(countries.first.states_count).to eq 2
+    end
+  end
+
   context "when a singular association id changes after preload" do
     let!(:contact_book1) { ContactBook.create(name: "The Yellow Pages") }
     let!(:contact_book2) { ContactBook.create(name: "The White Pages") }

spec/support/database.rb

@@ -9,6 +9,7 @@ def self.tables
       "CREATE TABLE phone_numbers (id INTEGER NOT NULL PRIMARY KEY, contact_id INTEGER NOT NULL, phone VARCHAR(10))",
       "CREATE TABLE countries (id INTEGER NOT NULL PRIMARY KEY, name VARCHAR(255))",
       "CREATE TABLE parents_children (id INTEGER NOT NULL PRIMARY KEY, parent_id INTEGER, child_id INTEGER)",
+      "CREATE TABLE states (id INTEGER NOT NULL PRIMARY KEY, name VARCHAR(255), region VARCHAR(255))",
     ]
   end
 

spec/support/models.rb

@@ -74,9 +74,15 @@ class Country < ActiveRecord::Base
   has_many :addresses
   has_many :contacts, through: :addresses
   has_many :contact_owners, through: :contacts, source_type: 'ContactOwner'
+  has_many :states, primary_key: 'name', foreign_key: "region"
 
   has_many_aggregate :contacts, :count, :count, "*"
   has_many_aggregate :contact_owners, :count, :count, "*"
+  has_many_aggregate :states, :count, :count, "*"
+end
+
+class State < ActiveRecord::Base
+  belongs_to :country, foreign_key: 'region'
 end
 
 class ContactOwner < ActiveRecord::Base