Make it possible to limit ancestors in a `GroupHierarchy`

Passing a parent_id will limit ancestors upto the specified parent if
it is found.

Using `ancestors` and `descendants` the `base` relation will not be included

2 files changed, 55 insertions, 4 deletions

diff --git a/lib/gitlab/group_hierarchy.rb b/lib/gitlab/group_hierarchy.rb
index 635f52131f9..d88c415e1e5 100644
--- a/lib/gitlab/group_hierarchy.rb
+++ b/lib/gitlab/group_hierarchy.rb
@@ -17,12 +17,33 @@ module Gitlab
       @model = ancestors_base.model
     end
 
+    # Returns the set of descendants of a given relation, but excluding the given
+    # relation
+    def descendants
+      base_and_descendants.where.not(id: descendants_base.select(:id))
+    end
+
+    # Returns the set of ancestors of a given relation, but excluding the given
+    # relation
+    #
+    # Passing an `upto` will stop the recursion once the specified parent_id is
+    # reached. So all ancestors *lower* than the specified ancestor will be
+    # included.
+    def ancestors(upto: nil)
+      read_only(base_and_ancestors_cte(upto).apply_to(model.all))
+        .where.not(id: ancestors_base.select(:id))
+    end
+
     # Returns a relation that includes the ancestors_base set of groups
     # and all their ancestors (recursively).
-    def base_and_ancestors
+    #
+    # Passing an `upto` will stop the recursion once the specified parent_id is
+    # reached. So all ancestors *lower* than the specified acestor will be
+    # included.
+    def base_and_ancestors(upto: nil)
       return ancestors_base unless Group.supports_nested_groups?
 
-      read_only(base_and_ancestors_cte.apply_to(model.all))
+      read_only(base_and_ancestors_cte(upto).apply_to(model.all))
     end
 
     # Returns a relation that includes the descendants_base set of groups
@@ -78,17 +99,19 @@ module Gitlab
 
     private
 
-    def base_and_ancestors_cte
+    def base_and_ancestors_cte(stop_id = nil)
       cte = SQL::RecursiveCTE.new(:base_and_ancestors)
 
       cte << ancestors_base.except(:order)
 
       # Recursively get all the ancestors of the base set.
-      cte << model
+      parent_query = model
         .from([groups_table, cte.table])
         .where(groups_table[:id].eq(cte.table[:parent_id]))
         .except(:order)
+      parent_query = parent_query.where(cte.table[:parent_id].not_eq(stop_id)) if stop_id
 
+      cte << parent_query
       cte
     end
 
diff --git a/spec/lib/gitlab/group_hierarchy_spec.rb b/spec/lib/gitlab/group_hierarchy_spec.rb
index 8dc83a6db7f..30686634af4 100644
--- a/spec/lib/gitlab/group_hierarchy_spec.rb
+++ b/spec/lib/gitlab/group_hierarchy_spec.rb
@@ -18,6 +18,12 @@ describe Gitlab::GroupHierarchy, :postgresql do
       expect(relation).to include(parent, child1)
     end
 
+    it 'can find ancestors upto a certain level' do
+      relation = described_class.new(Group.where(id: child2)).base_and_ancestors(upto: child1)
+
+      expect(relation).to contain_exactly(child2)
+    end
+
     it 'uses ancestors_base #initialize argument' do
       relation = described_class.new(Group.where(id: child2.id), Group.none).base_and_ancestors
 
@@ -55,6 +61,28 @@ describe Gitlab::GroupHierarchy, :postgresql do
     end
   end
 
+  describe '#descendants' do
+    it 'includes only the descendants' do
+      relation = described_class.new(Group.where(id: parent)).descendants
+
+      expect(relation).to contain_exactly(child1, child2)
+    end
+  end
+
+  describe '#ancestors' do
+    it 'includes only the ancestors' do
+      relation = described_class.new(Group.where(id: child2)).ancestors
+
+      expect(relation).to contain_exactly(child1, parent)
+    end
+
+    it 'can find ancestors upto a certain level' do
+      relation = described_class.new(Group.where(id: child2)).ancestors(upto: child1)
+
+      expect(relation).to be_empty
+    end
+  end
+
   describe '#all_groups' do
     let(:relation) do
       described_class.new(Group.where(id: child1.id)).all_groups