Add has_closure_tree_roots (plural) to allow has_many in related AR models

lib/closure_tree/has_closure_tree_root.rb

@@ -3,14 +3,27 @@ class MultipleRootError < StandardError; end
   class RootOrderingDisabledError < StandardError; end
 
   module HasClosureTreeRoot
-
     def has_closure_tree_root(assoc_name, options = {})
-       options[:class_name] ||= assoc_name.to_s.sub(/\Aroot_/, "").classify
+      options[:class_name] ||= assoc_name.to_s.sub(/\Aroot_/, "").classify
       options[:foreign_key] ||= self.name.underscore << "_id"
 
       has_one assoc_name, -> { where(parent: nil) }, **options
+      define_closure_tree_method(assoc_name, options, allow_multiple_roots: false)
+      connection_pool.release_connection
+    end
 
-      # Fetches the association, eager loading all children and given associations
+    def has_closure_tree_roots(assoc_name, options = {})
+      options[:class_name] ||= assoc_name.to_s.sub(/\Aroots_/, "").classify
+      options[:foreign_key] ||= self.name.underscore << "_id"
+
+      has_many assoc_name, -> { where(parent: nil) }, **options
+      define_closure_tree_method(assoc_name, options, allow_multiple_roots: true)
+      connection_pool.release_connection
+    end
+
+    private
+
+    def define_closure_tree_method(assoc_name, options, allow_multiple_roots: false)
       define_method("#{assoc_name}_including_tree") do |*args|
         reload = false
         reload = args.shift if args && (args.first == true || args.first == false)
@@ -28,25 +41,45 @@ def has_closure_tree_root(assoc_name, options = {})
 
         roots = options[:class_name].constantize.where(parent: nil, options[:foreign_key] => id).to_a
 
-        return nil if roots.empty?
+        if roots.empty?
+          return allow_multiple_roots ? [] : nil
+        end
 
-        if roots.size > 1
+        unless allow_multiple_roots || roots.size <= 1
           raise MultipleRootError.new("#{self.class.name}: has_closure_tree_root requires a single root")
         end
 
-        temp_root = roots.first
-        root = nil
-        id_hash = {}
-        parent_col_id = temp_root.class._ct.options[:parent_column_name]
-
         # Lookup inverse belongs_to association reflection on target class.
-        inverse = temp_root.class.reflections.values.detect do |r|
+        inverse = roots.first.class.reflections.values.detect do |r|
           r.macro == :belongs_to && r.klass == self.class
         end
 
         # Fetch all descendants in constant number of queries.
         # This is the last query-triggering statement in the method.
-        temp_root.self_and_descendants.includes(*assoc_map).each do |node|
+        nodes_to_process = if allow_multiple_roots && roots.size > 0
+          # For multiple roots, fetch all descendants at once to avoid N+1 queries
+          root_ids = roots.map(&:id)
+          klass = roots.first.class
+          hierarchy_table = klass._ct.quoted_hierarchy_table_name
+
+          # Get all descendants of all roots including the roots themselves
+          # (generations 0 = self, > 0 = descendants)
+          descendant_scope = klass.
+            joins("INNER JOIN #{hierarchy_table} ON #{hierarchy_table}.descendant_id = #{klass.quoted_table_name}.#{klass.primary_key}").
+            where("#{hierarchy_table}.ancestor_id IN (?)", root_ids).
+            includes(*assoc_map).
+            distinct
+
+          descendant_scope
+        else
+          roots.first.self_and_descendants.includes(*assoc_map)
+        end
+
+        id_hash = {}
+        parent_col_id = roots.first.class._ct.options[:parent_column_name]
+        root = nil
+
+        nodes_to_process.each do |node|
           id_hash[node.id] = node
           parent_node = id_hash[node[parent_col_id]]
 
@@ -62,8 +95,8 @@ def has_closure_tree_root(assoc_name, options = {})
 
           if parent_node
             parent_node.association(:children).target << node
-          else
-            # Capture the root we're going to use
+          elsif !allow_multiple_roots
+            # Capture the root we're going to use (only for single root case)
             root = node
           end
 
@@ -75,10 +108,12 @@ def has_closure_tree_root(assoc_name, options = {})
           end
         end
 
-        @closure_tree_roots[assoc_name][assoc_map] = root
-      end
+        result = allow_multiple_roots ? 
+          roots.map { |root| id_hash[root.id] } : 
+          root
 
-      connection_pool.release_connection
+        @closure_tree_roots[assoc_name][assoc_map] = result
+      end
     end
   end
 end

spec/closure_tree/has_closure_tree_roots_spec.rb

@@ -0,0 +1,223 @@
+require "spec_helper"
+
+RSpec.describe "has_closure_tree_roots" do
+  let!(:post) { Post.create!(title: "Test Post") }
+  let!(:post_reloaded) { post.class.find(post.id) } # Ensures we're starting fresh
+
+  before do
+    # Create a structure like this:
+    # Post
+    #  |- Comment1
+    #  |   |- Reply1-1
+    #  |   |- Reply1-2
+    #  |       |- Reply1-2-1
+    #  |- Comment2
+    #      |- Reply2-1
+
+    @comment1 = Comment.create!(body: "Top comment 1", post: post)
+    @comment2 = Comment.create!(body: "Top comment 2", post: post)
+
+    @reply1_1 = Comment.create!(body: "Reply 1-1", post: post, parent: @comment1)
+    @reply1_2 = Comment.create!(body: "Reply 1-2", post: post, parent: @comment1)
+    @reply2_1 = Comment.create!(body: "Reply 2-1", post: post, parent: @comment2)
+
+    @reply1_2_1 = Comment.create!(body: "Reply 1-2-1", post: post, parent: @reply1_2)
+  end
+
+  context "with basic config" do
+    it "loads all root comments in a constant number of queries" do
+      expect do
+        roots = post_reloaded.comments_including_tree
+        expect(roots.size).to eq 2
+        expect(roots[0].body).to eq "Top comment 1"
+        expect(roots[1].body).to eq "Top comment 2"
+        expect(roots[0].children[0].body).to eq "Reply 1-1"
+        expect(roots[0].children[1].body).to eq "Reply 1-2"
+        expect(roots[0].children[1].children[0].body).to eq "Reply 1-2-1"
+      end.to_not exceed_query_limit(2)
+    end
+
+    it "eager loads inverse association to post" do
+      expect do
+        roots = post_reloaded.comments_including_tree
+        expect(roots[0].post).to eq post
+        expect(roots[1].post).to eq post
+        expect(roots[0].children[0].post).to eq post
+        expect(roots[0].children[1].children[0].post).to eq post
+      end.to_not exceed_query_limit(2)
+    end
+
+    it "memoizes by assoc_map" do
+      post_reloaded.comments_including_tree.first.body = "changed1"
+      expect(post_reloaded.comments_including_tree.first.body).to eq "changed1"
+      expect(post_reloaded.comments_including_tree(true).first.body).to eq "Top comment 1"
+    end
+
+    it "works if true passed on first call" do
+      expect(post_reloaded.comments_including_tree(true).first.body).to eq "Top comment 1"
+    end
+
+    it "loads all nodes plus single association in a constant number of queries" do
+      # Add some attributes to test with - similar to contracts in the root spec
+      @comment1.update!(likes_count: 10)
+      @comment2.update!(likes_count: 5)
+      @reply1_1.update!(likes_count: 3)
+      @reply1_2.update!(likes_count: 7)
+      @reply2_1.update!(likes_count: 2)
+      @reply1_2_1.update!(likes_count: 4)
+
+      expect do
+        roots = post_reloaded.comments_including_tree
+        expect(roots.size).to eq 2
+        expect(roots[0].body).to eq "Top comment 1"
+        expect(roots[0].likes_count).to eq 10
+        expect(roots[0].children[1].likes_count).to eq 7
+        expect(roots[0].children[1].children[0].likes_count).to eq 4
+        expect(roots[1].children[0].body).to eq "Reply 2-1"
+      end.to_not exceed_query_limit(2)
+    end
+
+    it "loads all nodes and nested associations in a constant number of queries" do
+      # Create some nested associations to test with
+      user1 = User.create!(email: "[email protected]")
+      user2 = User.create!(email: "[email protected]")
+      user3 = User.create!(email: "[email protected]")
+
+      # Create comment_likes instead of using contracts
+      @comment1.comment_likes.create!(user: user1)
+      @comment2.comment_likes.create!(user: user2)
+      @reply1_1.comment_likes.create!(user: user3)
+
+      expect do
+        roots = post_reloaded.comments_including_tree(comment_likes: :user)
+        expect(roots.size).to eq 2
+        expect(roots[0].body).to eq "Top comment 1"
+        expect(roots[0].comment_likes.first.user.email).to eq "[email protected]"
+        expect(roots[1].comment_likes.first.user.email).to eq "[email protected]"
+        expect(roots[0].children[0].comment_likes.first.user.email).to eq "[email protected]"
+      end.to_not exceed_query_limit(4) # Without optimization, this would scale with number of nodes
+    end
+
+    context "with no comment roots" do
+      let(:empty_post) { Post.create!(title: "Empty Post") }
+
+      it "should return empty array" do
+        expect(empty_post.comments_including_tree).to eq([])
+      end
+    end
+
+    it "works if eager load association map is not given" do
+      expect do
+        roots = post_reloaded.comments_including_tree
+        expect(roots.size).to eq 2
+        expect(roots[0].body).to eq "Top comment 1"
+        expect(roots[0].children[1].children[0].body).to eq "Reply 1-2-1"
+      end.to_not exceed_query_limit(2)
+    end
+  end
+
+  context "when comment is destroyed" do
+    it "properly maintains the hierarchy" do
+      @comment1.destroy
+      roots = post_reloaded.comments_including_tree
+      expect(roots.size).to eq 1
+      expect(roots[0].body).to eq "Top comment 2"
+      expect(roots[0].children[0].body).to eq "Reply 2-1"
+    end
+  end
+
+  context "when comment is added after initial load" do
+    it "includes the new comment when reloaded" do
+      roots = post_reloaded.comments_including_tree
+      expect(roots.size).to eq 2
+
+      new_comment = Comment.create!(body: "New top comment", post: post)
+
+      # Should be memoized, so still 2
+      expect(post_reloaded.comments_including_tree.size).to eq 2
+
+      # With true, should reload and find 3
+      expect(post_reloaded.comments_including_tree(true).size).to eq 3
+    end
+  end
+
+  context "with no tree root" do
+    let(:empty_post) { Post.create!(title: "Empty Post") }
+
+    it "should return []" do
+      expect(empty_post.comments_including_tree).to eq([])
+    end
+  end
+
+  context "with explicit class_name and foreign_key" do
+    before do
+      # Create a model similar to Grouping in the models.rb file
+      class ForumPost < ApplicationRecord
+        self.table_name = "posts"
+        has_closure_tree_roots :thread_comments, class_name: 'Comment', foreign_key: 'post_id'
+      end
+
+      # Create the post and comments - reusing the same ones from above for simplicity
+      @post_collection = ForumPost.find(post.id)
+      @post_collection_reloaded = @post_collection.class.find(@post_collection.id)
+    end
+
+    after do
+      # Clean up our dynamically created class after the test
+      Object.send(:remove_const, :ForumPost) if Object.const_defined?(:ForumPost)
+    end
+
+    it "should still work" do
+      roots = @post_collection_reloaded.thread_comments_including_tree
+      expect(roots.size).to eq 2
+      expect(roots[0].body).to eq "Top comment 1"
+      expect(roots[0].children[1].body).to eq "Reply 1-2"
+    end
+  end
+
+  context "with bad class_name" do
+    before do
+      # Create a model with an invalid class_name
+      class BadClassPost < ApplicationRecord
+        self.table_name = "posts"
+        has_closure_tree_roots :invalid_comments, class_name: 'NonExistentComment'
+      end
+
+      @bad_class_post = BadClassPost.find(post.id)
+      @bad_class_post_reloaded = @bad_class_post.class.find(@bad_class_post.id)
+    end
+
+    after do
+      Object.send(:remove_const, :BadClassPost) if Object.const_defined?(:BadClassPost)
+    end
+
+    it "should error" do
+      expect do
+        @bad_class_post_reloaded.invalid_comments_including_tree
+      end.to raise_error(NameError)
+    end
+  end
+
+  context "with bad foreign_key" do
+    before do
+      # Create a model with an invalid foreign_key
+      class BadKeyPost < ApplicationRecord
+        self.table_name = "posts"
+        has_closure_tree_roots :broken_comments, class_name: 'Comment', foreign_key: 'nonexistent_id'
+      end
+
+      @bad_key_post = BadKeyPost.find(post.id)
+      @bad_key_post_reloaded = @bad_key_post.class.find(@bad_key_post.id)
+    end
+
+    after do
+      Object.send(:remove_const, :BadKeyPost) if Object.const_defined?(:BadKeyPost)
+    end
+
+    it "should error" do
+      expect do
+        @bad_key_post_reloaded.broken_comments_including_tree
+      end.to raise_error(ActiveRecord::StatementInvalid)
+    end
+  end
+end 

spec/support/models.rb

@@ -48,6 +48,28 @@ def add_destroyed_tag
 class DestroyedTag < ApplicationRecord
 end
 
+class Post < ApplicationRecord
+  # A post may have many comments, and each comment may have many replies. It makes little
+  # sense to have a "Root comment" in this context, so we use has_closure_tree_roots.
+  has_closure_tree_roots :comments
+end
+
+class Comment < ApplicationRecord
+  has_closure_tree dependent: :destroy
+  belongs_to :post
+  belongs_to :user, optional: true
+  has_many :comment_likes
+
+  # This is just for testing the eager loading of associations
+  attribute :likes_count, :integer
+end
+
+# To test eager loading of descendants with multiple roots, we need a model that has many
+class CommentLike < ApplicationRecord
+  belongs_to :comment
+  belongs_to :user
+end
+
 class Group < ApplicationRecord
   has_closure_tree_root :root_user
 end