Fix broken prefetch_associations of a polymorphic cache_belongs_to

Shopify · May 5, 2020 · a4eeff9 · a4eeff9
1 parent 95c3e3d
commit a4eeff9
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 17 deletions.
diff --git a/lib/identity_cache/cached/belongs_to.rb b/lib/identity_cache/cached/belongs_to.rb
@@ -37,28 +37,30 @@ def fetch(records)
 
       def fetch_async(load_strategy, records)
         if reflection.polymorphic?
-          cache_keys_to_associated_ids = {}
+          type_fetcher_to_db_ids_hash = {}
 
           records.each do |owner_record|
             associated_id = owner_record.send(reflection.foreign_key)
             next unless associated_id && !owner_record.instance_variable_defined?(records_variable_name)
-            associated_cache_key = Object.const_get(
+            foreign_type_fetcher = Object.const_get(
               owner_record.send(reflection.foreign_type)
             ).cached_model.cached_primary_index
-            unless cache_keys_to_associated_ids[associated_cache_key]
-              cache_keys_to_associated_ids[associated_cache_key] = {}
-            end
-            cache_keys_to_associated_ids[associated_cache_key][associated_id] = owner_record
+            (type_fetcher_to_db_ids_hash[foreign_type_fetcher] ||= []) << associated_id
           end
 
-          load_strategy.load_batch(cache_keys_to_associated_ids) do |associated_records_by_cache_key|
+          load_strategy.load_batch(type_fetcher_to_db_ids_hash) do |batch_load_result|
             batch_records = []
-            associated_records_by_cache_key.each do |cache_key, associated_records|
-              associated_records.keys.each do |id, associated_record|
-                owner_record = cache_keys_to_associated_ids.fetch(cache_key).fetch(id)
-                batch_records << owner_record
-                write(owner_record, associated_record)
-              end
+
+            records.each do |owner_record|
+              associated_id = owner_record.send(reflection.foreign_key)
+              next unless associated_id && !owner_record.instance_variable_defined?(records_variable_name)
+              foreign_type_fetcher = Object.const_get(
+                owner_record.send(reflection.foreign_type)
+              ).cached_model.cached_primary_index
+
+              associated_record = batch_load_result.fetch(foreign_type_fetcher).fetch(associated_id)
+              batch_records << owner_record
+              write(owner_record, associated_record)
             end
 
             yield batch_records

diff --git a/test/fetch_multi_test.rb b/test/fetch_multi_test.rb
@@ -292,12 +292,16 @@ def test_fetch_multi_with_polymorphic_has_one
     PolymorphicRecord.send(:cache_belongs_to, :owner)
 
     item = Item.create!
+    item2 = ItemTwo.create!
 
-    poly1 = item.polymorphic_record = PolymorphicRecord.create!
+    poly1 = PolymorphicRecord.create!(owner: item)
+    poly2 = PolymorphicRecord.create!(owner: item2)
 
-    assert_queries(2) do
-      PolymorphicRecord.fetch_multi(poly1.id, includes: :owner)
+    poly_records = assert_queries(3) do
+      PolymorphicRecord.fetch_multi([poly1.id, poly2.id], includes: :owner)
     end
+    assert_equal([poly1, poly2], poly_records)
+    assert_equal([item, item2], poly_records.map(&:fetch_owner))
   end
 
   def test_fetch_multi_with_polymorphic_has_many
@@ -311,9 +315,11 @@ def test_fetch_multi_with_polymorphic_has_many
     poly2 = item.polymorphic_records.create
     poly3 = item2.polymorphic_records.create
 
-    assert_queries(3) do
+    poly_records = assert_queries(3) do
       PolymorphicRecord.fetch_multi([poly1.id, poly2.id, poly3.id], includes: :owner)
     end
+    assert_equal([poly1, poly2, poly3], poly_records)
+    assert_equal([item, item, item2], poly_records.map(&:fetch_owner))
   end
 
   def test_fetch_multi_with_polymorphic_has_one_with_nil