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

sergioisidoro · sergioisidoro · commit ad3baf79cea1 · 2025-04-01T15:22:54.000+03:00
diff --git a/lib/closure_tree/has_closure_tree_root.rb b/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
diff --git a/spec/closure_tree/has_closure_tree_roots_spec.rb b/spec/closure_tree/has_closure_tree_roots_spec.rb
@@ -0,0 +1,154 @@
+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: "comment1@example.com")
+      user2 = User.create!(email: "comment2@example.com")
+      user3 = User.create!(email: "reply1_1@example.com")
+      
+      # 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 "comment1@example.com"
+        expect(roots[1].comment_likes.first.user.email).to eq "comment2@example.com"
+        expect(roots[0].children[0].comment_likes.first.user.email).to eq "reply1_1@example.com"
+      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
+  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 nested comment creation" do
+    it "properly builds the hierarchy" do
+      # Create a new root comment with nested children
+      new_comment = Comment.new(body: "New root", post: post)
+      reply1 = Comment.new(body: "New reply 1", post: post)
+      reply2 = Comment.new(body: "New reply 2", post: post)
+      
+      new_comment.children << reply1
+      new_comment.children << reply2
+      
+      new_comment.save!
+      
+      roots = post_reloaded.comments_including_tree(true)
+      new_root = roots.find { |r| r.body == "New root" }
+      
+      expect(new_root.children.size).to eq 2
+      expect(new_root.children.map(&:body)).to include("New reply 1", "New reply 2")
+    end
+  end
+end 
diff --git a/spec/support/models.rb b/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
diff --git a/spec/support/schema.rb b/spec/support/schema.rb
@@ -138,6 +138,33 @@
     t.integer 'generations', null: false
   end
 
+  create_table 'posts' do |t|
+    t.string 'title'
+    t.text 'content'
+    t.timestamps null: false
+  end
+
+  create_table 'comments' do |t|
+    t.text 'body'
+    t.references 'post'
+    t.references 'parent'
+    t.references 'user'
+    t.integer 'likes_count'
+    t.timestamps null: false
+  end
+
+  create_table 'comment_likes' do |t|
+    t.references 'comment', null: false
+    t.references 'user', null: false
+    t.timestamps null: false
+  end
+
+  create_table 'comment_hierarchies', id: false do |t|
+    t.references 'ancestor', null: false
+    t.references 'descendant', null: false
+    t.integer 'generations', null: false
+  end
+
   add_index 'label_hierarchies', %i[ancestor_id descendant_id generations], unique: true,
                                                                             name: 'lh_anc_desc_idx'
   add_index 'label_hierarchies', [:descendant_id], name: 'lh_desc_idx'
@@ -154,4 +181,9 @@
   add_foreign_key(:menu_item_hierarchies, :menu_items, column: 'descendant_id', on_delete: :cascade)
   add_foreign_key(:tag_hierarchies, :tags, column: 'ancestor_id', on_delete: :cascade)
   add_foreign_key(:tag_hierarchies, :tags, column: 'descendant_id', on_delete: :cascade)
+  add_foreign_key(:comments, :comments, column: 'parent_id', on_delete: :cascade)
+  add_foreign_key(:comment_hierarchies, :comments, column: 'ancestor_id', on_delete: :cascade)
+  add_foreign_key(:comment_hierarchies, :comments, column: 'descendant_id', on_delete: :cascade)
+  add_foreign_key(:comment_likes, :comments, column: 'comment_id', on_delete: :cascade)
+  add_foreign_key(:comment_likes, :users, column: 'user_id', on_delete: :cascade)
 end
