From e9880a9a6e57670779ba02b2d55a20abe8cd7c1b Mon Sep 17 00:00:00 2001
From: Tor Erik Linnerud <tel@jklm.no>
Date: Thu, 5 Dec 2024 00:11:36 -0500
Subject: [PATCH] Fix has_many_aggregate when association uses a custom primary
 key

Previously, if the primary key column for the assocation was not 'id',
has_many_aggregate would use the wrong key when looking up the aggregate
values in the preloaded_data hash.
---
 lib/jit_preloader/active_record/base.rb  |  3 ++-
 spec/lib/jit_preloader/preloader_spec.rb | 10 ++++++++++
 spec/support/database.rb                 |  1 +
 spec/support/models.rb                   |  6 ++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

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