Remove n+1 queries from GraphQL EpicType

Support hasChildren and hasIssues in a more optimised way.
Extend LazyEpicAggregate to receive a block to be executed.

ee/app/graphql/types/epic_type.rb

@@ -79,11 +79,9 @@ module Types
           description: 'Labels assigned to the epic'
 
     field :has_children, GraphQL::BOOLEAN_TYPE, null: false,
-          description: 'Indicates if the epic has children',
-          method: :has_children?
+          description: 'Indicates if the epic has children'
     field :has_issues, GraphQL::BOOLEAN_TYPE, null: false,
-          description: 'Indicates if the epic has direct issues',
-          method: :has_issues?
+          description: 'Indicates if the epic has direct issues'
 
     field :web_path, GraphQL::STRING_TYPE, null: false,
           description: 'Web path of the epic',
@@ -145,5 +143,24 @@ module Types
       resolve: -> (epic, args, ctx) do
         Epics::DescendantCountService.new(epic, ctx[:current_user])
       end
+
+    def has_children?
+      return object.has_children? unless Feature.enabled?(:unfiltered_epic_aggregates, object.group, default_enabled: true)
+
+      Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate.new(context, object.id, COUNT) do |node, _aggregate_object|
+        node.children.any?
+      end
+    end
+
+    def has_issues?
+      return object.has_issues? unless Feature.enabled?(:unfiltered_epic_aggregates, object.group, default_enabled: true)
+
+      Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate.new(context, object.id, COUNT) do |node, _aggregate_object|
+        node.has_issues?
+      end
+    end
+
+    alias_method :has_children, :has_children?
+    alias_method :has_issues, :has_issues?
   end
 end

ee/lib/gitlab/graphql/aggregations/epics/epic_node.rb

@@ -50,10 +50,10 @@ module Gitlab
 
           def to_s
             {
-                epic_id: @epic_id,
-                parent_id: @parent_id,
-                children: children,
-                object_id: object_id
+              epic_id: @epic_id,
+              parent_id: @parent_id,
+              children: children,
+              object_id: object_id
             }.to_s
           end
 
@@ -70,6 +70,12 @@ module Gitlab
             @sums[key] = direct_sum + sum_from_children
           end
 
+          def has_issues?
+            [CLOSED_ISSUE_STATE, OPENED_ISSUE_STATE].any? do |state|
+              value_from_records(COUNT, state, ISSUE_TYPE) > 0
+            end
+          end
+
           private
 
           def set_epic_attributes(record)

ee/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate.rb

@@ -13,7 +13,7 @@ module Gitlab
 
           # Because facets "count" and "weight_sum" share the same db query, but have a different graphql type object,
           # we can separate them and serve only the fields which are requested by the GraphQL query
-          def initialize(query_ctx, epic_id, aggregate_facet)
+          def initialize(query_ctx, epic_id, aggregate_facet, &block)
             @epic_id = epic_id
 
             error = validate_facet(aggregate_facet)
@@ -31,6 +31,8 @@ module Gitlab
             }
             # Register this ID to be loaded later:
             @lazy_state[:pending_ids] << epic_id
+
+            @block = block
           end
 
           # Return the loaded record, hitting the database if needed
@@ -41,7 +43,10 @@ module Gitlab
               load_records_into_tree
             end
 
-            aggregate_object(tree[@epic_id])
+            node = tree[@epic_id]
+            object = aggregate_object(node)
+
+            @block ? @block.call(node, object) : object
           end
 
           private

ee/spec/lib/gitlab/graphql/aggregations/epics/epic_node_spec.rb

@@ -39,6 +39,12 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
       allow(subject).to receive(:epic_info_flat_list).and_return(flat_info)
     end
 
+    shared_examples 'has_issues?' do |expected_result|
+      it "returns #{expected_result}" do
+        expect(subject.has_issues?).to eq expected_result
+      end
+    end
+
     context 'an epic with no child epics' do
       context 'with no child issues', :aggregate_results do
         let(:flat_info) { [] }
@@ -52,6 +58,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
+
+        it_behaves_like 'has_issues?', false
       end
 
       context 'with an issue with 0 weight', :aggregate_results do
@@ -70,9 +78,11 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
+
+        it_behaves_like 'has_issues?', true
       end
 
-      context 'with an issue with nonzero weight' do
+      context 'with an open issue with nonzero weight' do
         let(:flat_info) do
           [
               record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
@@ -88,6 +98,18 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
+
+        it_behaves_like 'has_issues?', true
+      end
+
+      context 'with a closed issue with nonzero weight' do
+        let(:flat_info) do
+          [
+              record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: CLOSED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
+          ]
+        end
+
+        it_behaves_like 'has_issues?', true
       end
     end
 
@@ -119,6 +141,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 1)
           expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
+
+        it_behaves_like 'has_issues?', false
       end
     end
   end

ee/spec/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate_spec.rb

@@ -92,6 +92,20 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
         expect(lazy_state[:pending_ids]).to be_empty
       end
 
+      context 'if a block is provided' do
+        let(:result) { 123 }
+
+        subject do
+          described_class.new(query_ctx, epic_id, COUNT) do |object|
+            result
+          end
+        end
+
+        it 'calls the block' do
+          expect(subject.epic_aggregate).to eq result
+        end
+      end
+
       it 'creates the parent-child associations', :aggregate_failures do
         subject.epic_aggregate
 

ee/spec/requests/api/graphql/group/epic/epic_aggregate_query_spec.rb

@@ -7,26 +7,26 @@ describe 'Epic aggregates (count and weight)' do
 
   let_it_be(:current_user) { create(:user) }
   let_it_be(:group) { create(:group, :public) }
-  let_it_be(:parent_epic) { create(:epic, id: 1, group: group, title: 'parent epic') }
+  let_it_be(:parent_epic) { create(:epic, group: group, title: 'parent epic') }
 
+  let(:target_epic) { parent_epic }
   let(:query) do
-    graphql_query_for('group', { fullPath: group.full_path }, query_graphql_field('epics', { iid: parent_epic.iid }, epic_aggregates_query))
+    graphql_query_for('group', { fullPath: target_epic.group.full_path }, query_graphql_field('epics', { iid: target_epic.iid }, epic_aggregates_query))
   end
-
   let(:epic_aggregates_query) do
     <<~QUERY
-    nodes {
-      descendantWeightSum {
-        openedIssues
-        closedIssues
-      }
-      descendantCounts {
-        openedEpics
-        closedEpics
-        openedIssues
-        closedIssues
+      nodes {
+        descendantWeightSum {
+          openedIssues
+          closedIssues
+        }
+        descendantCounts {
+          openedEpics
+          closedEpics
+          openedIssues
+          closedIssues
+        }
       }
-    }
     QUERY
   end
 
@@ -42,15 +42,17 @@ describe 'Epic aggregates (count and weight)' do
 
     let_it_be(:project) { create(:project, namespace: group) }
 
-    let_it_be(:epic_with_issues) { create(:epic, id: 2, group: subgroup, parent: parent_epic, title: 'epic with issues') }
-    let_it_be(:epic_without_issues) { create(:epic, :closed, id: 3, group: subgroup, parent: parent_epic, title: 'epic without issues') }
-    let_it_be(:closed_epic) { create(:epic, :closed, id: 4, group: subgroup, parent: parent_epic, title: 'closed epic') }
+    let_it_be(:epic_with_issues) { create(:epic, group: subgroup, parent: parent_epic, title: 'epic with issues') }
+    let_it_be(:epic_without_issues) { create(:epic, :closed, group: subgroup, parent: parent_epic, title: 'epic without issues') }
+    let_it_be(:closed_epic) { create(:epic, :closed, group: subgroup, parent: parent_epic, title: 'closed epic') }
 
     let_it_be(:issue1) { create(:issue, project: project, weight: 5, state: :opened) }
     let_it_be(:issue2) { create(:issue, project: project, weight: 7, state: :closed) }
+    let_it_be(:issue3) { create(:issue, project: project, weight: 0) }
 
     let_it_be(:epic_issue1) { create(:epic_issue, epic: epic_with_issues, issue: issue1) }
     let_it_be(:epic_issue2) { create(:epic_issue, epic: epic_with_issues, issue: issue2) }
+    let_it_be(:epic_issue3) { create(:epic_issue, epic: parent_epic, issue: issue3) }
 
     before do
       group.add_developer(current_user)
@@ -73,7 +75,7 @@ describe 'Epic aggregates (count and weight)' do
 
       it 'returns the issue counts' do
         issue_count_result = {
-          "openedIssues" => 1,
+          "openedIssues" => 2,
           "closedIssues" => 1
         }
 
@@ -83,6 +85,53 @@ describe 'Epic aggregates (count and weight)' do
       end
     end
 
+    shared_examples 'having correct values for' do |field_name|
+      let(:epic_aggregates_query) do
+        <<~QUERY
+          nodes {
+            #{field_name}
+          }
+        QUERY
+      end
+
+      it_behaves_like 'a working graphql query'
+
+      context 'when target epic has child epics or issues' do
+        let(:target_epic) { parent_epic }
+
+        it 'returns true' do
+          post_graphql(query, current_user: current_user)
+
+          expect(subject).to include(a_hash_including(field_name => true))
+        end
+      end
+
+      context 'when target epic has no child epics nor issues' do
+        let(:target_epic) { epic_without_issues }
+
+        it 'returns false' do
+          post_graphql(query, current_user: current_user)
+
+          expect(subject).to include(a_hash_including(field_name => false))
+        end
+      end
+    end
+
+    shared_examples 'efficient query' do
+      it 'does not result in N+1' do
+        control_count = ActiveRecord::QueryRecorder.new { post_graphql(query, current_user: current_user) }.count
+
+        query_with_multiple_epics = graphql_query_for('group', { fullPath: epic_with_issues.group.full_path }, query_graphql_field('epics', { iids: [epic_with_issues.iid, epic_without_issues.iid, parent_epic.iid] }, epic_aggregates_query))
+
+        # We still get multiple of these lines
+        # Group Load (0.6ms)  SELECT "namespaces".* FROM "namespaces" WHERE "namespaces"."type" = 'Group' AND "namespaces"."id" = x LIMIT 1
+        #   -> ee/app/policies/epic_policy.rb:4:in `block in <class:EpicPolicy>'
+        # So I'll add n to this number to take this into account.
+        # See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/27551#note_307716428
+        expect { post_graphql(query_with_multiple_epics, current_user: current_user) }.not_to exceed_query_limit(control_count + 2)
+      end
+    end
+
     context 'with feature flag enabled' do
       before do
         stub_feature_flags(unfiltered_epic_aggregates: true)
@@ -99,14 +148,40 @@ describe 'Epic aggregates (count and weight)' do
 
       it 'returns the weights' do
         descendant_weight_result = {
-            "openedIssues" => 5,
-            "closedIssues" => 7
+          "openedIssues" => 5,
+          "closedIssues" => 7
         }
 
         is_expected.to include(
           a_hash_including('descendantWeightSum' => a_hash_including(descendant_weight_result))
         )
       end
+
+      context 'when requesting has_issues' do
+        let(:epic_aggregates_query) do
+          <<~QUERY
+            nodes {
+              hasIssues
+            }
+          QUERY
+        end
+
+        it_behaves_like 'having correct values for', 'hasIssues'
+        it_behaves_like 'efficient query'
+      end
+
+      context 'when requesting has_children' do
+        let(:epic_aggregates_query) do
+          <<~QUERY
+            nodes {
+              hasChildren
+            }
+          QUERY
+        end
+
+        it_behaves_like 'having correct values for', 'hasChildren'
+        it_behaves_like 'efficient query'
+      end
     end
 
     context 'with feature flag disabled' do
@@ -114,17 +189,20 @@ describe 'Epic aggregates (count and weight)' do
         stub_feature_flags(unfiltered_epic_aggregates: false)
       end
 
+      it_behaves_like 'having correct values for', 'hasChildren'
+      it_behaves_like 'having correct values for', 'hasIssues'
+
       context 'when requesting counts' do
         let(:epic_aggregates_query) do
           <<~QUERY
-          nodes {
-            descendantCounts {
-              openedEpics
-              closedEpics
-              openedIssues
-              closedIssues
+            nodes {
+              descendantCounts {
+                openedEpics
+                closedEpics
+                openedIssues
+                closedIssues
+              }
             }
-          }
           QUERY
         end
 
@@ -140,12 +218,12 @@ describe 'Epic aggregates (count and weight)' do
       context 'when requesting weights' do
         let(:epic_aggregates_query) do
           <<~QUERY
-          nodes {
-            descendantWeightSum {
-              openedIssues
-              closedIssues
+            nodes {
+              descendantWeightSum {
+                openedIssues
+                closedIssues
+              }
             }
-          }
           QUERY
         end