Merge branch '337818-query-multiple-group-ancestors-at-once' into 'master'

Resolve "Query multiple group ancestors at once"

See merge request gitlab-org/gitlab!67652

app/models/namespaces/traversal/linear_scopes.rb

@@ -15,6 +15,28 @@ module Namespaces
           select('namespaces.traversal_ids[array_length(namespaces.traversal_ids, 1)] AS id')
         end
 
+        def self_and_ancestors(include_self: true, hierarchy_order: nil)
+          return super unless use_traversal_ids_for_ancestor_scopes?
+
+          records = unscoped
+            .without_sti_condition
+            .where(id: without_sti_condition.select('unnest(traversal_ids)'))
+            .order_by_depth(hierarchy_order)
+            .normal_select
+
+          if include_self
+            records
+          else
+            records.where.not(id: all.as_ids)
+          end
+        end
+
+        def self_and_ancestor_ids(include_self: true)
+          return super unless use_traversal_ids_for_ancestor_scopes?
+
+          self_and_ancestors(include_self: include_self).as_ids
+        end
+
         def self_and_descendants(include_self: true)
           return super unless use_traversal_ids?
 
@@ -22,11 +44,7 @@ module Namespaces
 
           distinct = records.select('DISTINCT on(namespaces.id) namespaces.*')
 
-          # Produce a query of the form: SELECT * FROM namespaces;
-          #
-          # When we have queries that break this SELECT * format we can run in to errors.
-          # For example `SELECT DISTINCT on(...)` will fail when we chain a `.count` c
-          unscoped.without_sti_condition.from(distinct, :namespaces)
+          distinct.normal_select
         end
 
         def self_and_descendant_ids(include_self: true)
@@ -42,12 +60,35 @@ module Namespaces
           unscope(where: :type)
         end
 
+        def order_by_depth(hierarchy_order)
+          return all unless hierarchy_order
+
+          depth_order = hierarchy_order == :asc ? :desc : :asc
+
+          all
+            .select(Arel.star, 'array_length(traversal_ids, 1) as depth')
+            .order(depth: depth_order, id: :asc)
+        end
+
+        # Produce a query of the form: SELECT * FROM namespaces;
+        #
+        # When we have queries that break this SELECT * format we can run in to errors.
+        # For example `SELECT DISTINCT on(...)` will fail when we chain a `.count` c
+        def normal_select
+          unscoped.without_sti_condition.from(all, :namespaces)
+        end
+
         private
 
         def use_traversal_ids?
           Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
         end
 
+        def use_traversal_ids_for_ancestor_scopes?
+          Feature.enabled?(:use_traversal_ids_for_ancestor_scopes, default_enabled: :yaml) &&
+          use_traversal_ids?
+        end
+
         def self_and_descendants_with_duplicates(include_self: true)
           base_ids = select(:id)
 

app/models/namespaces/traversal/recursive_scopes.rb

@@ -10,6 +10,22 @@ module Namespaces
           select('id')
         end
 
+        def self_and_ancestors(include_self: true, hierarchy_order: nil)
+          records = Gitlab::ObjectHierarchy.new(all).base_and_ancestors(hierarchy_order: hierarchy_order)
+
+          if include_self
+            records
+          else
+            records.where.not(id: all.as_ids)
+          end
+        end
+        alias_method :recursive_self_and_ancestors, :self_and_ancestors
+
+        def self_and_ancestor_ids(include_self: true)
+          self_and_ancestors(include_self: include_self).as_ids
+        end
+        alias_method :recursive_self_and_ancestor_ids, :self_and_ancestor_ids
+
         def descendant_ids
           recursive_descendants.as_ids
         end

config/feature_flags/development/use_traversal_ids_for_ancestor_scopes.yml

+---
+name: use_traversal_ids_for_ancestor_scopes
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67652
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/340159
+milestone: '14.3'
+type: development
+group: group::access
+default_enabled: false

spec/support/shared_examples/namespaces/traversal_scope_examples.rb

@@ -31,6 +31,125 @@ RSpec.shared_examples 'namespace traversal scopes' do
     it { expect(subject.where_values_hash).not_to have_key(:type) }
   end
 
+  describe '.order_by_depth' do
+    subject { described_class.where(id: [group_1, nested_group_1, deep_nested_group_1]).order_by_depth(direction) }
+
+    context 'ascending' do
+      let(:direction) { :asc }
+
+      it { is_expected.to eq [deep_nested_group_1, nested_group_1, group_1] }
+    end
+
+    context 'descending' do
+      let(:direction) { :desc }
+
+      it { is_expected.to eq [group_1, nested_group_1, deep_nested_group_1] }
+    end
+  end
+
+  describe '.normal_select' do
+    let(:query_result) { described_class.where(id: group_1).normal_select }
+
+    subject { query_result.column_names }
+
+    it { is_expected.to eq described_class.column_names }
+  end
+
+  shared_examples '.self_and_ancestors' do
+    subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestors }
+
+    it { is_expected.to contain_exactly(group_1, nested_group_1, group_2, nested_group_2) }
+
+    context 'when include_self is false' do
+      subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestors(include_self: false) }
+
+      it { is_expected.to contain_exactly(group_1, group_2) }
+    end
+
+    context 'when hierarchy_order is ascending' do
+      subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestors(hierarchy_order: :asc) }
+
+      it { is_expected.to eq [nested_group_1, nested_group_2, group_1, group_2] }
+    end
+
+    context 'when hierarchy_order is descending' do
+      subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestors(hierarchy_order: :desc) }
+
+      it { is_expected.to eq [group_1, group_2, nested_group_1, nested_group_2] }
+    end
+  end
+
+  describe '.self_and_ancestors' do
+    context "use_traversal_ids_ancestor_scopes feature flag is true" do
+      before do
+        stub_feature_flags(use_traversal_ids: true)
+        stub_feature_flags(use_traversal_ids_for_ancestor_scopes: true)
+      end
+
+      it_behaves_like '.self_and_ancestors'
+
+      it 'not make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).self_and_ancestors.load }.not_to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+
+    context "use_traversal_ids_ancestor_scopes feature flag is false" do
+      before do
+        stub_feature_flags(use_traversal_ids_for_ancestor_scopes: false)
+      end
+
+      it_behaves_like '.self_and_ancestors'
+
+      it 'make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).self_and_ancestors.load }.to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+  end
+
+  shared_examples '.self_and_ancestor_ids' do
+    subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestor_ids.pluck(:id) }
+
+    it { is_expected.to contain_exactly(group_1.id, nested_group_1.id, group_2.id, nested_group_2.id) }
+
+    context 'when include_self is false' do
+      subject do
+        described_class
+          .where(id: [nested_group_1, nested_group_2])
+          .self_and_ancestor_ids(include_self: false)
+          .pluck(:id)
+      end
+
+      it { is_expected.to contain_exactly(group_1.id, group_2.id) }
+    end
+  end
+
+  describe '.self_and_ancestor_ids' do
+    context "use_traversal_ids_ancestor_scopes feature flag is true" do
+      before do
+        stub_feature_flags(use_traversal_ids: true)
+        stub_feature_flags(use_traversal_ids_for_ancestor_scopes: true)
+      end
+
+      it_behaves_like '.self_and_ancestor_ids'
+
+      it 'make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).self_and_ancestor_ids.load }.not_to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+
+    context "use_traversal_ids_ancestor_scopes feature flag is false" do
+      before do
+        stub_feature_flags(use_traversal_ids_for_ancestor_scopes: false)
+      end
+
+      it_behaves_like '.self_and_ancestor_ids'
+
+      it 'make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).self_and_ancestor_ids.load }.to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+  end
+
   describe '.self_and_descendants' do
     subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendants }