From 02f9ea49c7dc3a42fc313d49d2d668b88dae7346 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8C=85=E5=AD=90?= <baozhecheng@foxmail.com>
Date: Tue, 1 Jun 2021 18:07:39 +0800
Subject: [PATCH] add cmd specify the version or branch #977 (#986)

* add cmd specify the version or branch
---
 cmd/kratos/internal/base/repo.go       | 36 ++++++++++++++++++++------
 cmd/kratos/internal/base/repo_test.go  |  2 +-
 cmd/kratos/internal/project/new.go     |  5 ++--
 cmd/kratos/internal/project/project.go |  4 ++-
 4 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/cmd/kratos/internal/base/repo.go b/cmd/kratos/internal/base/repo.go
index 8e23243d2..c00ba6678 100644
--- a/cmd/kratos/internal/base/repo.go
+++ b/cmd/kratos/internal/base/repo.go
@@ -1,6 +1,7 @@
 package base
 
 import (
+	"bytes"
 	"context"
 	"os"
 	"os/exec"
@@ -10,12 +11,13 @@ import (
 
 // Repo is git repository manager.
 type Repo struct {
-	url  string
-	home string
+	url    string
+	home   string
+	branch string
 }
 
 // NewRepo new a repository manager.
-func NewRepo(url string) *Repo {
+func NewRepo(url string, branch string) *Repo {
 	var start int
 	start = strings.Index(url, "//")
 	if start == -1 {
@@ -25,8 +27,9 @@ func NewRepo(url string) *Repo {
 	}
 	end := strings.LastIndex(url, "/")
 	return &Repo{
-		url:  url,
-		home: kratosHomeWithDir("repo/" + url[start:end]),
+		url:    url,
+		home:   kratosHomeWithDir("repo/" + url[start:end]),
+		branch: branch,
 	}
 }
 
@@ -37,15 +40,26 @@ func (r *Repo) Path() string {
 	if end == -1 {
 		end = len(r.url)
 	}
-	return path.Join(r.home, r.url[start+1:end])
+	branch := ""
+	if r.branch == "" {
+		branch = "@main"
+	} else {
+		branch = "@" + r.branch
+	}
+	return path.Join(r.home, r.url[start+1:end]+branch)
 }
 
 // Pull fetch the repository from remote url.
 func (r *Repo) Pull(ctx context.Context) error {
 	cmd := exec.Command("git", "pull")
 	cmd.Dir = r.Path()
-	cmd.Stderr = os.Stderr
+	var out bytes.Buffer
+	cmd.Stderr = &out
+	cmd.Stdout = os.Stdout
 	err := cmd.Run()
+	if strings.Contains(out.String(), "You are not currently on a branch.") {
+		return nil
+	}
 	return err
 }
 
@@ -54,8 +68,14 @@ func (r *Repo) Clone(ctx context.Context) error {
 	if _, err := os.Stat(r.Path()); !os.IsNotExist(err) {
 		return r.Pull(ctx)
 	}
-	cmd := exec.Command("git", "clone", r.url, r.Path())
+	cmd := &exec.Cmd{}
+	if r.branch == "" {
+		cmd = exec.Command("git", "clone", r.url, r.Path())
+	} else {
+		cmd = exec.Command("git", "clone", "-b", r.branch, r.url, r.Path())
+	}
 	cmd.Stderr = os.Stderr
+	cmd.Stdout = os.Stdout
 	err := cmd.Run()
 	return err
 }
diff --git a/cmd/kratos/internal/base/repo_test.go b/cmd/kratos/internal/base/repo_test.go
index 3ab488163..c3decaef7 100644
--- a/cmd/kratos/internal/base/repo_test.go
+++ b/cmd/kratos/internal/base/repo_test.go
@@ -6,7 +6,7 @@ import (
 )
 
 func TestRepo(t *testing.T) {
-	r := NewRepo("https://github.com/go-kratos/service-layout.git")
+	r := NewRepo("https://github.com/go-kratos/service-layout.git", "")
 	if err := r.Clone(context.Background()); err != nil {
 		t.Fatal(err)
 	}
diff --git a/cmd/kratos/internal/project/new.go b/cmd/kratos/internal/project/new.go
index daa4ddcb5..5c5197b6d 100644
--- a/cmd/kratos/internal/project/new.go
+++ b/cmd/kratos/internal/project/new.go
@@ -18,7 +18,8 @@ type Project struct {
 }
 
 // New new a project from remote repo.
-func (p *Project) New(ctx context.Context, dir string, layout string) error {
+func (p *Project) New(ctx context.Context, dir string, layout string, branch string) error {
+
 	to := path.Join(dir, p.Name)
 	if _, err := os.Stat(to); !os.IsNotExist(err) {
 		fmt.Printf("🚫 %s already exists\n", p.Name)
@@ -34,7 +35,7 @@ func (p *Project) New(ctx context.Context, dir string, layout string) error {
 		os.RemoveAll(to)
 	}
 	fmt.Printf("🚀 Creating service %s, layout repo is %s, please wait a moment.\n\n", p.Name, layout)
-	repo := base.NewRepo(layout)
+	repo := base.NewRepo(layout, branch)
 	if err := repo.CopyTo(ctx, to, p.Name, []string{".git", ".github"}); err != nil {
 		return err
 	}
diff --git a/cmd/kratos/internal/project/project.go b/cmd/kratos/internal/project/project.go
index f36be8cb5..01cf0c8af 100644
--- a/cmd/kratos/internal/project/project.go
+++ b/cmd/kratos/internal/project/project.go
@@ -19,12 +19,14 @@ var CmdNew = &cobra.Command{
 }
 
 var repoURL string
+var branch string
 
 func init() {
 	if repoURL = os.Getenv("KRATOS_LAYOUT_REPO"); repoURL == "" {
 		repoURL = "https://github.com/go-kratos/kratos-layout.git"
 	}
 	CmdNew.Flags().StringVarP(&repoURL, "-repo-url", "r", repoURL, "layout repo")
+	CmdNew.Flags().StringVarP(&branch, "-branch", "b", branch, "repo branch")
 }
 
 func run(cmd *cobra.Command, args []string) {
@@ -48,7 +50,7 @@ func run(cmd *cobra.Command, args []string) {
 		name = args[0]
 	}
 	p := &Project{Name: name}
-	if err := p.New(ctx, wd, repoURL); err != nil {
+	if err := p.New(ctx, wd, repoURL, branch); err != nil {
 		fmt.Fprintf(os.Stderr, "\033[31mERROR: %s\033[m\n", err)
 		return
 	}