Provisioning: Fix stage options for export (#108585)

* Add ref option to stage options

* Fix the issue with ref in export worker

* Add unit tests for export stage options

* Do not fail if ref is equal to the stage branch

* Format code again

* fix test

---------

Co-authored-by: Stephanie Hingtgen <stephanie.hingtgen@grafana.com>

pkg/registry/apis/provisioning/jobs/export/worker.go

@@ -62,9 +62,11 @@ func (r *ExportWorker) Process(ctx context.Context, repo repository.Repository,
 	}
 
 	cloneOptions := repository.StageOptions{
+		Ref:                   options.Branch,
+		Timeout:               10 * time.Minute,
+		PushOnWrites:          false,
 		Mode:                  repository.StageModeCommitOnlyOnce,
 		CommitOnlyOnceMessage: msg,
-		Timeout:               10 * time.Minute,
 	}
 
 	fn := func(repo repository.Repository, _ bool) error {

pkg/registry/apis/provisioning/jobs/export/worker_test.go

@@ -228,7 +228,9 @@ func TestExportWorker_ProcessStageOptions(t *testing.T) {
 	job := v0alpha1.Job{
 		Spec: v0alpha1.JobSpec{
 			Action: v0alpha1.JobActionPush,
-			Push:   &v0alpha1.ExportJobOptions{},
+			Push: &v0alpha1.ExportJobOptions{
+				Branch: "feature-branch",
+			},
 		},
 	}
 
@@ -239,7 +241,8 @@ func TestExportWorker_ProcessStageOptions(t *testing.T) {
 			Namespace: "test-namespace",
 		},
 		Spec: v0alpha1.RepositorySpec{
-			Workflows: []v0alpha1.Workflow{v0alpha1.WriteWorkflow},
+			Type:      v0alpha1.GitRepositoryType,
+			Workflows: []v0alpha1.Workflow{v0alpha1.WriteWorkflow, v0alpha1.BranchWorkflow},
 		},
 	})
 
@@ -258,9 +261,11 @@ func TestExportWorker_ProcessStageOptions(t *testing.T) {
 	mockExportFn.On("Execute", mock.Anything, "test-repo", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 	mockStageFn := NewMockWrapWithStageFn(t)
-	// Verify clone and push options
+	// Verify all stage options including Ref (branch), Timeout, and PushOnWrites
 	mockStageFn.On("Execute", mock.Anything, mockRepo, mock.MatchedBy(func(opts repository.StageOptions) bool {
-		return opts.Timeout == 10*time.Minute && opts.Mode == repository.StageModeCommitOnlyOnce
+		return opts.Ref == "feature-branch" &&
+			opts.Timeout == 10*time.Minute &&
+			!opts.PushOnWrites && opts.Mode == repository.StageModeCommitOnlyOnce
 	}), mock.Anything).Return(func(ctx context.Context, repo repository.Repository, stageOpts repository.StageOptions, fn func(repository.Repository, bool) error) error {
 		return fn(repo, true)
 	})
@@ -270,6 +275,89 @@ func TestExportWorker_ProcessStageOptions(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestExportWorker_ProcessStageOptionsWithBranch(t *testing.T) {
+	tests := []struct {
+		name        string
+		branch      string
+		expectedRef string
+		workflows   []v0alpha1.Workflow
+		repoType    v0alpha1.RepositoryType
+	}{
+		{
+			name:        "branch specified",
+			branch:      "develop",
+			expectedRef: "develop",
+			workflows:   []v0alpha1.Workflow{v0alpha1.WriteWorkflow, v0alpha1.BranchWorkflow},
+			repoType:    v0alpha1.GitRepositoryType,
+		},
+		{
+			name:        "empty branch",
+			branch:      "",
+			expectedRef: "",
+			workflows:   []v0alpha1.Workflow{v0alpha1.WriteWorkflow},
+			repoType:    v0alpha1.LocalRepositoryType,
+		},
+		{
+			name:        "main branch",
+			branch:      "main",
+			expectedRef: "main",
+			workflows:   []v0alpha1.Workflow{v0alpha1.WriteWorkflow, v0alpha1.BranchWorkflow},
+			repoType:    v0alpha1.GitRepositoryType,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			job := v0alpha1.Job{
+				Spec: v0alpha1.JobSpec{
+					Action: v0alpha1.JobActionPush,
+					Push: &v0alpha1.ExportJobOptions{
+						Branch: tt.branch,
+					},
+				},
+			}
+
+			mockRepo := repository.NewMockRepository(t)
+			mockRepo.On("Config").Return(&v0alpha1.Repository{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-repo",
+					Namespace: "test-namespace",
+				},
+				Spec: v0alpha1.RepositorySpec{
+					Type:      tt.repoType,
+					Workflows: tt.workflows,
+				},
+			})
+
+			mockProgress := jobs.NewMockJobProgressRecorder(t)
+			mockClients := resources.NewMockClientFactory(t)
+			mockResourceClients := resources.NewMockResourceClients(t)
+			mockClients.On("Clients", mock.Anything, "test-namespace").Return(mockResourceClients, nil)
+
+			mockRepoResources := resources.NewMockRepositoryResourcesFactory(t)
+			mockRepoResourcesClient := resources.NewMockRepositoryResources(t)
+			mockRepoResources.On("Client", mock.Anything, mock.Anything).Return(mockRepoResourcesClient, nil)
+
+			mockExportFn := NewMockExportFn(t)
+			mockExportFn.On("Execute", mock.Anything, "test-repo", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+			mockStageFn := NewMockWrapWithStageFn(t)
+			// Verify that the stage options contain the correct branch reference and other parameters
+			mockStageFn.On("Execute", mock.Anything, mockRepo, mock.MatchedBy(func(opts repository.StageOptions) bool {
+				return opts.Ref == tt.expectedRef &&
+					opts.Timeout == 10*time.Minute &&
+					!opts.PushOnWrites
+			}), mock.Anything).Return(func(ctx context.Context, repo repository.Repository, stageOpts repository.StageOptions, fn func(repository.Repository, bool) error) error {
+				return fn(repo, true)
+			})
+
+			r := NewExportWorker(mockClients, mockRepoResources, mockExportFn.Execute, mockStageFn.Execute)
+			err := r.Process(context.Background(), mockRepo, job, mockProgress)
+			require.NoError(t, err)
+		})
+	}
+}
+
 func TestExportWorker_ProcessExportFnError(t *testing.T) {
 	job := v0alpha1.Job{
 		Spec: v0alpha1.JobSpec{

pkg/registry/apis/provisioning/repository/git/staged.go

@@ -26,7 +26,12 @@ func NewStagedGitRepository(ctx context.Context, repo *gitRepository, opts repos
 		defer cancel()
 	}
 
-	ref, err := repo.client.GetRef(ctx, "refs/heads/"+repo.gitConfig.Branch)
+	branch := opts.Ref
+	if branch == "" {
+		branch = repo.gitConfig.Branch
+	}
+
+	ref, err := repo.client.GetRef(ctx, "refs/heads/"+branch)
 	if err != nil {
 		// TODO: opts.CreateIfNotExists doesn't make sense in the context of the staged repository
 		// because we only support the branch that is passed in.
@@ -47,8 +52,25 @@ func NewStagedGitRepository(ctx context.Context, repo *gitRepository, opts repos
 	}, nil
 }
 
+// isRefSupported checks if the given ref is supported for staged operations.
+// It returns true if ref is empty, equals the git config branch, or equals the staged options ref.
+func (r *stagedGitRepository) isRefSupported(ref string) bool {
+	if ref == "" {
+		return true
+	}
+	if ref == r.gitConfig.Branch {
+		return true
+	}
+	// Allow ref if it matches the staged options ref (the branch we're staging to)
+	stagingBranch := r.opts.Ref
+	if stagingBranch == "" {
+		stagingBranch = r.gitConfig.Branch
+	}
+	return ref == stagingBranch
+}
+
 func (r *stagedGitRepository) Read(ctx context.Context, path, ref string) (*repository.FileInfo, error) {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return nil, errors.New("ref is not supported for staged repository")
 	}
 
@@ -58,7 +80,7 @@ func (r *stagedGitRepository) Read(ctx context.Context, path, ref string) (*repo
 }
 
 func (r *stagedGitRepository) ReadTree(ctx context.Context, ref string) ([]repository.FileTreeEntry, error) {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return nil, errors.New("ref is not supported for staged repository")
 	}
 
@@ -89,7 +111,7 @@ func (r *stagedGitRepository) handleCommitAndPush(ctx context.Context, message s
 }
 
 func (r *stagedGitRepository) Create(ctx context.Context, path, ref string, data []byte, message string) error {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return errors.New("ref is not supported for staged repository")
 	}
 
@@ -108,7 +130,7 @@ func (r *stagedGitRepository) blobExists(ctx context.Context, path string) (bool
 }
 
 func (r *stagedGitRepository) Write(ctx context.Context, path, ref string, data []byte, message string) error {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return errors.New("ref is not supported for staged repository")
 	}
 
@@ -131,7 +153,7 @@ func (r *stagedGitRepository) Write(ctx context.Context, path, ref string, data
 }
 
 func (r *stagedGitRepository) Update(ctx context.Context, path, ref string, data []byte, message string) error {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return errors.New("ref is not supported for staged repository")
 	}
 
@@ -147,7 +169,7 @@ func (r *stagedGitRepository) Update(ctx context.Context, path, ref string, data
 }
 
 func (r *stagedGitRepository) Delete(ctx context.Context, path, ref, message string) error {
-	if ref != "" && ref != r.gitConfig.Branch {
+	if !r.isRefSupported(ref) {
 		return errors.New("ref is not supported for staged repository")
 	}
 

pkg/registry/apis/provisioning/repository/git/staged_test.go

@@ -17,10 +17,11 @@ import (
 
 func TestNewStagedGitRepository(t *testing.T) {
 	tests := []struct {
-		name      string
-		setupMock func(*mocks.FakeClient)
-		opts      repository.StageOptions
-		wantError error
+		name        string
+		setupMock   func(*mocks.FakeClient)
+		opts        repository.StageOptions
+		wantError   error
+		expectedRef string
 	}{
 		{
 			name: "succeeds with default options",
@@ -32,11 +33,29 @@ func TestNewStagedGitRepository(t *testing.T) {
 				mockWriter := &mocks.FakeStagedWriter{}
 				mockClient.NewStagedWriterReturns(mockWriter, nil)
 			},
+			expectedRef: "refs/heads/main",
 			opts: repository.StageOptions{
 				Mode: repository.StageModeCommitOnEach,
 			},
 			wantError: nil,
 		},
+		{
+			name: "succeeds with custom ref option",
+			setupMock: func(mockClient *mocks.FakeClient) {
+				mockClient.GetRefReturns(nanogit.Ref{
+					Name: "refs/heads/custom",
+					Hash: hash.Hash{1, 2, 3},
+				}, nil)
+				mockWriter := &mocks.FakeStagedWriter{}
+				mockClient.NewStagedWriterReturns(mockWriter, nil)
+			},
+			expectedRef: "refs/heads/custom",
+			opts: repository.StageOptions{
+				Ref:          "custom",
+				PushOnWrites: false,
+			},
+			wantError: nil,
+		},
 		{
 			name: "succeeds with BeforeFn",
 			setupMock: func(mockClient *mocks.FakeClient) {
@@ -50,7 +69,8 @@ func TestNewStagedGitRepository(t *testing.T) {
 			opts: repository.StageOptions{
 				Mode: repository.StageModeCommitOnEach,
 			},
-			wantError: nil,
+			wantError:   nil,
+			expectedRef: "refs/heads/main",
 		},
 		{
 			name: "succeeds with timeout",
@@ -66,7 +86,8 @@ func TestNewStagedGitRepository(t *testing.T) {
 				Mode:    repository.StageModeCommitOnEach,
 				Timeout: time.Second * 5,
 			},
-			wantError: nil,
+			expectedRef: "refs/heads/main",
+			wantError:   nil,
 		},
 		{
 			name: "succeeds with CommitOnlyOnce option",
@@ -82,7 +103,8 @@ func TestNewStagedGitRepository(t *testing.T) {
 				Mode:                  repository.StageModeCommitOnlyOnce,
 				CommitOnlyOnceMessage: "Custom commit message",
 			},
-			wantError: nil,
+			expectedRef: "refs/heads/main",
+			wantError:   nil,
 		},
 		{
 			name: "succeeds with CommitAndPushOnEach option",
@@ -97,7 +119,8 @@ func TestNewStagedGitRepository(t *testing.T) {
 			opts: repository.StageOptions{
 				Mode: repository.StageModeCommitAndPushOnEach,
 			},
-			wantError: nil,
+			expectedRef: "refs/heads/main",
+			wantError:   nil,
 		},
 		{
 			name: "fails with GetRef error",
@@ -121,7 +144,8 @@ func TestNewStagedGitRepository(t *testing.T) {
 			opts: repository.StageOptions{
 				Mode: repository.StageModeCommitOnEach,
 			},
-			wantError: errors.New("build staged writer: failed to create writer"),
+			wantError:   errors.New("build staged writer: failed to create writer"),
+			expectedRef: "refs/heads/main",
 		},
 	}
 
@@ -156,6 +180,10 @@ func TestNewStagedGitRepository(t *testing.T) {
 				require.Equal(t, tt.opts.Mode, actualOpts.Mode)
 				require.Equal(t, tt.opts.Timeout, actualOpts.Timeout)
 				require.Equal(t, tt.opts.CommitOnlyOnceMessage, actualOpts.CommitOnlyOnceMessage)
+
+				// Verify the expected ref
+				_, ref := mockClient.GetRefArgsForCall(0)
+				require.Equal(t, tt.expectedRef, ref)
 			}
 		})
 	}
@@ -207,6 +235,25 @@ func TestStagedGitRepository_Read(t *testing.T) {
 			ref:       "main",
 			wantError: nil,
 		},
+		{
+			name: "succeeds with ref matching stage options",
+			setupMock: func(mockClient *mocks.FakeClient) {
+				mockClient.GetRefReturns(nanogit.Ref{
+					Name: "refs/heads/feature",
+					Hash: hash.Hash{1, 2, 3},
+				}, nil)
+				mockClient.GetCommitReturns(&nanogit.Commit{
+					Tree: hash.Hash{4, 5, 6},
+				}, nil)
+				mockClient.GetBlobByPathReturns(&nanogit.Blob{
+					Content: []byte("file content"),
+					Hash:    hash.Hash{7, 8, 9},
+				}, nil)
+			},
+			path:      "test.yaml",
+			ref:       "feature",
+			wantError: nil,
+		},
 		{
 			name: "fails with unsupported ref",
 			setupMock: func(_ *mocks.FakeClient) {
@@ -223,7 +270,13 @@ func TestStagedGitRepository_Read(t *testing.T) {
 			mockClient := &mocks.FakeClient{}
 			tt.setupMock(mockClient)
 
-			stagedRepo := createTestStagedRepository(mockClient)
+			// Use stage options with ref "feature" for the specific test case
+			opts := repository.StageOptions{}
+			if tt.ref == "feature" {
+				opts.Ref = "feature"
+			}
+
+			stagedRepo := createTestStagedRepositoryWithWriter(&mocks.FakeStagedWriter{}, opts, mockClient)
 			fileInfo, err := stagedRepo.Read(context.Background(), tt.path, tt.ref)
 			if tt.wantError != nil {
 				require.EqualError(t, err, tt.wantError.Error())
@@ -986,6 +1039,75 @@ func TestStagedGitRepository_Remove(t *testing.T) {
 	})
 }
 
+func TestStagedGitRepository_isRefSupported(t *testing.T) {
+	tests := []struct {
+		name      string
+		stageOpts repository.StageOptions
+		gitBranch string
+		ref       string
+		expected  bool
+	}{
+		{
+			name:      "empty ref is supported",
+			stageOpts: repository.StageOptions{},
+			gitBranch: "main",
+			ref:       "",
+			expected:  true,
+		},
+		{
+			name:      "ref matches git config branch",
+			stageOpts: repository.StageOptions{},
+			gitBranch: "main",
+			ref:       "main",
+			expected:  true,
+		},
+		{
+			name:      "ref matches stage options ref",
+			stageOpts: repository.StageOptions{Ref: "feature"},
+			gitBranch: "main",
+			ref:       "feature",
+			expected:  true,
+		},
+		{
+			name:      "ref matches stage options ref when empty defaults to git branch",
+			stageOpts: repository.StageOptions{Ref: ""},
+			gitBranch: "main",
+			ref:       "main",
+			expected:  true,
+		},
+		{
+			name:      "unsupported ref",
+			stageOpts: repository.StageOptions{Ref: "feature"},
+			gitBranch: "main",
+			ref:       "other-branch",
+			expected:  false,
+		},
+		{
+			name:      "unsupported ref with empty stage options",
+			stageOpts: repository.StageOptions{},
+			gitBranch: "main",
+			ref:       "feature",
+			expected:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			stagedRepo := &stagedGitRepository{
+				gitRepository: &gitRepository{
+					gitConfig: RepositoryConfig{
+						Branch: tt.gitBranch,
+					},
+				},
+				opts: tt.stageOpts,
+			}
+
+			result := stagedRepo.isRefSupported(tt.ref)
+			require.Equal(t, tt.expected, result)
+		})
+	}
+}
+
 // Helper functions for creating test instances
 
 func createTestStagedRepository(mockClient *mocks.FakeClient) *stagedGitRepository {

pkg/registry/apis/provisioning/repository/staged.go

@@ -23,6 +23,10 @@ const (
 )
 
 type StageOptions struct {
+	// Ref custom ref
+	Ref string
+	// Push on every write
+	PushOnWrites bool
 	// Mode defines the staging and commit behavior
 	Mode StageMode
 	// Maximum time allowed for clone operation in seconds (0 means no limit)