Merge pull request #120 from adamryman/mrCheckoutTracking

[#104] Support setting the remote during `lab mr checkout`

cmd/mrCheckout.go

@@ -5,13 +5,20 @@ import (
 	"log"
 
 	"github.com/spf13/cobra"
+	"github.com/tcnksm/go-gitconfig"
 	"github.com/xanzy/go-gitlab"
 	"github.com/zaquestion/lab/internal/git"
 	lab "github.com/zaquestion/lab/internal/gitlab"
 )
 
-var (
+// mrCheckoutConfig holds configuration values for calls to lab mr checkout
+type mrCheckoutConfig struct {
 	branch string
+	track  bool
+}
+
+var (
+	mrCheckoutCfg mrCheckoutConfig
 )
 
 // listCmd represents the list command
@@ -36,26 +43,56 @@ var checkoutCmd = &cobra.Command{
 			fmt.Printf("MR #%d not found\n", mrID)
 			return
 		}
-		// https://docs.gitlab.com/ee/user/project/merge_requests/#checkout-merge-requests-locally
-		if branch == "" {
-			branch = mrs[0].SourceBranch
+
+		mr := mrs[0]
+		// If the config does not specify a branch, use the mr source branch name
+		if mrCheckoutCfg.branch == "" {
+			mrCheckoutCfg.branch = mr.SourceBranch
 		}
-		mr := fmt.Sprintf("refs/merge-requests/%d/head", mrID)
-		gitf := git.New("fetch", forkedFromRemote, fmt.Sprintf("%s:%s", mr, branch))
-		err = gitf.Run()
-		if err != nil {
+		// By default, fetch to configured branch
+		fetchToRef := mrCheckoutCfg.branch
+
+		// If track, make sure we have a remote for the mr author and then set
+		// the fetchToRef to the mr author/sourceBranch
+		if mrCheckoutCfg.track {
+			// Check if remote already exists
+			if _, err := gitconfig.Local("remote." + mr.Author.Username + ".url"); err != nil {
+				// Find and create remote
+				mrProject, err := lab.GetProject(mr.ProjectID)
+				if err != nil {
+					log.Fatal(err)
+				}
+				if err := git.RemoteAdd(mr.Author.Username, mrProject.SSHURLToRepo, "."); err != nil {
+					log.Fatal(err)
+				}
+			}
+			fetchToRef = fmt.Sprintf("refs/remotes/%s/%s", mr.Author.Username, mr.SourceBranch)
+		}
+
+		// https://docs.gitlab.com/ce/user/project/merge_requests/#checkout-merge-requests-locally
+		mrRef := fmt.Sprintf("refs/merge-requests/%d/head", mrID)
+		fetchRefSpec := fmt.Sprintf("%s:%s", mrRef, fetchToRef)
+		if err := git.New("fetch", forkedFromRemote, fetchRefSpec).Run(); err != nil {
 			log.Fatal(err)
 		}
 
-		gitc := git.New("checkout", branch)
-		err = gitc.Run()
-		if err != nil {
+		if mrCheckoutCfg.track {
+			// Create configured branch with tracking from fetchToRef
+			// git branch --flags <branchname> [<start-point>]
+			if err := git.New("branch", "--track", mrCheckoutCfg.branch, fetchToRef).Run(); err != nil {
+				log.Fatal(err)
+			}
+		}
+
+		// Check out branch
+		if err := git.New("checkout", mrCheckoutCfg.branch).Run(); err != nil {
 			log.Fatal(err)
 		}
 	},
 }
 
 func init() {
-	checkoutCmd.Flags().StringVarP(&branch, "branch", "b", "", "checkout merge request with <branch> name")
+	checkoutCmd.Flags().StringVarP(&mrCheckoutCfg.branch, "branch", "b", "", "checkout merge request with <branch> name")
+	checkoutCmd.Flags().BoolVarP(&mrCheckoutCfg.track, "track", "t", false, "set checked out branch to track mr author remote branch, adds remote if needed")
 	mrCmd.AddCommand(checkoutCmd)
 }

internal/gitlab/gitlab.go

@@ -81,6 +81,18 @@ var (
 	localProjects map[string]*gitlab.Project = make(map[string]*gitlab.Project)
 )
 
+// GetProject looks up a Gitlab project by ID.
+func GetProject(projectID int) (*gitlab.Project, error) {
+	target, resp, err := lab.Projects.GetProject(projectID)
+	if resp != nil && resp.StatusCode == http.StatusNotFound {
+		return nil, ErrProjectNotFound
+	}
+	if err != nil {
+		return nil, err
+	}
+	return target, nil
+}
+
 // FindProject looks up the Gitlab project. If the namespace is not provided in
 // the project string it will search for projects in the users namespace
 func FindProject(project string) (*gitlab.Project, error) {