From ab0cc6c4440fa89204b3a81ac0ce7fe30a81c82b Mon Sep 17 00:00:00 2001
From: Jose Diaz-Gonzalez <jose@seatgeek.com>
Date: Sat, 21 Nov 2020 20:37:54 -0500
Subject: [PATCH] fix: implement ps:restore

Also add all the niceties around parallelism that were introduced for the other commands.
---
 plugins/ps/ps.go                          | 41 +++++++++++++++++++++++
 plugins/ps/src/subcommands/subcommands.go |  4 ++-
 plugins/ps/subcommands.go                 | 25 ++++++++++++--
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/plugins/ps/ps.go b/plugins/ps/ps.go
index 81c52e8c2..0131fd368 100644
--- a/plugins/ps/ps.go
+++ b/plugins/ps/ps.go
@@ -92,6 +92,47 @@ func Restart(appName string) error {
 	return common.PlugnTrigger("release-and-deploy", []string{appName}...)
 }
 
+// Restore ensures an app that should be running is running on boot
+func Restore(appName string) error {
+	scheduler := common.GetAppScheduler(appName)
+	if err := common.PlugnTrigger("pre-restore", []string{scheduler, appName}...); err != nil {
+		return fmt.Errorf("Error running pre-restore: %s", err)
+	}
+
+	if err := common.PlugnTrigger("proxy-clear-config", []string{appName}...); err != nil {
+		return fmt.Errorf("Error clearing proxy config: %s", err)
+	}
+
+	if !common.IsDeployed(appName) {
+		common.LogWarn(fmt.Sprintf("App %s has not been deployed", appName))
+		return nil
+	}
+
+	b, _ := common.PlugnTriggerOutput("config-get", []string{appName, "DOKKU_APP_RESTORE"}...)
+	restore := strings.TrimSpace(string(b[:]))
+	if restore == "0" {
+		return nil
+	}
+
+	return Start(appName)
+}
+
+func RestorePrep() error {
+	apps, err := common.DokkuApps()
+	if err != nil {
+		common.LogWarn(err.Error())
+		return nil
+	}
+
+	for _, appName := range apps {
+		if err := common.PlugnTrigger("proxy-clear-config", []string{appName}...); err != nil {
+			return fmt.Errorf("Error clearing proxy config: %s", err)
+		}
+	}
+
+	return nil
+}
+
 // Start starts the app
 func Start(appName string) error {
 	imageTag, _ := common.GetRunningImageTag(appName)
diff --git a/plugins/ps/src/subcommands/subcommands.go b/plugins/ps/src/subcommands/subcommands.go
index a1041cbc7..6eec6e0ae 100644
--- a/plugins/ps/src/subcommands/subcommands.go
+++ b/plugins/ps/src/subcommands/subcommands.go
@@ -44,9 +44,11 @@ func main() {
 		err = ps.CommandRestart(appName, *allApps, *parallelCount)
 	case "restore":
 		args := flag.NewFlagSet("ps:restore", flag.ExitOnError)
+		allApps := args.Bool("all", false, "--all: restore all apps")
+		parallelCount := args.Int("parallel", ps.RunInSerial, "--parallel: number of apps to restore in parallel, -1 to match cpu count")
 		args.Parse(os.Args[2:])
 		appName := args.Arg(0)
-		err = ps.CommandRestore(appName)
+		err = ps.CommandRestore(appName, *allApps, *parallelCount)
 	case "retire":
 		args := flag.NewFlagSet("ps:retire", flag.ExitOnError)
 		args.Parse(os.Args[2:])
diff --git a/plugins/ps/subcommands.go b/plugins/ps/subcommands.go
index bda148ff8..0eb22e0c6 100644
--- a/plugins/ps/subcommands.go
+++ b/plugins/ps/subcommands.go
@@ -83,8 +83,29 @@ func CommandRestart(appName string, allApps bool, parallelCount int) error {
 
 // CommandRestore starts previously running apps e.g. after reboot
 // TODO: implement me
-func CommandRestore(appName string) error {
-	return nil
+func CommandRestore(appName string, allApps bool, parallelCount int) error {
+	if allApps {
+		if err := RestorePrep(); err != nil {
+			return err
+		}
+
+		return common.RunCommandAgainstAllApps(Restore, "restore", parallelCount)
+	}
+
+	if appName == "" {
+		common.LogWarn("Restore specified without app, assuming --all")
+
+		if err := RestorePrep(); err != nil {
+			return err
+		}
+		return common.RunCommandAgainstAllApps(Restore, "restore", parallelCount)
+	}
+
+	if err := common.VerifyAppName(appName); err != nil {
+		return err
+	}
+
+	return Restore(appName)
 }
 
 // CommandRetire ensures old containers are retired