Fix GitHub importer performance on branch existence check

2 files changed, 55 insertions, 7 deletions

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