Add latest changes from gitlab-org/gitlab@14-2-stable-eev14.2.0-rc42

4 files changed, 116 insertions, 18 deletions

diff --git a/app/models/namespaces/traversal/linear.rb b/app/models/namespaces/traversal/linear.rb
index 3d78f384634..33e8c3e5172 100644
--- a/app/models/namespaces/traversal/linear.rb
+++ b/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?
@@ -59,7 +52,7 @@ module Namespaces
       end
 
       def use_traversal_ids?
-        return false unless Feature.enabled?(:use_traversal_ids, root_ancestor, default_enabled: :yaml)
+        return false unless Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
 
         traversal_ids.present?
       end
@@ -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
@@ -190,7 +177,13 @@ module Namespaces
         if hierarchy_order
           depth_sql = "ABS(#{traversal_ids.count} - array_length(traversal_ids, 1))"
           skope = skope.select(skope.arel_table[Arel.star], "#{depth_sql} as depth")
-                       .order(depth: hierarchy_order)
+          # The SELECT includes an extra depth attribute. We wrap the SQL in a
+          # standard SELECT to avoid mismatched attribute errors when trying to
+          # chain future ActiveRelation commands, and retain the ordering.
+          skope = self.class
+            .without_sti_condition
+            .from(skope, self.class.table_name)
+            .order(depth: hierarchy_order)
         end
 
         skope
diff --git a/app/models/namespaces/traversal/linear_scopes.rb b/app/models/namespaces/traversal/linear_scopes.rb
new file mode 100644
index 00000000000..90fae8ef35d
--- /dev/null
+++ b/app/models/namespaces/traversal/linear_scopes.rb
@@ -0,0 +1,68 @@
+# 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(include_self: true)
+          return super unless use_traversal_ids?
+
+          records = self_and_descendants_with_duplicates(include_self: include_self)
+
+          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(include_self: true)
+          return super unless use_traversal_ids?
+
+          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.
+        # 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(include_self: true)
+          base_ids = select(:id)
+
+          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
+  end
+end
diff --git a/app/models/namespaces/traversal/recursive.rb b/app/models/namespaces/traversal/recursive.rb
index d9e8743aa50..c1ada715d6d 100644
--- a/app/models/namespaces/traversal/recursive.rb
+++ b/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?
diff --git a/app/models/namespaces/traversal/recursive_scopes.rb b/app/models/namespaces/traversal/recursive_scopes.rb
new file mode 100644
index 00000000000..be49d5d9d55
--- /dev/null
+++ b/app/models/namespaces/traversal/recursive_scopes.rb
@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Namespaces
+  module Traversal
+    module RecursiveScopes
+      extend ActiveSupport::Concern
+
+      class_methods do
+        def as_ids
+          select('id')
+        end
+
+        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(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
+    end
+  end
+end