Query ActiveRecord::Relation descendants

Enable linear traversal queries of multiple groups at the same time.
Improve the recursive query ObjectHierarchy interface.

Changelog: performance

app/models/namespaces/traversal/linear.rb

@@ -37,6 +37,7 @@ module Namespaces
   module Traversal
     module Linear
       extend ActiveSupport::Concern
+      include LinearScopes
 
       UnboundedSearch = Class.new(StandardError)
 
@@ -44,14 +45,6 @@ module Namespaces
         before_update :lock_both_roots, if: -> { sync_traversal_ids? && parent_id_changed? }
         after_create :sync_traversal_ids, if: -> { sync_traversal_ids? }
         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?
@@ -164,20 +157,14 @@ module Namespaces
         Namespace.lock.select(:id).where(id: roots).order(id: :asc).load
       end
 
-      # Make sure we drop the STI `type = 'Group'` condition for better performance.
-      # Logically equivalent so long as hierarchies remain homogeneous.
-      def without_sti_condition
-        self.class.unscope(where: :type)
-      end
-
       # Search this namespace's lineage. Bound inclusively by top node.
       def lineage(top: nil, bottom: nil, hierarchy_order: nil)
         raise UnboundedSearch, 'Must bound search by either top or bottom' unless top || bottom
 
-        skope = without_sti_condition
+        skope = self.class.without_sti_condition
 
         if top
-          skope = skope.traversal_ids_contains("{#{top.id}}")
+          skope = skope.where("traversal_ids @> ('{?}')", top.id)
         end
 
         if bottom

app/models/namespaces/traversal/linear_scopes.rb

+# frozen_string_literal: true
+
+module Namespaces
+  module Traversal
+    module LinearScopes
+      extend ActiveSupport::Concern
+
+      class_methods do
+        # When filtering namespaces by the traversal_ids column to compile a
+        # list of namespace IDs, it can be faster to reference the ID in
+        # traversal_ids than the primary key ID column.
+        def as_ids
+          return super unless use_traversal_ids?
+
+          select('namespaces.traversal_ids[array_length(namespaces.traversal_ids, 1)] AS id')
+        end
+
+        def self_and_descendants
+          return super unless use_traversal_ids?
+
+          without_dups = self_and_descendants_with_duplicates
+            .select('DISTINCT on(namespaces.id) namespaces.*')
+
+          # Wrap the `SELECT DISTINCT on(....)` with a normal query so we
+          # retain expected Rails behavior. Otherwise count and other
+          # aggregates won't work.
+          unscoped.without_sti_condition.from(without_dups, :namespaces)
+        end
+
+        def self_and_descendant_ids
+          return super unless use_traversal_ids?
+
+          self_and_descendants_with_duplicates.select('DISTINCT namespaces.id')
+        end
+
+        # Make sure we drop the STI `type = 'Group'` condition for better performance.
+        # Logically equivalent so long as hierarchies remain homogeneous.
+        def without_sti_condition
+          unscope(where: :type)
+        end
+
+        private
+
+        def use_traversal_ids?
+          Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
+        end
+
+        def self_and_descendants_with_duplicates
+          base_ids = select(:id)
+
+          unscoped
+            .without_sti_condition
+            .from("namespaces, (#{base_ids.to_sql}) base")
+            .where('namespaces.traversal_ids @> ARRAY[base.id]')
+        end
+      end
+    end
+  end
+end

app/models/namespaces/traversal/recursive.rb

@@ -4,6 +4,7 @@ module Namespaces
   module Traversal
     module Recursive
       extend ActiveSupport::Concern
+      include RecursiveScopes
 
       def root_ancestor
         return self if parent.nil?

app/models/namespaces/traversal/recursive_scopes.rb

+# frozen_string_literal: true
+
+module Namespaces
+  module Traversal
+    module RecursiveScopes
+      extend ActiveSupport::Concern
+
+      class_methods do
+        def as_ids
+          select('id')
+        end
+
+        def self_and_descendants
+          Gitlab::ObjectHierarchy.new(all).base_and_descendants
+        end
+        alias_method :recursive_self_and_descendants, :self_and_descendants
+
+        def self_and_descendant_ids
+          self_and_descendants.as_ids
+        end
+        alias_method :recursive_self_and_descendant_ids, :self_and_descendant_ids
+      end
+    end
+  end
+end

spec/models/namespace_spec.rb

@@ -207,9 +207,23 @@ RSpec.describe Namespace do
     it { is_expected.to include_module(Gitlab::VisibilityLevel) }
     it { is_expected.to include_module(Namespaces::Traversal::Recursive) }
     it { is_expected.to include_module(Namespaces::Traversal::Linear) }
+    it { is_expected.to include_module(Namespaces::Traversal::RecursiveScopes) }
+    it { is_expected.to include_module(Namespaces::Traversal::LinearScopes) }
   end
 
-  it_behaves_like 'linear namespace traversal'
+  context 'traversal scopes' do
+    context 'recursive' do
+      before do
+        stub_feature_flags(use_traversal_ids: false)
+      end
+
+      it_behaves_like 'namespace traversal scopes'
+    end
+
+    context 'linear' do
+      it_behaves_like 'namespace traversal scopes'
+    end
+  end
 
   context 'traversal_ids on create' do
     context 'default traversal_ids' do

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_scope_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'namespace traversal scopes' do
+  # Hierarchy 1
+  let_it_be(:group_1) { create(:group) }
+  let_it_be(:nested_group_1) { create(:group, parent: group_1) }
+  let_it_be(:deep_nested_group_1) { create(:group, parent: nested_group_1) }
+
+  # Hierarchy 2
+  let_it_be(:group_2) { create(:group) }
+  let_it_be(:nested_group_2) { create(:group, parent: group_2) }
+  let_it_be(:deep_nested_group_2) { create(:group, parent: nested_group_2) }
+
+  # All groups
+  let_it_be(:groups) do
+    [
+      group_1, nested_group_1, deep_nested_group_1,
+      group_2, nested_group_2, deep_nested_group_2
+    ]
+  end
+
+  describe '.as_ids' do
+    subject { described_class.where(id: [group_1, group_2]).as_ids.pluck(:id) }
+
+    it { is_expected.to contain_exactly(group_1.id, group_2.id) }
+  end
+
+  describe '.without_sti_condition' do
+    subject { described_class.without_sti_condition }
+
+    it { expect(subject.where_values_hash).not_to have_key(:type) }
+  end
+
+  describe '.self_and_descendants' do
+    subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendants }
+
+    it { is_expected.to contain_exactly(nested_group_1, deep_nested_group_1, nested_group_2, deep_nested_group_2) }
+
+    context 'with duplicate descendants' do
+      subject { described_class.where(id: [group_1, group_2, nested_group_1]).self_and_descendants }
+
+      it { is_expected.to match_array(groups) }
+    end
+  end
+
+  describe '.self_and_descendant_ids' do
+    subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendant_ids.pluck(:id) }
+
+    it { is_expected.to contain_exactly(nested_group_1.id, deep_nested_group_1.id, nested_group_2.id, deep_nested_group_2.id) }
+  end
+end