Reduce number of queries on Users table

app/models/concerns/cache_markdown_field.rb

@@ -163,9 +163,9 @@ module CacheMarkdownField
     refs = all_references(self.author)
 
     references = {}
-    references[:mentioned_users_ids] = refs.mentioned_users&.pluck(:id).presence
-    references[:mentioned_groups_ids] = refs.mentioned_groups&.pluck(:id).presence
-    references[:mentioned_projects_ids] = refs.mentioned_projects&.pluck(:id).presence
+    references[:mentioned_users_ids] = refs.mentioned_user_ids.presence
+    references[:mentioned_groups_ids] = refs.mentioned_group_ids.presence
+    references[:mentioned_projects_ids] = refs.mentioned_project_ids.presence
 
     # One retry is enough as next time `model_user_mention` should return the existing mention record,
     # that threw the `ActiveRecord::RecordNotUnique` exception in first place.

ee/spec/services/epic_issues/create_service_spec.rb

@@ -124,7 +124,7 @@ RSpec.describe EpicIssues::CreateService do
 
             include_examples 'returns success'
 
-            it 'does not perofrm N + 1 queries' do
+            it 'does not perform N + 1 queries' do
               allow(SystemNoteService).to receive(:epic_issue)
               allow(SystemNoteService).to receive(:issue_on_epic)
 
@@ -132,9 +132,9 @@ RSpec.describe EpicIssues::CreateService do
               extractor = double
               allow(Gitlab::ReferenceExtractor).to receive(:new).and_return(extractor)
               allow(extractor).to receive(:reset_memoized_values)
-              allow(extractor).to receive(:mentioned_users)
-              allow(extractor).to receive(:mentioned_groups)
-              allow(extractor).to receive(:mentioned_projects)
+              allow(extractor).to receive(:mentioned_user_ids)
+              allow(extractor).to receive(:mentioned_group_ids)
+              allow(extractor).to receive(:mentioned_project_ids)
               allow(extractor).to receive(:analyze)
               allow(extractor).to receive(:issues).and_return([issue])
 

lib/banzai/reference_extractor.rb

@@ -11,11 +11,11 @@ module Banzai
       @texts_and_contexts << { text: text, context: context }
     end
 
-    def references(type, project, current_user = nil)
+    def references(type, project, current_user, ids_only: false)
       context = RenderContext.new(project, current_user)
       processor = Banzai::ReferenceParser[type].new(context)
 
-      processor.process(html_documents)
+      processor.process(html_documents, ids_only: ids_only)
     end
 
     def reset_memoized_values

lib/banzai/reference_parser/base_parser.rb

@@ -76,9 +76,11 @@ module Banzai
       end
 
       # Returns an Array of objects referenced by any of the given HTML nodes.
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         ids = unique_attribute_values(nodes, self.class.data_attribute)
 
+        return ids if options.fetch(:ids_only, false)
+
         if ids.empty?
           references_relation.none
         else
@@ -194,7 +196,7 @@ module Banzai
 
       # Processes the list of HTML documents and returns an Array containing all
       # the references.
-      def process(documents)
+      def process(documents, ids_only: false)
         type = self.class.reference_type
         reference_options = self.class.reference_options
 
@@ -202,17 +204,17 @@ module Banzai
           Querying.css(document, "a[data-reference-type='#{type}'].gfm", reference_options).to_a
         end
 
-        gather_references(nodes)
+        gather_references(nodes, ids_only: ids_only)
       end
 
       # Gathers the references for the given HTML nodes.  Returns visible
       # references and a list of nodes which are not visible to the user
-      def gather_references(nodes)
+      def gather_references(nodes, ids_only: false)
         nodes = nodes_user_can_reference(current_user, nodes)
         visible = nodes_visible_to_user(current_user, nodes)
         not_visible = nodes - visible
 
-        { visible: referenced_by(visible), not_visible: not_visible }
+        { visible: referenced_by(visible, ids_only: ids_only), not_visible: not_visible }
       end
 
       # Returns a Hash containing the projects for a given list of HTML nodes.

lib/banzai/reference_parser/commit_parser.rb

@@ -5,7 +5,7 @@ module Banzai
     class CommitParser < BaseParser
       self.reference_type = :commit
 
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         commit_ids = commit_ids_per_project(nodes)
         projects = find_projects_for_hash_keys(commit_ids)
 

lib/banzai/reference_parser/commit_range_parser.rb

@@ -5,7 +5,7 @@ module Banzai
     class CommitRangeParser < BaseParser
       self.reference_type = :commit_range
 
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         range_ids = commit_range_ids_per_project(nodes)
         projects = find_projects_for_hash_keys(range_ids)
 

lib/banzai/reference_parser/external_issue_parser.rb

@@ -5,7 +5,7 @@ module Banzai
     class ExternalIssueParser < BaseParser
       self.reference_type = :external_issue
 
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         issue_ids = issue_ids_per_project(nodes)
         projects = find_projects_for_hash_keys(issue_ids)
         issues = []

lib/banzai/reference_parser/issuable_parser.rb

@@ -13,7 +13,7 @@ module Banzai
         end
       end
 
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         records = records_for_nodes(nodes)
 
         nodes.map { |node| records[node] }.compact.uniq

lib/banzai/reference_parser/user_parser.rb

@@ -5,7 +5,7 @@ module Banzai
     class UserParser < BaseParser
       self.reference_type = :user
 
-      def referenced_by(nodes)
+      def referenced_by(nodes, options = {})
         group_ids = []
         user_ids = []
         project_ids = []

lib/gitlab/reference_extractor.rb

@@ -24,8 +24,8 @@ module Gitlab
       super(text, context.merge(project: project))
     end
 
-    def references(type)
-      refs = super(type, project, current_user)
+    def references(type, ids_only: false)
+      refs = super(type, project, current_user, ids_only: ids_only)
       @stateful_not_visible_counter += refs[:not_visible].count
 
       refs[:visible]
@@ -41,6 +41,12 @@ module Gitlab
       define_method(type.to_s.pluralize) do
         @references[type] ||= references(type)
       end
+
+      if %w(mentioned_user mentioned_group mentioned_project).include?(type.to_s)
+        define_method("#{type}_ids") do
+          @references[type] ||= references(type, ids_only: true)
+        end
+      end
     end
 
     def issues

spec/lib/banzai/reference_parser/base_parser_spec.rb

@@ -78,6 +78,7 @@ RSpec.describe Banzai::ReferenceParser::BaseParser do
 
   describe '#referenced_by' do
     context 'when references_relation is implemented' do
+      context 'and ids_only is set to false' do
         it 'returns a collection of objects' do
           links = Nokogiri::HTML.fragment("<a data-foo='#{user.id}'></a>")
                     .children
@@ -87,6 +88,24 @@ RSpec.describe Banzai::ReferenceParser::BaseParser do
         end
       end
 
+      context 'and ids_only is set to true' do
+        it 'returns a collection of id values without performing a db query' do
+          links = Nokogiri::HTML.fragment("<a data-foo='1'></a><a data-foo='2'></a>").children
+
+          expect(subject).not_to receive(:references_relation)
+          expect(subject.referenced_by(links, ids_only: true)).to eq(%w(1 2))
+        end
+
+        context 'and the html fragment does not contain any attributes' do
+          it 'returns an empty array' do
+            links = Nokogiri::HTML.fragment("no links").children
+
+            expect(subject.referenced_by(links, ids_only: true)).to eq([])
+          end
+        end
+      end
+    end
+
     context 'when references_relation is not implemented' do
       it 'raises NotImplementedError' do
         links = Nokogiri::HTML.fragment('<a data-foo="1"></a>').children
@@ -188,7 +207,7 @@ RSpec.describe Banzai::ReferenceParser::BaseParser do
       dummy = Class.new(described_class) do
         self.reference_type = :test
 
-        def gather_references(nodes)
+        def gather_references(nodes, ids_only: false)
           nodes
         end
       end
@@ -222,7 +241,7 @@ RSpec.describe Banzai::ReferenceParser::BaseParser do
           nodes.select { |n| n.id > 5 }
         end
 
-        def referenced_by(nodes)
+        def referenced_by(nodes, ids_only: false)
           nodes.map(&:id)
         end
       end

spec/services/issues/update_service_spec.rb

@@ -488,6 +488,21 @@ RSpec.describe Issues::UpdateService, :mailer do
           end
         end
       end
+
+      it 'verifies the number of queries' do
+        update_issue(description: "- [ ] Task 1 #{user.to_reference}")
+
+        baseline = ActiveRecord::QueryRecorder.new do
+          update_issue(description: "- [x] Task 1 #{user.to_reference}")
+        end
+
+        recorded = ActiveRecord::QueryRecorder.new do
+          update_issue(description: "- [x] Task 1 #{user.to_reference}\n- [ ] Task 2 #{user.to_reference}")
+        end
+
+        expect(recorded.count).to eq(baseline.count - 1)
+        expect(recorded.cached_count).to eq(0)
+      end
     end
 
     context 'when description changed' do