Optimize query for reference parser

Contributes to https://gitlab.com/gitlab-org/gitlab/-/issues/271242

* Load only projects for merge requests
* Fix test about N+1 problem

config/feature_flags/development/optimize_merge_request_parser.yml

+---
+name: optimize_merge_request_parser
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/62490/
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/331893
+milestone: '14.0'
+type: development
+group: group::source code
+default_enabled: false

lib/banzai/reference_parser/merge_request_parser.rb

@@ -7,6 +7,19 @@ module Banzai
 
       self.reference_type = :merge_request
 
+      def nodes_visible_to_user(user, nodes)
+        return super if Feature.disabled?(:optimize_merge_request_parser, user, default_enabled: :yaml)
+
+        merge_request_nodes = nodes.select { |node| node.has_attribute?(self.class.data_attribute) }
+        records = projects_for_nodes(merge_request_nodes)
+
+        merge_request_nodes.select do |node|
+          project = records[node]
+
+          project && can_read_reference?(user, project)
+        end
+      end
+
       def records_for_nodes(nodes)
         @merge_requests_for_nodes ||= grouped_objects_for_nodes(
           nodes,
@@ -23,14 +36,19 @@ module Banzai
         )
       end
 
-      def can_read_reference?(user, merge_request)
+      def can_read_reference?(user, object)
         memo = strong_memoize(:can_read_reference) { {} }
 
-        project_id = merge_request.project_id
+        project_id = object.project_id
 
         return memo[project_id] if memo.key?(project_id)
 
-        memo[project_id] = can?(user, :read_merge_request_iid, merge_request.project)
+        memo[project_id] = can?(user, :read_merge_request_iid, object)
+      end
+
+      def projects_for_nodes(nodes)
+        @projects_for_nodes ||=
+          grouped_objects_for_nodes(nodes, Project.includes(:project_feature, :group, :namespace), 'data-project')
       end
     end
   end

spec/lib/banzai/reference_parser/commit_parser_spec.rb

@@ -130,11 +130,11 @@ RSpec.describe Banzai::ReferenceParser::CommitParser do
   end
 
   context 'when checking commits on another projects' do
-    let(:control_links) do
+    let!(:control_links) do
       [commit_link]
     end
 
-    let(:actual_links) do
+    let!(:actual_links) do
       control_links + [commit_link, commit_link]
     end
 

spec/lib/banzai/reference_parser/merge_request_parser_spec.rb

@@ -18,10 +18,19 @@ RSpec.describe Banzai::ReferenceParser::MergeRequestParser do
     context 'when the link has a data-issue attribute' do
       before do
         project.update_attribute(:visibility_level, Gitlab::VisibilityLevel::PUBLIC)
+        link['data-project'] = merge_request.project_id.to_s
         link['data-merge-request'] = merge_request.id.to_s
       end
 
       it_behaves_like "referenced feature visibility", "merge_requests"
+
+      context 'when optimize_merge_request_parser feature flag is off' do
+        before do
+          stub_feature_flags(optimize_merge_request_parser: false)
+        end
+
+        it_behaves_like "referenced feature visibility", "merge_requests"
+      end
     end
   end
 
@@ -29,6 +38,7 @@ RSpec.describe Banzai::ReferenceParser::MergeRequestParser do
     describe 'when the link has a data-merge-request attribute' do
       context 'using an existing merge request ID' do
         it 'returns an Array of merge requests' do
+          link['data-project'] = merge_request.project_id.to_s
           link['data-merge-request'] = merge_request.id.to_s
 
           expect(subject.referenced_by([link])).to eq([merge_request])
@@ -37,6 +47,7 @@ RSpec.describe Banzai::ReferenceParser::MergeRequestParser do
 
       context 'using a non-existing merge request ID' do
         it 'returns an empty Array' do
+          link['data-project'] = merge_request.project_id.to_s
           link['data-merge-request'] = ''
 
           expect(subject.referenced_by([link])).to eq([])
@@ -49,16 +60,16 @@ RSpec.describe Banzai::ReferenceParser::MergeRequestParser do
     let(:other_project) { create(:project, :public) }
     let(:other_merge_request) { create(:merge_request, source_project: other_project) }
 
-    let(:control_links) do
+    let!(:control_links) do
       [merge_request_link(other_merge_request)]
     end
 
-    let(:actual_links) do
+    let!(:actual_links) do
       control_links + [merge_request_link(create(:merge_request, :conflict, source_project: other_project))]
     end
 
     def merge_request_link(merge_request)
-      Nokogiri::HTML.fragment(%Q{<a data-merge-request="#{merge_request.id}"></a>}).children[0]
+      Nokogiri::HTML.fragment(%Q{<a data-project="#{merge_request.project_id}" data-merge-request="#{merge_request.id}"></a>}).children[0]
     end
 
     before do

spec/support/helpers/reference_parser_helpers.rb

@@ -11,25 +11,20 @@ module ReferenceParserHelpers
   end
 
   RSpec.shared_examples 'no project N+1 queries' do
-    it 'avoids N+1 queries in #nodes_visible_to_user', :request_store do
+    it 'avoids N+1 queries in #nodes_visible_to_user' do
       context = Banzai::RenderContext.new(project, user)
 
-      record_queries = lambda do |links|
-        ActiveRecord::QueryRecorder.new do
+      request = lambda do |links|
         described_class.new(context).nodes_visible_to_user(user, links)
       end
-      end
 
-      control = record_queries.call(control_links)
+      control = ActiveRecord::QueryRecorder.new { request.call(control_links) }
 
       create(:group_member, group: project.group) if project.group
       create(:project_member, project: project)
       create(:project_group_link, project: project)
 
-      actual = record_queries.call(actual_links)
-
-      expect(actual.count).to be <= control.count
-      expect(actual.cached_count).to be <= control.cached_count
+      expect { request.call(actual_links) }.not_to exceed_query_limit(control)
     end
   end