Refine linear queries in Namespace#all_projects

Changelog: performance

app/models/namespace.rb

@@ -271,7 +271,8 @@ class Namespace < ApplicationRecord
   # Includes projects from this namespace and projects from all subgroups
   # that belongs to this namespace
   def all_projects
-    namespace = user? ? self : self_and_descendants
+    namespace = user? ? self : self_and_descendant_ids
+
     Project.where(namespace: namespace)
   end
 

app/models/namespaces/traversal/linear.rb

@@ -46,6 +46,12 @@ module Namespaces
         after_update :sync_traversal_ids, if: -> { sync_traversal_ids? && saved_change_to_parent_id? }
 
         scope :traversal_ids_contains, ->(ids) { where("traversal_ids @> (?)", ids) }
+        # When filtering namespaces by the traversal_ids column to compile a
+        # list of namespace IDs, it's much faster to reference the ID in
+        # traversal_ids than the primary key ID column.
+        # WARNING This scope must be used behind a linear query feature flag
+        # such as `use_traversal_ids`.
+        scope :as_ids, -> { select('traversal_ids[array_length(traversal_ids, 1)] AS id') }
       end
 
       def sync_traversal_ids?
@@ -64,6 +70,12 @@ module Namespaces
         lineage(top: self)
       end
 
+      def self_and_descendant_ids
+        return super unless use_traversal_ids?
+
+        self_and_descendants.as_ids
+      end
+
       def descendants
         return super unless use_traversal_ids?
 

app/models/namespaces/traversal/recursive.rb

@@ -61,6 +61,11 @@ module Namespaces
       end
       alias_method :recursive_self_and_descendants, :self_and_descendants
 
+      def self_and_descendant_ids
+        self_and_descendants.select(:id)
+      end
+      alias_method :recursive_self_and_descendant_ids, :self_and_descendant_ids
+
       def object_hierarchy(ancestors_base)
         Gitlab::ObjectHierarchy.new(ancestors_base, options: { use_distinct: Feature.enabled?(:use_distinct_in_object_hierarchy, self) })
       end

spec/models/group_spec.rb

@@ -517,6 +517,10 @@ RSpec.describe Group do
         it { expect(group.self_and_descendants.to_sql).not_to include 'traversal_ids @>' }
       end
 
+      describe '#self_and_descendant_ids' do
+        it { expect(group.self_and_descendant_ids.to_sql).not_to include 'traversal_ids @>' }
+      end
+
       describe '#descendants' do
         it { expect(group.descendants.to_sql).not_to include 'traversal_ids @>' }
       end
@@ -533,6 +537,10 @@ RSpec.describe Group do
         it { expect(group.self_and_descendants.to_sql).to include 'traversal_ids @>' }
       end
 
+      describe '#self_and_descendant_ids' do
+        it { expect(group.self_and_descendant_ids.to_sql).to include 'traversal_ids @>' }
+      end
+
       describe '#descendants' do
         it { expect(group.descendants.to_sql).to include 'traversal_ids @>' }
       end

spec/models/namespace_spec.rb

@@ -5,6 +5,7 @@ require 'spec_helper'
 RSpec.describe Namespace do
   include ProjectForksHelper
   include GitHelpers
+  include ReloadHelpers
 
   let!(:namespace) { create(:namespace, :with_namespace_settings) }
   let(:gitlab_shell) { Gitlab::Shell.new }
@@ -199,6 +200,8 @@ RSpec.describe Namespace do
     it { is_expected.to include_module(Namespaces::Traversal::Linear) }
   end
 
+  it_behaves_like 'linear namespace traversal'
+
   context 'traversal_ids on create' do
     context 'default traversal_ids' do
       let(:namespace) { build(:namespace) }
@@ -1010,35 +1013,51 @@ RSpec.describe Namespace do
     end
   end
 
-  describe '#all_projects' do
+  shared_examples '#all_projects' do
     context 'when namespace is a group' do
-      let(:namespace) { create(:group) }
-      let(:child) { create(:group, parent: namespace) }
-      let!(:project1) { create(:project_empty_repo, namespace: namespace) }
-      let!(:project2) { create(:project_empty_repo, namespace: child) }
+      let_it_be(:namespace) { create(:group) }
+      let_it_be(:child) { create(:group, parent: namespace) }
+      let_it_be(:project1) { create(:project_empty_repo, namespace: namespace) }
+      let_it_be(:project2) { create(:project_empty_repo, namespace: child) }
+      let_it_be(:other_project) { create(:project_empty_repo) }
+
+      before do
+        reload_models(namespace, child)
+      end
 
       it { expect(namespace.all_projects.to_a).to match_array([project2, project1]) }
       it { expect(child.all_projects.to_a).to match_array([project2]) }
-
-      it 'queries for the namespace and its descendants' do
-        expect(Project).to receive(:where).with(namespace: [namespace, child])
-
-        namespace.all_projects
-      end
     end
 
     context 'when namespace is a user namespace' do
       let_it_be(:user) { create(:user) }
       let_it_be(:user_namespace) { create(:namespace, owner: user) }
       let_it_be(:project) { create(:project, namespace: user_namespace) }
+      let_it_be(:other_project) { create(:project_empty_repo) }
+
+      before do
+        reload_models(user_namespace)
+      end
 
       it { expect(user_namespace.all_projects.to_a).to match_array([project]) }
+    end
+  end
+
+  describe '#all_projects' do
+    context 'with use_traversal_ids feature flag enabled' do
+      before do
+        stub_feature_flags(use_traversal_ids: true)
+      end
 
-      it 'only queries for the namespace itself' do
-        expect(Project).to receive(:where).with(namespace: user_namespace)
+      include_examples '#all_projects'
+    end
 
-        user_namespace.all_projects
+    context 'with use_traversal_ids feature flag disabled' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
       end
+
+      include_examples '#all_projects'
     end
   end
 

spec/support/shared_examples/namespaces/linear_traversal_examples.rb

+# frozen_string_literal: true
+
+# Traversal examples common to linear and recursive methods are in
+# spec/support/shared_examples/namespaces/traversal_examples.rb
+
+RSpec.shared_examples 'linear namespace traversal' do
+  context 'when use_traversal_ids feature flag is enabled' do
+    before do
+      stub_feature_flags(use_traversal_ids: true)
+    end
+
+    context 'scopes' do
+      describe '.as_ids' do
+        let_it_be(:namespace1) { create(:group) }
+        let_it_be(:namespace2) { create(:group) }
+
+        subject { Namespace.where(id: [namespace1, namespace2]).as_ids.pluck(:id) }
+
+        it { is_expected.to contain_exactly(namespace1.id, namespace2.id) }
+      end
+    end
+  end
+end

spec/support/shared_examples/namespaces/traversal_examples.rb

@@ -122,4 +122,20 @@ RSpec.shared_examples 'namespace traversal' do
       it_behaves_like 'recursive version', :self_and_descendants
     end
   end
+
+  describe '#self_and_descendant_ids' do
+    let!(:group) { create(:group, path: 'git_lab') }
+    let!(:nested_group) { create(:group, parent: group) }
+    let!(:deep_nested_group) { create(:group, parent: nested_group) }
+
+    subject { group.self_and_descendant_ids.pluck(:id) }
+
+    it { is_expected.to contain_exactly(group.id, nested_group.id, deep_nested_group.id) }
+
+    describe '#recursive_self_and_descendant_ids' do
+      let(:groups) { [group, nested_group, deep_nested_group] }
+
+      it_behaves_like 'recursive version', :self_and_descendant_ids
+    end
+  end
 end