Poly helpers

lib/closure_tree/hierarchy_maintenance.rb

@@ -66,10 +66,10 @@ def rebuild!(called_by_rebuild = false)
         unless root?
           _ct.connection.execute <<-SQL.squish
             INSERT INTO #{_ct.quoted_hierarchy_table_name}
-              (ancestor_id, descendant_id, generations)
-            SELECT x.ancestor_id, #{_ct.quote(_ct_id)}, x.generations + 1
+              (ancestor_id, ancestor_type, descendant_id, descendant_type, generations)
+            SELECT x.ancestor_id, x.ancestor_type, #{_ct.quote(_ct_id)}, #{_ct.quote(_ct.model_class.to_s)}, x.generations + 1
             FROM #{_ct.quoted_hierarchy_table_name} x
-            WHERE x.descendant_id = #{_ct.quote(_ct_parent_id)}
+            WHERE x.descendant_id = #{_ct.quote(_ct_parent_id)} AND x.descendant_type = #{_ct.quote(parent.class.to_s)}
           SQL
         end
 
@@ -97,8 +97,8 @@ def delete_hierarchy_references
             SELECT DISTINCT descendant_id
             FROM (SELECT descendant_id
               FROM #{_ct.quoted_hierarchy_table_name}
-              WHERE ancestor_id = #{_ct.quote(id)}
-                 OR descendant_id = #{_ct.quote(id)}
+              WHERE (ancestor_id = #{_ct.quote(id)} AND ancestor_type = #{_ct.quote(self.class.to_s)})
+                 OR (descendant_id = #{_ct.quote(id)} AND descendant_type = #{_ct.quote(self.class.to_s)})
             ) #{ _ct.t_alias_keyword } x )
         SQL
       end
@@ -119,13 +119,15 @@ def cleanup!
         hierarchy_table = hierarchy_class.arel_table
 
         [:descendant_id, :ancestor_id].each do |foreign_key|
-          alias_name = foreign_key.to_s.split('_').first + "s"
+          key_prefix = foreign_key.to_s.split('_')
+          key_as_type = "#{key_prefix}_type".to_sym
+          alias_name = key_prefix.first + "s"
           alias_table = Arel::Table.new(table_name).alias(alias_name)
           arel_join = hierarchy_table.join(alias_table, Arel::Nodes::OuterJoin)
                                      .on(alias_table[primary_key].eq(hierarchy_table[foreign_key]))
                                      .join_sources
 
-          lonely_childs = hierarchy_class.joins(arel_join).where(alias_table[primary_key].eq(nil))
+          lonely_childs = hierarchy_class.joins(arel_join).where(alias_table[primary_key].eq(nil)).where(key_as_type => self.class.to_s)
           ids = lonely_childs.pluck(foreign_key)
 
           hierarchy_class.where(hierarchy_table[foreign_key].in(ids)).delete_all

lib/closure_tree/model.rb

@@ -5,17 +5,19 @@ module Model
     extend ActiveSupport::Concern
 
     included do
-
       belongs_to :parent, nil,
         class_name: _ct.model_class.to_s,
         foreign_key: _ct.parent_column_name,
         inverse_of: :children,
         touch: _ct.options[:touch],
-        optional: true
+        optional: true,
+        polymorphic: true
 
+      where_for_ancestors = { ancestor_type: _ct.model_class.to_s }
+      where_for_descendants = { descendant_type: _ct.model_class.to_s }
       order_by_generations = -> { Arel.sql("#{_ct.quoted_hierarchy_table_name}.generations ASC") }
 
-      has_many :children, *_ct.has_many_order_with_option, **{
+      has_many :children, *_ct.has_many_order_with_option_and_where({ parent_type: _ct.model_class.to_s }), **{
         class_name: _ct.model_class.to_s,
         foreign_key: _ct.parent_column_name,
         dependent: _ct.options[:dependent],
@@ -28,21 +30,60 @@ def hash_tree(options = {})
           end
         end
 
-      has_many :ancestor_hierarchies, *_ct.has_many_order_without_option(order_by_generations),
+      has_many :ancestor_hierarchies, *_ct.has_many_order_without_option_and_where(order_by_generations, where_for_ancestors),
         class_name: _ct.hierarchy_class_name,
         foreign_key: 'descendant_id'
 
       has_many :self_and_ancestors, *_ct.has_many_order_without_option(order_by_generations),
         through: :ancestor_hierarchies,
         source: :ancestor
 
-      has_many :descendant_hierarchies, *_ct.has_many_order_without_option(order_by_generations),
+      has_many :descendant_hierarchies, *_ct.has_many_order_without_option_and_where(order_by_generations, where_for_descendants),
         class_name: _ct.hierarchy_class_name,
         foreign_key: 'ancestor_id'
 
       has_many :self_and_descendants, *_ct.has_many_order_with_option(order_by_generations),
         through: :descendant_hierarchies,
-        source: :descendant
+        source: :descendant, source_type: _ct.model_class.to_s
+    end
+
+    def poly_children
+      scope = hierarchy_class.where(ancestor_id: id, ancestor_type: _ct.model_class.to_s, generations: 1)
+      map_and_group_relation(scope, :descendants)
+    end
+
+    def poly_self_and_descendants
+      scope = hierarchy_class.where(ancestor: self)
+      map_and_group_relation(scope, :descendants)
+    end
+
+    def poly_descendants
+      poly_self_and_descendants - [self]
+    end
+
+    def poly_self_and_ancestors
+      scope = hierarchy_class.where(descendant: self)
+      map_and_group_relation(scope, :ancestors)
+    end
+
+    def poly_ancestors
+      poly_self_and_ancestors - [self]
+    end
+
+    def map_and_group_relation(scope, group)
+      scope
+        .map { |h_data| group_by_map(h_data, group) }
+        .group_by { |h_data| h_data[:type] }
+        .map { |type, h_data_array| type.constantize.where(id: h_data_array.map { |h_data| h_data[:id] }) }
+        .flatten
+    end
+
+    def group_by_map(h_data, group)
+      if group == :ancestors
+        { type: h_data.ancestor_type, id: h_data.ancestor_id }
+      elsif group == :descendants
+        { type: h_data.descendant_type, id: h_data.descendant_id }
+      end
     end
 
     # Delegate to the Support instance on the class:
@@ -72,6 +113,20 @@ def root
       self_and_ancestors.where(_ct.parent_column_name.to_sym => nil).first
     end
 
+    def polymorphic_root
+      polymorphic_ancestor_lookup = {}
+
+      hierarchy_class.where(descendant_id: self.id).pluck(:ancestor_type, :ancestor_id).each do |ancestor_type, ancestor_id|
+        polymorphic_ancestor_lookup[ancestor_type] ||= []
+        polymorphic_ancestor_lookup[ancestor_type] << ancestor_id
+      end
+
+      polymorphic_ancestor_lookup.keys.each do |ancestor_class|
+        no_parent_ancestor = ancestor_class.constantize.where(id: polymorphic_ancestor_lookup[ancestor_class], parent_id: nil).first
+        return no_parent_ancestor if no_parent_ancestor.present?
+      end
+    end
+
     def leaves
       self_and_descendants.leaves
     end

lib/closure_tree/support.rb

@@ -32,34 +32,34 @@ def initialize(model_class, options)
     end
 
     def hierarchy_class_for_model
-      parent_class = model_class.module_parent
-      hierarchy_class = parent_class.const_set(short_hierarchy_class_name, Class.new(model_class.superclass))
-      use_attr_accessible = use_attr_accessible?
-      include_forbidden_attributes_protection = include_forbidden_attributes_protection?
-      model_class_name = model_class.to_s
-      hierarchy_class.class_eval do
-        include ActiveModel::ForbiddenAttributesProtection if include_forbidden_attributes_protection
-        belongs_to :ancestor, class_name: model_class_name
-        belongs_to :descendant, class_name: model_class_name
-        attr_accessible :ancestor, :descendant, :generations if use_attr_accessible
-        def ==(other)
-          self.class == other.class && ancestor_id == other.ancestor_id && descendant_id == other.descendant_id
-        end
-        alias :eql? :==
-        def hash
-          ancestor_id.hash << 31 ^ descendant_id.hash
+      @@hierarchy_class_for_model ||= begin
+        parent_class = model_class.module_parent
+        hierarchy_class = parent_class.const_set(short_hierarchy_class_name, Class.new(model_class.superclass))
+        use_attr_accessible = use_attr_accessible?
+        include_forbidden_attributes_protection = include_forbidden_attributes_protection?
+        hierarchy_class.class_eval do
+          include ActiveModel::ForbiddenAttributesProtection if include_forbidden_attributes_protection
+          belongs_to :ancestor, polymorphic: true
+          belongs_to :descendant, polymorphic: true
+          attr_accessible :ancestor, :descendant, :generations if use_attr_accessible
+          def ==(other)
+            self.class == other.class && ancestor_id == other.ancestor_id && descendant_id == other.descendant_id
+          end
+          alias :eql? :==
+          def hash
+            ancestor_id.hash << 31 ^ descendant_id.hash
+          end
         end
+        hierarchy_class.table_name = hierarchy_table_name
+        hierarchy_class
       end
-      hierarchy_class.table_name = hierarchy_table_name
-      hierarchy_class
     end
 
     def hierarchy_table_name
       # We need to use the table_name, not something like ct_class.to_s.demodulize + "_hierarchies",
       # because they may have overridden the table name, which is what we want to be consistent with
       # in order for the schema to make sense.
-      tablename = options[:hierarchy_table_name] ||
-        remove_prefix_and_suffix(table_name).singularize + "_hierarchies"
+      tablename = "hierarchies"
 
       ActiveRecord::Base.table_name_prefix + tablename + ActiveRecord::Base.table_name_suffix
     end
@@ -84,6 +84,10 @@ def has_many_order_without_option(order_by_opt)
       [lambda { order(order_by_opt.call) }]
     end
 
+    def has_many_order_without_option_and_where(order_by_opt, where_clause)
+      [lambda { where(where_clause).order(order_by_opt.call) }]
+    end
+
     def has_many_order_with_option(order_by_opt=nil)
       order_options = [order_by_opt, order_by].compact
       [lambda {
@@ -92,6 +96,14 @@ def has_many_order_with_option(order_by_opt=nil)
       }]
     end
 
+    def has_many_order_with_option_and_where(order_by_opt=nil, where_clause)
+      order_options = [order_by_opt, order_by].compact
+      [lambda {
+        order_options = order_options.map { |o| o.is_a?(Proc) ? o.call : o }
+        where(where_clause).order(order_options)
+      }]
+    end
+
     def ids_from(scope)
       scope.pluck(model_class.primary_key)
     end

lib/closure_tree/support_attributes.rb

@@ -17,7 +17,7 @@ def quoted_value(value)
     end
 
     def hierarchy_class_name
-      options[:hierarchy_class_name] || model_class.to_s + "Hierarchy"
+      "Hierarchy"
     end
 
     def primary_key_column

lib/generators/closure_tree/templates/create_hierarchies_table.rb.erb

@@ -1,16 +1,18 @@
 class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
   def change
-    create_table :<%= migration_name %>, id: false do |t|
+    create_table :hierarchies, id: false do |t|
       t.<%= primary_key_type %> :ancestor_id, null: false
+      t.string :ancestor_type, null: false
       t.<%= primary_key_type %> :descendant_id, null: false
+      t.string :descendant_type, null: false
       t.integer :generations, null: false
     end
 
-    add_index :<%= migration_name %>, [:ancestor_id, :descendant_id, :generations],
+    add_index :hierarchies, [:ancestor_id, :ancestor_type, :descendant_id, :descendant_type, :generations],
       unique: true,
-      name: "<%= file_name %>_anc_desc_idx"
+      name: "anc_desc_idx"
 
-    add_index :<%= migration_name -%>, [:descendant_id],
-      name: "<%= file_name %>_desc_idx"
+    add_index :hierarchies, [:descendant_id],
+      name: "desc_idx"
   end
 end

spec/polymorphic_spec.rb

@@ -0,0 +1,69 @@
+require 'spec_helper'
+
+RSpec.describe 'Polymorphic methods' do
+  it 'can have polymorphic parents' do
+    parent = Project.create!(name: 'parent')
+    child = Project.create!(name: 'child', parent: parent)
+    expect(child.parent).to eq(parent)
+
+    expect(parent.descendants).to include(child)
+  end
+
+  it 'can have different type children #poly_children' do
+    root_parent = Project.create!(name: 'root parent')
+    parent = Project.create!(name: 'parent', parent: root_parent)
+    child = Task.create!(name: 'child', parent: parent)
+    another_child = Task.create!(name: 'child', parent: parent)
+    sub_task = Task.create!(name: sub_task, parent: child)
+
+    expect(child.parent).to eq(parent)
+    expect(parent.poly_children).to include(child)
+    expect(parent.poly_children).to include(another_child)
+    expect(root_parent.poly_children).to include(parent)
+    expect(child.poly_children).to include(sub_task)
+  end
+
+  it 'can have poly_self_and_descendants and poly_descendants' do
+    root_parent = Project.create!(name: 'root parent')
+    parent = Project.create!(name: 'parent', parent: root_parent)
+    child = Task.create!(name: 'child', parent: parent)
+    another_child = Task.create!(name: 'child', parent: parent)
+    sub_task = Task.create!(name: sub_task, parent: child)
+
+    expect(root_parent.poly_self_and_descendants).to match_array([
+      root_parent,
+      parent,
+      child,
+      another_child,
+      sub_task,
+    ])
+
+    expect(root_parent.poly_descendants).to match_array([
+      parent,
+      child,
+      another_child,
+      sub_task,
+    ])
+
+    expect(child.poly_descendants).to match_array([
+      sub_task,
+    ])
+  end
+
+  it 'can have poly_self_and_ancestors' do
+    root_parent = Project.create!(name: 'root parent')
+    parent = Project.create!(name: 'parent', parent: root_parent)
+    child = Task.create!(name: 'child', parent: parent)
+
+    expect(child.poly_self_and_ancestors).to match_array([
+      child,
+      parent,
+      root_parent,
+    ])
+
+    expect(child.poly_ancestors).to match_array([
+      parent,
+      root_parent,
+    ])
+  end
+end

spec/spec_helper.rb

@@ -80,7 +80,7 @@ def sqlite?
     end
 
     config.after(:suite) do
-      FileUtils.remove_entry_secure ENV['FLOCK_DIR']
+      FileUtils.remove_entry_secure(ENV['FLOCK_DIR'], true)
     end
   end
 end
@@ -103,3 +103,15 @@ def sqlite?
 require_relative 'support/helpers'
 require_relative 'support/exceed_query_limit'
 require_relative 'support/query_counter'
+
+class Organziation < ActiveRecord::Base
+  has_closure_tree
+end
+
+class Project < ActiveRecord::Base
+  has_closure_tree
+end
+
+class Task < ActiveRecord::Base
+  has_closure_tree
+end