Implement OperationService.UserAddBranch Gitaly RPCgitaly-add-branch

2 files changed, 47 insertions, 7 deletions

diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 9ea5ed04b13..22b735c6f7b 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -656,13 +656,13 @@ module Gitlab
       end
 
       def add_branch(branch_name, user:, target:)
-        target_object = Ref.dereference_object(lookup(target))
-        raise InvalidRef.new("target not found: #{target}") unless target_object
-
-        OperationService.new(user, self).add_branch(branch_name, target_object.oid)
-        find_branch(branch_name)
-      rescue Rugged::ReferenceError => ex
-        raise InvalidRef, ex
+        gitaly_migrate(:operation_user_create_branch) do |is_enabled|
+          if is_enabled
+            gitaly_add_branch(branch_name, user, target)
+          else
+            rugged_add_branch(branch_name, user, target)
+          end
+        end
       end
 
       def add_tag(tag_name, user:, target:, message: nil)
@@ -1081,6 +1081,10 @@ module Gitlab
         @gitaly_repository_client ||= Gitlab::GitalyClient::RepositoryService.new(self)
       end
 
+      def gitaly_operation_client
+        @gitaly_operation_client ||= Gitlab::GitalyClient::OperationService.new(self)
+      end
+
       def gitaly_migrate(method, status: Gitlab::GitalyClient::MigrationStatus::OPT_IN, &block)
         Gitlab::GitalyClient.migrate(method, status: status, &block)
       rescue GRPC::NotFound => e
@@ -1472,6 +1476,22 @@ module Gitlab
           file.write(gitattributes_content)
         end
       end
+
+      def gitaly_add_branch(branch_name, user, target)
+        gitaly_operation_client.user_create_branch(branch_name, user, target)
+      rescue GRPC::FailedPrecondition => ex
+        raise InvalidRef, ex
+      end
+
+      def rugged_add_branch(branch_name, user, target)
+        target_object = Ref.dereference_object(lookup(target))
+        raise InvalidRef.new("target not found: #{target}") unless target_object
+
+        OperationService.new(user, self).add_branch(branch_name, target_object.oid)
+        find_branch(branch_name)
+      rescue Rugged::ReferenceError
+        raise InvalidRef, ex
+      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/operation_service.rb b/lib/gitlab/gitaly_client/operation_service.rb
index 2d5440e7ea8..46bd5c18603 100644
--- a/lib/gitlab/gitaly_client/operation_service.rb
+++ b/lib/gitlab/gitaly_client/operation_service.rb
@@ -40,6 +40,26 @@ module Gitlab
       rescue GRPC::FailedPrecondition => e
         raise Gitlab::Git::Repository::InvalidRef, e
       end
+
+      def user_create_branch(branch_name, user, start_point)
+        request = Gitaly::UserCreateBranchRequest.new(
+          repository: @gitaly_repo,
+          branch_name: GitalyClient.encode(branch_name),
+          user: Util.gitaly_user(user),
+          start_point: GitalyClient.encode(start_point)
+        )
+        response = GitalyClient.call(@repository.storage, :operation_service,
+          :user_create_branch, request)
+        if response.pre_receive_error.present?
+          raise Gitlab::Git::HooksService::PreReceiveError.new(response.pre_receive_error)
+        end
+
+        branch = response.branch
+        return nil unless branch
+
+        target_commit = Gitlab::Git::Commit.decorate(@repository, branch.target_commit)
+        Gitlab::Git::Branch.new(@repository, branch.name, target_commit.id, target_commit)
+      end
     end
   end
 end