From cb58cbbc2c95166b565c72b3cef87404c2b3426a Mon Sep 17 00:00:00 2001
From: Krzysztof Kotlarek <kotlarek.krzysztof@gmail.com>
Date: Mon, 14 Sep 2020 11:58:28 +1000
Subject: [PATCH] FEATURE: allow to extend topic_eager_loads in Search (#10625)

This additional interface is required by encrypt plugin
---
 lib/plugin/instance.rb  | 11 +++++++++++
 lib/search.rb           | 14 ++++++++++++--
 spec/lib/search_spec.rb | 30 ++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/lib/plugin/instance.rb b/lib/plugin/instance.rb
index 4729223a9e5..882c6184419 100644
--- a/lib/plugin/instance.rb
+++ b/lib/plugin/instance.rb
@@ -206,6 +206,17 @@ class Plugin::Instance
     Search.advanced_filter(trigger, &block)
   end
 
+  # Allow to eager load additional tables in Search. Useful to avoid N+1 performance problems.
+  # Example usage:
+  #   register_search_topic_eager_load do |opts|
+  #     %i(example_table)
+  #   end
+  # OR
+  #   register_search_topic_eager_load(%i(example_table))
+  def register_search_topic_eager_load(tables = nil, &block)
+    Search.custom_topic_eager_load(tables, &block)
+  end
+
   # Request a new size for topic thumbnails
   # Will respect plugin enabled setting is enabled
   # Size should be an array with two elements [max_width, max_height]
diff --git a/lib/search.rb b/lib/search.rb
index caa5edca206..f2ded80c485 100644
--- a/lib/search.rb
+++ b/lib/search.rb
@@ -297,6 +297,14 @@ class Search
     @advanced_filters
   end
 
+  def self.custom_topic_eager_load(tables = nil, &block)
+    (@custom_topic_eager_loads ||= []) << (tables || block)
+  end
+
+  def self.custom_topic_eager_loads
+    Array.wrap(@custom_topic_eager_loads)
+  end
+
   advanced_filter(/^in:personal-direct$/) do |posts|
     if @guardian.user
       posts
@@ -1194,11 +1202,13 @@ class Search
       topic_eager_loads << :tags
     end
 
+    Search.custom_topic_eager_loads.each do |custom_loads|
+      topic_eager_loads.concat(custom_loads.is_a?(Array) ? custom_loads : custom_loads.call(search_pms: @search_pms).to_a)
+    end
+
     query.includes(topic: topic_eager_loads)
   end
 
-  private
-
   # Limited for performance reasons since `TS_HEADLINE` is slow when the text
   # document is too long.
   MAX_LENGTH_FOR_HEADLINE = 2500
diff --git a/spec/lib/search_spec.rb b/spec/lib/search_spec.rb
index 784b4f8c544..1aaf6edb49e 100644
--- a/spec/lib/search_spec.rb
+++ b/spec/lib/search_spec.rb
@@ -94,4 +94,34 @@ describe Search do
       end
     end
   end
+
+  context "custom_eager_load" do
+    fab!(:topic) { Fabricate(:topic) }
+    fab!(:post) { Fabricate(:post, topic: topic) }
+
+    before do
+      SearchIndexer.enable
+      SearchIndexer.index(topic, force: true)
+    end
+
+    it "includes custom tables" do
+      begin
+        expect(Search.execute("test").posts[0].topic.association(:category).loaded?).to be true
+        expect(Search.execute("test").posts[0].topic.association(:tags).loaded?).to be false
+
+        SiteSetting.tagging_enabled = true
+        Search.custom_topic_eager_load([:topic_users])
+        Search.custom_topic_eager_load() do
+          [:bookmarks]
+        end
+
+        expect(Search.execute("test").posts[0].topic.association(:tags).loaded?).to be true
+        expect(Search.execute("test").posts[0].topic.association(:topic_users).loaded?).to be true
+        expect(Search.execute("test").posts[0].topic.association(:bookmarks).loaded?).to be true
+      ensure
+        SiteSetting.tagging_enabled = false
+        Search.instance_variable_set(:@custom_topic_eager_loads, [])
+      end
+    end
+  end
 end