Added preloading of inverse association

Author: Tom Smyth

lib/closure_tree/has_closure_tree_root.rb

@@ -29,16 +29,23 @@ def has_closure_tree_root(assoc_name, options = {})
         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|
+          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|
           id_hash[node.id] = node
           parent_node = id_hash[node[parent_col_id]]
 
-          # Preload parent association
+          # Pre-assign parent association
           parent_assoc = node.association(:parent)
           parent_assoc.loaded!
           parent_assoc.target = parent_node
 
-          # Preload children association as empty for now,
+          # Pre-assign children association as empty for now,
           # children will be added in subsequent loop iterations
           children_assoc = node.association(:children)
           children_assoc.loaded!
@@ -49,6 +56,13 @@ def has_closure_tree_root(assoc_name, options = {})
             # Capture the root we're going to use
             root = node
           end
+
+          # Pre-assign inverse association back to this class, if it exists on target class.
+          if inverse
+            inverse_assoc = node.association(inverse.name)
+            inverse_assoc.loaded!
+            inverse_assoc.target = self
+          end
         end
 
         root

spec/db/models.rb

@@ -58,6 +58,7 @@ class User < ActiveRecord::Base
     :hierarchy_table_name => 'referral_hierarchies'
 
   has_many :contracts, inverse_of: :user
+  belongs_to :group # Can't use and don't need inverse_of here when using has_closure_tree_root.
 
   def indirect_contracts
     Contract.where(:user_id => descendant_ids)

spec/has_closure_tree_root_spec.rb

@@ -9,6 +9,7 @@
   let!(:user4) { User.create!(email: "[email protected]", group_id: group.id) }
   let!(:user5) { User.create!(email: "[email protected]", group_id: group.id) }
   let!(:user6) { User.create!(email: "[email protected]", group_id: group.id) }
+  let!(:group_reloaded) { group.class.find(group.id) } # Ensures were starting fresh.
 
   before do
     # The tree (contract types in parens)
@@ -37,13 +38,9 @@
   context "with basic config" do
     let!(:group) { Group.create!(name: "TheGroup") }
 
-    before do
-      group = Group.first # Ensure we're starting fresh
-    end
-
     it "loads all nodes and associations in a constant number of queries" do
       expect do
-        root = group.root_user_including_tree(contracts: :contract_type)
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
         expect(root.children[0].email).to eq "[email protected]"
         expect(root.children[0].parent.children[1].email).to eq "[email protected]"
         expect(root.children[1].contracts.map(&:contract_type).map(&:name)).to eq %w(Type1 Type2)
@@ -53,9 +50,17 @@
       end.to_not exceed_query_limit(4) # Without this feature, this is 15, and scales with number of nodes.
     end
 
+    it "eager loads inverse association to group" do
+      expect do
+        root = group_reloaded.root_user_including_tree
+        expect(root.group).to eq group
+        expect(root.children[0].group).to eq group
+      end.to_not exceed_query_limit(2)
+    end
+
     it "works if eager load association map is not given" do
       expect do
-        root = group.root_user_including_tree
+        root = group_reloaded.root_user_including_tree
         expect(root.children[0].email).to eq "[email protected]"
         expect(root.children[0].parent.children[1].children[0].email).to eq "[email protected]"
       end.to_not exceed_query_limit(2)
@@ -74,7 +79,7 @@
 
       it "should error" do
         expect do
-          root = group.root_user_including_tree(contracts: :contract_type)
+          root = group_reloaded.root_user_including_tree(contracts: :contract_type)
         end.to raise_error(ClosureTree::MultipleRootError)
       end
     end
@@ -84,7 +89,7 @@
     let(:group) { Grouping.create!(name: "TheGrouping") }
 
     it "should still work" do
-      root = group.root_person_including_tree(contracts: :contract_type)
+      root = group_reloaded.root_person_including_tree(contracts: :contract_type)
       expect(root.children[0].email).to eq "[email protected]"
     end
   end
@@ -94,7 +99,7 @@
 
     it "should error" do
       expect do
-        root = group.root_user_including_tree(contracts: :contract_type)
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
       end.to raise_error(NameError)
     end
   end
@@ -104,7 +109,7 @@
 
     it "should error" do
       expect do
-        root = group.root_user_including_tree(contracts: :contract_type)
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
       end.to raise_error(ActiveRecord::StatementInvalid)
     end
   end