Include subgroups in GraphQL field for vulnerabilityGrades for Groups

This change adds ability to include subgroups when fetching
vulnerabilityGrades for a group.

doc/api/graphql/reference/gitlab_schema.graphql

@@ -8725,7 +8725,12 @@ type Group {
   """
   Represents vulnerable project counts for each grade
   """
-  vulnerabilityGrades: [VulnerableProjectsByGrade!]!
+  vulnerabilityGrades(
+    """
+    Include grades belonging to subgroups
+    """
+    includeSubgroups: Boolean = false
+  ): [VulnerableProjectsByGrade!]!
 
   """
   Vulnerability scanners reported on the project vulnerabilties of the group and its subgroups

doc/api/graphql/reference/gitlab_schema.json

@@ -23634,7 +23634,16 @@
               "name": "vulnerabilityGrades",
               "description": "Represents vulnerable project counts for each grade",
               "args": [
-
+                {
+                  "name": "includeSubgroups",
+                  "description": "Include grades belonging to subgroups",
+                  "type": {
+                    "kind": "SCALAR",
+                    "name": "Boolean",
+                    "ofType": null
+                  },
+                  "defaultValue": "false"
+                }
               ],
               "type": {
                 "kind": "NON_NULL",

ee/app/assets/javascripts/security_dashboard/graphql/group_vulnerability_grades.query.graphql

@@ -3,7 +3,7 @@
 
 query groupVulnerabilityGrades($fullPath: ID!) {
   group(fullPath: $fullPath) {
-    vulnerabilityGrades {
+    vulnerabilityGrades(includeSubgroups: true) {
       grade
       projects {
         nodes {

ee/app/graphql/ee/types/group_type.rb

@@ -65,9 +65,7 @@ module EE
               [::Types::VulnerableProjectsByGradeType],
               null: false,
               description: 'Represents vulnerable project counts for each grade',
-              resolve: -> (obj, _args, ctx) {
-                ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate.new(ctx, obj)
-              }
+              resolver: ::Resolvers::VulnerabilitiesGradeResolver
       end
     end
   end

ee/app/graphql/resolvers/vulnerabilities_grade_resolver.rb

+# frozen_string_literal: true
+
+module Resolvers
+  class VulnerabilitiesGradeResolver < VulnerabilitiesBaseResolver
+    type [::Types::VulnerableProjectsByGradeType], null: true
+
+    argument :include_subgroups, GraphQL::BOOLEAN_TYPE,
+              required: false,
+              default_value: false,
+              description: 'Include grades belonging to subgroups'
+
+    def resolve(**args)
+      ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate
+        .new(context, vulnerable, include_subgroups: args[:include_subgroups])
+    end
+  end
+end

ee/app/models/vulnerabilities/projects_grade.rb

@@ -19,9 +19,11 @@ module Vulnerabilities
       vulnerable.projects.with_vulnerability_statistics.inc_routes.where(id: project_ids)
     end
 
-    def self.grades_for(vulnerables)
+    def self.grades_for(vulnerables, include_subgroups: false)
+      projects = include_subgroups ? vulnerables.map(&:all_projects) : vulnerables.map(&:projects)
+
       ::Vulnerabilities::Statistic
-        .for_project(vulnerables.map(&:projects).reduce(&:or))
+        .for_project(projects.reduce(&:or))
         .group(:letter_grade)
         .select(:letter_grade, 'array_agg(project_id) project_ids')
         .then do |statistics|

ee/changelogs/unreleased/243755-add-include-subgroups-to-group-vulnerability-grades-graphql.yml

+---
+title: Include subgroups in GraphQL field for vulnerabilityGrades for Groups
+merge_request: 45518
+author:
+type: changed

ee/lib/gitlab/graphql/aggregations/vulnerability_statistics/lazy_aggregate.rb

@@ -7,27 +7,27 @@ module Gitlab
         class LazyAggregate
           attr_reader :vulnerable, :lazy_state
 
-          def initialize(query_ctx, vulnerable)
+          def initialize(query_ctx, vulnerable, include_subgroups: false)
             @vulnerable = vulnerable.respond_to?(:sync) ? vulnerable.sync : vulnerable
-
+            @include_subgroups = include_subgroups
             # Initialize the loading state for this query,
             # or get the previously-initiated state
             @lazy_state = query_ctx[:lazy_vulnerability_statistics_aggregate] ||= {
-              pending_vulnerables: Set.new,
-              loaded_objects: {}
+              pending_vulnerables: { true => Set.new, false => Set.new },
+              loaded_objects: { true => {}, false => {} }
             }
             # Register this ID to be loaded later:
-            @lazy_state[:pending_vulnerables] << vulnerable
+            @lazy_state[:pending_vulnerables][@include_subgroups] << vulnerable
           end
 
           # Return the loaded record, hitting the database if needed
           def execute
             # Check if the record was already loaded
-            if @lazy_state[:pending_vulnerables].present?
+            if @lazy_state[:pending_vulnerables][@include_subgroups].present?
               load_records_into_loaded_objects
             end
 
-            @lazy_state[:loaded_objects][@vulnerable]
+            @lazy_state[:loaded_objects][@include_subgroups][@vulnerable]
           end
 
           private
@@ -35,14 +35,17 @@ module Gitlab
           def load_records_into_loaded_objects
             # The record hasn't been loaded yet, so
             # hit the database with all pending IDs to prevent N+1
-            pending_vulnerables = @lazy_state[:pending_vulnerables].to_a
-            grades = ::Vulnerabilities::ProjectsGrade.grades_for(pending_vulnerables)
+            @lazy_state[:pending_vulnerables].each do |include_subgroups, pending_vulnerables|
+              next if pending_vulnerables.blank?
 
-            pending_vulnerables.each do |vulnerable|
-              @lazy_state[:loaded_objects][vulnerable] = grades[vulnerable]
-            end
+              grades = ::Vulnerabilities::ProjectsGrade.grades_for(pending_vulnerables, include_subgroups: include_subgroups)
 
-            @lazy_state[:pending_vulnerables].clear
+              pending_vulnerables.each do |vulnerable|
+                @lazy_state[:loaded_objects][include_subgroups][vulnerable] = grades[vulnerable]
+              end
+
+              @lazy_state[:pending_vulnerables][include_subgroups].clear
+            end
           end
         end
       end

ee/spec/graphql/resolvers/vulnerabilities_grade_resolver_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Resolvers::VulnerabilitiesGradeResolver do
+  include GraphqlHelpers
+
+  subject { resolve(described_class, obj: group, args: args, ctx: { current_user: user }).execute }
+
+  let_it_be(:group) { create(:group) }
+  let_it_be(:subgroup) { create(:group, parent: group) }
+  let_it_be(:project) { create(:project, namespace: group) }
+  let_it_be(:project_in_subgroup) { create(:project, namespace: subgroup) }
+  let_it_be(:user) { create(:user) }
+
+  let_it_be(:vulnerability_statistic_1) { create(:vulnerability_statistic, :grade_f, project: project) }
+  let_it_be(:vulnerability_statistic_2) { create(:vulnerability_statistic, :grade_d, project: project_in_subgroup) }
+
+  describe '#resolve' do
+    let(:args) { { include_subgroups: include_subgroups } }
+
+    context 'when include_subgroups is set to true' do
+      let(:include_subgroups) { true }
+
+      it 'returns project grades for projects in group and its subgroups' do
+        expect(subject.map(&:grade)).to match_array(%w[d f])
+      end
+    end
+
+    context 'when include_subgroups is set to true' do
+      let(:include_subgroups) { false }
+
+      it 'returns project grades for projects in group only' do
+        expect(subject.map(&:grade)).to match_array(%w[f])
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/graphql/aggregations/vulnerability_statistics/lazy_aggregate_spec.rb

@@ -7,29 +7,37 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
     {}
   end
 
-  let(:vulnerable) { create(:group) }
-  let(:blocks_vulnerable) { 18 }
-  let(:blocking_vulnerable) { 38 }
+  let_it_be(:vulnerable) { create(:group) }
+  let_it_be(:other_vulnerable) { create(:group) }
+  let(:include_subgroups) { true }
 
   describe '#initialize' do
     it 'adds the vulnerable to the lazy state' do
-      subject = described_class.new(query_ctx, vulnerable)
+      subject = described_class.new(query_ctx, vulnerable, include_subgroups: true)
 
-      expect(subject.lazy_state[:pending_vulnerables]).to match [vulnerable]
+      expect(subject.lazy_state[:pending_vulnerables]).to eq(true => Set.new([vulnerable]), false => Set.new)
       expect(subject.vulnerable).to match vulnerable
     end
   end
 
   describe '#execute' do
-    subject { described_class.new(query_ctx, vulnerable) }
+    let(:fake_data) do
+      {
+        vulnerable => [::Vulnerabilities::ProjectsGrade.new(vulnerable, 'a', [])],
+        other_vulnerable => [::Vulnerabilities::ProjectsGrade.new(other_vulnerable, 'b', [])]
+      }
+    end
+
+    subject { described_class.new(query_ctx, vulnerable, include_subgroups: include_subgroups) }
 
     before do
       subject.instance_variable_set(:@lazy_state, fake_state)
+      allow(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).and_return(fake_data)
     end
 
     context 'if the record has already been loaded' do
       let(:fake_state) do
-        { pending_vulnerables: Set.new, loaded_objects: { vulnerable => [::Vulnerabilities::ProjectsGrade.new(vulnerable, 'a', [])] } }
+        { pending_vulnerables: { true => Set.new, false => Set.new }, loaded_objects: { true => { vulnerable => [::Vulnerabilities::ProjectsGrade.new(vulnerable, 'a', [])] }, false => {} } }
       end
 
       it 'does not make the query again' do
@@ -40,24 +48,14 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
     end
 
     context 'if the record has not been loaded' do
-      let(:other_vulnerable) { create(:group) }
+      let(:include_subgroups) { true }
       let(:fake_state) do
-        { pending_vulnerables: Set.new([vulnerable, other_vulnerable]), loaded_objects: {} }
-      end
-
-      let(:fake_data) do
-        {
-          vulnerable => [::Vulnerabilities::ProjectsGrade.new(vulnerable, 'a', [])],
-          other_vulnerable => [::Vulnerabilities::ProjectsGrade.new(other_vulnerable, 'b', [])]
-        }
-      end
-
-      before do
-        allow(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).and_return(fake_data)
+        { pending_vulnerables: { true => Set.new([vulnerable]), false => Set.new([other_vulnerable]) }, loaded_objects: { true => {}, false => {} } }
       end
 
       it 'makes the query' do
-        expect(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).with([vulnerable, other_vulnerable])
+        expect(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).with([vulnerable], include_subgroups: true)
+        expect(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).with([other_vulnerable], include_subgroups: false)
 
         subject.execute
       end
@@ -65,7 +63,7 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
       it 'clears the pending IDs' do
         subject.execute
 
-        expect(subject.lazy_state[:pending_vulnerables]).to be_empty
+        expect(subject.lazy_state[:pending_vulnerables].values).to all(be_empty)
       end
     end
   end

ee/spec/models/vulnerabilities/projects_grade_spec.rb

@@ -4,38 +4,64 @@ require 'spec_helper'
 
 RSpec.describe Vulnerabilities::ProjectsGrade do
   let_it_be(:group) { create(:group) }
+  let_it_be(:subgroup) { create(:group, parent: group) }
   let_it_be(:project_1) { create(:project, group: group) }
   let_it_be(:project_2) { create(:project, group: group) }
   let_it_be(:project_3) { create(:project, group: group) }
   let_it_be(:project_4) { create(:project, group: group) }
   let_it_be(:project_5) { create(:project, group: group) }
+  let_it_be(:project_6) { create(:project, group: subgroup) }
 
   let_it_be(:vulnerability_statistic_1) { create(:vulnerability_statistic, :grade_a, project: project_1) }
   let_it_be(:vulnerability_statistic_2) { create(:vulnerability_statistic, :grade_b, project: project_2) }
   let_it_be(:vulnerability_statistic_3) { create(:vulnerability_statistic, :grade_b, project: project_3) }
   let_it_be(:vulnerability_statistic_4) { create(:vulnerability_statistic, :grade_c, project: project_4) }
   let_it_be(:vulnerability_statistic_5) { create(:vulnerability_statistic, :grade_f, project: project_5) }
+  let_it_be(:vulnerability_statistic_6) { create(:vulnerability_statistic, :grade_d, project: project_6) }
 
   describe '.grades_for' do
     let(:compare_key) { ->(projects_grade) { [projects_grade.grade, projects_grade.project_ids.sort] } }
+    let(:include_subgroups) { false }
 
-    subject(:projects_grades) { described_class.grades_for([vulnerable]) }
+    subject(:projects_grades) { described_class.grades_for([vulnerable], include_subgroups: include_subgroups) }
 
     context 'when the given vulnerable is a Group' do
       let(:vulnerable) { group }
-      let(:expected_projects_grades) do
-        {
-          vulnerable => [
-            described_class.new(vulnerable, 'a', [project_1.id]),
-            described_class.new(vulnerable, 'b', [project_2.id, project_3.id]),
-            described_class.new(vulnerable, 'c', [project_4.id]),
-            described_class.new(vulnerable, 'f', [project_5.id])
-          ]
-        }
+
+      context 'when subgroups are not included' do
+        let(:expected_projects_grades) do
+          {
+            vulnerable => [
+              described_class.new(vulnerable, 'a', [project_1.id]),
+              described_class.new(vulnerable, 'b', [project_2.id, project_3.id]),
+              described_class.new(vulnerable, 'c', [project_4.id]),
+              described_class.new(vulnerable, 'f', [project_5.id])
+            ]
+          }
+        end
+
+        it 'returns the letter grades for given vulnerable' do
+          expect(projects_grades[vulnerable].map(&compare_key)).to match_array(expected_projects_grades[vulnerable].map(&compare_key))
+        end
       end
 
-      it 'returns the letter grades for given vulnerable' do
-        expect(projects_grades[vulnerable].map(&compare_key)).to match_array(expected_projects_grades[vulnerable].map(&compare_key))
+      context 'when subgroups are included' do
+        let(:include_subgroups) { true }
+        let(:expected_projects_grades) do
+          {
+            vulnerable => [
+              described_class.new(vulnerable, 'a', [project_1.id]),
+              described_class.new(vulnerable, 'b', [project_2.id, project_3.id]),
+              described_class.new(vulnerable, 'c', [project_4.id]),
+              described_class.new(vulnerable, 'd', [project_6.id]),
+              described_class.new(vulnerable, 'f', [project_5.id])
+            ]
+          }
+        end
+
+        it 'returns the letter grades for given vulnerable' do
+          expect(projects_grades[vulnerable].map(&compare_key)).to match_array(expected_projects_grades[vulnerable].map(&compare_key))
+        end
       end
     end