Incorporate Gitaly's local_branches operation into repo codelocal-branches-gitaly

9 files changed, 140 insertions, 23 deletions

diff --git a/GITALY_SERVER_VERSION b/GITALY_SERVER_VERSION
index a918a2aa18d..a3df0a6959e 100644
--- a/GITALY_SERVER_VERSION
+++ b/GITALY_SERVER_VERSION
@@ -1 +1 @@
-0.6.0
+0.8.0
diff --git a/app/models/repository.rb b/app/models/repository.rb
index 7bb874d7744..75bd252b6c6 100644
--- a/app/models/repository.rb
+++ b/app/models/repository.rb
@@ -644,22 +644,8 @@ class Repository
     "#{name}-#{highest_branch_id + 1}"
   end
 
-  # Remove archives older than 2 hours
   def branches_sorted_by(value)
-    case value
-    when 'name'
-      branches.sort_by(&:name)
-    when 'updated_desc'
-      branches.sort do |a, b|
-        commit(b.dereferenced_target).committed_date <=> commit(a.dereferenced_target).committed_date
-      end
-    when 'updated_asc'
-      branches.sort do |a, b|
-        commit(a.dereferenced_target).committed_date <=> commit(b.dereferenced_target).committed_date
-      end
-    else
-      branches
-    end
+    raw_repository.local_branches(value)
   end
 
   def tags_sorted_by(value)
diff --git a/changelogs/unreleased/gitaly-local-branches.yml b/changelogs/unreleased/gitaly-local-branches.yml
new file mode 100644
index 00000000000..adcc0fa6280
--- /dev/null
+++ b/changelogs/unreleased/gitaly-local-branches.yml
@@ -0,0 +1,4 @@
+---
+title: Add suport for find_local_branches GRPC from Gitaly
+merge_request: 10059
+author:
diff --git a/lib/gitlab/git/commit.rb b/lib/gitlab/git/commit.rb
index 3a73697dc5d..957119bc5e3 100644
--- a/lib/gitlab/git/commit.rb
+++ b/lib/gitlab/git/commit.rb
@@ -55,6 +55,7 @@ module Gitlab
         #   Commit.find(repo, 'master')
         #
         def find(repo, commit_id = "HEAD")
+          return commit_id if commit_id.is_a?(Gitlab::Git::Commit)
           return decorate(commit_id) if commit_id.is_a?(Rugged::Commit)
 
           obj = if commit_id.is_a?(String)
diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index d7dac9f6149..df0cbeb0422 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -86,14 +86,16 @@ module Gitlab
       end
 
       # Returns an Array of Branches
-      def branches
-        rugged.branches.map do |rugged_ref|
+      def branches(filter: nil, sort_by: nil)
+        branches = rugged.branches.each(filter).map do |rugged_ref|
           begin
             Gitlab::Git::Branch.new(self, rugged_ref.name, rugged_ref.target)
           rescue Rugged::ReferenceError
             # Omit invalid branch
           end
-        end.compact.sort_by(&:name)
+        end.compact
+
+        sort_branches(branches, sort_by)
       end
 
       def reload_rugged
@@ -114,10 +116,21 @@ module Gitlab
         Gitlab::Git::Branch.new(self, rugged_ref.name, rugged_ref.target) if rugged_ref
       end
 
-      def local_branches
-        rugged.branches.each(:local).map do |branch|
-          Gitlab::Git::Branch.new(self, branch.name, branch.target)
+      def local_branches(sort_by = nil)
+        Gitlab::GitalyClient.migrate(:local_branches) do |is_enabled|
+          if is_enabled
+            gitaly_ref_client.local_branches(sort_by).map do |branch|
+              target = branch_target_from_gitaly_response(branch)
+              Gitlab::Git::Branch.new(self, branch.name, target)
+            end
+          else
+            branches(filter: :local, sort_by: sort_by)
+          end
         end
+      rescue GRPC::NotFound => e
+        raise NoRepository.new(e)
+      rescue GRPC::BadStatus => e
+        raise CommandError.new(e)
       end
 
       # Returns the number of valid branches
@@ -1270,9 +1283,53 @@ module Gitlab
         diff.each_patch
       end
 
+      def sort_branches(branches, sort_by)
+        case sort_by
+        when 'name'
+          branches.sort_by(&:name)
+        when 'updated_desc'
+          branches.sort do |a, b|
+            b.dereferenced_target.committed_date <=> a.dereferenced_target.committed_date
+          end
+        when 'updated_asc'
+          branches.sort do |a, b|
+            a.dereferenced_target.committed_date <=> b.dereferenced_target.committed_date
+          end
+        else
+          branches
+        end
+      end
+
       def gitaly_ref_client
         @gitaly_ref_client ||= Gitlab::GitalyClient::Ref.new(self)
       end
+
+      def branch_target_from_gitaly_response(response)
+        # Git messages have no encoding enforcements. However, in the UI we only
+        # handle UTF-8, so basically we cross our fingers that the message force
+        # encoded to UTF-8 is readable.
+        message = response.commit_subject.dup.force_encoding('UTF-8')
+
+        # NOTE: For ease of parsing in Gitaly, we have only the subject of
+        # the commit and not the full message. This is ok, since all the
+        # code that uses `local_branches` only cares at most about the
+        # commit message.
+        # TODO: Once gitaly "takes over" Rugged consider separating the
+        # subject from the message to make it clearer when there's one
+        # available but not the other.
+        hash = {
+          id: response.commit_id,
+          message: message,
+          authored_date: Time.at(response.commit_author.date.seconds),
+          author_name: response.commit_author.name,
+          author_email: response.commit_author.email,
+          committed_date: Time.at(response.commit_committer.date.seconds),
+          committer_name: response.commit_committer.name,
+          committer_email: response.commit_committer.email,
+        }
+
+        Gitlab::Git::Commit.decorate(hash)
+      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/ref.rb b/lib/gitlab/gitaly_client/ref.rb
index d3c0743db4e..71ec6532ef4 100644
--- a/lib/gitlab/gitaly_client/ref.rb
+++ b/lib/gitlab/gitaly_client/ref.rb
@@ -34,6 +34,12 @@ module Gitlab
         stub.find_ref_name(request).name
       end
 
+      def local_branches(sort_by = nil)
+        request = Gitaly::FindLocalBranchesRequest.new(repository: @gitaly_repo)
+        request.sort_by = sort_by_param(sort_by) if sort_by
+        consume_branches_response(stub.find_local_branches(request))
+      end
+
       private
 
       def consume_refs_response(response, prefix:)
@@ -41,6 +47,16 @@ module Gitlab
           r.names.map { |name| name.sub(/\A#{Regexp.escape(prefix)}/, '') }
         end
       end
+
+      def sort_by_param(sort_by)
+        enum_value = Gitaly::FindLocalBranchesRequest::SortBy.resolve(sort_by.upcase.to_sym)
+        raise ArgumentError, "Invalid sort_by key `#{sort_by}`" unless enum_value
+        enum_value
+      end
+
+      def consume_branches_response(response)
+        response.flat_map { |r| r.branches }
+      end
     end
   end
 end
diff --git a/spec/lib/gitlab/git/repository_spec.rb b/spec/lib/gitlab/git/repository_spec.rb
index 3d6d7292b42..a81370eba4f 100644
--- a/spec/lib/gitlab/git/repository_spec.rb
+++ b/spec/lib/gitlab/git/repository_spec.rb
@@ -1036,7 +1036,9 @@ describe Gitlab::Git::Repository, seed_helper: true do
       ref = double()
       allow(ref).to receive(:name) { 'bad-branch' }
       allow(ref).to receive(:target) { raise Rugged::ReferenceError }
-      allow(repository.rugged).to receive(:branches) { [ref] }
+      branches = double()
+      allow(branches).to receive(:each) { [ref].each }
+      allow(repository.rugged).to receive(:branches) { branches }
     end
 
     it 'should return empty branches' do
@@ -1247,6 +1249,28 @@ describe Gitlab::Git::Repository, seed_helper: true do
       expect(@repo.local_branches.any? { |branch| branch.name == 'remote_branch' }).to eq(false)
       expect(@repo.local_branches.any? { |branch| branch.name == 'local_branch' }).to eq(true)
     end
+
+    context 'with gitaly enabled' do
+      before { stub_gitaly }
+
+      it 'gets the branches from GitalyClient' do
+        expect_any_instance_of(Gitlab::GitalyClient::Ref).to receive(:local_branches).
+          and_return([])
+        @repo.local_branches
+      end
+
+      it 'wraps GRPC not found' do
+        expect_any_instance_of(Gitlab::GitalyClient::Ref).to receive(:local_branches).
+          and_raise(GRPC::NotFound)
+        expect { @repo.local_branches }.to raise_error(Gitlab::Git::Repository::NoRepository)
+      end
+
+      it 'wraps GRPC exceptions' do
+        expect_any_instance_of(Gitlab::GitalyClient::Ref).to receive(:local_branches).
+          and_raise(GRPC::Unknown)
+        expect { @repo.local_branches }.to raise_error(Gitlab::Git::CommandError)
+      end
+    end
   end
 
   def create_remote_branch(remote_name, branch_name, source_branch_name)
diff --git a/spec/lib/gitlab/gitaly_client/ref_spec.rb b/spec/lib/gitlab/gitaly_client/ref_spec.rb
index 5405eafd281..8769a604260 100644
--- a/spec/lib/gitlab/gitaly_client/ref_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/ref_spec.rb
@@ -38,4 +38,27 @@ describe Gitlab::GitalyClient::Ref do
       client.default_branch_name
     end
   end
+
+  describe '#local_branches' do
+    it 'sends a find_local_branches message' do
+      expect_any_instance_of(Gitaly::Ref::Stub).
+        to receive(:find_local_branches).with(gitaly_request_with_repo_path(repo_path)).
+          and_return([])
+
+      client.local_branches
+    end
+
+    it 'parses and sends the sort parameter' do
+      expect_any_instance_of(Gitaly::Ref::Stub).
+        to receive(:find_local_branches).
+          with(gitaly_request_with_params(sort_by: :UPDATED_DESC)).
+          and_return([])
+
+      client.local_branches('updated_desc')
+    end
+
+    it 'raises an argument error if an invalid sort_by parameter is passed' do
+      expect { client.local_branches('invalid_sort') }.to raise_error(ArgumentError)
+    end
+  end
 end
diff --git a/spec/support/matchers/gitaly_matchers.rb b/spec/support/matchers/gitaly_matchers.rb
index 65dbc01f6e4..ed14bcec9f2 100644
--- a/spec/support/matchers/gitaly_matchers.rb
+++ b/spec/support/matchers/gitaly_matchers.rb
@@ -1,3 +1,9 @@
 RSpec::Matchers.define :gitaly_request_with_repo_path do |path|
   match { |actual| actual.repository.path == path }
 end
+
+RSpec::Matchers.define :gitaly_request_with_params do |params|
+  match do |actual|
+    params.reduce(true) { |r, (key, val)| r && actual.send(key) == val }
+  end
+end