Merge branch 'osw-handle-diff-unfolding-in-batches' into 'master'

Handle diffs unfolding for batch load endpoint

See merge request gitlab-org/gitlab!18153

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -25,6 +25,9 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
     return render_404 unless diffable
 
     diffs = diffable.diffs_in_batch(params[:page], params[:per_page], diff_options: diff_options)
+    positions = @merge_request.note_positions_for_paths(diffs.diff_file_paths, current_user)
+
+    diffs.unfold_diff_files(positions.unfoldable)
 
     options = {
       merge_request: @merge_request,

app/models/merge_request.rb

@@ -454,6 +454,15 @@ class MergeRequest < ApplicationRecord
     merge_request_diffs.where.not(id: merge_request_diff.id)
   end
 
+  # Overwritten in EE
+  def note_positions_for_paths(paths, _user = nil)
+    positions = notes.new_diff_notes.joins(:note_diff_file)
+      .where('note_diff_files.old_path IN (?) OR note_diff_files.new_path IN (?)', paths, paths)
+      .positions
+
+    Gitlab::Diff::PositionCollection.new(positions, diff_head_sha)
+  end
+
   def preloads_discussion_diff_highlighting?
     true
   end

app/models/note.rb

@@ -193,6 +193,12 @@ class Note < ApplicationRecord
       groups
     end
 
+    def positions
+      where.not(position: nil)
+        .select(:id, :type, :position) # ActiveRecord needs id and type for typecasting.
+        .map(&:position)
+    end
+
     def count_for_collection(ids, type)
       user.select('noteable_id', 'COUNT(*) as count')
         .group(:noteable_id)

ee/app/models/draft_note.rb

@@ -28,6 +28,12 @@ class DraftNote < ApplicationRecord
 
   scope :authored_by, ->(u) { where(author_id: u.id) }
 
+  def self.positions
+    where.not(position: nil)
+      .select(:position)
+      .map(&:position)
+  end
+
   def project
     merge_request.target_project
   end

ee/app/models/ee/merge_request.rb

@@ -106,6 +106,18 @@ module EE
       hidden.count
     end
 
+    override :note_positions_for_paths
+    def note_positions_for_paths(file_paths, user = nil)
+      return super unless user
+
+      positions = draft_notes
+        .authored_by(user)
+        .positions
+        .select { |pos| file_paths.include?(pos.file_path) }
+
+      super.concat(positions)
+    end
+
     def participant_approvers
       strong_memoize(:participant_approvers) do
         next [] unless approval_needed?

ee/spec/models/merge_request_spec.rb

@@ -50,6 +50,48 @@ describe MergeRequest do
     end
   end
 
+  describe '#note_positions_for_paths' do
+    let(:user) { create(:user) }
+    let(:merge_request) { create(:merge_request, :with_diffs) }
+    let(:project) { merge_request.project }
+    let!(:diff_note) do
+      create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+    end
+    let!(:draft_note) do
+      create(:draft_note_on_text_diff, author: user, merge_request: merge_request)
+    end
+
+    let(:file_paths) { merge_request.diffs.diff_files.map(&:file_path) }
+
+    subject do
+      merge_request.note_positions_for_paths(file_paths)
+    end
+
+    it 'returns a Gitlab::Diff::PositionCollection' do
+      expect(subject).to be_a(Gitlab::Diff::PositionCollection)
+    end
+
+    context 'when user is given' do
+      subject do
+        merge_request.note_positions_for_paths(file_paths, user)
+      end
+
+      it 'returns notes and draft notes positions' do
+        expect(subject).to match_array([draft_note.position, diff_note.position])
+      end
+    end
+
+    context 'when user is not given' do
+      subject do
+        merge_request.note_positions_for_paths(file_paths)
+      end
+
+      it 'returns notes positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+  end
+
   describe '#participant_approvers' do
     let(:approvers) { create_list(:user, 2) }
     let(:code_owners) { create_list(:user, 2) }

lib/gitlab/diff/file_collection/merge_request_diff_batch.rb

@@ -29,6 +29,10 @@ module Gitlab
           }
         end
 
+        def diff_file_paths
+          diff_files.map(&:file_path)
+        end
+
         override :diffs
         def diffs
           strong_memoize(:diffs) do

lib/gitlab/diff/lines_unfolder.rb

@@ -54,7 +54,7 @@ module Gitlab
       def unfold_required?
         strong_memoize(:unfold_required) do
           next false unless @diff_file.text?
-          next false unless @position.on_text? && @position.unchanged?
+          next false unless @position.unfoldable?
           next false if @diff_file.new_file? || @diff_file.deleted_file?
           next false unless @position.old_line
           # Invalid position (MR import scenario)

lib/gitlab/diff/position.rb

@@ -79,6 +79,10 @@ module Gitlab
         formatter.line_age
       end
 
+      def unfoldable?
+        on_text? && unchanged?
+      end
+
       def unchanged?
         type.nil?
       end

lib/gitlab/diff/position_collection.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Diff
+    class PositionCollection
+      include Enumerable
+
+      # collection - An array of Gitlab::Diff::Position
+      def initialize(collection, diff_head_sha)
+        @collection = collection
+        @diff_head_sha = diff_head_sha
+      end
+
+      def each(&block)
+        @collection.each(&block)
+      end
+
+      def concat(positions)
+        tap { @collection.concat(positions) }
+      end
+
+      # Doing a lightweight filter in-memory given we're not prepared for querying
+      # positions (https://gitlab.com/gitlab-org/gitlab/issues/33271).
+      def unfoldable
+        select do |position|
+          position.unfoldable? && position.head_sha == @diff_head_sha
+        end
+      end
+    end
+  end
+end

spec/controllers/projects/merge_requests/diffs_controller_spec.rb

@@ -258,5 +258,26 @@ describe Projects::MergeRequests::DiffsController do
 
     it_behaves_like 'forked project with submodules'
     it_behaves_like 'persisted preferred diff view cookie'
+
+    context 'diff unfolding' do
+      let!(:unfoldable_diff_note) do
+        create(:diff_note_on_merge_request, :folded_position, project: project, noteable: merge_request)
+      end
+
+      let!(:diff_note) do
+        create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+      end
+
+      it 'unfolds correct diff file positions' do
+        expect_next_instance_of(Gitlab::Diff::FileCollection::MergeRequestDiffBatch) do |instance|
+          expect(instance)
+            .to receive(:unfold_diff_files)
+            .with([unfoldable_diff_note.position])
+            .and_call_original
+        end
+
+        go
+      end
+    end
   end
 end

spec/factories/notes.rb

@@ -62,6 +62,18 @@ FactoryBot.define do
         )
       end
 
+      trait :folded_position do
+        position do
+          Gitlab::Diff::Position.new(
+            old_path: "files/ruby/popen.rb",
+            new_path: "files/ruby/popen.rb",
+            old_line: 1,
+            new_line: 1,
+            diff_refs: diff_refs
+          )
+        end
+      end
+
       trait :resolved do
         resolved_at { Time.now }
         resolved_by { create(:user) }

spec/lib/gitlab/diff/position_collection_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Diff::PositionCollection do
+  let(:merge_request) { build(:merge_request) }
+
+  def build_text_position(attrs = {})
+    attributes = {
+      old_path: "files/ruby/popen.rb",
+      new_path: "files/ruby/popen.rb",
+      old_line: nil,
+      new_line: 14,
+      diff_refs: merge_request.diff_refs
+    }.merge(attrs)
+
+    Gitlab::Diff::Position.new(attributes)
+  end
+
+  def build_image_position(attrs = {})
+    attributes = {
+      old_path: "files/images/any_image.png",
+      new_path: "files/images/any_image.png",
+      width: 10,
+      height: 10,
+      x: 1,
+      y: 1,
+      diff_refs: merge_request.diff_refs,
+      position_type: "image"
+    }.merge(attrs)
+
+    Gitlab::Diff::Position.new(attributes)
+  end
+
+  let(:text_position) { build_text_position }
+  let(:folded_text_position) { build_text_position(old_line: 1, new_line: 1) }
+  let(:image_position) { build_image_position }
+  let(:head_sha) { merge_request.diff_head_sha }
+
+  let(:collection) do
+    described_class.new([text_position, folded_text_position, image_position], head_sha)
+  end
+
+  describe '#to_a' do
+    it 'returns all positions' do
+      expect(collection.to_a).to eq([text_position, folded_text_position, image_position])
+    end
+  end
+
+  describe '#unfoldable' do
+    it 'returns unfoldable diff positions' do
+      expect(collection.unfoldable).to eq([folded_text_position])
+    end
+
+    context 'when given head_sha does not match with positions head_sha' do
+      let(:head_sha) { 'unknown' }
+
+      it 'returns no position' do
+        expect(collection.unfoldable).to be_empty
+      end
+    end
+  end
+
+  describe '#concat' do
+    let(:new_text_position) { build_text_position(old_line: 1, new_line: 1) }
+
+    it 'returns a Gitlab::Diff::Position' do
+      expect(collection.concat([new_text_position])).to be_a(described_class)
+    end
+
+    it 'concatenates the new position to the collection' do
+      collection.concat([new_text_position])
+
+      expect(collection.to_a).to eq([text_position, folded_text_position, image_position, new_text_position])
+    end
+  end
+end

spec/models/merge_request_spec.rb

@@ -650,6 +650,46 @@ describe MergeRequest do
     end
   end
 
+  describe '#note_positions_for_paths' do
+    let(:merge_request) { create(:merge_request, :with_diffs) }
+    let(:project) { merge_request.project }
+    let!(:diff_note) do
+      create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+    end
+
+    let(:file_paths) { merge_request.diffs.diff_files.map(&:file_path) }
+
+    subject do
+      merge_request.note_positions_for_paths(file_paths)
+    end
+
+    it 'returns a Gitlab::Diff::PositionCollection' do
+      expect(subject).to be_a(Gitlab::Diff::PositionCollection)
+    end
+
+    context 'within all diff files' do
+      it 'returns correct positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+
+    context 'within specific diff file' do
+      let(:file_paths) { [diff_note.position.file_path] }
+
+      it 'returns correct positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+
+    context 'within no diff files' do
+      let(:file_paths) { [] }
+
+      it 'returns no positions' do
+        expect(subject.to_a).to be_empty
+      end
+    end
+  end
+
   describe '#discussions_diffs' do
     let(:merge_request) { create(:merge_request) }