Avoid N + 1 queries when fetching a list of groups via GraphQL

GraphQL node authorization will make a lot of queries
to the DB. This preloads groups feature availability

app/models/preloaders/group_policy_preloader.rb

@@ -8,15 +8,12 @@ module Preloaders
     end
 
     def execute
-      Preloaders::UserMaxAccessLevelInGroupsPreloader.new(@groups, @current_user).execute
-      Preloaders::GroupRootAncestorPreloader.new(@groups, root_ancestor_preloads).execute
+      Preloaders::UserMaxAccessLevelInGroupsPreloader.new(groups, current_user).execute
     end
 
     private
 
-    def root_ancestor_preloads
-      []
-    end
+    attr_reader :groups, :current_user
   end
 end
 

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

@@ -5,11 +5,18 @@ module EE
     module GroupPolicyPreloader
       extend ::Gitlab::Utils::Override
 
+      override :execute
+      def execute
+        super
+
+        ::Preloaders::GroupRootAncestorPreloader.new(groups, root_ancestor_preloads).execute
+        ::Gitlab::GroupPlansPreloader.new.preload(groups) if ::Gitlab::CurrentSettings.should_check_namespace_plan?
+      end
+
       private
 
-      override :root_ancestor_preloads
       def root_ancestor_preloads
-        [*super, :ip_restrictions, :saml_provider]
+        [:ip_restrictions, :saml_provider]
       end
     end
   end

ee/lib/gitlab/group_plans_preloader.rb

@@ -15,7 +15,9 @@ module Gitlab
     # plans.
     # rubocop: disable CodeReuse/ActiveRecord
     def preload(groups)
-      groups_and_ancestors = groups_and_ancestors_for(groups)
+      groups_and_ancestors = groups_and_ancestors_for(
+        activerecord_relation(groups)
+      )
       # A Hash mapping group IDs to their corresponding Group instances.
       groups_map = groups_and_ancestors.each_with_object({}) do |group, hash|
         hash[group.id] = group
@@ -52,7 +54,6 @@ module Gitlab
         group.memoized_plans = plans_map[group.id].map { |id| plans[id] }
       end
     end
-    # rubocop: enable CodeReuse/ActiveRecord
 
     # Returns an ActiveRecord::Relation that includes the given groups, and all
     # their (recursive) ancestors.
@@ -62,5 +63,16 @@ module Gitlab
        .join_gitlab_subscription
        .select('namespaces.id', 'namespaces.parent_id', 'gitlab_subscriptions.hosted_plan_id')
     end
+
+    private
+
+    def activerecord_relation(groups)
+      if groups.is_a?(ActiveRecord::Relation)
+        groups
+      else
+        Group.where(id: groups)
+      end
+    end
+    # rubocop: enable CodeReuse/ActiveRecord
   end
 end

ee/spec/lib/gitlab/group_plans_preloader_spec.rb

@@ -4,16 +4,12 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::GroupPlansPreloader, :saas do
   describe '#preload' do
-    let!(:plan1) { create(:free_plan, name: 'plan-1') }
-    let!(:plan2) { create(:free_plan, name: 'plan-2') }
+    let_it_be(:plan1) { create(:free_plan, name: 'plan-1') }
+    let_it_be(:plan2) { create(:free_plan, name: 'plan-2') }
 
-    let(:preloaded_groups) do
-      # We don't use the factory objects here because they might have the plan
-      # loaded already (as we specify the plan when creating them).
-      described_class.new.preload(Group.order(id: :asc))
-    end
+    let(:preloaded_groups) { described_class.new.preload(pristine_groups) }
 
-    before do
+    before_all do
       group1 = create(:group, name: 'group-1')
       create(:gitlab_subscription, namespace: group1, hosted_plan_id: plan1.id)
 
@@ -51,6 +47,16 @@ RSpec.describe Gitlab::GroupPlansPreloader, :saas do
       end
     end
 
-    it_behaves_like 'preloading cases'
+    context 'when an ActiveRecord relationship is provided' do
+      let(:pristine_groups) { Group.order(id: :asc) }
+
+      it_behaves_like 'preloading cases'
+    end
+
+    context 'when an array of groups is provided' do
+      let(:pristine_groups) { Group.order(id: :asc).to_a }
+
+      it_behaves_like 'preloading cases'
+    end
   end
 end

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

@@ -11,6 +11,7 @@ RSpec.describe Preloaders::GroupPolicyPreloader do
   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] }
+  let(:should_check_namespace_plan) { false }
 
   before_all do
     guest_group.add_guest(user)
@@ -19,6 +20,10 @@ RSpec.describe Preloaders::GroupPolicyPreloader do
     public_maintainer_group.add_maintainer(user)
   end
 
+  before do
+    allow(::Gitlab::CurrentSettings).to receive(:should_check_namespace_plan?).and_return(should_check_namespace_plan)
+  end
+
   context 'when ip_restrictions feature is enabled' do
     before do
       stub_licensed_features(group_ip_restriction: true)
@@ -41,6 +46,24 @@ RSpec.describe Preloaders::GroupPolicyPreloader do
     end
   end
 
+  context 'when check_namespace_plan setting is disabled' do
+    it 'does not preload group plans' do
+      expect(::Gitlab::GroupPlansPreloader).not_to receive(:new)
+
+      preload_groups_for_policy(user)
+    end
+  end
+
+  context 'when check_namespace_plan setting is enabled' do
+    let(:should_check_namespace_plan) { true }
+
+    it 'preloads group plans' do
+      expect(::Gitlab::GroupPlansPreloader).to receive_message_chain(:new, :preload)
+
+      preload_groups_for_policy(user)
+    end
+  end
+
   def authorize_all_groups(current_user, group_list = base_groups)
     group_list.each { |group| current_user.can?(:read_group, group) }
   end

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

@@ -32,6 +32,20 @@ RSpec.describe 'Query current user groups' do
     public_maintainer_group.add_maintainer(user)
   end
 
+  shared_examples 'no N + 1 DB queries' do
+    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
+
   context 'when permission_scope is CREATE_PROJECTS' do
     let(:group_arguments) { { permission_scope: :CREATE_PROJECTS } }
 
@@ -46,16 +60,20 @@ RSpec.describe 'Query current user groups' 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) }
+      context 'when check_namespace_plan setting is enabled' do
+        before do
+          stub_application_setting(check_namespace_plan: true)
+        end
 
-        create(:group, :private).tap { |group| group.add_maintainer(current_user) }
-        create(:group, :private, parent: private_maintainer_group)
+        it_behaves_like 'no N + 1 DB queries'
+      end
 
-        another_root = create(:group, :private, name: 'root-3', path: 'root-3')
-        create(:group, :private, parent: another_root).tap { |group| group.add_maintainer(current_user) }
+      context 'when check_namespace_plan setting is disabled' do
+        before do
+          stub_application_setting(check_namespace_plan: false)
+        end
 
-        expect { post_graphql(query, current_user: current_user) }.not_to exceed_query_limit(control)
+        it_behaves_like 'no N + 1 DB queries'
       end
     end
   end