Merge branch 'manifest_default_branch' into 'master'

Set default branch from manifest on restore

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/6482

Merged-by: Toon Claes <toon@gitlab.com>
Approved-by: Toon Claes <toon@gitlab.com>
Co-authored-by: James Fargher <jfargher@gitlab.com>

3 files changed, 143 insertions, 37 deletions

diff --git a/internal/backup/backup.go b/internal/backup/backup.go
index fa15cce94..0ced2c4c1 100644
--- a/internal/backup/backup.go
+++ b/internal/backup/backup.go
@@ -49,6 +49,8 @@ type Backup struct {
 	Steps []Step `toml:"steps"`
 	// ObjectFormat is the name of the object hash used by the repository.
 	ObjectFormat string `toml:"object_format"`
+	// HeadReference is the reference that HEAD points to.
+	HeadReference string `toml:"head_reference,omitempty"`
 }
 
 // Step represents an incremental step that makes up a complete backup for a repository
@@ -98,14 +100,16 @@ type Repository interface {
 	// repository cannot be found.
 	Remove(ctx context.Context) error
 	// Create creates the repository.
-	Create(ctx context.Context, hash git.ObjectHash) error
+	Create(ctx context.Context, hash git.ObjectHash, defaultBranch string) error
 	// FetchBundle fetches references from a bundle. Refs will be mirrored to
 	// the repository.
-	FetchBundle(ctx context.Context, reader io.Reader) error
+	FetchBundle(ctx context.Context, reader io.Reader, updateHead bool) error
 	// SetCustomHooks updates the custom hooks for the repository.
 	SetCustomHooks(ctx context.Context, reader io.Reader) error
 	// ObjectHash detects the object hash used by the repository.
 	ObjectHash(ctx context.Context) (git.ObjectHash, error)
+	// HeadReference fetches the reference pointed to by HEAD.
+	HeadReference(ctx context.Context) (git.ReferenceName, error)
 }
 
 // ResolveLocator returns a locator implementation based on a locator identifier.
@@ -235,6 +239,12 @@ func (mgr *Manager) Create(ctx context.Context, req *CreateRequest) error {
 	}
 	backup.ObjectFormat = hash.Format
 
+	headRef, err := repo.HeadReference(ctx)
+	if err != nil {
+		return fmt.Errorf("manager: %w", err)
+	}
+	backup.HeadReference = headRef.String()
+
 	refs, err := repo.ListRefs(ctx)
 	if err != nil {
 		return fmt.Errorf("manager: %w", err)
@@ -293,7 +303,9 @@ func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
 		return fmt.Errorf("manager: %w", err)
 	}
 
-	if err := repo.Create(ctx, hash); err != nil {
+	defaultBranch, defaultBranchKnown := git.ReferenceName(backup.HeadReference).Branch()
+
+	if err := repo.Create(ctx, hash, defaultBranch); err != nil {
 		return fmt.Errorf("manager: %w", err)
 	}
 
@@ -323,7 +335,7 @@ func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
 		// Git bundles can not be created for empty repositories. Since empty
 		// repository backups do not contain a bundle, skip bundle restoration.
 		if len(refs) > 0 {
-			if err := mgr.restoreBundle(ctx, repo, step.BundlePath); err != nil {
+			if err := mgr.restoreBundle(ctx, repo, step.BundlePath, !defaultBranchKnown); err != nil {
 				return fmt.Errorf("manager: %w", err)
 			}
 		}
@@ -470,14 +482,14 @@ func (mgr *Manager) readRefs(ctx context.Context, path string) ([]git.Reference,
 	return refs, nil
 }
 
-func (mgr *Manager) restoreBundle(ctx context.Context, repo Repository, path string) error {
+func (mgr *Manager) restoreBundle(ctx context.Context, repo Repository, path string, updateHead bool) error {
 	reader, err := mgr.sink.GetReader(ctx, path)
 	if err != nil {
 		return fmt.Errorf("restore bundle: %q: %w", path, err)
 	}
 	defer reader.Close()
 
-	if err := repo.FetchBundle(ctx, reader); err != nil {
+	if err := repo.FetchBundle(ctx, reader, updateHead); err != nil {
 		return fmt.Errorf("restore bundle: %q: %w", path, err)
 	}
 	return nil
diff --git a/internal/backup/backup_test.go b/internal/backup/backup_test.go
index ab4acdd08..14341ec6a 100644
--- a/internal/backup/backup_test.go
+++ b/internal/backup/backup_test.go
@@ -134,9 +134,10 @@ func TestManager_Create(t *testing.T) {
 						repo:     repo,
 						repoPath: repoPath,
 						expectedBackup: &backup.Backup{
-							ID:           backupID,
-							Repository:   vanityRepo,
-							ObjectFormat: gittest.DefaultObjectHash.Format,
+							ID:            backupID,
+							Repository:    vanityRepo,
+							ObjectFormat:  gittest.DefaultObjectHash.Format,
+							HeadReference: git.DefaultRef.String(),
 							Steps: []backup.Step{
 								{
 									BundlePath:      joinBackupPath(t, "", vanityRepo, backupID, "001.bundle"),
@@ -163,9 +164,10 @@ func TestManager_Create(t *testing.T) {
 						repo:     repo,
 						repoPath: repoPath,
 						expectedBackup: &backup.Backup{
-							ID:           backupID,
-							Repository:   vanityRepo,
-							ObjectFormat: gittest.DefaultObjectHash.Format,
+							ID:            backupID,
+							Repository:    vanityRepo,
+							ObjectFormat:  gittest.DefaultObjectHash.Format,
+							HeadReference: git.DefaultRef.String(),
 							Steps: []backup.Step{
 								{
 									BundlePath:      joinBackupPath(t, "", vanityRepo, backupID, "001.bundle"),
@@ -189,9 +191,10 @@ func TestManager_Create(t *testing.T) {
 						repo:     emptyRepo,
 						repoPath: repoPath,
 						expectedBackup: &backup.Backup{
-							ID:           backupID,
-							Repository:   vanityRepo,
-							ObjectFormat: gittest.DefaultObjectHash.Format,
+							ID:            backupID,
+							Repository:    vanityRepo,
+							ObjectFormat:  gittest.DefaultObjectHash.Format,
+							HeadReference: git.DefaultRef.String(),
 							Steps: []backup.Step{
 								{
 									BundlePath:      joinBackupPath(t, "", vanityRepo, backupID, "001.bundle"),
@@ -799,8 +802,6 @@ func TestManager_Restore_latest(t *testing.T) {
 }
 
 func TestManager_Restore_specific(t *testing.T) {
-	gittest.SkipWithSHA256(t)
-
 	t.Parallel()
 
 	const backupID = "abc123"
@@ -865,17 +866,18 @@ func TestManager_Restore_specific(t *testing.T) {
 			backupRoot := testhelper.TempDir(t)
 
 			for _, tc := range []struct {
-				desc          string
-				setup         func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
-				alwaysCreate  bool
-				expectExists  bool
-				expectedPaths []string
-				expectedErrAs error
+				desc                  string
+				setup                 func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
+				alwaysCreate          bool
+				expectExists          bool
+				expectedPaths         []string
+				expectedErrAs         error
+				expectedHeadReference git.ReferenceName
 			}{
 				{
 					desc: "missing backup",
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						repo, _ := gittest.CreateRepository(tb, ctx, cfg)
 
 						return repo, nil
 					},
@@ -884,7 +886,9 @@ func TestManager_Restore_specific(t *testing.T) {
 				{
 					desc: "single incremental",
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						gittest.SkipWithSHA256(tb) // sha256 only works with manifest files
+
+						repo, _ := gittest.CreateRepository(tb, ctx, cfg)
 
 						relativePath := stripRelativePath(tb, repo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
@@ -901,9 +905,11 @@ func TestManager_Restore_specific(t *testing.T) {
 				{
 					desc: "many incrementals",
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						_, expectedRepoPath := gittest.CreateRepository(t, ctx, cfg)
+						gittest.SkipWithSHA256(tb) // sha256 only works with manifest files
 
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						_, expectedRepoPath := gittest.CreateRepository(tb, ctx, cfg)
+
+						repo, _ := gittest.CreateRepository(tb, ctx, cfg)
 
 						root := gittest.WriteCommit(tb, cfg, expectedRepoPath,
 							gittest.WithBranch("master"),
@@ -922,7 +928,7 @@ func TestManager_Restore_specific(t *testing.T) {
 							"refs/heads/master",
 							"refs/heads/other",
 						)
-						refs1 := gittest.Exec(t, cfg, "-C", expectedRepoPath, "show-ref", "--head")
+						refs1 := gittest.Exec(tb, cfg, "-C", expectedRepoPath, "show-ref", "--head")
 
 						master2 := gittest.WriteCommit(tb, cfg, expectedRepoPath,
 							gittest.WithBranch("master"),
@@ -935,7 +941,7 @@ func TestManager_Restore_specific(t *testing.T) {
 							"refs/heads/master",
 							"refs/heads/other",
 						)
-						refs2 := gittest.Exec(t, cfg, "-C", expectedRepoPath, "show-ref", "--head")
+						refs2 := gittest.Exec(tb, cfg, "-C", expectedRepoPath, "show-ref", "--head")
 
 						relativePath := stripRelativePath(tb, repo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
@@ -956,6 +962,59 @@ func TestManager_Restore_specific(t *testing.T) {
 					},
 					expectExists: true,
 				},
+				{
+					desc: "manifest, empty backup",
+					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
+						repo, _ := gittest.CreateRepository(tb, ctx, cfg)
+
+						testhelper.WriteFiles(tb, backupRoot, map[string]any{
+							filepath.Join("manifests", repo.GetStorageName(), repo.GetRelativePath(), backupID+".toml"): fmt.Sprintf(
+								`object_format = %q
+head_reference = 'refs/heads/banana'
+
+[[steps]]
+bundle_path = 'repo.bundle'
+ref_path = 'repo.refs'
+custom_hooks_path = 'custom_hooks.tar'
+`, gittest.DefaultObjectHash.Format),
+							"repo.refs": "",
+						})
+
+						return repo, new(git.Checksum)
+					},
+					expectExists:          true,
+					expectedHeadReference: "refs/heads/banana",
+				},
+				{
+					desc: "manifest",
+					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
+						repo, _ := gittest.CreateRepository(tb, ctx, cfg)
+
+						testhelper.WriteFiles(tb, backupRoot, map[string]any{
+							filepath.Join("manifests", repo.GetStorageName(), repo.GetRelativePath(), backupID+".toml"): fmt.Sprintf(
+								`object_format = %q
+head_reference = 'refs/heads/banana'
+
+[[steps]]
+bundle_path = 'repo.bundle'
+ref_path = 'repo.refs'
+custom_hooks_path = 'custom_hooks.tar'
+`, gittest.DefaultObjectHash.Format),
+							"repo.bundle": repoBundle,
+							"repo.refs":   repoRefs,
+						})
+
+						checksum := gittest.ChecksumRepo(tb, cfg, repoPath)
+						// Negate off the default branch since the manifest is
+						// explicitly setting a different unborn branch that
+						// will not be part of the checksum.
+						checksum.Add(git.NewReference("HEAD", commitID))
+
+						return repo, checksum
+					},
+					expectExists:          true,
+					expectedHeadReference: "refs/heads/banana",
+				},
 			} {
 				t.Run(tc.desc, func(t *testing.T) {
 					repo, expectedChecksum := tc.setup(t)
@@ -1000,10 +1059,18 @@ func TestManager_Restore_specific(t *testing.T) {
 						// the repository through Praefect. In order to get to the correct disk paths, we need
 						// to get the replica path of the rewritten repository.
 						repoPath := filepath.Join(cfg.Storages[0].Path, gittest.GetReplicaPath(t, ctx, cfg, repo))
+
 						for _, p := range tc.expectedPaths {
 							require.FileExists(t, filepath.Join(repoPath, p))
 						}
 					}
+
+					if len(tc.expectedHeadReference) > 0 {
+						repoPath := filepath.Join(cfg.Storages[0].Path, gittest.GetReplicaPath(t, ctx, cfg, repo))
+
+						ref := gittest.GetSymbolicRef(t, cfg, repoPath, "HEAD")
+						require.Equal(t, tc.expectedHeadReference, git.ReferenceName(ref.Target))
+					}
 				})
 			}
 		})
diff --git a/internal/backup/repository.go b/internal/backup/repository.go
index 1f527234f..beef0c3d8 100644
--- a/internal/backup/repository.go
+++ b/internal/backup/repository.go
@@ -209,11 +209,12 @@ func (rr *remoteRepository) Remove(ctx context.Context) error {
 }
 
 // Create creates the repository.
-func (rr *remoteRepository) Create(ctx context.Context, hash git.ObjectHash) error {
+func (rr *remoteRepository) Create(ctx context.Context, hash git.ObjectHash, defaultBranch string) error {
 	repoClient := rr.newRepoClient()
 	if _, err := repoClient.CreateRepository(ctx, &gitalypb.CreateRepositoryRequest{
-		Repository:   rr.repo,
-		ObjectFormat: hash.ProtoFormat,
+		Repository:    rr.repo,
+		ObjectFormat:  hash.ProtoFormat,
+		DefaultBranch: []byte(defaultBranch),
 	}); err != nil {
 		return fmt.Errorf("remote repository: create: %w", err)
 	}
@@ -222,13 +223,13 @@ func (rr *remoteRepository) Create(ctx context.Context, hash git.ObjectHash) err
 
 // FetchBundle fetches references from a bundle. Refs will be mirrored to the
 // repository.
-func (rr *remoteRepository) FetchBundle(ctx context.Context, reader io.Reader) error {
+func (rr *remoteRepository) FetchBundle(ctx context.Context, reader io.Reader, updateHead bool) error {
 	repoClient := rr.newRepoClient()
 	stream, err := repoClient.FetchBundle(ctx)
 	if err != nil {
 		return fmt.Errorf("remote repository: fetch bundle: %w", err)
 	}
-	request := &gitalypb.FetchBundleRequest{Repository: rr.repo, UpdateHead: true}
+	request := &gitalypb.FetchBundleRequest{Repository: rr.repo, UpdateHead: updateHead}
 	bundle := streamio.NewWriter(func(p []byte) error {
 		request.Data = p
 		if err := stream.Send(request); err != nil {
@@ -292,6 +293,21 @@ func (rr *remoteRepository) ObjectHash(ctx context.Context) (git.ObjectHash, err
 	return git.ObjectHashByProto(response.Format)
 }
 
+// HeadReference returns the current value of HEAD.
+func (rr *remoteRepository) HeadReference(ctx context.Context) (git.ReferenceName, error) {
+	refClient := rr.newRefClient()
+
+	response, err := refClient.FindDefaultBranchName(ctx, &gitalypb.FindDefaultBranchNameRequest{
+		Repository: rr.repo,
+		HeadOnly:   true,
+	})
+	if err != nil {
+		return "", fmt.Errorf("remote repository: head reference: %w", err)
+	}
+
+	return git.ReferenceName(response.Name), nil
+}
+
 func (rr *remoteRepository) newRepoClient() gitalypb.RepositoryServiceClient {
 	return gitalypb.NewRepositoryServiceClient(rr.conn)
 }
@@ -388,7 +404,7 @@ func (r *localRepository) Remove(ctx context.Context) error {
 }
 
 // Create creates the repository.
-func (r *localRepository) Create(ctx context.Context, hash git.ObjectHash) error {
+func (r *localRepository) Create(ctx context.Context, hash git.ObjectHash, defaultBranch string) error {
 	if err := repoutil.Create(
 		ctx,
 		r.logger,
@@ -399,6 +415,7 @@ func (r *localRepository) Create(ctx context.Context, hash git.ObjectHash) error
 		r.repo,
 		func(repository *gitalypb.Repository) error { return nil },
 		repoutil.WithObjectHash(hash),
+		repoutil.WithBranchName(defaultBranch),
 	); err != nil {
 		return fmt.Errorf("local repository: create: %w", err)
 	}
@@ -407,9 +424,9 @@ func (r *localRepository) Create(ctx context.Context, hash git.ObjectHash) error
 
 // FetchBundle fetches references from a bundle. Refs will be mirrored to the
 // repository.
-func (r *localRepository) FetchBundle(ctx context.Context, reader io.Reader) error {
+func (r *localRepository) FetchBundle(ctx context.Context, reader io.Reader, updateHead bool) error {
 	err := r.repo.FetchBundle(ctx, r.txManager, reader, &localrepo.FetchBundleOpts{
-		UpdateHead: true,
+		UpdateHead: updateHead,
 	})
 	if err != nil {
 		return fmt.Errorf("local repository: fetch bundle: %w", err)
@@ -441,3 +458,13 @@ func (r *localRepository) ObjectHash(ctx context.Context) (git.ObjectHash, error
 
 	return hash, nil
 }
+
+// HeadReference returns the current value of HEAD.
+func (r *localRepository) HeadReference(ctx context.Context) (git.ReferenceName, error) {
+	head, err := r.repo.HeadReference(ctx)
+	if err != nil {
+		return "", fmt.Errorf("local repository: head reference: %w", err)
+	}
+
+	return head, nil
+}