1 files changed, 88 insertions, 16 deletions

diff --git a/lib/gitlab/github_import/representation/diff_note.rb b/lib/gitlab/github_import/representation/diff_note.rb
index a3dcd2e380c..fecff0644c2 100644
--- a/lib/gitlab/github_import/representation/diff_note.rb
+++ b/lib/gitlab/github_import/representation/diff_note.rb
@@ -7,13 +7,14 @@ module Gitlab
         include ToHash
         include ExposeAttribute
 
-        attr_reader :attributes
-
-        expose_attribute :noteable_type, :noteable_id, :commit_id, :file_path,
-                         :diff_hunk, :author, :note, :created_at, :updated_at,
-                         :original_commit_id, :note_id
-
+        NOTEABLE_TYPE = 'MergeRequest'
         NOTEABLE_ID_REGEX = %r{/pull/(?<iid>\d+)}i.freeze
+        DISCUSSION_CACHE_KEY = 'github-importer/discussion-id-map/%{project_id}/%{noteable_id}/%{original_note_id}'
+
+        expose_attribute :noteable_id, :commit_id, :file_path,
+          :diff_hunk, :author, :created_at, :updated_at,
+          :original_commit_id, :note_id, :end_line, :start_line,
+          :side, :in_reply_to_id
 
         # Builds a diff note from a GitHub API response.
         #
@@ -30,7 +31,6 @@ module Gitlab
 
           user = Representation::User.from_api_response(note.user) if note.user
           hash = {
-            noteable_type: 'MergeRequest',
             noteable_id: matches[:iid].to_i,
             file_path: note.path,
             commit_id: note.commit_id,
@@ -42,7 +42,9 @@ module Gitlab
             updated_at: note.updated_at,
             note_id: note.id,
             end_line: note.line,
-            start_line: note.start_line
+            start_line: note.start_line,
+            side: note.side,
+            in_reply_to_id: note.in_reply_to_id
           }
 
           new(hash)
@@ -56,21 +58,41 @@ module Gitlab
           new(hash)
         end
 
+        attr_accessor :merge_request, :project
+
         # attributes - A Hash containing the raw note details. The keys of this
         #              Hash must be Symbols.
         def initialize(attributes)
           @attributes = attributes
+
+          @note_formatter = DiffNotes::SuggestionFormatter.new(
+            note: attributes[:note],
+            start_line: attributes[:start_line],
+            end_line: attributes[:end_line]
+          )
+        end
+
+        def noteable_type
+          NOTEABLE_TYPE
+        end
+
+        def contains_suggestion?
+          @note_formatter.contains_suggestion?
+        end
+
+        def note
+          @note_formatter.formatted_note
         end
 
         def line_code
           diff_line = Gitlab::Diff::Parser.new.parse(diff_hunk.lines).to_a.last
 
-          Gitlab::Git
-            .diff_line_code(file_path, diff_line.new_pos, diff_line.old_pos)
+          Gitlab::Git.diff_line_code(file_path, diff_line.new_pos, diff_line.old_pos)
         end
 
         # Returns a Hash that can be used to populate `notes.st_diff`, removing
         # the need for requesting Git data for every diff note.
+        # Used when importing with LegacyDiffNote
         def diff_hash
           {
             diff: diff_hunk,
@@ -85,12 +107,15 @@ module Gitlab
           }
         end
 
-        def note
-          @note ||= DiffNotes::SuggestionFormatter.formatted_note_for(
-            note: attributes[:note],
-            start_line: attributes[:start_line],
-            end_line: attributes[:end_line]
-          )
+        # Used when importing with DiffNote
+        def diff_position
+          position_params = {
+            diff_refs: merge_request.diff_refs,
+            old_path: file_path,
+            new_path: file_path
+          }
+
+          Gitlab::Diff::Position.new(position_params.merge(diff_line_params))
         end
 
         def github_identifiers
@@ -100,6 +125,53 @@ module Gitlab
             noteable_type: noteable_type
           }
         end
+
+        def discussion_id
+          if in_reply_to_id.present?
+            current_discussion_id
+          else
+            Discussion.discussion_id(
+              Struct
+              .new(:noteable_id, :noteable_type)
+              .new(merge_request.id, NOTEABLE_TYPE)
+            ).tap do |discussion_id|
+              cache_discussion_id(discussion_id)
+            end
+          end
+        end
+
+        private
+
+        # Required by ExposeAttribute
+        attr_reader :attributes
+
+        def diff_line_params
+          if addition?
+            { new_line: end_line, old_line: nil }
+          else
+            { new_line: nil, old_line: end_line }
+          end
+        end
+
+        def addition?
+          side == 'RIGHT'
+        end
+
+        def cache_discussion_id(discussion_id)
+          Gitlab::Cache::Import::Caching.write(discussion_id_cache_key(note_id), discussion_id)
+        end
+
+        def current_discussion_id
+          Gitlab::Cache::Import::Caching.read(discussion_id_cache_key(in_reply_to_id))
+        end
+
+        def discussion_id_cache_key(id)
+          DISCUSSION_CACHE_KEY % {
+            project_id: project.id,
+            noteable_id: merge_request.id,
+            original_note_id: id
+          }
+        end
       end
     end
   end