Adds max_descendants_depth to ObjectHierarchy

CE-port of 10546-fix-epic-depth-validation

2 files changed, 73 insertions, 10 deletions

diff --git a/lib/gitlab/object_hierarchy.rb b/lib/gitlab/object_hierarchy.rb
index f2772c733c7..65bcf5e6ec4 100644
--- a/lib/gitlab/object_hierarchy.rb
+++ b/lib/gitlab/object_hierarchy.rb
@@ -5,6 +5,8 @@ module Gitlab
   #
   # This class uses recursive CTEs and as a result will only work on PostgreSQL.
   class ObjectHierarchy
+    DEPTH_COLUMN = :depth
+
     attr_reader :ancestors_base, :descendants_base, :model
 
     # ancestors_base - An instance of ActiveRecord::Relation for which to
@@ -27,6 +29,17 @@ module Gitlab
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
+    # Returns the maximum depth starting from the base
+    # A base object with no children has a maximum depth of `1`
+    def max_descendants_depth
+      unless hierarchy_supported?
+        # This makes the return value consistent with the case where hierarchy is supported
+        return descendants_base.exists? ? 1 : nil
+      end
+
+      base_and_descendants(with_depth: true).maximum(DEPTH_COLUMN)
+    end
+
     # Returns the set of ancestors of a given relation, but excluding the given
     # relation
     #
@@ -64,10 +77,15 @@ module Gitlab
 
     # Returns a relation that includes the descendants_base set of objects
     # and all their descendants (recursively).
-    def base_and_descendants
-      return descendants_base unless hierarchy_supported?
-
-      read_only(base_and_descendants_cte.apply_to(model.all))
+    #
+    # When `with_depth` is `true`, a `depth` column is included where it starts with `1` for the base objects
+    # and incremented as we go down the descendant tree
+    def base_and_descendants(with_depth: false)
+      unless hierarchy_supported?
+        return with_depth ? descendants_base.select("1 as #{DEPTH_COLUMN}", objects_table[Arel.star]) : descendants_base
+      end
+
+      read_only(base_and_descendants_cte(with_depth: with_depth).apply_to(model.all))
     end
 
     # Returns a relation that includes the base objects, their ancestors,
@@ -124,10 +142,9 @@ module Gitlab
     # rubocop: disable CodeReuse/ActiveRecord
     def base_and_ancestors_cte(stop_id = nil, hierarchy_order = nil)
       cte = SQL::RecursiveCTE.new(:base_and_ancestors)
-      depth_column = :depth
 
       base_query = ancestors_base.except(:order)
-      base_query = base_query.select("1 as #{depth_column}", objects_table[Arel.star]) if hierarchy_order
+      base_query = base_query.select("1 as #{DEPTH_COLUMN}", objects_table[Arel.star]) if hierarchy_order
 
       cte << base_query
 
@@ -137,7 +154,7 @@ module Gitlab
         .where(objects_table[:id].eq(cte.table[:parent_id]))
         .except(:order)
 
-      parent_query = parent_query.select(cte.table[depth_column] + 1, objects_table[Arel.star]) if hierarchy_order
+      parent_query = parent_query.select(cte.table[DEPTH_COLUMN] + 1, objects_table[Arel.star]) if hierarchy_order
       parent_query = parent_query.where(cte.table[:parent_id].not_eq(stop_id)) if stop_id
 
       cte << parent_query
@@ -146,17 +163,23 @@ module Gitlab
     # rubocop: enable CodeReuse/ActiveRecord
 
     # rubocop: disable CodeReuse/ActiveRecord
-    def base_and_descendants_cte
+    def base_and_descendants_cte(with_depth: false)
       cte = SQL::RecursiveCTE.new(:base_and_descendants)
 
-      cte << descendants_base.except(:order)
+      base_query = descendants_base.except(:order)
+      base_query = base_query.select("1 as #{DEPTH_COLUMN}", objects_table[Arel.star]) if with_depth
+
+      cte << base_query
 
       # Recursively get all the descendants of the base set.
-      cte << model
+      descendants_query = model
         .from([objects_table, cte.table])
         .where(objects_table[:parent_id].eq(cte.table[:id]))
         .except(:order)
 
+      descendants_query = descendants_query.select(cte.table[DEPTH_COLUMN] + 1, objects_table[Arel.star]) if with_depth
+
+      cte << descendants_query
       cte
     end
     # rubocop: enable CodeReuse/ActiveRecord
diff --git a/spec/lib/gitlab/object_hierarchy_spec.rb b/spec/lib/gitlab/object_hierarchy_spec.rb
index 4700a7ad2e1..e6e9ae3223e 100644
--- a/spec/lib/gitlab/object_hierarchy_spec.rb
+++ b/spec/lib/gitlab/object_hierarchy_spec.rb
@@ -81,6 +81,24 @@ describe Gitlab::ObjectHierarchy, :postgresql do
       expect { relation.update_all(share_with_group_lock: false) }
         .to raise_error(ActiveRecord::ReadOnlyRecord)
     end
+
+    context 'when with_depth is true' do
+      let(:relation) do
+        described_class.new(Group.where(id: parent.id)).base_and_descendants(with_depth: true)
+      end
+
+      it 'includes depth in the results' do
+        object_depths = {
+          parent.id => 1,
+          child1.id => 2,
+          child2.id => 3
+        }
+
+        relation.each do |object|
+          expect(object.depth).to eq(object_depths[object.id])
+        end
+      end
+    end
   end
 
   describe '#descendants' do
@@ -91,6 +109,28 @@ describe Gitlab::ObjectHierarchy, :postgresql do
     end
   end
 
+  describe '#max_descendants_depth' do
+    subject { described_class.new(base_relation).max_descendants_depth }
+
+    context 'when base relation is empty' do
+      let(:base_relation) { Group.where(id: nil) }
+
+      it { expect(subject).to be_nil }
+    end
+
+    context 'when base has no children' do
+      let(:base_relation) { Group.where(id: child2) }
+
+      it { expect(subject).to eq(1) }
+    end
+
+    context 'when base has grandchildren' do
+      let(:base_relation) { Group.where(id: parent) }
+
+      it { expect(subject).to eq(3) }
+    end
+  end
+
   describe '#ancestors' do
     it 'includes only the ancestors' do
       relation = described_class.new(Group.where(id: child2)).ancestors