Re-use queries in reference parsersreuse-queries-in-reference-parsers

This caches various queries to ensure that multiple reference extraction
runs re-use any objects queried in previous runs.

4 files changed, 113 insertions, 4 deletions

diff --git a/CHANGELOG b/CHANGELOG
index b44ca54a39c..25274fa8a6a 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -20,6 +20,7 @@ v 8.10.0 (unreleased)
   - Add Spring EmojiOne updates.
   - Fix viewing notification settings when a project is pending deletion
   - Fix pagination when sorting by columns with lots of ties (like priority)
+  - The Markdown reference parsers now re-use query results to prevent running the same queries multiple times !5020
   - Updated project header design
   - Exclude email check from the standard health check
   - Updated layout for Projects, Groups, Users on Admin area !4424
diff --git a/lib/banzai/reference_parser/base_parser.rb b/lib/banzai/reference_parser/base_parser.rb
index 3d7b9c4a024..6cf218aaa0d 100644
--- a/lib/banzai/reference_parser/base_parser.rb
+++ b/lib/banzai/reference_parser/base_parser.rb
@@ -133,8 +133,9 @@ module Banzai
         return {} if nodes.empty?
 
         ids = unique_attribute_values(nodes, attribute)
+        rows = collection_objects_for_ids(collection, ids)
 
-        collection.where(id: ids).each_with_object({}) do |row, hash|
+        rows.each_with_object({}) do |row, hash|
           hash[row.id] = row
         end
       end
@@ -153,6 +154,31 @@ module Banzai
         values.to_a
       end
 
+      # Queries the collection for the objects with the given IDs.
+      #
+      # If the RequestStore module is enabled this method will only query any
+      # objects that have not yet been queried. For objects that have already
+      # been queried the object is returned from the cache.
+      def collection_objects_for_ids(collection, ids)
+        if RequestStore.active?
+          cache = collection_cache[collection_cache_key(collection)]
+          to_query = ids.map(&:to_i) - cache.keys
+
+          unless to_query.empty?
+            collection.where(id: to_query).each { |row| cache[row.id] = row }
+          end
+
+          cache.values
+        else
+          collection.where(id: ids)
+        end
+      end
+
+      # Returns the cache key to use for a collection.
+      def collection_cache_key(collection)
+        collection.respond_to?(:model) ? collection.model : collection
+      end
+
       # Processes the list of HTML documents and returns an Array containing all
       # the references.
       def process(documents)
@@ -189,7 +215,7 @@ module Banzai
       end
 
       def find_projects_for_hash_keys(hash)
-        Project.where(id: hash.keys)
+        collection_objects_for_ids(Project, hash.keys)
       end
 
       private
@@ -199,6 +225,12 @@ module Banzai
       def lazy(&block)
         Gitlab::Lazy.new(&block)
       end
+
+      def collection_cache
+        RequestStore[:banzai_collection_cache] ||= Hash.new do |hash, key|
+          hash[key] = {}
+        end
+      end
     end
   end
 end
diff --git a/lib/banzai/reference_parser/user_parser.rb b/lib/banzai/reference_parser/user_parser.rb
index a12b0d19560..863f5725d3b 100644
--- a/lib/banzai/reference_parser/user_parser.rb
+++ b/lib/banzai/reference_parser/user_parser.rb
@@ -73,7 +73,7 @@ module Banzai
       def find_users(ids)
         return [] if ids.empty?
 
-        User.where(id: ids).to_a
+        collection_objects_for_ids(User, ids)
       end
 
       def find_users_for_groups(ids)
@@ -85,7 +85,8 @@ module Banzai
       def find_users_for_projects(ids)
         return [] if ids.empty?
 
-        Project.where(id: ids).flat_map { |p| p.team.members.to_a }
+        collection_objects_for_ids(Project, ids).
+          flat_map { |p| p.team.members.to_a }
       end
     end
   end
diff --git a/spec/lib/banzai/reference_parser/base_parser_spec.rb b/spec/lib/banzai/reference_parser/base_parser_spec.rb
index 543b4786d84..ac9c66e2663 100644
--- a/spec/lib/banzai/reference_parser/base_parser_spec.rb
+++ b/spec/lib/banzai/reference_parser/base_parser_spec.rb
@@ -234,4 +234,79 @@ describe Banzai::ReferenceParser::BaseParser, lib: true do
         to eq([project])
     end
   end
+
+  describe '#collection_objects_for_ids' do
+    context 'with RequestStore disabled' do
+      it 'queries the collection directly' do
+        collection = User.all
+
+        expect(collection).to receive(:where).twice.and_call_original
+
+        2.times do
+          expect(subject.collection_objects_for_ids(collection, [user.id])).
+            to eq([user])
+        end
+      end
+    end
+
+    context 'with RequestStore enabled' do
+      before do
+        cache = Hash.new { |hash, key| hash[key] = {} }
+
+        allow(RequestStore).to receive(:active?).and_return(true)
+        allow(subject).to receive(:collection_cache).and_return(cache)
+      end
+
+      it 'queries the collection on the first call' do
+        expect(subject.collection_objects_for_ids(User, [user.id])).
+          to eq([user])
+      end
+
+      it 'does not query previously queried objects' do
+        collection = User.all
+
+        expect(collection).to receive(:where).once.and_call_original
+
+        2.times do
+          expect(subject.collection_objects_for_ids(collection, [user.id])).
+            to eq([user])
+        end
+      end
+
+      it 'casts String based IDs to Fixnums before querying objects' do
+        2.times do
+          expect(subject.collection_objects_for_ids(User, [user.id.to_s])).
+            to eq([user])
+        end
+      end
+
+      it 'queries any additional objects after the first call' do
+        other_user = create(:user)
+
+        expect(subject.collection_objects_for_ids(User, [user.id])).
+          to eq([user])
+
+        expect(subject.collection_objects_for_ids(User, [user.id, other_user.id])).
+          to eq([user, other_user])
+      end
+
+      it 'caches objects on a per collection class basis' do
+        expect(subject.collection_objects_for_ids(User, [user.id])).
+          to eq([user])
+
+        expect(subject.collection_objects_for_ids(Project, [project.id])).
+          to eq([project])
+      end
+    end
+  end
+
+  describe '#collection_cache_key' do
+    it 'returns the cache key for a Class' do
+      expect(subject.collection_cache_key(Project)).to eq(Project)
+    end
+
+    it 'returns the cache key for an ActiveRecord::Relation' do
+      expect(subject.collection_cache_key(Project.all)).to eq(Project)
+    end
+  end
 end