From a918a3d763fad6e11f5b5bd92289c8d26c53ab65 Mon Sep 17 00:00:00 2001
From: lijianguo <lijianguo1991@outlook.com>
Date: Fri, 24 May 2024 15:15:21 +0800
Subject: [PATCH] chore(ps): restart process types

---
 cmd/cmd.go        |  2 +-
 cmd/ps.go         | 20 +++++++++++++++++---
 cmd/ps_test.go    | 15 +++++++++------
 go.mod            |  2 +-
 go.sum            |  4 ++--
 parser/ps.go      | 10 ++++++----
 parser/ps_test.go |  2 +-
 7 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 49a6f4d9..8c11f4ff 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -87,7 +87,7 @@ type Commander interface {
 	PsLogs(string, string, int, bool, string) error
 	PsExec(string, string, bool, bool, []string) error
 	PsScale(string, []string) error
-	PsRestart(string, string) error
+	PsRestart(string, []string, string) error
 	PsDescribe(string, string) error
 	RegistryList(string) error
 	RegistrySet(string, []string) error
diff --git a/cmd/ps.go b/cmd/ps.go
index 3a5284f4..dd45e02e 100644
--- a/cmd/ps.go
+++ b/cmd/ps.go
@@ -138,17 +138,31 @@ func (d *DryccCmd) PsScale(appID string, targets []string) error {
 }
 
 // PsRestart restarts an app's processes.
-func (d *DryccCmd) PsRestart(appID, target string) error {
+func (d *DryccCmd) PsRestart(appID string, targets []string, confirm string) error {
 	s, appID, err := load(d.ConfigFile, appID)
 	if err != nil {
 		return err
 	}
+	if len(targets) == 0 && confirm == "" {
+		d.Printf(` !    WARNING: Potentially Restart Action
+ !    This command will restart all processes of the application
+ !    To proceed, type "yes" !
 
+> `)
+
+		fmt.Scanln(&confirm)
+		if confirm != "yes" {
+			return fmt.Errorf("cancel the restart action")
+		}
+	}
 	d.Printf("Restarting processes... but first, %s!\n", drinkOfChoice())
 	startTime := time.Now()
 	quit := progress(d.WOut)
-
-	err = ps.Restart(s.Client, appID, target)
+	ptypes := strings.Join(targets, ",")
+	targetMap := map[string]string{
+		"types": ptypes,
+	}
+	err = ps.Restart(s.Client, appID, targetMap)
 	quit <- true
 	<-quit
 	if err != nil {
diff --git a/cmd/ps_test.go b/cmd/ps_test.go
index e2da889d..d8097d17 100644
--- a/cmd/ps_test.go
+++ b/cmd/ps_test.go
@@ -221,33 +221,36 @@ func TestPsRestart(t *testing.T) {
 	var b bytes.Buffer
 	cmdr := DryccCmd{WOut: &b, ConfigFile: cf}
 
-	server.Mux.HandleFunc("/v2/apps/foo/pods/restart/", func(w http.ResponseWriter, _ *http.Request) {
+	server.Mux.HandleFunc("/v2/apps/foo/pods/restart/", func(w http.ResponseWriter, r *http.Request) {
+		testutil.AssertBody(t, map[string]string{"types": ""}, r)
 		testutil.SetHeaders(w)
 		w.WriteHeader(http.StatusNoContent)
 	})
 
 	b.Reset()
-	err = cmdr.PsRestart("foo", "")
+	err = cmdr.PsRestart("foo", []string{""}, "yes")
 	assert.NoError(t, err)
 
-	server.Mux.HandleFunc("/v2/apps/coolapp/pods/restart/", func(w http.ResponseWriter, _ *http.Request) {
+	server.Mux.HandleFunc("/v2/apps/coolapp/pods/restart/", func(w http.ResponseWriter, r *http.Request) {
+		testutil.AssertBody(t, map[string]string{"types": "web"}, r)
 		testutil.SetHeaders(w)
 		w.WriteHeader(http.StatusNoContent)
 	})
 
 	b.Reset()
 
-	err = cmdr.PsRestart("coolapp", "")
+	err = cmdr.PsRestart("coolapp", []string{"web"}, "")
 	assert.NoError(t, err)
 
-	server.Mux.HandleFunc("/v2/apps/testapp/pods/web/restart/", func(w http.ResponseWriter, _ *http.Request) {
+	server.Mux.HandleFunc("/v2/apps/testapp/pods/restart/", func(w http.ResponseWriter, r *http.Request) {
+		testutil.AssertBody(t, map[string]string{"types": "web,worker"}, r)
 		testutil.SetHeaders(w)
 		w.WriteHeader(http.StatusNoContent)
 	})
 
 	b.Reset()
 
-	err = cmdr.PsRestart("testapp", "web")
+	err = cmdr.PsRestart("testapp", []string{"web", "worker"}, "")
 	assert.NoError(t, err)
 }
 
diff --git a/go.mod b/go.mod
index 538c887e..50cc6e93 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.22
 require (
 	github.com/containerd/console v1.0.4
 	github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815
-	github.com/drycc/controller-sdk-go v0.0.0-20240523052610-656dcb6b9db3
+	github.com/drycc/controller-sdk-go v0.0.0-20240524075643-99bda554d60c
 	github.com/drycc/pkg v0.0.0-20240225112316-78fc9239f51f
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/stretchr/testify v1.9.0
diff --git a/go.sum b/go.sum
index bb36be61..7803d436 100644
--- a/go.sum
+++ b/go.sum
@@ -4,8 +4,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815 h1:bWDMxwH3px2JBh6AyO7hdCn/PkvCZXii8TGj7sbtEbQ=
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
-github.com/drycc/controller-sdk-go v0.0.0-20240523052610-656dcb6b9db3 h1:eITOqdzcBH/QQVDD/zwHyq5JY1+xXr3MsPg066uR4fk=
-github.com/drycc/controller-sdk-go v0.0.0-20240523052610-656dcb6b9db3/go.mod h1:n6eQe1irJqjwLo/7t9+Dhdv6faSESQN+ATnZRBP3/Uc=
+github.com/drycc/controller-sdk-go v0.0.0-20240524075643-99bda554d60c h1:wxbZlPUsQfs7a3TnfAhdJ37kS3Gy56WkQ1k0keJJs24=
+github.com/drycc/controller-sdk-go v0.0.0-20240524075643-99bda554d60c/go.mod h1:n6eQe1irJqjwLo/7t9+Dhdv6faSESQN+ATnZRBP3/Uc=
 github.com/drycc/pkg v0.0.0-20240225112316-78fc9239f51f h1:kgjvUQJeAszDoU1Vo4vTTE92KI8Av3JPb6Qn890niXg=
 github.com/drycc/pkg v0.0.0-20240225112316-78fc9239f51f/go.mod h1:n+QxGif6ha9CEoxVnlipxb9IdmerybcUSzTEDFkvjiA=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
diff --git a/parser/ps.go b/parser/ps.go
index 4c5e1f46..175f5025 100644
--- a/parser/ps.go
+++ b/parser/ps.go
@@ -138,15 +138,17 @@ func psRestart(argv []string, cmdr cmd.Commander) error {
 	usage := `
 Restart an application or process type.
 
-Usage: drycc ps:restart [<type>] [options]
+Usage: drycc ps:restart [<type>...] [options]
 
 Arguments:
   <type>
-    the process name as defined in your Procfile, such as 'web' or 'worker'.
+    the process name as defined in your Procfile, such as 'web' or 'web worker'.
 
 Options:
   -a --app=<app>
     the uniquely identifiable name for the application.
+  --confirm=yes
+    To proceed, type "yes".
 `
 
 	args, err := docopt.ParseArgs(usage, argv, "")
@@ -156,8 +158,8 @@ Options:
 	}
 
 	apps := safeGetString(args, "--app")
-	tp := safeGetString(args, "<type>")
-	return cmdr.PsRestart(apps, tp)
+	confirm := safeGetString(args, "--confirm")
+	return cmdr.PsRestart(apps, args["<type>"].([]string), confirm)
 }
 
 func psScale(argv []string, cmdr cmd.Commander) error {
diff --git a/parser/ps_test.go b/parser/ps_test.go
index fff4e654..e700b2f0 100644
--- a/parser/ps_test.go
+++ b/parser/ps_test.go
@@ -28,7 +28,7 @@ func (d FakeDryccCmd) PsScale(string, []string) error {
 	return errors.New("ps:scale")
 }
 
-func (d FakeDryccCmd) PsRestart(string, string) error {
+func (d FakeDryccCmd) PsRestart(string, []string, string) error {
 	return errors.New("ps:restart")
 }
 
