Merge pull request #413 from Shopify/fetch_batch

Prefetch with batch operation objects.

lib/identity_cache.rb

@@ -5,8 +5,11 @@
 
 require "identity_cache/version"
 require "identity_cache/encoder"
+require "identity_cache/prefetch"
+require "identity_cache/prefetch/batch"
+require "identity_cache/prefetch/segment"
+require "identity_cache/prefetch/operation"
 require "identity_cache/cached"
-require "identity_cache/cached/prefetcher"
 require "identity_cache/cached/embedded_fetching"
 require "identity_cache/cached/association"
 require "identity_cache/cached/belongs_to"

lib/identity_cache/prefetch.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module IdentityCache
+  module Cached
+    module Prefetch
+    end
+
+    private_constant :Prefetch
+  end
+end

lib/identity_cache/prefetch/batch.rb

@@ -0,0 +1,22 @@
+module IdentityCache
+  module Prefetch
+    class Batch
+      def initialize(operation)
+        @operation = operation
+        @segments   = []
+      end
+
+      attr_reader :segments
+
+      def add(cached_association, parent)
+        Segment.new(self, cached_association, parent).tap do |segment|
+          segments << segment
+        end
+      end
+
+      def load
+        segments.map(&:load)
+      end
+    end
+  end
+end

lib/identity_cache/prefetch/operation.rb

@@ -0,0 +1,54 @@
+module IdentityCache
+  module Prefetch
+    class Operation
+      def initialize(klass, associations, records)
+        @records = records.to_a
+        @batches = {}
+
+        build(klass, associations)
+      end
+
+      attr_reader :batches, :records
+
+      def load
+        batches.each_value(&:load)
+      end
+
+      private
+
+      def build(klass, associations, parent: self, level: 0)
+        return if records.empty?
+
+        batch = batches[level] ||= Batch.new(self)
+
+        Array.wrap(associations).each do |association|
+          case association
+          when Symbol
+            batch.add(
+              klass.cached_association(association),
+              parent
+            )
+          when Hash
+            association.each do |parent_association, nested_associations|
+              segment = batch.add(
+                klass.cached_association(parent_association),
+                parent
+              )
+
+              nested_klass = klass.reflect_on_association(parent_association).klass
+
+              build(
+                nested_klass,
+                nested_associations,
+                parent: segment,
+                level: level.next
+              )
+            end
+          else
+            raise TypeError, "Invalid association class #{association.class}"
+          end
+        end
+      end
+    end
+  end
+end

lib/identity_cache/prefetch/segment.rb

@@ -0,0 +1,34 @@
+module IdentityCache
+  module Prefetch
+    class Segment
+      ASSOCIATION_FETCH_EVENT = "association_fetch.identity_cache"
+
+      def initialize(batch, cached_association, provider)
+        @batch = batch
+        @cached_association = cached_association
+        @provider = provider
+      end
+
+      def load
+        ActiveSupport::Notifications.instrument(ASSOCIATION_FETCH_EVENT, association: association) do
+          records
+        end
+      end
+
+      def records
+        @records ||= if @cached_association.reflection.active_record.should_use_cache?
+          @cached_association.fetch(@provider.records)
+        else
+          ActiveRecord::Associations::Preloader.new.preload(@provider.records, association)
+          @provider.records
+        end
+      end
+
+      private
+
+      def association
+        @cached_association.name
+      end
+    end
+  end
+end

lib/identity_cache/query_api.rb

@@ -91,7 +91,7 @@ def fetch_multi(*ids, includes: nil)
 
       # Prefetches cached associations on a collection of records
       def prefetch_associations(includes, records)
-        Cached::Prefetcher.prefetch(self, includes, records)
+        Prefetch::Operation.new(self, includes, records).load
       end
 
       # Invalidates the primary cache index for the associated record. Will not invalidate cached attributes.

test/prefetch/batch_test.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+require "test_helper"
+
+module IdentityCache
+  module Prefetch
+    class BatchTest < IdentityCache::TestCase
+      def setup
+        super
+
+        AssociatedRecord.send(:cache_belongs_to, :item)
+
+        @record    = AssociatedRecord.create!
+        @operation = Operation.new(AssociatedRecord, [], AssociatedRecord.all)
+        @batch = add_batch(operation, 0)
+      end
+
+      attr_reader :record, :operation, :batch
+
+      def test_load
+        batch.segments.each do |segment|
+          segment.expects(:load)
+        end
+
+        batch.load
+      end
+
+      def test_add
+        batch.add(AssociatedRecord.cached_association(:item), operation)
+
+        assert_equal(1, batch.segments.count)
+      end
+
+      private
+
+      def add_batch(operation, level)
+        operation.batches[level] = Batch.new(operation)
+      end
+    end
+  end
+end

test/prefetch/operation_test.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+require "test_helper"
+
+module IdentityCache
+  module Prefetch
+    class OperationTest < IdentityCache::TestCase
+      def setup
+        super
+
+        AssociatedRecord.send(:cache_belongs_to, :item)
+        AssociatedRecord.send(:cache_has_one, :deeply_associated, embed: :id)
+        AssociatedRecord.send(:cache_has_many, :deeply_associated_records)
+        DeeplyAssociatedRecord.send(:cache_belongs_to, :item)
+        Item.send(:cache_has_one, :associated, embed: :id)
+
+        @record = AssociatedRecord.create!
+      end
+
+      attr_reader :record
+
+      def test_load
+        operation = Operation.new(AssociatedRecord, :item, [record])
+
+        operation.batches.each_value do |batch|
+          batch.expects(:load)
+        end
+
+        operation.load
+      end
+
+      def test_records
+        operation = Operation.new(AssociatedRecord, :item, AssociatedRecord.all)
+
+        assert_equal([record], operation.records)
+        assert_instance_of(Array, operation.records)
+      end
+
+      def test_batches
+        operation = Operation.new(
+          AssociatedRecord,
+          [
+            :item,
+            { deeply_associated: :item },
+            { deeply_associated_records: { item: :associated } },
+          ],
+          [record],
+        )
+
+        assert_equal(3, operation.batches.count)
+        assert_equal(3, operation.batches.values.first.segments.count)
+        assert_equal(2, operation.batches.values.second.segments.count)
+        assert_equal(1, operation.batches.values.third.segments.count)
+      end
+    end
+  end
+end

test/prefetch/segment_test.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+require "test_helper"
+
+module IdentityCache
+  module Prefetch
+    class SegmentTest < IdentityCache::TestCase
+      def setup
+        super
+
+        AssociatedRecord.send(:cache_belongs_to, :item)
+
+        @record = AssociatedRecord.create!
+        @cached_association = AssociatedRecord.cached_association(:item)
+        @operation = Operation.new(AssociatedRecord, [], AssociatedRecord.all)
+        @batch = add_batch(operation, 0)
+      end
+
+      attr_reader :record, :cached_association, :operation, :batch
+
+      def test_load
+        segment = batch.add(cached_association, operation)
+
+        cached_association.expects(:fetch).with([record])
+
+        segment.load
+      end
+
+      def test_nested_load
+        Item.send(:cache_has_many, :associated_records)
+
+        next_batch = add_batch(operation, 1)
+        nested_cached_association = Item.cached_association(:associated_records)
+        nested_record = Item.create!(title: "Rocket Shoes")
+        record.update!(item: nested_record)
+
+        segment = batch.add(cached_association, operation)
+        nested_segment = next_batch.add(nested_cached_association, segment)
+
+        cached_association.expects(:fetch).with([record]).returns([nested_record])
+        nested_cached_association.expects(:fetch).with([nested_record])
+
+
+        nested_segment.load
+      end
+
+      def test_records
+        nested_record = Item.create!(title: "Invisible Ink")
+        record.update!(item: nested_record)
+
+        cached_association.expects(:fetch).with([record]).returns([nested_record])
+
+        segment = batch.add(cached_association, operation)
+
+        assert_equal([nested_record], segment.records)
+      end
+
+      private
+
+      def add_batch(operation, level)
+        operation.batches[level] = Batch.new(operation)
+      end
+    end
+  end
+end

test/prefetch_associations_test.rb

@@ -435,7 +435,7 @@ def test_prefetch_associations
     private
 
     def prefetch(klass, includes, records)
-      Cached::Prefetcher.prefetch(klass, includes, records)
+      Prefetch::Operation.new(klass, includes, records).load
     end
 
     def setup_has_many_children_and_grandchildren(*parents)