has_closure_tree_root - Better tree eager loading

lib/closure_tree.rb

@@ -4,6 +4,7 @@ module ClosureTree
   extend ActiveSupport::Autoload
 
   autoload :HasClosureTree
+  autoload :HasClosureTreeRoot
   autoload :Support
   autoload :HierarchyMaintenance
   autoload :Model
@@ -25,4 +26,5 @@ def self.configuration
 
 ActiveSupport.on_load :active_record do
   ActiveRecord::Base.send :extend, ClosureTree::HasClosureTree
+  ActiveRecord::Base.send :extend, ClosureTree::HasClosureTreeRoot
 end

lib/closure_tree/has_closure_tree_root.rb

@@ -0,0 +1,88 @@
+module ClosureTree
+  class MultipleRootError < StandardError; end
+
+  module HasClosureTreeRoot
+
+    def has_closure_tree_root(assoc_name, options = {})
+      options.assert_valid_keys(
+        :class_name,
+        :foreign_key
+      )
+
+      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
+
+      # Fetches the association, eager loading all children and given associations
+      define_method("#{assoc_name}_including_tree") do |assoc_map_or_reload = nil, assoc_map = nil|
+        reload = false
+        if assoc_map_or_reload.is_a?(::Hash)
+          assoc_map = assoc_map_or_reload
+        else
+          reload = assoc_map_or_reload
+        end
+
+        unless reload
+          # Memoize
+          @closure_tree_roots ||= {}
+          @closure_tree_roots[assoc_name] ||= {}
+          if @closure_tree_roots[assoc_name].has_key?(assoc_map)
+            return @closure_tree_roots[assoc_name][assoc_map]
+          end
+        end
+
+        roots = options[:class_name].constantize.where(parent: nil, options[:foreign_key] => id).to_a
+
+        return nil if roots.empty?
+
+        if 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|
+          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]]
+
+          # Pre-assign parent association
+          parent_assoc = node.association(:parent)
+          parent_assoc.loaded!
+          parent_assoc.target = parent_node
+
+          # Pre-assign children association as empty for now,
+          # children will be added in subsequent loop iterations
+          children_assoc = node.association(:children)
+          children_assoc.loaded!
+
+          if parent_node
+            parent_node.association(:children).target << node
+          else
+            # 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
+
+        @closure_tree_roots[assoc_name][assoc_map] = root
+      end
+    end
+  end
+end

spec/db/models.rb

@@ -35,13 +35,30 @@ def add_destroyed_tag
 class DestroyedTag < ActiveRecord::Base
 end
 
+class Group < ActiveRecord::Base
+  has_closure_tree_root :root_user
+end
+
+class Grouping < ActiveRecord::Base
+  has_closure_tree_root :root_person, class_name: "User", foreign_key: :group_id
+end
+
+class UserSet < ActiveRecord::Base
+  has_closure_tree_root :root_user, class_name: "Useur"
+end
+
+class Team < ActiveRecord::Base
+  has_closure_tree_root :root_user, class_name: "User", foreign_key: :grp_id
+end
+
 class User < ActiveRecord::Base
   acts_as_tree :parent_column_name => "referrer_id",
     :name_column => 'email',
     :hierarchy_class_name => 'ReferralHierarchy',
     :hierarchy_table_name => 'referral_hierarchies'
 
-  has_many :contracts
+  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)
@@ -53,7 +70,12 @@ def to_s
 end
 
 class Contract < ActiveRecord::Base
-  belongs_to :user
+  belongs_to :user, inverse_of: :contracts
+  belongs_to :contract_type, inverse_of: :contracts
+end
+
+class ContractType < ActiveRecord::Base
+  has_many :contracts, inverse_of: :contract_type
 end
 
 class Label < ActiveRecord::Base

spec/db/schema.rb

@@ -43,16 +43,38 @@
   add_index "tag_hierarchies", [:ancestor_id, :descendant_id, :generations], :unique => true, :name => "tag_anc_desc_idx"
   add_index "tag_hierarchies", [:descendant_id], :name => "tag_desc_idx"
 
+  create_table "groups" do |t|
+    t.string "name", null: false
+  end
+
+  create_table "groupings" do |t|
+    t.string "name", null: false
+  end
+
+  create_table "user_sets" do |t|
+    t.string "name", null: false
+  end
+
+  create_table "teams" do |t|
+    t.string "name", null: false
+  end
+
   create_table "users" do |t|
     t.string "email"
     t.integer "referrer_id"
+    t.integer "group_id"
     t.timestamps null: false
   end
 
   add_foreign_key(:users, :users, :column => 'referrer_id')
 
   create_table "contracts" do |t|
     t.integer "user_id", :null => false
+    t.integer "contract_type_id"
+  end
+
+  create_table "contract_types" do |t|
+    t.string "name", :null => false
   end
 
   create_table "referral_hierarchies", :id => false do |t|

spec/has_closure_tree_root_spec.rb

@@ -0,0 +1,132 @@
+require "spec_helper"
+
+describe "has_closure_tree_root" do
+  let!(:ct1) { ContractType.create!(name: "Type1") }
+  let!(:ct2) { ContractType.create!(name: "Type2") }
+  let!(:user1) { User.create!(email: "1@example.com", group_id: group.id) }
+  let!(:user2) { User.create!(email: "2@example.com", group_id: group.id) }
+  let!(:user3) { User.create!(email: "3@example.com", group_id: group.id) }
+  let!(:user4) { User.create!(email: "4@example.com", group_id: group.id) }
+  let!(:user5) { User.create!(email: "5@example.com", group_id: group.id) }
+  let!(:user6) { User.create!(email: "6@example.com", group_id: group.id) }
+  let!(:group_reloaded) { group.class.find(group.id) } # Ensures were starting fresh.
+
+  before do
+    # The tree (contract types in parens)
+    #
+    #                   U1(1)
+    #                  /    \
+    #              U2(1)   U3(1&2)
+    #             /        /     \
+    #         U4(2)      U5(1)   U6(2)
+
+    user1.children << user2
+    user1.children << user3
+    user2.children << user4
+    user3.children << user5
+    user3.children << user6
+
+    user1.contracts.create!(contract_type: ct1)
+    user2.contracts.create!(contract_type: ct1)
+    user3.contracts.create!(contract_type: ct1)
+    user3.contracts.create!(contract_type: ct2)
+    user4.contracts.create!(contract_type: ct2)
+    user5.contracts.create!(contract_type: ct1)
+    user6.contracts.create!(contract_type: ct2)
+  end
+
+  context "with basic config" do
+    let!(:group) { Group.create!(name: "TheGroup") }
+
+    it "loads all nodes and associations in a constant number of queries" do
+      expect do
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
+        expect(root.children[0].email).to eq "2@example.com"
+        expect(root.children[0].parent.children[1].email).to eq "3@example.com"
+        expect(root.children[1].contracts.map(&:contract_type).map(&:name)).to eq %w(Type1 Type2)
+        expect(root.children[1].children[0].contracts[0].contract_type.name).to eq "Type1"
+        expect(root.children[0].children[0].contracts[0].user.
+          parent.parent.children[1].children[1].contracts[0].contract_type.name).to eq "Type2"
+      end.to_not exceed_query_limit(4) # Without this feature, this is 15, and scales with number of nodes.
+    end
+
+    it "memoizes by assoc_map" do
+      group_reloaded.root_user_including_tree.email = "x"
+      expect(group_reloaded.root_user_including_tree.email).to eq "x"
+      group_reloaded.root_user_including_tree(contracts: :contract_type).email = "y"
+      expect(group_reloaded.root_user_including_tree(contracts: :contract_type).email).to eq "y"
+      expect(group_reloaded.root_user_including_tree.email).to eq "x"
+    end
+
+    it "doesn't memoize if true argument passed" do
+      group_reloaded.root_user_including_tree.email = "x"
+      expect(group_reloaded.root_user_including_tree(true).email).to eq "1@example.com"
+      group_reloaded.root_user_including_tree(contracts: :contract_type).email = "y"
+      expect(group_reloaded.root_user_including_tree(true, contracts: :contract_type).email).
+        to eq "1@example.com"
+    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_reloaded.root_user_including_tree
+        expect(root.children[0].email).to eq "2@example.com"
+        expect(root.children[0].parent.children[1].children[0].email).to eq "5@example.com"
+      end.to_not exceed_query_limit(2)
+    end
+
+    context "with no tree root" do
+      let(:group2) { Group.create!(name: "OtherGroup") }
+
+      it "should return nil" do
+        expect(group2.root_user_including_tree(contracts: :contract_type)).to be_nil
+      end
+    end
+
+    context "with multiple tree roots" do
+      let!(:other_root) { User.create!(email: "10@example.com", group_id: group.id) }
+
+      it "should error" do
+        expect do
+          root = group_reloaded.root_user_including_tree(contracts: :contract_type)
+        end.to raise_error(ClosureTree::MultipleRootError)
+      end
+    end
+  end
+
+  context "with explicit class_name and foreign_key" do
+    let(:group) { Grouping.create!(name: "TheGrouping") }
+
+    it "should still work" do
+      root = group_reloaded.root_person_including_tree(contracts: :contract_type)
+      expect(root.children[0].email).to eq "2@example.com"
+    end
+  end
+
+  context "with bad class_name" do
+    let(:group) { UserSet.create!(name: "TheUserSet") }
+
+    it "should error" do
+      expect do
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
+      end.to raise_error(NameError)
+    end
+  end
+
+  context "with bad foreign_key" do
+    let(:group) { Team.create!(name: "TheTeam") }
+
+    it "should error" do
+      expect do
+        root = group_reloaded.root_user_including_tree(contracts: :contract_type)
+      end.to raise_error(ActiveRecord::StatementInvalid)
+    end
+  end
+end

spec/support/exceed_query_limit.rb

@@ -0,0 +1,18 @@
+# Derived from http://stackoverflow.com/a/13423584/153896. Updated for RSpec 3.
+RSpec::Matchers.define :exceed_query_limit do |expected|
+  supports_block_expectations
+
+  match do |block|
+    query_count(&block) > expected
+  end
+
+  failure_message_when_negated do |actual|
+    "Expected to run maximum #{expected} queries, got #{@counter.query_count}"
+  end
+
+  def query_count(&block)
+    @counter = ActiveRecord::QueryCounter.new
+    ActiveSupport::Notifications.subscribed(@counter.to_proc, 'sql.active_record', &block)
+    @counter.query_count
+  end
+end

spec/support/query_counter.rb

@@ -0,0 +1,18 @@
+# From http://stackoverflow.com/a/13423584/153896
+module ActiveRecord
+  class QueryCounter
+    attr_reader :query_count
+
+    def initialize
+      @query_count = 0
+    end
+
+    def to_proc
+      lambda(&method(:callback))
+    end
+
+    def callback(name, start, finish, message_id, values)
+      @query_count += 1 unless %w(CACHE SCHEMA).include?(values[:name])
+    end
+  end
+end