Implement custom preloads (#129)

README.md

@@ -220,6 +220,9 @@ Associations with any of the following options cannot be eager loaded:
 
 Goldiloader detects associations with any of these options and disables automatic eager loading on them.
 
+It might still be possible to eager load these with Goldiloader by using [custom preloads](#custom-preloads).
+
+
 ### Eager Loading Limitation Workarounds
 
 Most of the Rails limitations with eager loading can be worked around by pushing the problematic SQL into the database via database views. Consider the following example with associations that can't be eager loaded due to SQL limits:
@@ -275,6 +278,64 @@ class RecentPostReference < ActiveRecord::Base
 end
 ```
 
+## Custom Preloads
+
+In addition to preloading relations, you can also define custom preloads by yourself in your model. The only requirement is that you need to be able to perform a lookup for multiple records/ids and return a single Hash with the ids as keys.
+If that's the case, these preloads can nearly be anything. Some examples could be:
+
+- simple aggregations (count, sum, maximum, etc.)
+- more complex custom SQL queries
+- external API requests (ElasticSearch, Redis, etc.)
+- relations with primary keys stored in a `jsonb` column
+
+Here's how:
+
+```ruby
+class Blog < ActiveRecord::Base
+  has_many :posts
+
+  def posts_count
+    goldiload do |ids|
+      # By default, `ids` will be an array of `Blog#id`s
+      Post
+        .where(blog_id: ids)
+        .group(:blog_id)
+        .count
+    end
+  end
+end
+```
+
+The first time you call the `posts_count` method, it will call the block with all model ids from the current context and reuse the result from the block for all other models in the context.
+
+A more complex example might use a custom primary key instead of `id`, use a non ActiveRecord API and have more complex return values than just scalar values:
+
+
+```ruby
+class Post < ActiveRecord::Base
+  def main_translator_reference
+    json_payload[:main_translator_reference]
+  end
+
+  def main_translator
+    goldiload(key: :main_translator_reference) do |references|
+      # `references` will be an array of `Post#main_translator_reference`s
+      SomeExternalApi.fetch_translators(
+        id: references
+      ).index_by(&:id)
+    end
+  end
+end
+```
+
+**Note:** The `goldiload` method will use the `source_location` of the given block as a cache name to distinguish between multiple defined preloads. If this causes an issue for you, you can also pass a cache name explicitly as the first argument to the `goldiload` method.
+
+
+### Gotchas
+
+Even though the methods in the examples above (`posts_count`, `main_translator`) are actually instance methods, the block passed to `goldiload` should not contain any references to these instances, as this could break the internal lookup/caching mechanism. We prevent this for the `self` keyword, so you'll get a `NoMethodError`. If you get this, you might want to think about the implementation rather than just trying to work around the exception.
+
+
 ## Upgrading
 
 ### From 0.x, 1.x

lib/goldiloader.rb

@@ -3,6 +3,7 @@
 require 'active_support/all'
 require 'active_record'
 require 'goldiloader/compatibility'
+require 'goldiloader/custom_preloads'
 require 'goldiloader/auto_include_context'
 require 'goldiloader/scope_info'
 require 'goldiloader/association_options'

lib/goldiloader/active_record_patches.rb

@@ -25,6 +25,11 @@ def reload(*)
       @auto_include_context = nil
       super
     end
+
+    def goldiload(cache_name = nil, key: self.class.primary_key, &block)
+      cache_name ||= block.source_location.join(':')
+      auto_include_context.preloaded(self, cache_name: cache_name, key: key, &block)
+    end
   end
   ::ActiveRecord::Base.include(::Goldiloader::BasePatch)
 

lib/goldiloader/auto_include_context.rb

@@ -51,5 +51,7 @@ def register_models(models)
     end
 
     alias_method :register_model, :register_models
+
+    prepend Goldiloader::CustomPreloads
   end
 end

lib/goldiloader/custom_preloads.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Goldiloader
+  module CustomPreloads
+    def initialize
+      super
+      @custom_preloads = nil
+    end
+
+    def preloaded(model, cache_name:, key:, &block)
+      unless preloaded?(cache_name)
+        ids = models.map do |record|
+          record.public_send(key)
+        end
+
+        # We're using instance_exec instead of a simple yield to make sure that the
+        # given block does not have any references to the model instance as this might
+        # lead to unexpected results. The block will be executed in the context of the
+        # class of the model instead.
+        block_context = models.first.class
+        preloaded_hash = block_context.instance_exec(ids, &block)
+        store_preloaded(cache_name, preloaded_hash)
+      end
+      fetch_preloaded(cache_name, model, key: key)
+    end
+
+    private
+
+    def store_preloaded(cache_name, preloaded_hash)
+      @custom_preloads ||= {}
+      @custom_preloads[cache_name] = preloaded_hash
+    end
+
+    def fetch_preloaded(cache_name, instance, key:)
+      @custom_preloads&.dig(cache_name, instance.public_send(key))
+    end
+
+    def preloaded?(cache_name)
+      @custom_preloads&.key?(cache_name)
+    end
+  end
+end

spec/db/schema.rb

@@ -105,6 +105,32 @@ class Blog < ActiveRecord::Base
 
   has_one :post_with_order, -> { order(:id) }, class_name: 'Post'
 
+  def posts_count
+    goldiload do |ids|
+      Post.where(blog_id: ids).group(:blog_id).count
+    end
+  end
+
+  def tags_count
+    goldiload do |ids|
+      Tag
+        .joins(:posts)
+        .where(posts: { blog_id: ids })
+        .group(:blog_id)
+        .count
+    end
+  end
+
+  # rubocop:disable Style/RedundantSelf
+  def custom_preload_with_self_reference
+    goldiload do |ids|
+      ids.to_h do |id|
+        [id, self.posts.count]
+      end
+    end
+  end
+  # rubocop:enable Style/RedundantSelf
+
   def posts_overridden
     'boom'
   end
@@ -128,6 +154,18 @@ class Post < ActiveRecord::Base
 
   after_destroy :after_post_destroy
 
+  def author_global_id
+    author&.to_gid&.to_s
+  end
+
+  def author_via_global_id
+    goldiload key: :author_global_id do |gids|
+      GlobalID::Locator.locate_many(gids).index_by do |author|
+        author.to_gid.to_s
+      end
+    end
+  end
+
   def after_post_destroy
     # Hook for tests
   end

spec/goldiloader/goldiloader_spec.rb

@@ -978,6 +978,81 @@
     end
   end
 
+  context "custom preloads" do
+    before do
+      # create some additional records to make sure we actually have different counts
+      blog1.posts.create!(title: 'another-post') do |post|
+        post.tags << Tag.create!(name: 'some tag')
+      end
+    end
+
+    let(:blogs) { Blog.order(:name).to_a }
+
+    it "returns custom preloads" do
+      expected_post_counts = blogs.map do |blog|
+        blog.posts.count
+      end
+
+      expected_tag_counts = blogs.map do |blog|
+        blog.posts.sum { |post| post.tags.count }
+      end
+
+      expect do
+        expect(blogs.map(&:posts_count)).to eq expected_post_counts
+        expect(blogs.map(&:tags_count)).to eq expected_tag_counts
+      end.to execute_queries(Post => 1, Tag => 1)
+    end
+
+    it "works without a collection" do
+      expect(blog1.posts_count).to eq blog1.posts.count
+      expect(blog2.posts_count).to eq blog2.posts.count
+    end
+
+    it "prevents self references to the model inside the block" do
+      expect do
+        blog1.custom_preload_with_self_reference
+      end.to raise_error(NoMethodError)
+    end
+
+    it "uses different caches for different blocks" do
+      result1 = blog1.goldiload do |ids|
+        ids.to_h { |id| [id, 42] }
+      end
+      expect(result1).to eq 42
+
+      result2 = blog1.goldiload do |ids|
+        ids.to_h { |id| [id, 666] }
+      end
+      expect(result2).to eq 666
+    end
+
+    it "can use an explicit cache_name" do
+      # Define explicit cache key :random_cache_key
+      blog1.goldiload(:random_cache_key) do |ids|
+        ids.to_h { |id| [id, 42] }
+      end
+
+      # Another blog for the same key
+      result = blog1.goldiload(:random_cache_key) do |ids|
+        # :nocov:
+        ids.to_h { |id| [id, 666] }
+        # :nocov:
+      end
+
+      # First block should be used
+      expect(result).to eq 42
+    end
+
+    it "can preload with a custom key" do
+      posts = Post.all.order(id: :asc)
+      expected_authors = posts.map(&:author)
+
+      expect do
+        expect(posts.map(&:author_via_global_id)).to eq expected_authors
+      end.to execute_queries(User => 1)
+    end
+  end
+
   describe "#globally_enabled" do
     context "enabled" do
       it "allows setting per thread" do

spec/spec_helper.rb

@@ -39,7 +39,7 @@
 
     expected_counts_by_table = expected_counts.transform_keys(&:table_name)
 
-    table_extractor = /SELECT .* FROM "(.+)" WHERE/
+    table_extractor = /SELECT .* FROM "(\w+)" (WHERE|INNER JOIN)/
     actual_counts_by_table = @actual_queries.group_by do |query|
       table_extractor.match(query)[1]
     end.transform_values(&:size)