Fix GitHub importer performance on branch existence check

3 files changed, 57 insertions, 35 deletions

diff --git a/lib/github/import.rb b/lib/github/import.rb
index b20614b3060..ff5d7db2705 100644
--- a/lib/github/import.rb
+++ b/lib/github/import.rb
@@ -172,7 +172,7 @@ module Github
           next unless merge_request.new_record? && pull_request.valid?
 
           begin
-            restore_branches(pull_request)
+            pull_request.restore_branches!
 
             author_id   = user_id(pull_request.author, project.creator_id)
             description = format_description(pull_request.description, pull_request.author)
@@ -208,7 +208,7 @@ module Github
           rescue => e
             error(:pull_request, pull_request.url, e.message)
           ensure
-            clean_up_restored_branches(pull_request)
+            pull_request.remove_restored_branches!
           end
         end
 
@@ -325,32 +325,6 @@ module Github
       end
     end
 
-    def restore_branches(pull_request)
-      restore_source_branch(pull_request) unless pull_request.source_branch_exists?
-      restore_target_branch(pull_request) unless pull_request.target_branch_exists?
-    end
-
-    def restore_source_branch(pull_request)
-      repository.create_branch(pull_request.source_branch_name, pull_request.source_branch_sha)
-    end
-
-    def restore_target_branch(pull_request)
-      repository.create_branch(pull_request.target_branch_name, pull_request.target_branch_sha)
-    end
-
-    def remove_branch(name)
-      repository.delete_branch(name)
-    rescue Rugged::ReferenceError
-      errors << { type: :branch, url: nil, error: "Could not clean up restored branch: #{name}" }
-    end
-
-    def clean_up_restored_branches(pull_request)
-      return if pull_request.opened?
-
-      remove_branch(pull_request.source_branch_name) unless pull_request.source_branch_exists?
-      remove_branch(pull_request.target_branch_name) unless pull_request.target_branch_exists?
-    end
-
     def label_ids(labels)
       labels.map { |attrs| cached[:label_ids][attrs.fetch('name')] }.compact
     end
diff --git a/lib/github/representation/branch.rb b/lib/github/representation/branch.rb
index d1dac6944f0..c6fa928d565 100644
--- a/lib/github/representation/branch.rb
+++ b/lib/github/representation/branch.rb
@@ -26,13 +26,25 @@ module Github
       end
 
       def exists?
-        branch_exists? && commit_exists?
+        @exists ||= branch_exists? && commit_exists?
       end
 
       def valid?
         sha.present? && ref.present?
       end
 
+      def restore!(name)
+        repository.create_branch(name, sha)
+      rescue Gitlab::Git::Repository::InvalidRef => e
+        Rails.logger.error("#{self.class.name}: Could not restore branch #{name}: #{e}")
+      end
+
+      def remove!(name)
+        repository.delete_branch(name)
+      rescue Rugged::ReferenceError => e
+        Rails.logger.error("#{self.class.name}: Could not remove branch #{name}: #{e}")
+      end
+
       private
 
       def branch_exists?
diff --git a/lib/github/representation/pull_request.rb b/lib/github/representation/pull_request.rb
index ac9c8283b4b..61d5e9455d9 100644
--- a/lib/github/representation/pull_request.rb
+++ b/lib/github/representation/pull_request.rb
@@ -1,8 +1,6 @@
 module Github
   module Representation
     class PullRequest < Representation::Issuable
-      attr_reader :project
-
       delegate :user, :repo, :ref, :sha, to: :source_branch, prefix: true
       delegate :user, :exists?, :repo, :ref, :sha, :short_sha, to: :target_branch, prefix: true
 
@@ -10,10 +8,6 @@ module Github
         project
       end
 
-      def source_branch_exists?
-        !cross_project? && source_branch.exists?
-      end
-
       def source_branch_name
         @source_branch_name ||=
           if cross_project? || !source_branch_exists?
@@ -23,6 +17,12 @@ module Github
           end
       end
 
+      def source_branch_exists?
+        return if defined?(@source_branch_exists)
+
+        @source_branch_exists = !cross_project? && source_branch.exists?
+      end
+
       def target_project
         project
       end
@@ -31,6 +31,10 @@ module Github
         @target_branch_name ||= target_branch_exists? ? target_branch_ref : target_branch_name_prefixed
       end
 
+      def target_branch_exists?
+        @target_branch_exists ||= target_branch.exists?
+      end
+
       def state
         return 'merged' if raw['state'] == 'closed' && raw['merged_at'].present?
         return 'closed' if raw['state'] == 'closed'
@@ -46,6 +50,18 @@ module Github
         source_branch.valid? && target_branch.valid?
       end
 
+      def restore_branches!
+        restore_source_branch!
+        restore_target_branch!
+      end
+
+      def remove_restored_branches!
+        return if opened?
+
+        remove_source_branch!
+        remove_target_branch!
+      end
+
       private
 
       def project
@@ -73,6 +89,26 @@ module Github
 
         source_branch_repo.id != target_branch_repo.id
       end
+
+      def restore_source_branch!
+        return if source_branch_exists?
+
+        source_branch.restore!(source_branch_name)
+      end
+
+      def restore_target_branch!
+        return if target_branch_exists?
+
+        target_branch.restore!(target_branch_name)
+      end
+
+      def remove_source_branch!
+        source_branch.remove!(source_branch_name) unless source_branch_exists?
+      end
+
+      def remove_target_branch!
+        target_branch.remove!(target_branch_name) unless target_branch_exists?
+      end
     end
   end
 end