Allow passing custom key to BatchLoader#batch

README.md

@@ -18,8 +18,10 @@ This gem provides a generic lazy batching mechanism to avoid N+1 DB queries, HTT
   * [RESTful API example](#restful-api-example)
   * [GraphQL example](#graphql-example)
   * [Loading multiple items](#loading-multiple-items)
+  * [Batch key](#batch-key)
   * [Caching](#caching)
 * [Installation](#installation)
+* [API](#api)
 * [Implementation details](#implementation-details)
 * [Development](#development)
 * [Contributing](#contributing)
@@ -300,6 +302,30 @@ BatchLoader.for(user.id).batch(default_value: []) do |comment_ids, loader|
 end
 ```
 
+### Batch key
+
+It's possible to reuse the same `BatchLoader#batch` block for loading different types of data by specifying a unique `key`.
+For example, with polymorphic associations:
+
+```ruby
+def lazy_association(post)
+  id = post.association_id
+  key = post.association_type
+
+  BatchLoader.for(id).batch(key: key) do |ids, loader, args|
+    model = Object.const_get(args[:key])
+    model.where(id: ids).each { |record| record.call(record.id, record) }
+  end
+end
+post1 = Post.save(association_id: 1, association_type: 'Tag')
+post2 = Post.save(association_id: 1, association_type: 'Category')
+
+lazy_association(post1) # SELECT * FROM tags WHERE id IN (1)
+lazy_association(post2) # SELECT * FROM categories WHERE id IN (1)
+```
+
+It's also required to pass custom `key` when using `BatchLoader` with metaprogramming (e.g. `eval`).
+
 ### Caching
 
 By default `BatchLoader` caches the loaded values. You can test it by running something like:
@@ -364,6 +390,24 @@ Or install it yourself as:
 
     $ gem install batch-loader
 
+## API
+
+```ruby
+BatchLoader.for(item).batch(default_value: default_value, cache: cache, key: key) do |items, loader, args|
+  # ...
+end
+```
+
+| Argument Key    | Default                                       | Description                                                   |
+| --------------- | --------------------------------------------- | ------------------------------------------------------------- |
+| `item`          | -                                             | Item which will be collected and used for batching.           |
+| `default_value` | `nil`                                         | Value returned by default after batching.                     |
+| `cache`         | `true`                                        | Set `false` to disable caching between the same executions.   |
+| `key`           | `nil`                                         | Pass custom key to uniquely identify the batch block.         |
+| `items`         | -                                             | List of collected items for batching.                         |
+| `loader`        | -                                             | Lambda which should be called to load values loaded in batch. |
+| `args`          | `{default_value: nil, cache: true, key: nil}` | Arguments passed to the `batch` method.                       |
+
 ## Implementation details
 
 See the [slides](https://speakerdeck.com/exaspark/batching-a-powerful-way-to-solve-n-plus-1-queries) [37-42].

lib/batch_loader.rb

@@ -23,9 +23,10 @@ def initialize(item:, executor_proxy: nil)
     @__executor_proxy = executor_proxy
   end
 
-  def batch(default_value: nil, cache: true, &batch_block)
+  def batch(default_value: nil, cache: true, key: nil, &batch_block)
     @default_value = default_value
     @cache = cache
+    @key = key
     @batch_block = batch_block
     __executor_proxy.add(item: @item)
 
@@ -78,7 +79,8 @@ def __ensure_batched
 
     items = __executor_proxy.list_items
     loader = __loader
-    @batch_block.call(items, loader)
+    args = {default_value: @default_value, cache: @cache, key: @key}
+    @batch_block.call(items, loader, args)
     items.each do |item|
       next if __executor_proxy.value_loaded?(item: item)
       loader.call(item, @default_value)
@@ -126,7 +128,7 @@ def __purge_cache
   def __executor_proxy
     @__executor_proxy ||= begin
       raise NoBatchError.new("Please provide a batch block first") unless @batch_block
-      BatchLoader::ExecutorProxy.new(@default_value, &@batch_block)
+      BatchLoader::ExecutorProxy.new(@default_value, @key, &@batch_block)
     end
   end
 

lib/batch_loader/executor_proxy.rb

@@ -6,10 +6,10 @@ class BatchLoader
   class ExecutorProxy
     attr_reader :default_value, :block, :global_executor
 
-    def initialize(default_value, &block)
+    def initialize(default_value, key, &block)
       @default_value = default_value
       @block = block
-      @block_hash_key = block.source_location
+      @block_hash_key = "#{key}#{block.source_location}"
       @global_executor = BatchLoader::Executor.ensure_current
     end
 

spec/batch_loader_spec.rb

@@ -87,6 +87,26 @@
     end
   end
 
+  context 'with custom key' do
+    it 'batches multiple items by key' do
+      author = Author.save(id: 1)
+      reader = Reader.save(id: 2)
+      batch_loader = ->(type, id) do
+        BatchLoader.for(id).batch(key: type) do |ids, loader, args|
+          args[:key].where(id: ids).each { |user| loader.call(user.id, user) }
+        end
+      end
+
+      laoded_author = batch_loader.call(Author, 1)
+      loader_reader = batch_loader.call(Reader, 2)
+
+      expect(Author).to receive(:where).with(id: [1]).once.and_call_original
+      expect(laoded_author).to eq(author)
+      expect(Reader).to receive(:where).with(id: [2]).once.and_call_original
+      expect(loader_reader).to eq(reader)
+    end
+  end
+
   context 'loader' do
     it 'loads the data even in a separate thread' do
       lazy = BatchLoader.for(1).batch do |nums, loader|
@@ -107,7 +127,7 @@
           thread.join
         end
       end
-      slow_executor_proxy = SlowExecutorProxy.new([], &batch_block)
+      slow_executor_proxy = SlowExecutorProxy.new([], nil, &batch_block)
       lazy = BatchLoader.new(item: 1, executor_proxy: slow_executor_proxy).batch(default_value: [], &batch_block)
 
       expect(lazy).to match_array([1, 2])

spec/fixtures/models.rb

@@ -38,22 +38,23 @@ class User
   class << self
     def save(id:)
       ensure_init_store
-      @users[id] = new(id: id)
+      @store[self][id] = new(id: id)
     end
 
     def where(id:)
       ensure_init_store
-      @users.each_with_object([]) { |(k, v), memo| memo << v if id.include?(k) }
+      @store[self].each_with_object([]) { |(k, v), memo| memo << v if id.include?(k) }
     end
 
     def destroy_all
-      @users = {}
+      ensure_init_store
+      @store[self] = {}
     end
 
     private
 
     def ensure_init_store
-      @users ||= {}
+      @store ||= Hash.new { |h, k| h[k] = {} }
     end
   end
 
@@ -73,3 +74,9 @@ def some_private_method
     :some_private_method
   end
 end
+
+class Author < User
+end
+
+class Reader < User
+end