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.

2 files changed, 37 insertions, 4 deletions

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