Query multiple group descendants at once

Changelog: performance

app/models/namespaces/traversal/linear_scopes.rb

@@ -15,22 +15,25 @@ module Namespaces
           select('namespaces.traversal_ids[array_length(namespaces.traversal_ids, 1)] AS id')
         end
 
-        def self_and_descendants
+        def self_and_descendants(include_self: true)
           return super unless use_traversal_ids?
 
-          without_dups = self_and_descendants_with_duplicates
-            .select('DISTINCT on(namespaces.id) namespaces.*')
+          records = self_and_descendants_with_duplicates(include_self: include_self)
 
-          # 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)
+          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)
         end
 
-        def self_and_descendant_ids
+        def self_and_descendant_ids(include_self: true)
           return super unless use_traversal_ids?
 
-          self_and_descendants_with_duplicates.select('DISTINCT namespaces.id')
+          self_and_descendants_with_duplicates(include_self: include_self)
+            .select('DISTINCT namespaces.id')
         end
 
         # Make sure we drop the STI `type = 'Group'` condition for better performance.
@@ -45,13 +48,19 @@ module Namespaces
           Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
         end
 
-        def self_and_descendants_with_duplicates
+        def self_and_descendants_with_duplicates(include_self: true)
           base_ids = select(:id)
 
-          unscoped
+          records = unscoped
             .without_sti_condition
             .from("namespaces, (#{base_ids.to_sql}) base")
             .where('namespaces.traversal_ids @> ARRAY[base.id]')
+
+          if include_self
+            records
+          else
+            records.where('namespaces.id <> base.id')
+          end
         end
       end
     end

app/models/namespaces/traversal/recursive_scopes.rb

@@ -10,13 +10,24 @@ module Namespaces
           select('id')
         end
 
-        def self_and_descendants
-          Gitlab::ObjectHierarchy.new(all).base_and_descendants
+        def descendant_ids
+          recursive_descendants.as_ids
+        end
+        alias_method :recursive_descendant_ids, :descendant_ids
+
+        def self_and_descendants(include_self: true)
+          base = if include_self
+                   unscoped.where(id: all.as_ids)
+                 else
+                   unscoped.where(parent_id: all.as_ids)
+                 end
+
+          Gitlab::ObjectHierarchy.new(base).base_and_descendants
         end
         alias_method :recursive_self_and_descendants, :self_and_descendants
 
-        def self_and_descendant_ids
-          self_and_descendants.as_ids
+        def self_and_descendant_ids(include_self: true)
+          self_and_descendants(include_self: include_self).as_ids
         end
         alias_method :recursive_self_and_descendant_ids, :self_and_descendant_ids
       end

spec/support/shared_examples/namespaces/traversal_scope_examples.rb

@@ -41,11 +41,28 @@ RSpec.shared_examples 'namespace traversal scopes' do
 
       it { is_expected.to match_array(groups) }
     end
+
+    context 'when include_self is false' do
+      subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendants(include_self: false) }
+
+      it { is_expected.to contain_exactly(deep_nested_group_1, deep_nested_group_2) }
+    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) }
+
+    context 'when include_self is false' do
+      subject do
+        described_class
+          .where(id: [nested_group_1, nested_group_2])
+          .self_and_descendant_ids(include_self: false)
+          .pluck(:id)
+      end
+
+      it { is_expected.to contain_exactly(deep_nested_group_1.id, deep_nested_group_2.id) }
+    end
   end
 end