Import pull requests from GitHub where the source branch was removed

2 files changed, 47 insertions, 27 deletions

diff --git a/lib/gitlab/github_import/importer.rb b/lib/gitlab/github_import/importer.rb
index 0f9e3ee14ee..a3f27891784 100644
--- a/lib/gitlab/github_import/importer.rb
+++ b/lib/gitlab/github_import/importer.rb
@@ -3,12 +3,15 @@ module Gitlab
     class Importer
       include Gitlab::ShellAdapter
 
-      attr_reader :project, :client
+      attr_reader :client, :project, :repo, :repo_url
 
       def initialize(project)
         @project = project
-        if import_data_credentials
-          @client = Client.new(import_data_credentials[:user])
+        @repo = project.import_source
+        @repo_url = project.import_url
+
+        if credentials
+          @client = Client.new(credentials[:user])
           @formatter = Gitlab::ImportFormatter.new
         else
           raise Projects::ImportService::Error, "Unable to find project import data credentials for project ID: #{@project.id}"
@@ -22,8 +25,8 @@ module Gitlab
 
       private
 
-      def import_data_credentials
-        @import_data_credentials ||= project.import_data.credentials if project.import_data
+      def credentials
+        @credentials ||= project.import_data.credentials if project.import_data
       end
 
       def import_labels
@@ -68,22 +71,31 @@ module Gitlab
       end
 
       def import_pull_requests
-        client.pull_requests(project.import_source, state: :all,
-                                                    sort: :created,
-                                                    direction: :asc).each do |raw_data|
-          pull_request = PullRequestFormatter.new(project, raw_data)
-
-          if pull_request.valid?
-            merge_request = MergeRequest.new(pull_request.attributes)
-
-            if merge_request.save
-              apply_labels(pull_request.number, merge_request)
-              import_comments(pull_request.number, merge_request)
-              import_comments_on_diff(pull_request.number, merge_request)
-            end
+        pull_requests = client.pull_requests(repo, state: :all, sort: :created, direction: :asc)
+                              .map { |raw| PullRequestFormatter.new(project, raw) }
+                              .reject(&:cross_project?)
+
+        source_branches_removed = pull_requests.reject(&:source_branch_exists?)
+        source_branches_removed.each do |pull_request|
+          client.create_ref(repo, "refs/heads/#{pull_request.source_branch}", pull_request.source_sha)
+        end
+
+        project.repository.fetch_ref(repo_url, '+refs/heads/*', 'refs/heads/*')
+
+        pull_requests.each do |pull_request|
+          merge_request = MergeRequest.new(pull_request.attributes)
+
+          if merge_request.save
+            apply_labels(pull_request.number, merge_request)
+            import_comments(pull_request.number, merge_request)
+            import_comments_on_diff(pull_request.number, merge_request)
           end
         end
 
+        source_branches_removed.each do |pull_request|
+          client.delete_ref(repo, "heads/#{pull_request.source_branch}")
+        end
+
         true
       rescue ActiveRecord::RecordInvalid => e
         raise Projects::ImportService::Error, e.message
diff --git a/lib/gitlab/github_import/pull_request_formatter.rb b/lib/gitlab/github_import/pull_request_formatter.rb
index d21b942ad4b..be970f724c6 100644
--- a/lib/gitlab/github_import/pull_request_formatter.rb
+++ b/lib/gitlab/github_import/pull_request_formatter.rb
@@ -7,7 +7,7 @@ module Gitlab
           title: raw_data.title,
           description: description,
           source_project: source_project,
-          source_branch: source_branch.name,
+          source_branch: source_branch,
           target_project: target_project,
           target_branch: target_branch.name,
           state: state,
@@ -27,6 +27,22 @@ module Gitlab
         !cross_project? && source_branch.present? && target_branch.present?
       end
 
+      def cross_project?
+        source_repo.present? && target_repo.present? && source_repo.id != target_repo.id
+      end
+
+      def source_branch_exists?
+        source_project.repository.branch_names.include?(source_branch)
+      end
+
+      def source_branch
+        raw_data.head.ref
+      end
+
+      def source_sha
+        raw_data.head.sha
+      end
+
       private
 
       def assigned?
@@ -51,10 +67,6 @@ module Gitlab
         raw_data.body || ""
       end
 
-      def cross_project?
-        source_repo.present? && target_repo.present? && source_repo.id != target_repo.id
-      end
-
       def description
         formatter.author_line(author) + body
       end
@@ -73,10 +85,6 @@ module Gitlab
         raw_data.head.repo
       end
 
-      def source_branch
-        source_project.repository.find_branch(raw_data.head.ref)
-      end
-
       def target_project
         project
       end