Add batch loaded associations in API::Entities::MergeRequestBasic

With this change, whenever the Grape::Entity serializer is sent a list
of entities, all the associations that are batch loaded will be loaded
in a single SQL query.

app/models/label_link.rb

@@ -9,4 +9,9 @@ class LabelLink < ApplicationRecord
 
   validates :target, presence: true, unless: :importing?
   validates :label, presence: true, unless: :importing?
+
+  scope :preloaded, -> { preload(:label) }
+
+  scope :for_targets, ->(type:, scope:) { where(target_type: type, target_id: scope) }
+  scope :for_merge_requests, ->(merge_request_scope) { for_targets(type: MergeRequest, scope: merge_request_scope) }
 end

app/models/merge_request.rb

@@ -254,16 +254,11 @@ class MergeRequest < ApplicationRecord
   scope :join_project, -> { joins(:target_project) }
   scope :references_project, -> { references(:target_project) }
 
-  PROJECT_ROUTE_AND_NAMESPACE_ROUTE = [
+  scope :with_api_entity_associations, -> {
+    preload(
       target_project: [:route, { namespace: :route }],
       source_project: [:route, { namespace: :route }]
-  ].freeze
-
-  scope :with_api_entity_associations, -> {
-    preload(:assignees, :author, :unresolved_notes, :labels, :milestone,
-            :timelogs, :latest_merge_request_diff,
-            *PROJECT_ROUTE_AND_NAMESPACE_ROUTE,
-            metrics: [:latest_closed_by, :merged_by])
+    )
   }
   scope :by_target_branch_wildcard, ->(wildcard_branch_name) do
     where("target_branch LIKE ?", ApplicationRecord.sanitize_sql_like(wildcard_branch_name).tr('*', '%'))

app/models/merge_request/metrics.rb

@@ -5,6 +5,11 @@ class MergeRequest::Metrics < ApplicationRecord
   belongs_to :pipeline, class_name: 'Ci::Pipeline', foreign_key: :pipeline_id
   belongs_to :latest_closed_by, class_name: 'User'
   belongs_to :merged_by, class_name: 'User'
+
+  scope :preloaded, -> { preload(:merged_by, :latest_closed_by) }
+  scope :for_merge_requests, -> (merge_request_scope) do
+    where(merge_request: merge_request_scope)
+  end
 end
 
 MergeRequest::Metrics.prepend_if_ee('EE::MergeRequest::Metrics')

app/models/merge_request_assignee.rb

@@ -6,5 +6,11 @@ class MergeRequestAssignee < ApplicationRecord
 
   validates :assignee, uniqueness: { scope: :merge_request_id }
 
+  scope :preloaded, -> { preload(:assignee) }
+
   scope :in_projects, ->(project_ids) { joins(:merge_request).where("merge_requests.target_project_id in (?)", project_ids) }
+
+  scope :for_merge_requests, ->(merge_request_scope) do
+    where(merge_request: merge_request_scope)
+  end
 end

app/models/merge_request_diff.rb

@@ -100,6 +100,10 @@ class MergeRequestDiff < ApplicationRecord
     joins(merge_request: :metrics).where(condition)
   end
 
+  scope :for_merge_requests, -> (merge_request_scope) do
+    where(merge_request: merge_request_scope)
+  end
+
   def self.ids_for_external_storage_migration(limit:)
     # No point doing any work unless the feature is enabled
     return [] unless Gitlab.config.external_diffs.enabled

app/models/timelog.rb

@@ -20,6 +20,10 @@ class Timelog < ApplicationRecord
     where('spent_at BETWEEN ? AND ?', start_time, end_time)
   end
 
+  scope :for_merge_requests, -> (merge_request_scope) do
+    where(merge_request: merge_request_scope)
+  end
+
   def issuable
     issue || merge_request
   end

ee/lib/ee/api/releases.rb

@@ -9,7 +9,7 @@ module EE
         resource :projects, requirements: ::API::API::NAMESPACE_OR_PROJECT_REQUIREMENTS do
           desc 'Create Evidence for a Release' do
             detail 'This feature was introduced in GitLab 12.10.'
-            success Entities::Release
+            success ::API::Entities::Release
           end
           params do
             requires :tag_name, type: String, desc: 'The name of the tag', as: :tag

lib/api/entities/issuable_entity.rb

@@ -8,14 +8,26 @@ module API
       expose :title, :description
       expose :state, :created_at, :updated_at
 
-      # Avoids an N+1 query when metadata is included
-      def issuable_metadata(subject, options, method, args = nil)
-        cached_subject = options.dig(:issuable_metadata, subject.id)
+      def presented
+        lazy_issuable_metadata
 
-        if cached_subject
-          cached_subject[method]
-        else
-          subject.public_send(method, *args) # rubocop: disable GitlabSecurity/PublicSend
+        super
+      end
+
+      def issuable_metadata
+        lazy_issuable_metadata
+      end
+
+      protected
+
+      def lazy_issuable_metadata
+        BatchLoader.for(object).batch(key: :issuable_metadata) do |models, loader|
+          issuable_metadata = Gitlab::IssuableMetadata.new(nil, models)
+          metadata_by_id = issuable_metadata.data
+
+          models.each do |issuable|
+            loader.call(issuable, metadata_by_id[issuable.id])
+          end
         end
       end
     end

lib/api/entities/issuable_time_stats.rb

@@ -7,6 +7,12 @@ module API
         Gitlab::TimeTrackingFormatter.output(time_spent)
       end
 
+      def presented
+        lazy_timelogs
+
+        super
+      end
+
       expose :time_estimate
       expose :total_time_spent
       expose :human_time_estimate
@@ -15,12 +21,19 @@ module API
         expose :total_time_spent, as: :human_total_time_spent
       end
 
-      # rubocop: disable CodeReuse/ActiveRecord
+      private
+
+      def lazy_timelogs
+        BatchLoader.for(object.id).batch(key: :timelogs, default_value: []) do |ids, loader|
+          Timelog.for_merge_requests(ids).find_each do |timelog|
+            loader.call(timelog.merge_request_id) { |acc| acc << timelog }
+          end
+        end
+      end
+
       def total_time_spent
-        # Avoids an N+1 query since timelogs are preloaded
-        object.timelogs.map(&:time_spent).sum
+        lazy_timelogs.sum(&:time_spend) # rubocop:disable CodeReuse/ActiveRecord
       end
-      # rubocop: enable CodeReuse/ActiveRecord
     end
   end
 end

lib/api/entities/merge_request_basic.rb

@@ -3,17 +3,35 @@
 module API
   module Entities
     class MergeRequestBasic < IssuableEntity
+      # Evaluate the lazy exposures to trigger the BatchLoader
+      # before any object is serialized.
+      def presented
+        lazy_merge_request_metrics
+        lazy_diff
+        lazy_assignees
+        lazy_author
+        lazy_milestone
+        lazy_labels
+
+        # TODO: we could have a `:batch` exposure option to automatically scan
+        # exposures and evaluate the block so that the BatchLoader is primed
+        time_stats = self.class.find_exposure(:time_stats)
+        Object.const_get(time_stats.using_class_name, false).new(object).presented
+
+        super
+      end
+
       expose :merged_by, using: Entities::UserBasic do |merge_request, _options|
-        merge_request.metrics&.merged_by
+        lazy_merge_request_metrics&.merged_by
       end
       expose :merged_at do |merge_request, _options|
-        merge_request.metrics&.merged_at
+        lazy_merge_request_metrics&.merged_at
       end
       expose :closed_by, using: Entities::UserBasic do |merge_request, _options|
-        merge_request.metrics&.latest_closed_by
+        lazy_merge_request_metrics&.latest_closed_by
       end
       expose :closed_at do |merge_request, _options|
-        merge_request.metrics&.latest_closed_at
+        lazy_merge_request_metrics&.latest_closed_at
       end
       expose :title_html, if: -> (_, options) { options[:render_html] } do |entity|
         MarkupHelper.markdown_field(entity, :title)
@@ -22,23 +40,30 @@ module API
         MarkupHelper.markdown_field(entity, :description)
       end
       expose :target_branch, :source_branch
-      expose(:user_notes_count) { |merge_request, options| issuable_metadata(merge_request, options, :user_notes_count) }
-      expose(:upvotes)          { |merge_request, options| issuable_metadata(merge_request, options, :upvotes) }
-      expose(:downvotes)        { |merge_request, options| issuable_metadata(merge_request, options, :downvotes) }
-      expose :assignee, using: ::API::Entities::UserBasic do |merge_request|
-        merge_request.assignee
+      expose(:user_notes_count) { |merge_request, options| issuable_metadata.user_notes_count }
+      expose(:upvotes)          { |merge_request, options| issuable_metadata.upvotes }
+      expose(:downvotes)        { |merge_request, options| issuable_metadata.downvotes }
+
+      with_options using: Entities::UserBasic do
+        expose :lazy_author, as: :author
+        expose :lazy_assignees, as: :assignees
+        expose :lazy_assignee, as: :assignee do |merge_request, options|
+          lazy_assignees.first
         end
-      expose :author, :assignees, using: Entities::UserBasic
+      end
+
       expose :source_project_id, :target_project_id
       expose :labels do |merge_request, options|
+        lazy_labels do |label|
           if options[:with_labels_details]
-          ::API::Entities::LabelBasic.represent(merge_request.labels.sort_by(&:title))
+            Entities::LabelBasic.new(label)
           else
-          merge_request.labels.map(&:title).sort
+            label.title
+          end
         end
       end
       expose :work_in_progress?, as: :work_in_progress
-      expose :milestone, using: Entities::Milestone
+      expose :lazy_milestone, as: :milestone, using: Entities::Milestone
       expose :merge_when_pipeline_succeeds
 
       # Ideally we should deprecate `MergeRequest#merge_status` exposure and
@@ -51,15 +76,20 @@ module API
         merge_request.check_mergeability(async: true) unless options[:skip_merge_status_recheck]
         merge_request.public_merge_status
       end
-      expose :diff_head_sha, as: :sha
+      expose :diff_head_sha, as: :sha do |_, options|
+        lazy_diff.read_attribute(:head_commit_sha)
+      end
       expose :merge_commit_sha
       expose :squash_commit_sha
       expose :discussion_locked
       expose :should_remove_source_branch?, as: :should_remove_source_branch
       expose :force_remove_source_branch?, as: :force_remove_source_branch
-      expose :allow_collaboration, if: -> (merge_request, _) { merge_request.for_fork? }
+
+      with_options if: -> (merge_request, _) { merge_request.for_fork? } do
+        expose :allow_collaboration
         # Deprecated
-      expose :allow_collaboration, as: :allow_maintainer_to_push, if: -> (merge_request, _) { merge_request.for_fork? }
+        expose :allow_collaboration, as: :allow_maintainer_to_push
+      end
 
       # reference is deprecated in favour of references
       # Introduced [Gitlab 12.6](https://gitlab.com/gitlab-org/gitlab/merge_requests/20354)
@@ -83,6 +113,72 @@ module API
       expose :task_completion_status
       expose :cannot_be_merged?, as: :has_conflicts
       expose :mergeable_discussions_state?, as: :blocking_discussions_resolved
+
+      private
+
+      def lazy_merge_request_metrics
+        BatchLoader.for(object.id).batch(key: :merge_request_metrics) do |models, loader|
+          ::MergeRequest::Metrics
+            .preloaded
+            .for_merge_requests(models)
+            .find_each do |metric|
+            loader.call(metric.merge_request_id, metric)
+          end
+        end
+      end
+
+      def lazy_diff
+        BatchLoader.for(object.id).batch(key: :merge_request_diff) do |ids, loader|
+          ::MergeRequestDiff
+            .for_merge_requests(ids)
+            .find_each do |diff|
+            loader.call(diff.merge_request_id, diff)
+          end
+        end
+      end
+
+      def lazy_assignees
+        BatchLoader.for(object.id).batch(key: :assignees, default_value: []) do |ids, loader|
+          ::MergeRequestAssignee
+            .preloaded
+            .for_merge_requests(ids)
+            .find_each do |assignment|
+            loader.call(assignment.merge_request_id) { |acc| acc << assignment.assignee }
+          end
+        end
+      end
+
+      def lazy_author
+        BatchLoader.for(object.author_id).batch(key: :author) do |ids, loader|
+          ::User.id_in(ids).find_each do |author|
+            loader.call(author.id, author)
+          end
+        end
+      end
+
+      def lazy_milestone
+        BatchLoader.for(object.milestone_id).batch(key: :milestone) do |ids, loader|
+          ::Milestone
+            .with_api_entity_associations
+            .id_in(ids)
+            .find_each do |milestone|
+            loader.call(milestone.id, milestone)
+          end
+        end
+      end
+
+      def lazy_labels(&block)
+        BatchLoader.for(object.id).batch(key: :labels, default_value: []) do |ids, loader|
+          ::LabelLink
+            .preloaded
+            .for_merge_requests(ids)
+            .find_each do |link|
+            loader.call(link.target_id) do |memo|
+              memo << yield(link.label)
+            end
+          end
+        end
+      end
     end
   end
 end

spec/lib/api/entities/merge_request_basic_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::API::Entities::MergeRequestBasic do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project) { create(:project, :public) }
+  let_it_be(:merge_request) { create(:merge_request) }
+  let_it_be(:labels) { create_list(:label, 3) }
+  let_it_be(:merge_requests) { create_list(:labeled_merge_request, 10, :unique_branches, :with_diffs, labels: labels) }
+
+  let(:scope) { MergeRequest.with_api_entity_associations }
+
+  # This mimics the behavior of the `Grape::Entity` serializer
+  def present(obj)
+    described_class.new(obj).presented
+  end
+
+  describe "#with_api_entity_associations" do
+    it "avoids N+1 queries" do
+      query = scope.find(merge_request.id)
+
+      control = ActiveRecord::QueryRecorder.new do
+        present(query).to_json
+      end
+
+      query = scope.all
+      batch = ActiveRecord::QueryRecorder.new do
+        entities = query.map(&method(:present))
+
+        entities.to_json
+      end
+
+      # The current threshold is 3 query per entity maximum.
+      expect(batch.count).to be_within(3 * query.count).of(control.count)
+    end
+  end
+end