1 files changed, 54 insertions, 0 deletions
diff --git a/app/models/namespaces/traversal/cached.rb b/app/models/namespaces/traversal/cached.rb
index 55eaaa4667e..b962038d039 100644
--- a/app/models/namespaces/traversal/cached.rb
+++ b/app/models/namespaces/traversal/cached.rb
@@ -10,8 +10,62 @@ module Namespaces
         after_destroy :invalidate_descendants_cache
       end
 
+      override :self_and_descendant_ids
+      def self_and_descendant_ids
+        return super unless attempt_to_use_cached_data?
+
+        scope_with_cached_ids(
+          super,
+          self.class,
+          Namespaces::Descendants.arel_table[:self_and_descendant_group_ids]
+        )
+      end
+
+      override :all_project_ids
+      def all_project_ids
+        return super unless attempt_to_use_cached_data?
+
+        scope_with_cached_ids(
+          all_projects.select(:id),
+          Project,
+          Namespaces::Descendants.arel_table[:all_project_ids]
+        )
+      end
+
       private
 
+      # This method implements an OR based cache lookup using COALESCE, similar what you would do in Ruby:
+      # return cheap_cached_data || expensive_uncached_data
+      def scope_with_cached_ids(consistent_ids_scope, model, cached_ids_column)
+        # Look up the cached ids and unnest them into rows if the cache is up to date.
+        cache_lookup_query = Namespaces::Descendants
+          .where(outdated_at: nil, namespace_id: id)
+          .select(cached_ids_column.as('ids'))
+
+        # Invoke the consistent lookup query and collect the ids as a single array value
+        consistent_descendant_ids_scope = model
+          .from(consistent_ids_scope.arel.as(model.table_name))
+          .reselect(Arel::Nodes::NamedFunction.new('ARRAY_AGG', [model.arel_table[:id]]).as('ids'))
+          .unscope(where: :type)
+
+        from = <<~SQL
+        UNNEST(
+          COALESCE(
+            (SELECT ids FROM (#{cache_lookup_query.to_sql}) cached_query),
+            (SELECT ids FROM (#{consistent_descendant_ids_scope.to_sql}) consistent_query))
+        ) AS #{model.table_name}(id)
+        SQL
+
+        model
+          .from(from)
+          .unscope(where: :type)
+          .select(:id)
+      end
+
+      def attempt_to_use_cached_data?
+        Feature.enabled?(:group_hierarchy_optimization, self, type: :beta)
+      end
+
       override :sync_traversal_ids
       def sync_traversal_ids
         super