From 4a4338fcc60dcee33b0075213213ae51e8395a6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alejandro=20Rodr=C3=ADguez?= <alejorro70@gmail.com>
Date: Wed, 27 Jun 2018 17:55:40 -0400
Subject: Use Gitaly's OperationService.UserUpdateBranch RPC

---
 lib/gitlab/git/repository.rb                       |  8 ++++-
 lib/gitlab/gitaly_client/operation_service.rb      | 16 +++++++++
 .../gitlab/gitaly_client/operation_service_spec.rb | 41 ++++++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 29b3663a52a..6fb1af4d057 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -651,7 +651,13 @@ module Gitlab
       end
 
       def update_branch(branch_name, user:, newrev:, oldrev:)
-        OperationService.new(user, self).update_branch(branch_name, newrev, oldrev)
+        gitaly_migrate(:operation_user_update_branch) do |is_enabled|
+          if is_enabled
+            gitaly_operations_client.user_update_branch(branch_name, user, newrev, oldrev)
+          else
+            OperationService.new(user, self).update_branch(branch_name, newrev, oldrev)
+          end
+        end
       end
 
       def rm_branch(branch_name, user:)
diff --git a/lib/gitlab/gitaly_client/operation_service.rb b/lib/gitlab/gitaly_client/operation_service.rb
index ab2c61f6782..555733d1834 100644
--- a/lib/gitlab/gitaly_client/operation_service.rb
+++ b/lib/gitlab/gitaly_client/operation_service.rb
@@ -68,6 +68,22 @@ module Gitlab
         raise Gitlab::Git::Repository::InvalidRef, ex
       end
 
+      def user_update_branch(branch_name, user, newrev, oldrev)
+        request = Gitaly::UserUpdateBranchRequest.new(
+          repository: @gitaly_repo,
+          branch_name: encode_binary(branch_name),
+          user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
+          newrev: encode_binary(newrev),
+          oldrev: encode_binary(oldrev)
+        )
+
+        response = GitalyClient.call(@repository.storage, :operation_service, :user_update_branch, request)
+
+        if pre_receive_error = response.pre_receive_error.presence
+          raise Gitlab::Git::PreReceiveError, pre_receive_error
+        end
+      end
+
       def user_delete_branch(branch_name, user)
         request = Gitaly::UserDeleteBranchRequest.new(
           repository: @gitaly_repo,
diff --git a/spec/lib/gitlab/gitaly_client/operation_service_spec.rb b/spec/lib/gitlab/gitaly_client/operation_service_spec.rb
index 9709f1f5646..031d1e87dc1 100644
--- a/spec/lib/gitlab/gitaly_client/operation_service_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/operation_service_spec.rb
@@ -53,6 +53,47 @@ describe Gitlab::GitalyClient::OperationService do
     end
   end
 
+  describe '#user_update_branch' do
+    let(:branch_name) { 'my-branch' }
+    let(:newrev) { '01e' }
+    let(:oldrev) { '01d' }
+    let(:request) do
+      Gitaly::UserUpdateBranchRequest.new(
+        repository: repository.gitaly_repository,
+        branch_name: branch_name,
+        newrev: newrev,
+        oldrev: oldrev,
+        user: gitaly_user
+      )
+    end
+    let(:response) { Gitaly::UserUpdateBranchResponse.new }
+
+    subject { client.user_update_branch(branch_name, user, newrev, oldrev) }
+
+    it 'sends a user_update_branch message' do
+      expect_any_instance_of(Gitaly::OperationService::Stub)
+        .to receive(:user_update_branch).with(request, kind_of(Hash))
+        .and_return(response)
+
+      subject
+    end
+
+    context "when pre_receive_error is present" do
+      let(:response) do
+        Gitaly::UserUpdateBranchResponse.new(pre_receive_error: "something failed")
+      end
+
+      it "throws a PreReceive exception" do
+        expect_any_instance_of(Gitaly::OperationService::Stub)
+          .to receive(:user_update_branch).with(request, kind_of(Hash))
+          .and_return(response)
+
+        expect { subject }.to raise_error(
+          Gitlab::Git::PreReceiveError, "something failed")
+      end
+    end
+  end
+
   describe '#user_delete_branch' do
     let(:branch_name) { 'my-branch' }
     let(:request) do
-- 
cgit v1.2.1