Allow prefetch_associations on association that was prefetched before

Shopify · Feb 16, 2024 · ccb2fd7 · ccb2fd7
1 parent d1e53f7
commit ccb2fd7
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 9 deletions.
diff --git a/lib/identity_cache/cached/belongs_to.rb b/lib/identity_cache/cached/belongs_to.rb
@@ -77,16 +77,20 @@ def fetch_async(load_strategy, records)
             end
           end
 
-          load_strategy.load_multi(
-            reflection.klass.cached_primary_index,
-            ids_to_owner_record.keys
-          ) do |associated_records_by_id|
-            associated_records_by_id.each do |id, associated_record|
-              owner_record = ids_to_owner_record.fetch(id)
-              write(owner_record, associated_record)
-            end
+          if ids_to_owner_record.any?
+            load_strategy.load_multi(
+              reflection.klass.cached_primary_index,
+              ids_to_owner_record.keys
+            ) do |associated_records_by_id|
+              associated_records_by_id.each do |id, associated_record|
+                owner_record = ids_to_owner_record.fetch(id)
+                write(owner_record, associated_record)
+              end
 
-            yield associated_records_by_id.values.compact
+              yield associated_records_by_id.values.compact
+            end
+          else
+            yield records.filter_map { |record| record.instance_variable_get(records_variable_name) }
           end
         end
       end

diff --git a/test/prefetch_associations_test.rb b/test/prefetch_associations_test.rb
@@ -270,6 +270,35 @@ def test_fetch_multi_batch_fetches_first_level_associations_who_dont_include_ide
       end
     end
 
+    def test_fetch_multi_batch_fetches_non_embedded_second_level_belongs_to_associations_in_repeating_prefetch
+      Item.send(:cache_belongs_to, :item)
+      Item.send(:cache_has_many, :associated_records, embed: :ids)
+
+      bob_item = @bob.create_item!(title: "bob's item")
+      @bob.save!
+      fred_item = @fred.create_item!(title: "fred's item")
+      @fred.save!
+
+      bob_associated_record = bob_item.associated_records.create!(name: "bob child")
+      fred_associated_record = fred_item.associated_records.create!(name: "fred child")
+
+      cached_bob, cached_fred = Item.fetch_multi(
+        [@bob.id, @fred.id], includes: [:item]
+      )
+
+      Item.prefetch_associations({ item: { associated_records: [] } }, [cached_bob, cached_fred])
+
+      assert_queries(0) do
+        assert_memcache_operations(0) do
+          assert_equal(bob_item,  cached_bob.fetch_item)
+          assert_equal(fred_item, cached_fred.fetch_item)
+
+          assert_equal([bob_associated_record], cached_bob.fetch_item.fetch_associated_records)
+          assert_equal([fred_associated_record], cached_fred.fetch_item.fetch_associated_records)
+        end
+      end
+    end
+
     def test_fetch_multi_batch_fetches_non_embedded_second_level_has_many_associations
       Item.send(:cache_has_many, :associated_records, embed: :ids)
       AssociatedRecord.send(:cache_has_many, :deeply_associated_records, embed: :ids)