Implement preload_associations_lazily on relations (#36)

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## master
 
+- [PR [#36](https://github.com/DmitryTsepelev/ar_lazy_preload/pull/36)] Feature preload_associations_lazily, allow turn on automatic preload only for given ActiveRecord::Relation ([@mpospelov][])
+
 ## 0.3.2 (2020-07-21)
 
 - [PR [#33](https://github.com/DmitryTsepelev/ar_lazy_preload/pull/34)] Feature skip_preload, allow turn automatic preload off ([@OuYangJinTing][])

README.md

@@ -50,6 +50,16 @@ If you want to turn automatic preload off for a specific record, you can call `.
 users.first.skip_preload.posts # => SELECT * FROM posts WHERE user_id = ?
 ```
 
+### Relation auto preloading
+
+Another alternative for auto preloading is using relation `#preload_associations_lazily` method
+
+```ruby
+posts = User.preload_associations_lazily.flat_map(&:posts)
+# => SELECT * FROM users LIMIT 10
+# => SELECT * FROM posts WHERE user_id in (...)
+```
+
 ## Installation
 
 Add this line to your application's Gemfile, and you're all set:

lib/ar_lazy_preload/active_record/base.rb

@@ -5,6 +5,7 @@ module ArLazyPreload
   module Base
     def self.included(base)
       base.class.delegate :lazy_preload, to: :all
+      base.class.delegate :preload_associations_lazily, to: :all
     end
 
     attr_accessor :lazy_preload_context

lib/ar_lazy_preload/active_record/relation.rb

@@ -5,6 +5,8 @@
 module ArLazyPreload
   # ActiveRecord::Relation patch with lazy preloading support
   module Relation
+    attr_writer :preload_associations_lazily_setting
+
     # Enhanced #load method will check if association has not been loaded yet and add a context
     # for lazy preloading to loaded each record
     def load
@@ -13,12 +15,30 @@ def load
       if need_context
         Context.register(
           records: ar_lazy_preload_records,
-          association_tree: lazy_preload_values
+          association_tree: lazy_preload_values,
+          auto_preload: preload_associations_lazily_setting
         )
       end
       result
     end
 
+    # Lazily autoloads all association
+    # example:
+    #
+    #   users = User.preload_associations_lazily
+    #   users.each do |user|
+    #     user.posts.flat_map {|post| post.comments.map(&:id)}
+    #   end
+    #
+    # Same effect can be achieved by User.lazy_preload(posts: :comments)
+    def preload_associations_lazily
+      spawn.tap { |relation| relation.preload_associations_lazily_setting = true }
+    end
+
+    def preload_associations_lazily_setting
+      @preload_associations_lazily_setting ||= false
+    end
+
     # Specify relationships to be loaded lazily when association is loaded for the first time. For
     # example:
     #

lib/ar_lazy_preload/associated_context_builder.rb

@@ -23,7 +23,7 @@ def initialize(parent_context:, association_name:)
 
     # Takes all the associated records for the records, attached to the :parent_context and creates
     # a preloading context for them
-    def perform
+    def perform # rubocop:disable Metrics/MethodLength
       associated_records = parent_context.records.flat_map do |record|
         next if record.nil?
 
@@ -32,14 +32,18 @@ def perform
         reflection.collection? ? record_association.target : record_association.reader
       end
 
-      Context.register(records: associated_records, association_tree: child_association_tree)
+      Context.register(
+        records: associated_records,
+        association_tree: child_association_tree,
+        auto_preload: parent_context.auto_preload?
+      )
     end
 
     private
 
     def child_association_tree
       # `association_tree` is unnecessary when auto preload is enabled
-      return nil if ArLazyPreload.config.auto_preload?
+      return nil if parent_context.auto_preload?
 
       AssociationTreeBuilder.new(parent_context.association_tree).subtree_for(association_name)
     end

lib/ar_lazy_preload/context.rb

@@ -7,10 +7,10 @@
 module ArLazyPreload
   class Context
     # Initiates lazy preload context for given records
-    def self.register(records:, association_tree:)
+    def self.register(records:, association_tree:, auto_preload: false)
       return if records.empty?
 
-      if ArLazyPreload.config.auto_preload?
+      if ArLazyPreload.config.auto_preload? || auto_preload
         Contexts::AutoPreloadContext.new(records: records)
       elsif association_tree.any?
         Contexts::LazyPreloadContext.new(

lib/ar_lazy_preload/contexts/auto_preload_context.rb

@@ -4,6 +4,10 @@ module ArLazyPreload
   module Contexts
     # This class is responsible for automatic association preloading
     class AutoPreloadContext < BaseContext
+      def auto_preload?
+        true
+      end
+
       protected
 
       def association_needs_preload?(_association_name)

lib/ar_lazy_preload/contexts/base_context.rb

@@ -28,12 +28,17 @@ def try_preload_lazily(association_name)
         perform_preloading(association_name)
       end
 
+      def auto_preload?
+        false
+      end
+
       protected
 
       def association_needs_preload?(_association_name)
         raise NotImplementedError
       end
 
+
       private
 
       def perform_preloading(association_name)

spec/ar_lazy_preload/preload_associations_lazily_spec.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+describe "ActiveRecord::Relation.preload_associations_lazily" do
+  let!(:user1) { create(:user) }
+  let!(:user2) { create(:user) }
+
+  let!(:post) { create(:post, user: user1) }
+  let!(:comment1) { create(:comment, user: user1, post: post) }
+  let!(:comment1) { create(:comment, user: user2, post: post) }
+
+  describe "auto preloading" do
+    subject { Comment.preload_associations_lazily }
+
+    # SELECT "comments".* FROM "comments"
+    # SELECT "users".* FROM "users" WHERE "users"."id" IN (...)
+    it "loads association automatically" do
+      expect { subject.each { |comment| comment.user&.id } }.to make_database_queries(count: 2)
+    end
+  end
+end