From a16cdda893e6224839fec7ead4179a789b2e6318 Mon Sep 17 00:00:00 2001
From: Richard Penner <richard.penner@shopify.com>
Date: Fri, 23 Feb 2024 18:55:04 +0000
Subject: [PATCH] Adds support for models using compound primary keys

---
 lib/identity_cache/cached/attribute.rb | 22 +++++++++++++++++++---
 test/custom_primary_keys_test.rb       | 12 ++++++++++++
 test/helpers/active_record_objects.rb  |  2 ++
 test/helpers/models.rb                 | 11 +++++++++++
 4 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/lib/identity_cache/cached/attribute.rb b/lib/identity_cache/cached/attribute.rb
index c02dd9e8..2257c082 100644
--- a/lib/identity_cache/cached/attribute.rb
+++ b/lib/identity_cache/cached/attribute.rb
@@ -19,7 +19,14 @@ def initialize(model, attribute_or_proc, alias_name, key_fields, unique)
       end
 
       def attribute
-        @attribute ||= @attribute_proc.call.to_sym
+        @attribute ||= begin
+          res = @attribute_proc.call
+          if res.is_a?(Array)
+            res.map(&:to_sym)
+          else
+            res.to_sym
+          end
+        end
       end
 
       def fetch(db_key)
@@ -59,7 +66,11 @@ def cache_key(index_key)
       def load_one_from_db(key)
         query = model.reorder(nil).where(load_from_db_where_conditions(key))
         query = query.limit(1) if unique
-        results = query.pluck(attribute)
+        results = if attribute.is_a?(Array)
+          query.pluck(*attribute)
+        else
+          query.pluck(attribute)
+        end
         unique ? results.first : results
       end
 
@@ -138,10 +149,15 @@ def field_types
 
       def cache_key_prefix
         @cache_key_prefix ||= begin
+          attribute_key_prefix = if attribute.is_a?(Array)
+            attribute.join("/")
+          else
+            attribute.to_s
+          end
           unique_indicator = unique ? "" : "s"
           "attr#{unique_indicator}" \
             ":#{model.base_class.name}" \
-            ":#{attribute}" \
+            ":#{attribute_key_prefix}" \
             ":#{key_fields.join("/")}:"
         end
       end
diff --git a/test/custom_primary_keys_test.rb b/test/custom_primary_keys_test.rb
index 03ceffe6..1ed236ec 100644
--- a/test/custom_primary_keys_test.rb
+++ b/test/custom_primary_keys_test.rb
@@ -10,6 +10,11 @@ def setup
     @parent_record = CustomParentRecord.create!(parent_primary_key: 1)
     @child_record_1 = CustomChildRecord.create!(custom_parent_record: @parent_record, child_primary_key: 1)
     @child_record_2 = CustomChildRecord.create!(custom_parent_record: @parent_record, child_primary_key: 2)
+
+    CompositePrimaryKeyRecord.cache_has_many(:custom_child_records)
+    CPKReference.cache_belongs_to(:composite_primary_key_record)
+    @composite_record = CompositePrimaryKeyRecord.create!(key_part_one: 1, key_part_two: 2)
+    @cpk_reference = CPKReference.create!(composite_primary_key_record: @composite_record)
   end
 
   def test_fetch_parent
@@ -17,4 +22,11 @@ def test_fetch_parent
       CustomParentRecord.fetch(@parent_record.parent_primary_key)
     end
   end
+
+  def test_fetch_composite_primary_key_record
+    assert_nothing_raised do
+      cpk_record = CompositePrimaryKeyRecord.fetch([@composite_record.key_part_one, @composite_record.key_part_two])
+      refute_nil cpk_record
+    end
+  end
 end
diff --git a/test/helpers/active_record_objects.rb b/test/helpers/active_record_objects.rb
index cd38e2c2..296859e3 100644
--- a/test/helpers/active_record_objects.rb
+++ b/test/helpers/active_record_objects.rb
@@ -50,6 +50,8 @@ def teardown_models
     Object.send(:remove_const, "Deeply")
     Object.send(:remove_const, "CustomParentRecord")
     Object.send(:remove_const, "CustomChildRecord")
+    Object.send(:remove_const, "CompositePrimaryKeyRecord")
+    Object.send(:remove_const, "CPKReference")
     IdentityCache.const_get(:ParentModelExpiration).send(:lazy_hooks).clear
   end
 end
diff --git a/test/helpers/models.rb b/test/helpers/models.rb
index 668c5028..70a9610a 100644
--- a/test/helpers/models.rb
+++ b/test/helpers/models.rb
@@ -96,3 +96,14 @@ class CustomChildRecord < ActiveRecord::Base
   belongs_to :custom_parent_record, foreign_key: :parent_id
   self.primary_key = "child_primary_key"
 end
+
+class CompositePrimaryKeyRecord < ActiveRecord::Base
+  include IdentityCache
+  has_many :cpk_references, foreign_key: [:key_part_one, :key_part_two]
+  self.primary_key = [:key_part_one, :key_part_two]
+end
+
+class CPKReference < ActiveRecord::Base
+  include IdentityCache
+  belongs_to :composite_primary_key_record, foreign_key: [:key_part_one, :key_part_two]
+end