feat: add --branch option to clone or checkout different dotfiles branch (#8331)

Emyrk · web-flow · commit 9f5bc7c10bb2 · 2023-07-06T20:24:04.000Z
* feat: --branch option to clone different dotfiles branch
* chore: checkout specified branch if dotfiles already exist
diff --git a/cli/dotfiles.go b/cli/dotfiles.go
@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -18,6 +19,8 @@ import (
 
 func (r *RootCmd) dotfiles() *clibase.Cmd {
 	var symlinkDir string
+	var gitbranch string
+
 	cmd := &clibase.Cmd{
 		Use:        "dotfiles <git_repo_url>",
 		Middleware: clibase.RequireNArgs(1),
@@ -102,6 +105,9 @@ func (r *RootCmd) dotfiles() *clibase.Cmd {
 				}
 				gitCmdDir = cfgDir
 				subcommands = []string{"clone", inv.Args[0], dotfilesRepoDir}
+				if gitbranch != "" {
+					subcommands = append(subcommands, "--branch", gitbranch)
+				}
 				promptText = fmt.Sprintf("Cloning %s into directory %s.\n\n  Continue?", gitRepo, dotfilesDir)
 			}
 
@@ -140,6 +146,23 @@ func (r *RootCmd) dotfiles() *clibase.Cmd {
 				_, _ = fmt.Fprintln(inv.Stdout, cliui.DefaultStyles.Error.Render("Failed to update repo, continuing..."))
 			}
 
+			if dotfilesExists && gitbranch != "" {
+				// If the repo exists and the git-branch is specified, we need to check out the branch. We do this after
+				// git pull to make sure the branch was pulled down locally. If we do this before the pull, we could be
+				// trying to checkout a branch that does not yet exist locally and get a git error.
+				_, _ = fmt.Fprintf(inv.Stdout, "Dotfiles git branch %q specified\n", gitbranch)
+				err := ensureCorrectGitBranch(inv, ensureCorrectGitBranchParams{
+					repoDir:       dotfilesDir,
+					gitSSHCommand: gitsshCmd,
+					gitBranch:     gitbranch,
+				})
+				if err != nil {
+					// Do not block on this error, just log it and continue
+					_, _ = fmt.Fprintln(inv.Stdout,
+						cliui.DefaultStyles.Error.Render(fmt.Sprintf("Failed to use branch %q (%s), continuing...", err.Error(), gitbranch)))
+				}
+			}
+
 			// save git repo url so we can detect changes next time
 			err = cfg.DotfilesURL().Write(gitRepo)
 			if err != nil {
@@ -246,11 +269,59 @@ func (r *RootCmd) dotfiles() *clibase.Cmd {
 			Description: "Specifies the directory for the dotfiles symlink destinations. If empty, will use $HOME.",
 			Value:       clibase.StringOf(&symlinkDir),
 		},
+		{
+			Flag:          "branch",
+			FlagShorthand: "b",
+			Description: "Specifies which branch to clone. " +
+				"If empty, will default to cloning the default branch or using the existing branch in the cloned repo on disk.",
+			Value: clibase.StringOf(&gitbranch),
+		},
 		cliui.SkipPromptOption(),
 	}
 	return cmd
 }
 
+type ensureCorrectGitBranchParams struct {
+	repoDir       string
+	gitSSHCommand string
+	gitBranch     string
+}
+
+func ensureCorrectGitBranch(baseInv *clibase.Invocation, params ensureCorrectGitBranchParams) error {
+	dotfileCmd := func(cmd string, args ...string) *exec.Cmd {
+		c := exec.CommandContext(baseInv.Context(), cmd, args...)
+		c.Dir = params.repoDir
+		c.Env = append(baseInv.Environ.ToOS(), fmt.Sprintf(`GIT_SSH_COMMAND=%s -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no`, params.gitSSHCommand))
+		c.Stdout = baseInv.Stdout
+		c.Stderr = baseInv.Stderr
+		return c
+	}
+	c := dotfileCmd("git", "branch", "--show-current")
+	// Save the output
+	var out bytes.Buffer
+	c.Stdout = &out
+	err := c.Run()
+	if err != nil {
+		return xerrors.Errorf("getting current git branch: %w", err)
+	}
+
+	if strings.TrimSpace(out.String()) != params.gitBranch {
+		// Checkout and pull the branch
+		c := dotfileCmd("git", "checkout", params.gitBranch)
+		err := c.Run()
+		if err != nil {
+			return xerrors.Errorf("checkout git branch %q: %w", params.gitBranch, err)
+		}
+
+		c = dotfileCmd("git", "pull", "--ff-only")
+		err = c.Run()
+		if err != nil {
+			return xerrors.Errorf("pull git branch %q: %w", params.gitBranch, err)
+		}
+	}
+	return nil
+}
+
 // dirExists checks if the path exists and is a directory.
 func dirExists(name string) (bool, error) {
 	fi, err := os.Stat(name)
diff --git a/cli/dotfiles_test.go b/cli/dotfiles_test.go
@@ -80,6 +80,52 @@ func TestDotfiles(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, string(b), "wow\n")
 	})
+	t.Run("InstallScriptChangeBranch", func(t *testing.T) {
+		t.Parallel()
+		if runtime.GOOS == "windows" {
+			t.Skip("install scripts on windows require sh and aren't very practical")
+		}
+		_, root := clitest.New(t)
+		testRepo := testGitRepo(t, root)
+
+		// We need an initial commit to start the `main` branch
+		c := exec.Command("git", "commit", "--allow-empty", "-m", `"initial commit"`)
+		c.Dir = testRepo
+		err := c.Run()
+		require.NoError(t, err)
+
+		// nolint:gosec
+		err = os.WriteFile(filepath.Join(testRepo, "install.sh"), []byte("#!/bin/bash\necho wow > "+filepath.Join(string(root), ".bashrc")), 0o750)
+		require.NoError(t, err)
+
+		c = exec.Command("git", "checkout", "-b", "other_branch")
+		c.Dir = testRepo
+		err = c.Run()
+		require.NoError(t, err)
+
+		c = exec.Command("git", "add", "install.sh")
+		c.Dir = testRepo
+		err = c.Run()
+		require.NoError(t, err)
+
+		c = exec.Command("git", "commit", "-m", `"add install.sh"`)
+		c.Dir = testRepo
+		err = c.Run()
+		require.NoError(t, err)
+
+		c = exec.Command("git", "checkout", "main")
+		c.Dir = testRepo
+		err = c.Run()
+		require.NoError(t, err)
+
+		inv, _ := clitest.New(t, "dotfiles", "--global-config", string(root), "--symlink-dir", string(root), "-y", testRepo, "-b", "other_branch")
+		err = inv.Run()
+		require.NoError(t, err)
+
+		b, err := os.ReadFile(filepath.Join(string(root), ".bashrc"))
+		require.NoError(t, err)
+		require.Equal(t, string(b), "wow\n")
+	})
 	t.Run("SymlinkBackup", func(t *testing.T) {
 		t.Parallel()
 		_, root := clitest.New(t)
@@ -152,5 +198,10 @@ func testGitRepo(t *testing.T, root config.Root) string {
 	err = c.Run()
 	require.NoError(t, err)
 
+	c = exec.Command("git", "checkout", "-b", "main")
+	c.Dir = dir
+	err = c.Run()
+	require.NoError(t, err)
+
 	return dir
 }
diff --git a/cli/testdata/coder_dotfiles_--help.golden b/cli/testdata/coder_dotfiles_--help.golden
@@ -7,6 +7,11 @@ Personalize your workspace by applying a canonical dotfiles repository
      [40m [0m[91;40m$ coder dotfiles --yes git@github.com:example/dotfiles.git[0m[40m [0m
 
 [1mOptions[0m
+  -b, --branch string
+          Specifies which branch to clone. If empty, will default to cloning the
+          default branch or using the existing branch in the cloned repo on
+          disk.
+
       --symlink-dir string, $CODER_SYMLINK_DIR
           Specifies the directory for the dotfiles symlink destinations. If
           empty, will use $HOME.
diff --git a/docs/cli/dotfiles.md b/docs/cli/dotfiles.md
@@ -20,6 +20,14 @@ coder dotfiles [flags] <git_repo_url>
 
 ## Options
 
+### -b, --branch
+
+|      |                     |
+| ---- | ------------------- |
+| Type | <code>string</code> |
+
+Specifies which branch to clone. If empty, will default to cloning the default branch or using the existing branch in the cloned repo on disk.
+
 ### --symlink-dir
 
 |             |                                 |
