diff --git a/lib/jit_preloader/active_record/base.rb b/lib/jit_preloader/active_record/base.rb index c9ea1b1..0fac855 100644 --- a/lib/jit_preloader/active_record/base.rb +++ b/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 diff --git a/spec/lib/jit_preloader/preloader_spec.rb b/spec/lib/jit_preloader/preloader_spec.rb index 92e5c11..354c8b8 100644 --- a/spec/lib/jit_preloader/preloader_spec.rb +++ b/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") } diff --git a/spec/support/database.rb b/spec/support/database.rb index 7f09eb6..c1a600a 100644 --- a/spec/support/database.rb +++ b/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 diff --git a/spec/support/models.rb b/spec/support/models.rb index 05f132b..28841cb 100644 --- a/spec/support/models.rb +++ b/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