Make ES results more efficient with eager loading

ee/lib/gitlab/elastic/search_results.rb

@@ -25,13 +25,13 @@ module Gitlab
       def objects(scope, page = nil)
         case scope
         when 'projects'
-          projects.page(page).per(per_page).records
+          eager_load(projects, page, eager: [:route, :namespace])
         when 'issues'
-          issues.page(page).per(per_page).records
+          eager_load(issues, page, eager: { project: [:route, :namespace] })
         when 'merge_requests'
-          merge_requests.page(page).per(per_page).records
+          eager_load(merge_requests, page, eager: { target_project: [:route, :namespace] })
         when 'milestones'
-          milestones.page(page).per(per_page).records
+          eager_load(milestones, page, eager: { project: [:route, :namespace] })
         when 'blobs'
           blobs.page(page).per(per_page)
         when 'wiki_blobs'
@@ -45,6 +45,21 @@ module Gitlab
         end
       end
 
+      # Apply some eager loading to the `records` of an ES result object without
+      # losing pagination information
+      def eager_load(es_result, page, eager:)
+        page ||= 1
+        paginated_base = es_result.page(page).per(per_page)
+        relation = paginated_base.records.includes(eager) # rubocop:disable CodeReuse/ActiveRecord
+
+        Kaminari.paginate_array(
+          relation,
+          total_count: paginated_base.total_count,
+          limit: per_page,
+          offset: per_page * (page - 1)
+        )
+      end
+
       def generic_search_results
         @generic_search_results ||= Gitlab::SearchResults.new(current_user, limit_projects, query)
       end

ee/spec/features/search/elastic/global_search_spec.rb

@@ -15,7 +15,7 @@ describe 'Global elastic search', :elastic do
     stub_ee_application_setting(elasticsearch_search: false, elasticsearch_indexing: false)
   end
 
-  shared_examples 'a pure Elasticsearch result' do
+  shared_examples 'an efficient database result' do
     it 'avoids N+1 database queries' do
       create(object, creation_args)
       Gitlab::Elastic::Helper.refresh_index
@@ -38,7 +38,7 @@ describe 'Global elastic search', :elastic do
       let(:path) { search_path(search: 'initial', scope: 'issues') }
       let(:query_count_multiplier) { 0 }
 
-      it_behaves_like 'a pure Elasticsearch result'
+      it_behaves_like 'an efficient database result'
     end
 
     context 'searching projects' do
@@ -48,25 +48,16 @@ describe 'Global elastic search', :elastic do
       # Each Project requires 4 extra queries: one for each "count" (forks, open MRs, open Issues) and one for access level
       let(:query_count_multiplier) { 4 }
 
-      it_behaves_like 'a pure Elasticsearch result'
+      it_behaves_like 'an efficient database result'
     end
 
     context 'searching merge requests' do
-      it 'avoids N+1 database queries' do
-        path = search_path(search: 'initial', scope: 'merge_requests')
-
-        create(:merge_request, title: 'initial', source_project: project)
-        Gitlab::Elastic::Helper.refresh_index
-
-        control_count = ActiveRecord::QueryRecorder.new { visit path }.count
-
-        merge_requests = create_list(:merge_request, 10, title: 'initial')
-        merge_requests.each { |mr| mr.target_project.add_maintainer(user) }
-        Gitlab::Elastic::Helper.refresh_index
+      let(:object) { :merge_request }
+      let(:creation_args) { { title: 'initial' } }
+      let(:path) { search_path(search: 'initial*', scope: 'merge_requests') }
+      let(:query_count_multiplier) { 0 }
 
-        # Each MR loaded has a Route load that has been tricky to track down
-        expect { visit path }.not_to exceed_query_limit(control_count + 11)
-      end
+      it_behaves_like 'an efficient database result'
     end
 
     context 'searching milestones' do
@@ -75,7 +66,7 @@ describe 'Global elastic search', :elastic do
       let(:path) { search_path(search: 'milestone*', scope: 'milestones') }
       let(:query_count_multiplier) { 0 }
 
-      it_behaves_like 'a pure Elasticsearch result'
+      it_behaves_like 'an efficient database result'
     end
   end