Migrate Git::Repository#write_ref to Gitaly

2 files changed, 32 insertions, 4 deletions

diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 3a7930154e5..e65719e7f92 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -1105,10 +1105,14 @@ module Gitlab
       end
 
       def write_ref(ref_path, ref, old_ref: nil, shell: true)
-        if shell
-          shell_write_ref(ref_path, ref, old_ref)
-        else
-          rugged_write_ref(ref_path, ref)
+        ref_path = "#{Gitlab::Git::BRANCH_REF_PREFIX}#{ref_path}" unless ref_path.start_with?("refs/") || ref_path == "HEAD"
+
+        gitaly_migrate(:write_ref) do |is_enabled|
+          if is_enabled
+            gitaly_repository_client.write_ref(ref_path, ref, old_ref, shell)
+          else
+            local_write_ref(ref_path, ref, old_ref: old_ref, shell: shell)
+          end
         end
       end
 
@@ -1419,6 +1423,14 @@ module Gitlab
 
       private
 
+      def local_write_ref(ref_path, ref, old_ref: nil, shell: true)
+        if shell
+          shell_write_ref(ref_path, ref, old_ref)
+        else
+          rugged_write_ref(ref_path, ref)
+        end
+      end
+
       def shell_write_ref(ref_path, ref, old_ref)
         raise ArgumentError, "invalid ref_path #{ref_path.inspect}" if ref_path.include?(' ')
         raise ArgumentError, "invalid ref #{ref.inspect}" if ref.include?("\x00")
diff --git a/lib/gitlab/gitaly_client/repository_service.rb b/lib/gitlab/gitaly_client/repository_service.rb
index b0dbaf11598..7adf32af209 100644
--- a/lib/gitlab/gitaly_client/repository_service.rb
+++ b/lib/gitlab/gitaly_client/repository_service.rb
@@ -203,6 +203,22 @@ module Gitlab
           timeout: GitalyClient.default_timeout
         )
       end
+
+      def write_ref(ref_path, ref, old_ref, shell)
+        request = Gitaly::WriteRefRequest.new(
+          repository: @gitaly_repo,
+          ref: ref_path.b,
+          revision: ref.b,
+          shell: shell
+        )
+        request.old_revision = old_ref.b unless old_ref.nil?
+
+        response = GitalyClient.call(@storage, :repository_service, :write_ref, request)
+
+        raise Gitlab::Git::CommandError, encode!(response.error) if response.error.present?
+
+        true
+      end
     end
   end
 end