operations: Let UserUpdateBranch hook into transactions

This commit adjusts the `UserUpdateBranch` RPC to hook into transactions
in case the corresponding feature flag is enabled.

5 files changed, 37 insertions, 26 deletions

diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 274ea6419..9ce300417 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -159,6 +159,7 @@ func (c *Coordinator) accessorStreamParameters(ctx context.Context, call grpcCal
 
 var transactionRPCs = map[string]featureflag.FeatureFlag{
 	"/gitaly.OperationService/UserCreateBranch": featureflag.ReferenceTransactionsOperationService,
+	"/gitaly.OperationService/UserUpdateBranch": featureflag.ReferenceTransactionsOperationService,
 	"/gitaly.SSHService/SSHReceivePack":         featureflag.ReferenceTransactionsSSHService,
 	"/gitaly.SmartHTTPService/PostReceivePack":  featureflag.ReferenceTransactionsSmartHTTPService,
 }
diff --git a/internal/service/operations/update_branches_test.go b/internal/service/operations/update_branches_test.go
index f151d9f69..542d6cff2 100644
--- a/internal/service/operations/update_branches_test.go
+++ b/internal/service/operations/update_branches_test.go
@@ -21,35 +21,44 @@ var (
 )
 
 func TestSuccessfulUserUpdateBranchRequest(t *testing.T) {
-	ctx, cancel := testhelper.Context()
-	defer cancel()
+	featureSets, err := testhelper.NewFeatureSets([]featureflag.FeatureFlag{featureflag.ReferenceTransactions})
+	require.NoError(t, err)
 
-	testRepo, _, cleanupFn := testhelper.NewTestRepo(t)
-	defer cleanupFn()
+	for _, featureSet := range featureSets {
+		t.Run(featureSet.String(), func(t *testing.T) {
+			ctx, cancel := testhelper.Context()
+			defer cancel()
 
-	serverSocketPath, stop := runOperationServiceServer(t)
-	defer stop()
+			ctx = featureSet.WithParent(ctx)
 
-	client, conn := newOperationClient(t, serverSocketPath)
-	defer conn.Close()
+			testRepo, _, cleanupFn := testhelper.NewTestRepo(t)
+			defer cleanupFn()
 
-	request := &gitalypb.UserUpdateBranchRequest{
-		Repository: testRepo,
-		BranchName: []byte(updateBranchName),
-		Newrev:     newrev,
-		Oldrev:     oldrev,
-		User:       testhelper.TestUser,
-	}
+			serverSocketPath, stop := runOperationServiceServer(t)
+			defer stop()
 
-	response, err := client.UserUpdateBranch(ctx, request)
+			client, conn := newOperationClient(t, serverSocketPath)
+			defer conn.Close()
 
-	require.NoError(t, err)
-	require.Empty(t, response.PreReceiveError)
+			request := &gitalypb.UserUpdateBranchRequest{
+				Repository: testRepo,
+				BranchName: []byte(updateBranchName),
+				Newrev:     newrev,
+				Oldrev:     oldrev,
+				User:       testhelper.TestUser,
+			}
 
-	branchCommit, err := log.GetCommit(ctx, testRepo, updateBranchName)
+			response, err := client.UserUpdateBranch(ctx, request)
 
-	require.NoError(t, err)
-	require.Equal(t, string(newrev), branchCommit.Id)
+			require.NoError(t, err)
+			require.Empty(t, response.PreReceiveError)
+
+			branchCommit, err := log.GetCommit(ctx, testRepo, updateBranchName)
+
+			require.NoError(t, err)
+			require.Equal(t, string(newrev), branchCommit.Id)
+		})
+	}
 }
 
 func TestSuccessfulGitHooksForUserUpdateBranchRequest(t *testing.T) {
diff --git a/ruby/lib/gitaly_server/operations_service.rb b/ruby/lib/gitaly_server/operations_service.rb
index 0f4db6c7d..5d63ef1e1 100644
--- a/ruby/lib/gitaly_server/operations_service.rb
+++ b/ruby/lib/gitaly_server/operations_service.rb
@@ -72,9 +72,10 @@ module GitalyServer
       newrev = get_param!(request, :newrev)
       oldrev = get_param!(request, :oldrev)
       gitaly_user = get_param!(request, :user)
+      transaction = Praefect::Transaction.from_metadata(call.metadata)
 
       user = Gitlab::Git::User.from_gitaly(gitaly_user)
-      repo.update_branch(branch_name, user: user, newrev: newrev, oldrev: oldrev)
+      repo.update_branch(branch_name, user: user, newrev: newrev, oldrev: oldrev, transaction: transaction)
 
       Gitaly::UserUpdateBranchResponse.new
     rescue Gitlab::Git::Repository::InvalidRef, Gitlab::Git::CommitError => ex
diff --git a/ruby/lib/gitlab/git/operation_service.rb b/ruby/lib/gitlab/git/operation_service.rb
index e6025b673..abc489384 100644
--- a/ruby/lib/gitlab/git/operation_service.rb
+++ b/ruby/lib/gitlab/git/operation_service.rb
@@ -96,9 +96,9 @@ module Gitlab
         end
       end
 
-      def update_branch(branch_name, newrev, oldrev, push_options: nil)
+      def update_branch(branch_name, newrev, oldrev, push_options: nil, transaction: nil)
         ref = Gitlab::Git::BRANCH_REF_PREFIX + branch_name
-        update_ref_in_hooks(ref, newrev, oldrev, push_options: push_options)
+        update_ref_in_hooks(ref, newrev, oldrev, push_options: push_options, transaction: transaction)
       end
 
       # Yields the given block (which should return a commit) and
diff --git a/ruby/lib/gitlab/git/repository.rb b/ruby/lib/gitlab/git/repository.rb
index 675f4b2e2..d9d337a23 100644
--- a/ruby/lib/gitlab/git/repository.rb
+++ b/ruby/lib/gitlab/git/repository.rb
@@ -251,8 +251,8 @@ module Gitlab
         end
       end
 
-      def update_branch(branch_name, user:, newrev:, oldrev:, push_options: nil)
-        OperationService.new(user, self).update_branch(branch_name, newrev, oldrev, push_options: push_options)
+      def update_branch(branch_name, user:, newrev:, oldrev:, push_options: nil, transaction: nil)
+        OperationService.new(user, self).update_branch(branch_name, newrev, oldrev, push_options: push_options, transaction: transaction)
       end
 
       def rm_branch(branch_name, user:)