1 files changed, 37 insertions, 6 deletions
diff --git a/lib/gitlab/group_hierarchy.rb b/lib/gitlab/group_hierarchy.rb
index 5a31e56cb30..42ded7c286f 100644
--- a/lib/gitlab/group_hierarchy.rb
+++ b/lib/gitlab/group_hierarchy.rb
@@ -17,12 +17,32 @@ 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)
+      base_and_ancestors(upto: upto).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?
 
-      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
@@ -30,7 +50,7 @@ module Gitlab
     def base_and_descendants
       return descendants_base unless Group.supports_nested_groups?
 
-      base_and_descendants_cte.apply_to(model.all)
+      read_only(base_and_descendants_cte.apply_to(model.all))
     end
 
     # Returns a relation that includes the base groups, their ancestors,
@@ -67,26 +87,30 @@ module Gitlab
       union = SQL::Union.new([model.unscoped.from(ancestors_table),
                               model.unscoped.from(descendants_table)])
 
-      model
+      relation = model
         .unscoped
         .with
         .recursive(ancestors.to_arel, descendants.to_arel)
         .from("(#{union.to_sql}) #{model.table_name}")
+
+      read_only(relation)
     end
 
     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
 
@@ -107,5 +131,12 @@ module Gitlab
     def groups_table
       model.arel_table
     end
+
+    def read_only(relation)
+      # relations using a CTE are not safe to use with update_all as it will
+      # throw away the CTE, hence we mark them as read-only.
+      relation.extend(Gitlab::Database::ReadOnlyRelation)
+      relation
+    end
   end
 end