GraphQL user.groups node authorization DB query count optimization

GraphQL API authorizes each Group node in a list.
This MR optimizes DB query count by preloading some of the required
information for Group authorization.

Changelog: performance

app/graphql/resolvers/users/groups_resolver.rb

@@ -20,7 +20,7 @@ module Resolvers
                description: 'Filter by permissions the user has on groups.'
 
       before_connection_authorization do |nodes, current_user|
-        Preloaders::UserMaxAccessLevelInGroupsPreloader.new(nodes, current_user).execute
+        Preloaders::GroupPolicyPreloader.new(nodes, current_user).execute
       end
 
       def resolve_with_lookahead(**args)

app/models/preloaders/group_policy_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  class GroupPolicyPreloader
+    def initialize(groups, current_user)
+      @groups = groups
+      @current_user = current_user
+    end
+
+    def execute
+      Preloaders::UserMaxAccessLevelInGroupsPreloader.new(@groups, @current_user).execute
+      Preloaders::GroupRootAncestorPreloader.new(@groups, root_ancestor_preloads).execute
+    end
+
+    private
+
+    def root_ancestor_preloads
+      []
+    end
+  end
+end
+
+Preloaders::GroupPolicyPreloader.prepend_mod_with('Preloaders::GroupPolicyPreloader')

app/models/preloaders/group_root_ancestor_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  class GroupRootAncestorPreloader
+    def initialize(groups, root_ancestor_preloads = [])
+      @groups = groups
+      @root_ancestor_preloads = root_ancestor_preloads
+    end
+
+    def execute
+      return unless ::Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
+
+      # type == 'Group' condition located on subquery to prevent a filter in the query
+      root_query = Namespace.joins("INNER JOIN (#{join_sql}) as root_query ON root_query.root_id = namespaces.id")
+                        .select('namespaces.*, root_query.id as source_id')
+
+      root_query = root_query.preload(*@root_ancestor_preloads) if @root_ancestor_preloads.any?
+
+      root_ancestors_by_id = root_query.group_by(&:source_id)
+
+      @groups.each do |group|
+        group.root_ancestor = root_ancestors_by_id[group.id].first
+      end
+    end
+
+    private
+
+    def join_sql
+      Group.select('id, traversal_ids[1] as root_id').where(id: @groups.map(&:id)).to_sql
+    end
+  end
+end

app/models/preloaders/user_max_access_level_in_groups_preloader.rb

@@ -3,7 +3,6 @@
 module Preloaders
   # This class preloads the max access level (role) for the user within the given groups and
   # stores the values in requests store.
-  # Will only be able to preload max access level for groups where the user is a direct member
   class UserMaxAccessLevelInGroupsPreloader
     include BulkMemberAccessLoad
 
@@ -13,8 +12,17 @@ module Preloaders
     end
 
     def execute
+      if ::Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
+        preload_with_traversal_ids
+      else
+        preload_direct_memberships
+      end
+    end
+
+    private
+
+    def preload_direct_memberships
       group_memberships = GroupMember.active_without_invites_and_requests
-                                     .non_minimal_access
                                      .where(user: @user, source_id: @groups)
                                      .group(:source_id)
                                      .maximum(:access_level)
@@ -23,5 +31,22 @@ module Preloaders
         merge_value_to_request_store(User, @user.id, group_id, max_access_level)
       end
     end
+
+    def preload_with_traversal_ids
+      max_access_levels = GroupMember.active_without_invites_and_requests
+                                     .where(user: @user)
+                                     .joins("INNER JOIN (#{traversal_join_sql}) as hierarchy ON members.source_id = hierarchy.traversal_id")
+                                     .group('hierarchy.id')
+                                     .maximum(:access_level)
+
+      @groups.each do |group|
+        max_access_level = max_access_levels[group.id] || Gitlab::Access::NO_ACCESS
+        merge_value_to_request_store(User, @user.id, group.id, max_access_level)
+      end
+    end
+
+    def traversal_join_sql
+      Namespace.select('id, unnest(traversal_ids) as traversal_id').where(id: @groups.map(&:id)).to_sql
+    end
   end
 end

ee/app/graphql/ee/resolvers/users/groups_resolver.rb

@@ -11,7 +11,7 @@ module EE
 
         override :unconditional_includes
         def unconditional_includes
-          [:saml_provider, *super]
+          [:ip_restrictions, *super]
         end
       end
     end

ee/app/models/ee/preloaders/group_policy_preloader.rb

+# frozen_string_literal: true
+
+module EE
+  module Preloaders
+    module GroupPolicyPreloader
+      extend ::Gitlab::Utils::Override
+
+      private
+
+      override :root_ancestor_preloads
+      def root_ancestor_preloads
+        [*super, :ip_restrictions, :saml_provider]
+      end
+    end
+  end
+end

ee/spec/models/ee/preloaders/group_policy_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::GroupPolicyPreloader do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:root_parent) { create(:group, :private, name: 'root-1', path: 'root-1') }
+  let_it_be(:guest_group) { create(:group, name: 'public guest', path: 'public-guest') }
+  let_it_be(:private_maintainer_group) { create(:group, :private, name: 'b private maintainer', path: 'b-private-maintainer', parent: root_parent) }
+  let_it_be(:private_developer_group) { create(:group, :private, project_creation_level: nil, name: 'c public developer', path: 'c-public-developer') }
+  let_it_be(:public_maintainer_group) { create(:group, :private, name: 'a public maintainer', path: 'a-public-maintainer') }
+
+  let(:base_groups) { [guest_group, private_maintainer_group, private_developer_group, public_maintainer_group] }
+
+  before_all do
+    guest_group.add_guest(user)
+    private_maintainer_group.add_maintainer(user)
+    private_developer_group.add_developer(user)
+    public_maintainer_group.add_maintainer(user)
+  end
+
+  context 'when ip_restrictions feature is enabled' do
+    before do
+      stub_licensed_features(group_ip_restriction: true)
+    end
+
+    it 'avoids N+1 queries when authorizing a list of groups', :request_store do
+      preload_groups_for_policy(user)
+      control = ActiveRecord::QueryRecorder.new { authorize_all_groups(user) }
+
+      new_group1 = create(:group, :private).tap { |group| group.add_maintainer(user) }
+      new_group2 = create(:group, :private, parent: private_maintainer_group)
+
+      another_root = create(:group, :private, name: 'root-3', path: 'root-3')
+      new_group3 = create(:group, :private, parent: another_root).tap { |group| group.add_maintainer(user) }
+
+      pristine_groups = Group.where(id: base_groups + [new_group1, new_group2, new_group3]).to_a
+
+      preload_groups_for_policy(user, pristine_groups)
+      expect { authorize_all_groups(user, pristine_groups) }.not_to exceed_query_limit(control)
+    end
+  end
+
+  def authorize_all_groups(current_user, group_list = base_groups)
+    group_list.each { |group| current_user.can?(:read_group, group) }
+  end
+
+  def preload_groups_for_policy(current_user, group_list = base_groups)
+    described_class.new(group_list, current_user).execute
+  end
+end

ee/spec/requests/api/graphql/current_user/groups_query_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Query current user groups' do
+  include GraphqlHelpers
+
+  let_it_be(:user) { create(:user) }
+  let_it_be(:root_parent) { create(:group, :private, name: 'root-1', path: 'root-1') }
+  let_it_be(:guest_group) { create(:group, name: 'public guest', path: 'public-guest') }
+  let_it_be(:private_maintainer_group) { create(:group, :private, name: 'b private maintainer', path: 'b-private-maintainer', parent: root_parent) }
+  let_it_be(:private_developer_group) { create(:group, :private, project_creation_level: nil, name: 'c public developer', path: 'c-public-developer') }
+  let_it_be(:public_maintainer_group) { create(:group, :private, name: 'a public maintainer', path: 'a-public-maintainer') }
+
+  let(:group_arguments) { {} }
+  let(:current_user) { user }
+
+  let(:fields) do
+    <<~GRAPHQL
+      nodes { id path fullPath name }
+    GRAPHQL
+  end
+
+  let(:query) do
+    graphql_query_for('currentUser', {}, query_graphql_field('groups', group_arguments, fields))
+  end
+
+  before_all do
+    guest_group.add_guest(user)
+    private_maintainer_group.add_maintainer(user)
+    private_developer_group.add_developer(user)
+    public_maintainer_group.add_maintainer(user)
+  end
+
+  context 'when permission_scope is CREATE_PROJECTS' do
+    let(:group_arguments) { { permission_scope: :CREATE_PROJECTS } }
+
+    before do
+      post_graphql(query, current_user: current_user)
+    end
+
+    it_behaves_like 'a working graphql query'
+
+    context 'when ip_restrictions feature is enabled' do
+      before do
+        stub_licensed_features(group_ip_restriction: true)
+      end
+
+      it 'avoids N+1 queries', :request_store do
+        control = ActiveRecord::QueryRecorder.new { post_graphql(query, current_user: current_user) }
+
+        create(:group, :private).tap { |group| group.add_maintainer(current_user) }
+        create(:group, :private, parent: private_maintainer_group)
+
+        another_root = create(:group, :private, name: 'root-3', path: 'root-3')
+        create(:group, :private, parent: another_root).tap { |group| group.add_maintainer(current_user) }
+
+        expect { post_graphql(query, current_user: current_user) }.not_to exceed_query_limit(control)
+      end
+    end
+  end
+end

spec/models/preloaders/group_policy_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::GroupPolicyPreloader do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:root_parent) { create(:group, :private, name: 'root-1', path: 'root-1') }
+  let_it_be(:guest_group) { create(:group, name: 'public guest', path: 'public-guest') }
+  let_it_be(:private_maintainer_group) { create(:group, :private, name: 'b private maintainer', path: 'b-private-maintainer', parent: root_parent) }
+  let_it_be(:private_developer_group) { create(:group, :private, project_creation_level: nil, name: 'c public developer', path: 'c-public-developer') }
+  let_it_be(:public_maintainer_group) { create(:group, :private, name: 'a public maintainer', path: 'a-public-maintainer') }
+
+  let(:base_groups) { [guest_group, private_maintainer_group, private_developer_group, public_maintainer_group] }
+
+  before_all do
+    guest_group.add_guest(user)
+    private_maintainer_group.add_maintainer(user)
+    private_developer_group.add_developer(user)
+    public_maintainer_group.add_maintainer(user)
+  end
+
+  it 'avoids N+1 queries when authorizing a list of groups', :request_store do
+    preload_groups_for_policy(user)
+    control = ActiveRecord::QueryRecorder.new { authorize_all_groups(user) }
+
+    new_group1 = create(:group, :private).tap { |group| group.add_maintainer(user) }
+    new_group2 = create(:group, :private, parent: private_maintainer_group)
+
+    another_root = create(:group, :private, name: 'root-3', path: 'root-3')
+    new_group3 = create(:group, :private, parent: another_root).tap { |group| group.add_maintainer(user) }
+
+    pristine_groups = Group.where(id: base_groups + [new_group1, new_group2, new_group3])
+
+    preload_groups_for_policy(user, pristine_groups)
+    expect { authorize_all_groups(user, pristine_groups) }.not_to exceed_query_limit(control)
+  end
+
+  def authorize_all_groups(current_user, group_list = base_groups)
+    group_list.each { |group| current_user.can?(:read_group, group) }
+  end
+
+  def preload_groups_for_policy(current_user, group_list = base_groups)
+    described_class.new(group_list, current_user).execute
+  end
+end

spec/models/preloaders/group_root_ancestor_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::GroupRootAncestorPreloader do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:root_parent1) { create(:group, :private, name: 'root-1', path: 'root-1') }
+  let_it_be(:root_parent2) { create(:group, :private, name: 'root-2', path: 'root-2') }
+  let_it_be(:guest_group) { create(:group, name: 'public guest', path: 'public-guest') }
+  let_it_be(:private_maintainer_group) { create(:group, :private, name: 'b private maintainer', path: 'b-private-maintainer', parent: root_parent1) }
+  let_it_be(:private_developer_group) { create(:group, :private, project_creation_level: nil, name: 'c public developer', path: 'c-public-developer') }
+  let_it_be(:public_maintainer_group) { create(:group, :private, name: 'a public maintainer', path: 'a-public-maintainer', parent: root_parent2) }
+
+  let(:root_query_regex) { /\ASELECT.+FROM "namespaces" WHERE "namespaces"."id" = \d+/ }
+  let(:additional_preloads) { [] }
+  let(:groups) { [guest_group, private_maintainer_group, private_developer_group, public_maintainer_group] }
+  let(:pristine_groups) { Group.where(id: groups) }
+
+  shared_examples 'executes N matching DB queries' do |expected_query_count, query_method = nil|
+    it 'executes the specified root_ancestor queries' do
+      expect do
+        pristine_groups.each do |group|
+          root_ancestor = group.root_ancestor
+
+          root_ancestor.public_send(query_method) if query_method.present?
+        end
+      end.to make_queries_matching(root_query_regex, expected_query_count)
+    end
+
+    it 'strong_memoizes the correct root_ancestor' do
+      pristine_groups.each do |group|
+        expected_parent_id = group.root_ancestor.id == group.id ? nil : group.root_ancestor.id
+
+        expect(group.parent_id).to eq(expected_parent_id)
+      end
+    end
+  end
+
+  context 'when the preloader is used' do
+    before do
+      preload_ancestors
+    end
+
+    context 'when no additional preloads are provided' do
+      it_behaves_like 'executes N matching DB queries', 0
+    end
+
+    context 'when additional preloads are provided' do
+      let(:additional_preloads) { [:route] }
+      let(:root_query_regex) { /\ASELECT.+FROM "routes" WHERE "routes"."source_id" = \d+/ }
+
+      it_behaves_like 'executes N matching DB queries', 0, :full_path
+    end
+  end
+
+  context 'when the preloader is not used' do
+    it_behaves_like 'executes N matching DB queries', 2
+  end
+
+  def preload_ancestors
+    described_class.new(pristine_groups, additional_preloads).execute
+  end
+end

spec/models/preloaders/user_max_access_level_in_groups_preloader_spec.rb

@@ -13,32 +13,47 @@ RSpec.describe Preloaders::UserMaxAccessLevelInGroupsPreloader do
 
   shared_examples 'executes N max member permission queries to the DB' do
     it 'executes the specified max membership queries' do
-      queries = ActiveRecord::QueryRecorder.new do
-        groups.each { |group| user.can?(:read_group, group) }
-      end
+      expect { groups.each { |group| user.can?(:read_group, group) } }.to make_queries_matching(max_query_regex, expected_query_count)
+    end
 
-      max_queries = queries.log.grep(max_query_regex)
+    it 'caches the correct access_level for each group' do
+      groups.each do |group|
+        access_level_from_db = group.members_with_parents.where(user_id: user.id).group(:user_id).maximum(:access_level)[user.id] || Gitlab::Access::NO_ACCESS
+        cached_access_level = group.max_member_access_for_user(user)
 
-      expect(max_queries.count).to eq(expected_query_count)
+        expect(cached_access_level).to eq(access_level_from_db)
+      end
     end
   end
 
   context 'when the preloader is used', :request_store do
-    before do
-      described_class.new(groups, user).execute
-    end
+    context 'when user has indirect access to groups' do
+      let_it_be(:child_maintainer) { create(:group, :private, parent: group1).tap {|g| g.add_maintainer(user)} }
+      let_it_be(:child_indirect_access) { create(:group, :private, parent: group1) }
 
-    it_behaves_like 'executes N max member permission queries to the DB' do
-      # Will query all groups where the user is not already a member
-      let(:expected_query_count) { 1 }
-    end
+      let(:groups) { [group1, group2, group3, child_maintainer, child_indirect_access] }
+
+      context 'when traversal_ids feature flag is disabled' do
+        it_behaves_like 'executes N max member permission queries to the DB' do
+          before do
+            stub_feature_flags(use_traversal_ids: false)
+            described_class.new(groups, user).execute
+          end
+
+          # One query for group with no access and another one per group where the user is not a direct member
+          let(:expected_query_count) { 2 }
+        end
+      end
 
-    context 'when user has access but is not a direct member of the group' do
-      let(:groups) { [group1, group2, group3, create(:group, :private, parent: group1)] }
+      context 'when traversal_ids feature flag is enabled' do
+        it_behaves_like 'executes N max member permission queries to the DB' do
+          before do
+            stub_feature_flags(use_traversal_ids: true)
+            described_class.new(groups, user).execute
+          end
 
-      it_behaves_like 'executes N max member permission queries to the DB' do
-        # One query for group with no access and another one where the user is not a direct member
-        let(:expected_query_count) { 2 }
+          let(:expected_query_count) { 0 }
+        end
       end
     end
   end