Preload all associations in Vulnerability GraphQL API

This change adds preloading associations in Vulnerability GraphQL to
prevent N+1 queries. It also solves problem with N+1 query when loading
Vulnerability Findings.

ee/app/graphql/ee/gitlab_schema.rb

@@ -8,6 +8,7 @@ module EE
       lazy_resolve ::Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate, :epic_aggregate
       lazy_resolve ::Gitlab::Graphql::Aggregations::Issues::LazyBlockAggregate, :block_aggregate
       lazy_resolve ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate, :aggregate
+      lazy_resolve ::Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate, :user_notes_count_aggregate
     end
   end
 end

ee/app/graphql/resolvers/vulnerabilities/issue_links_resolver.rb

@@ -9,10 +9,10 @@ module Resolvers
                required: false,
                description: 'Filter issue links by link type'
 
-      delegate :issue_links, :finding, to: :object, private: true
+      delegate :issue_links, :finding, :created_issue_links, to: :object, private: true
 
       def resolve(link_type: nil, **)
-        links = issue_links.by_link_type(link_type)
+        links = issue_links_by_link_type(link_type)
         return links if links.present? || link_type != 'created'
 
         issue_feedback = finding.issue_feedback
@@ -20,6 +20,17 @@ module Resolvers
 
         issue_links.build(issue_id: issue_feedback.id, link_type: :created)
       end
+
+      private
+
+      def issue_links_by_link_type(link_type)
+        case link_type.to_s
+        when Types::Vulnerability::IssueLinkTypeEnum.enum['created']
+          created_issue_links
+        else
+          issue_links.by_link_type(link_type)
+        end
+      end
     end
   end
 end

ee/app/graphql/resolvers/vulnerabilities_resolver.rb

@@ -29,7 +29,10 @@ module Resolvers
     def resolve(**args)
       return Vulnerability.none unless vulnerable
 
-      vulnerabilities(args).with_findings_and_scanner.ordered
+      vulnerabilities(args)
+        .with_findings_scanner_and_identifiers
+        .with_created_issue_links_and_issues
+        .ordered
     end
 
     private

ee/app/graphql/types/vulnerability_type.rb

@@ -31,7 +31,10 @@ module Types
           description: "Indicates whether the vulnerability is fixed on the default branch or not"
 
     field :user_notes_count, GraphQL::INT_TYPE, null: false,
-          description: 'Number of user notes attached to the vulnerability'
+          description: 'Number of user notes attached to the vulnerability',
+          resolve: -> (obj, _args, ctx) {
+            ::Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate.new(ctx, obj)
+          }
 
     field :vulnerability_path, GraphQL::STRING_TYPE, null: true,
           description: "URL to the vulnerability's details page",

ee/app/models/ee/note.rb

@@ -13,6 +13,11 @@ module EE
       scope :searchable, -> { where(system: false).includes(:noteable) }
       scope :by_humans, -> { user.joins(:author).merge(::User.humans) }
       scope :with_suggestions, -> { joins(:suggestions) }
+      scope :count_for_vulnerability_id, ->(vulnerability_id) do
+        where(noteable_type: Vulnerability.name, noteable_id: vulnerability_id)
+          .group(:noteable_id)
+          .count
+      end
     end
 
     # Original method in Elastic::ApplicationSearch

ee/app/models/vulnerabilities/finding.rb

@@ -292,6 +292,12 @@ module Vulnerabilities
 
     # Array.difference (-) method uses hash and eql? methods to do comparison
     def hash
+      # This is causing N+1 queries whenever we are calling findings, ActiveRecord uses #hash method to make sure the
+      # array with findings is uniq before preloading. This method is used only in Gitlab::Ci::Reports::Security::VulnerabilityReportsComparer
+      # where we are normalizing security report findings into instances of Vulnerabilities::Finding, this is why we are using original implementation
+      # when Finding is persisted and identifiers are not preloaded.
+      return super if persisted? && !identifiers.loaded?
+
       report_type.hash ^ location_fingerprint.hash ^ first_fingerprint.hash
     end
 

ee/app/models/vulnerability.rb

@@ -35,6 +35,7 @@ class Vulnerability < ApplicationRecord
 
   has_many :findings, class_name: 'Vulnerabilities::Finding', inverse_of: :vulnerability
   has_many :issue_links, class_name: 'Vulnerabilities::IssueLink', inverse_of: :vulnerability
+  has_many :created_issue_links, -> { created }, class_name: 'Vulnerabilities::IssueLink', inverse_of: :vulnerability
   has_many :related_issues, through: :issue_links, source: :issue do
     def with_vulnerability_links
       select('issues.*, vulnerability_issue_links.id AS vulnerability_link_id, '\
@@ -62,6 +63,8 @@ class Vulnerability < ApplicationRecord
 
   scope :with_findings, -> { includes(:findings) }
   scope :with_findings_and_scanner, -> { includes(findings: :scanner) }
+  scope :with_findings_scanner_and_identifiers, -> { includes(findings: [:scanner, :identifiers, finding_identifiers: :identifier]) }
+  scope :with_created_issue_links_and_issues, -> { includes(created_issue_links: :issue) }
 
   scope :for_projects, -> (project_ids) { where(project_id: project_ids) }
   scope :with_report_types, -> (report_types) { where(report_type: report_types) }

ee/changelogs/unreleased/233017-preload-associations-in-graphql-vulnerability-type.yml

+---
+title: Preload all associations in Vulnerability GraphQL API
+merge_request: 38556
+author:
+type: performance

ee/lib/gitlab/graphql/aggregations/vulnerabilities/lazy_user_notes_count_aggregate.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Aggregations
+      module Vulnerabilities
+        class LazyUserNotesCountAggregate
+          attr_reader :vulnerability, :lazy_state
+
+          def initialize(query_ctx, vulnerability)
+            @vulnerability = vulnerability.respond_to?(:sync) ? vulnerability.sync : vulnerability
+
+            # Initialize the loading state for this query,
+            # or get the previously-initiated state
+            @lazy_state = query_ctx[:lazy_aggregate] ||= {
+              pending_vulnerability_ids: Set.new,
+              loaded_objects: {}
+            }
+            # Register this ID to be loaded later:
+            @lazy_state[:pending_vulnerability_ids] << vulnerability.id
+          end
+
+          # Return the loaded record, hitting the database if needed
+          def user_notes_count_aggregate
+            # Check if the record was already loaded
+            if @lazy_state[:pending_vulnerability_ids].present?
+              load_records_into_loaded_objects
+            end
+
+            @lazy_state[:loaded_objects][@vulnerability.id]
+          end
+
+          private
+
+          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_vulnerability_ids = @lazy_state[:pending_vulnerability_ids].to_a
+            counts = ::Note.user.count_for_vulnerability_id(pending_vulnerability_ids)
+
+            pending_vulnerability_ids.each do |vulnerability_id|
+              @lazy_state[:loaded_objects][vulnerability_id] = counts[vulnerability_id].to_i
+            end
+
+            @lazy_state[:pending_vulnerability_ids].clear
+          end
+        end
+      end
+    end
+  end
+end

ee/spec/graphql/resolvers/vulnerabilities_resolver_spec.rb

@@ -80,6 +80,20 @@ RSpec.describe Resolvers::VulnerabilitiesResolver do
       it 'only returns vulnerabilities belonging to the given projects' do
         is_expected.to contain_exactly(project2_vulnerability)
       end
+
+      context 'with multiple project IDs' do
+        let(:filters) { { project_id: [project.id, project2.id] } }
+
+        it 'avoids N+1 queries' do
+          control_count = ActiveRecord::QueryRecorder.new do
+            resolve(described_class, obj: vulnerable, args: { project_id: [project2.id] }, ctx: { current_user: current_user })
+          end.count
+
+          expect do
+            subject
+          end.not_to exceed_query_limit(control_count)
+        end
+      end
     end
 
     context 'when resolving vulnerabilities for a project' do

ee/spec/lib/gitlab/graphql/aggregations/vulnerabilities/lazy_user_notes_count_aggregate_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate do
+  let(:query_ctx) do
+    {}
+  end
+
+  let(:vulnerability) { create(:vulnerability) }
+
+  describe '#initialize' do
+    it 'adds the vulnerability to the lazy state' do
+      subject = described_class.new(query_ctx, vulnerability)
+
+      expect(subject.lazy_state[:pending_vulnerability_ids]).to match [vulnerability.id]
+      expect(subject.vulnerability).to match vulnerability
+    end
+  end
+
+  describe '#user_notes_count_aggregate' do
+    subject { described_class.new(query_ctx, vulnerability) }
+
+    before do
+      subject.instance_variable_set(:@lazy_state, fake_state)
+    end
+
+    context 'if the record has already been loaded' do
+      let(:fake_state) do
+        { pending_vulnerability_ids: Set.new, loaded_objects: { vulnerability.id => 10 } }
+      end
+
+      it 'does not make the query again' do
+        expect(::Note).not_to receive(:user)
+
+        subject.user_notes_count_aggregate
+      end
+    end
+
+    context 'if the record has not been loaded' do
+      let(:other_vulnerability) { create(:vulnerability) }
+      let(:fake_state) do
+        { pending_vulnerability_ids: Set.new([vulnerability.id, other_vulnerability.id]), loaded_objects: {} }
+      end
+
+      let(:fake_data) do
+        {
+          vulnerability.id => 10,
+          other_vulnerability.id => 14
+        }
+      end
+
+      before do
+        allow(::Note).to receive_message_chain(:user, :count_for_vulnerability_id).and_return(fake_data)
+      end
+
+      it 'makes the query' do
+        expect(::Note).to receive_message_chain(:user, :count_for_vulnerability_id).with([vulnerability.id, other_vulnerability.id])
+
+        subject.user_notes_count_aggregate
+      end
+
+      it 'clears the pending IDs' do
+        subject.user_notes_count_aggregate
+
+        expect(subject.lazy_state[:pending_vulnerability_ids]).to be_empty
+      end
+    end
+  end
+end

ee/spec/models/note_spec.rb

@@ -145,4 +145,17 @@ RSpec.describe Note do
       expect(described_class.by_humans).to match_array([user_note])
     end
   end
+
+  describe '.count_for_vulnerability_id' do
+    it 'counts notes by vulnerability id' do
+      vulnerability_1 = create(:vulnerability)
+      vulnerability_2 = create(:vulnerability)
+
+      create(:note, noteable: vulnerability_1, project: vulnerability_1.project)
+      create(:note, noteable: vulnerability_2, project: vulnerability_2.project)
+      create(:note, noteable: vulnerability_2, project: vulnerability_2.project)
+
+      expect(described_class.count_for_vulnerability_id([vulnerability_1.id, vulnerability_2.id])).to eq(vulnerability_1.id => 1, vulnerability_2.id => 2)
+    end
+  end
 end

ee/spec/models/vulnerability_spec.rb

@@ -34,6 +34,7 @@ RSpec.describe Vulnerability do
     it { is_expected.to belong_to(:epic) }
     it { is_expected.to have_many(:findings).class_name('Vulnerabilities::Finding').inverse_of(:vulnerability) }
     it { is_expected.to have_many(:issue_links).class_name('Vulnerabilities::IssueLink').inverse_of(:vulnerability) }
+    it { is_expected.to have_many(:created_issue_links).class_name('Vulnerabilities::IssueLink').inverse_of(:vulnerability).conditions(link_type: Vulnerabilities::IssueLink.link_types['created']) }
     it { is_expected.to have_many(:related_issues).through(:issue_links).source(:issue) }
     it { is_expected.to belong_to(:author).class_name('User') }
     it { is_expected.to belong_to(:updated_by).class_name('User') }