Linear version of Namespace#ancestors_upto

Reduce query complexity and improve performance.

app/models/namespaces/traversal/linear.rb

@@ -64,6 +64,13 @@ module Namespaces
         traversal_ids.present?
       end
 
+      def use_traversal_ids_for_ancestors_upto?
+        return false unless use_traversal_ids?
+        return false unless Feature.enabled?(:use_traversal_ids_for_ancestors_upto, root_ancestor, default_enabled: :yaml)
+
+        traversal_ids.present?
+      end
+
       def use_traversal_ids_for_root_ancestor?
         return false unless Feature.enabled?(:use_traversal_ids_for_root_ancestor, default_enabled: :yaml)
 
@@ -114,6 +121,35 @@ module Namespaces
         hierarchy_order == :desc ? traversal_ids[0..-2] : traversal_ids[0..-2].reverse
       end
 
+      # Returns all ancestors upto but excluding the top.
+      # When no top is given, all ancestors are returned.
+      # When top is not found, returns all ancestors.
+      #
+      # This copies the behavior of the recursive method. We will deprecate
+      # this behavior soon.
+      def ancestors_upto(top = nil, hierarchy_order: nil)
+        return super unless use_traversal_ids_for_ancestors_upto?
+
+        # We can't use a default value in the method definition above because
+        # we need to preserve those specific parameters for super.
+        hierarchy_order ||= :desc
+
+        # Get all ancestor IDs inclusively between top and our parent.
+        top_index = top ? traversal_ids.find_index(top.id) : 0
+        ids = traversal_ids[top_index...-1]
+        ids_string = ids.map { |id| Integer(id) }.join(',')
+
+        # WITH ORDINALITY lets us order the result to match traversal_ids order.
+        from_sql = <<~SQL
+          unnest(ARRAY[#{ids_string}]::bigint[]) WITH ORDINALITY AS ancestors(id, ord)
+          INNER JOIN namespaces ON namespaces.id = ancestors.id
+        SQL
+
+        self.class
+          .from(Arel.sql(from_sql))
+          .order('ancestors.ord': hierarchy_order)
+      end
+
       def self_and_ancestors(hierarchy_order: nil)
         return super unless use_traversal_ids_for_ancestors?
 

app/models/namespaces/traversal/recursive.rb

@@ -46,6 +46,7 @@ module Namespaces
         object_hierarchy(self.class.where(id: id))
           .ancestors(upto: top, hierarchy_order: hierarchy_order)
       end
+      alias_method :recursive_ancestors_upto, :ancestors_upto
 
       def self_and_ancestors(hierarchy_order: nil)
         return self.class.where(id: id) unless parent_id

config/feature_flags/development/use_traversal_ids_for_ancestors_upto.yml

+---
+name: use_traversal_ids_for_ancestors_upto
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/72662
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/343619
+milestone: '14.6'
+type: development
+group: group::access
+default_enabled: false

spec/models/concerns/group_descendant_spec.rb

@@ -19,14 +19,16 @@ RSpec.describe GroupDescendant do
 
         query_count = ActiveRecord::QueryRecorder.new { test_group.hierarchy }.count
 
-        expect(query_count).to eq(1)
+        # use_traversal_ids_for_ancestors_upto actor based feature flag check adds an extra query.
+        expect(query_count).to eq(2)
       end
 
       it 'only queries once for the ancestors when a top is given' do
         test_group = create(:group, parent: subsub_group).reload
 
         recorder = ActiveRecord::QueryRecorder.new { test_group.hierarchy(subgroup) }
-        expect(recorder.count).to eq(1)
+        # use_traversal_ids_for_ancestors_upto actor based feature flag check adds an extra query.
+        expect(recorder.count).to eq(2)
       end
 
       it 'builds a hierarchy for a group' do

spec/models/group_spec.rb

@@ -533,6 +533,10 @@ RSpec.describe Group do
       describe '#ancestors' do
         it { expect(group.ancestors.to_sql).not_to include 'traversal_ids <@' }
       end
+
+      describe '#ancestors_upto' do
+        it { expect(group.ancestors_upto.to_sql).not_to include "WITH ORDINALITY" }
+      end
     end
 
     context 'linear' do
@@ -566,6 +570,10 @@ RSpec.describe Group do
         end
       end
 
+      describe '#ancestors_upto' do
+        it { expect(group.ancestors_upto.to_sql).to include "WITH ORDINALITY" }
+      end
+
       context 'when project namespace exists in the group' do
         let!(:project) { create(:project, group: group) }
         let!(:project_namespace) { project.project_namespace }

spec/models/namespace_spec.rb

@@ -700,20 +700,6 @@ RSpec.describe Namespace do
     end
   end
 
-  describe '#ancestors_upto' do
-    let(:parent) { create(:group) }
-    let(:child) { create(:group, parent: parent) }
-    let(:child2) { create(:group, parent: child) }
-
-    it 'returns all ancestors when no namespace is given' do
-      expect(child2.ancestors_upto).to contain_exactly(child, parent)
-    end
-
-    it 'includes ancestors upto but excluding the given ancestor' do
-      expect(child2.ancestors_upto(parent)).to contain_exactly(child)
-    end
-  end
-
   describe '#move_dir', :request_store do
     shared_examples "namespace restrictions" do
       context "when any project has container images" do
@@ -1274,6 +1260,38 @@ RSpec.describe Namespace do
     end
   end
 
+  describe '#use_traversal_ids_for_ancestors_upto?' do
+    let_it_be(:namespace, reload: true) { create(:namespace) }
+
+    subject { namespace.use_traversal_ids_for_ancestors_upto? }
+
+    context 'when use_traversal_ids_for_ancestors_upto feature flag is true' do
+      before do
+        stub_feature_flags(use_traversal_ids_for_ancestors_upto: true)
+      end
+
+      it { is_expected.to eq true }
+
+      it_behaves_like 'disabled feature flag when traversal_ids is blank'
+    end
+
+    context 'when use_traversal_ids_for_ancestors_upto feature flag is false' do
+      before do
+        stub_feature_flags(use_traversal_ids_for_ancestors_upto: false)
+      end
+
+      it { is_expected.to eq false }
+    end
+
+    context 'when use_traversal_ids? feature flag is false' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
+      end
+
+      it { is_expected.to eq false }
+    end
+  end
+
   describe '#users_with_descendants' do
     let(:user_a) { create(:user) }
     let(:user_b) { create(:user) }

spec/support/shared_examples/namespaces/traversal_examples.rb

@@ -205,6 +205,58 @@ RSpec.shared_examples 'namespace traversal' do
     end
   end
 
+  shared_examples '#ancestors_upto' do
+    let(:parent) { create(:group) }
+    let(:child) { create(:group, parent: parent) }
+    let(:child2) { create(:group, parent: child) }
+
+    it 'returns all ancestors when no namespace is given' do
+      expect(child2.ancestors_upto).to contain_exactly(child, parent)
+    end
+
+    it 'includes ancestors upto but excluding the given ancestor' do
+      expect(child2.ancestors_upto(parent)).to contain_exactly(child)
+    end
+
+    context 'with asc hierarchy_order' do
+      it 'returns the correct ancestor ids' do
+        expect(child2.ancestors_upto(hierarchy_order: :asc)).to eq([child, parent])
+      end
+    end
+
+    context 'with desc hierarchy_order' do
+      it 'returns the correct ancestor ids' do
+        expect(child2.ancestors_upto(hierarchy_order: :desc)).to eq([parent, child])
+      end
+    end
+
+    describe '#recursive_self_and_ancestor_ids' do
+      it 'is equivalent to ancestors_upto' do
+        recursive_result = child2.recursive_ancestors_upto(parent)
+        linear_result = child2.ancestors_upto(parent)
+        expect(linear_result).to match_array recursive_result
+      end
+
+      it 'makes a recursive query' do
+        expect { child2.recursive_ancestors_upto.try(:load) }.to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+  end
+
+  describe '#ancestors_upto' do
+    context 'with use_traversal_ids_for_ancestors_upto enabled' do
+      include_examples '#ancestors_upto'
+    end
+
+    context 'with use_traversal_ids_for_ancestors_upto disabled' do
+      before do
+        stub_feature_flags(use_traversal_ids_for_ancestors_upto: false)
+      end
+
+      include_examples '#ancestors_upto'
+    end
+  end
+
   describe '#descendants' do
     let!(:another_group) { create(:group) }
     let!(:another_group_nested) { create(:group, parent: another_group) }