实现API节点远程升级

internal/rpc/services/service_api_node.go

@@ -1,7 +1,6 @@
 package services
 
 import (
-	"bytes"
 	"compress/gzip"
 	"context"
 	"crypto/md5"
@@ -10,12 +9,13 @@ import (
 	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
 	rpcutils "github.com/TeaOSLab/EdgeAPI/internal/rpc/utils"
+	executils "github.com/TeaOSLab/EdgeAPI/internal/utils/exec"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/dbs"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"io"
 	"os"
-	"os/exec"
 	"path/filepath"
 	"runtime"
 )
@@ -437,16 +437,18 @@ func (this *APINodeService) UploadAPINodeFile(ctx context.Context, req *pb.Uploa
 		}
 
 		// 检查文件是否可执行
-		var versionCmd = exec.Command(targetFile, "-V")
-		var versionBuf = &bytes.Buffer{}
-		versionCmd.Stdout = versionBuf
+		var versionCmd = executils.NewCmd(targetFile, "-V").WithStdout().WithStderr()
 		err = versionCmd.Run()
 		if err != nil {
-			return nil, errors.New("test file failed: " + err.Error())
+			return nil, errors.New("test file failed: " + versionCmd.Stderr())
+		}
+		var newVersion = versionCmd.Stdout()
+		if len(newVersion) == 0 {
+			return nil, errors.New("test file failed, new version should not be empty")
 		}
 
 		// 检查版本
-		if stringutil.VersionCompare(versionCmd.String(), teaconst.Version) >= 0 {
+		if stringutil.VersionCompare(newVersion, teaconst.Version) <= 0 {
 			return &pb.UploadAPINodeFileResponse{}, nil
 		}
 
@@ -460,11 +462,20 @@ func (this *APINodeService) UploadAPINodeFile(ctx context.Context, req *pb.Uploa
 			return nil, errors.New("rename file failed: " + err.Error())
 		}
 
-		// 重启
-		var cmd = exec.Command(exe, "restart")
-		err = cmd.Start()
+		// 执行升级
+		if !Tea.IsTesting() { // 开发环境下防止破坏本地数据库
+			var upgradeCmd = executils.NewCmd(exe, "upgrade").WithStderr()
+			err = upgradeCmd.Run()
 			if err != nil {
-			return nil, errors.New("start new process failed: " + err.Error())
+				return nil, errors.New("execute 'upgrade' command failed: " + upgradeCmd.Stderr())
+			}
+		}
+
+		// 重启
+		var restartCmd = executils.NewCmd(exe, "restart").WithStderr()
+		err = restartCmd.Start()
+		if err != nil {
+			return nil, errors.New("start new process failed: " + restartCmd.Stderr())
 		}
 	}
 

internal/utils/exec/cmd.go

@@ -0,0 +1,162 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package executils
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"os/exec"
+	"strings"
+	"time"
+)
+
+type Cmd struct {
+	name string
+	args []string
+	env  []string
+	dir  string
+
+	ctx        context.Context
+	timeout    time.Duration
+	cancelFunc func()
+
+	captureStdout bool
+	captureStderr bool
+
+	stdout *bytes.Buffer
+	stderr *bytes.Buffer
+
+	rawCmd *exec.Cmd
+}
+
+func NewCmd(name string, args ...string) *Cmd {
+	return &Cmd{
+		name: name,
+		args: args,
+	}
+}
+
+func NewTimeoutCmd(timeout time.Duration, name string, args ...string) *Cmd {
+	return (&Cmd{
+		name: name,
+		args: args,
+	}).WithTimeout(timeout)
+}
+
+func (this *Cmd) WithTimeout(timeout time.Duration) *Cmd {
+	this.timeout = timeout
+
+	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
+	this.ctx = ctx
+	this.cancelFunc = cancelFunc
+
+	return this
+}
+
+func (this *Cmd) WithStdout() *Cmd {
+	this.captureStdout = true
+	return this
+}
+
+func (this *Cmd) WithStderr() *Cmd {
+	this.captureStderr = true
+	return this
+}
+
+func (this *Cmd) WithEnv(env []string) *Cmd {
+	this.env = env
+	return this
+}
+
+func (this *Cmd) WithDir(dir string) *Cmd {
+	this.dir = dir
+	return this
+}
+
+func (this *Cmd) Start() error {
+	var cmd = this.compose()
+	return cmd.Start()
+}
+
+func (this *Cmd) Wait() error {
+	var cmd = this.compose()
+	return cmd.Wait()
+}
+
+func (this *Cmd) Run() error {
+	if this.cancelFunc != nil {
+		defer this.cancelFunc()
+	}
+
+	var cmd = this.compose()
+	return cmd.Run()
+}
+
+func (this *Cmd) RawStdout() string {
+	if this.stdout != nil {
+		return this.stdout.String()
+	}
+	return ""
+}
+
+func (this *Cmd) Stdout() string {
+	return strings.TrimSpace(this.RawStdout())
+}
+
+func (this *Cmd) RawStderr() string {
+	if this.stderr != nil {
+		return this.stderr.String()
+	}
+	return ""
+}
+
+func (this *Cmd) Stderr() string {
+	return strings.TrimSpace(this.RawStderr())
+}
+
+func (this *Cmd) String() string {
+	if this.rawCmd != nil {
+		return this.rawCmd.String()
+	}
+	var newCmd = exec.Command(this.name, this.args...)
+	return newCmd.String()
+}
+
+func (this *Cmd) Process() *os.Process {
+	if this.rawCmd != nil {
+		return this.rawCmd.Process
+	}
+	return nil
+}
+
+func (this *Cmd) compose() *exec.Cmd {
+	if this.rawCmd != nil {
+		return this.rawCmd
+	}
+
+	if this.ctx != nil {
+		this.rawCmd = exec.CommandContext(this.ctx, this.name, this.args...)
+	} else {
+		this.rawCmd = exec.Command(this.name, this.args...)
+	}
+
+	if this.env != nil {
+		this.rawCmd.Env = this.env
+	}
+
+	if len(this.dir) > 0 {
+		this.rawCmd.Dir = this.dir
+	}
+
+	if this.captureStdout {
+		this.stdout = &bytes.Buffer{}
+		this.rawCmd.Stdout = this.stdout
+	}
+	if this.captureStderr {
+		this.stderr = &bytes.Buffer{}
+		this.rawCmd.Stderr = this.stderr
+	}
+
+	return this.rawCmd
+}

internal/utils/exec/cmd_test.go

@@ -0,0 +1,61 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package executils_test
+
+import (
+	executils "github.com/TeaOSLab/EdgeAPI/internal/utils/exec"
+	"testing"
+	"time"
+)
+
+func TestNewTimeoutCmd_Sleep(t *testing.T) {
+	var cmd = executils.NewTimeoutCmd(1*time.Second, "sleep", "3")
+	cmd.WithStdout()
+	cmd.WithStderr()
+	err := cmd.Run()
+	t.Log("error:", err)
+	t.Log("stdout:", cmd.Stdout())
+	t.Log("stderr:", cmd.Stderr())
+}
+
+func TestNewTimeoutCmd_Echo(t *testing.T) {
+	var cmd = executils.NewTimeoutCmd(10*time.Second, "echo", "-n", "hello")
+	cmd.WithStdout()
+	cmd.WithStderr()
+	err := cmd.Run()
+	t.Log("error:", err)
+	t.Log("stdout:", cmd.Stdout())
+	t.Log("stderr:", cmd.Stderr())
+}
+
+func TestNewTimeoutCmd_Echo2(t *testing.T) {
+	var cmd = executils.NewCmd("echo", "hello")
+	cmd.WithStdout()
+	cmd.WithStderr()
+	err := cmd.Run()
+	t.Log("error:", err)
+	t.Log("stdout:", cmd.Stdout())
+	t.Log("raw stdout:", cmd.RawStdout())
+	t.Log("stderr:", cmd.Stderr())
+	t.Log("raw stderr:", cmd.RawStderr())
+}
+
+func TestNewTimeoutCmd_Echo3(t *testing.T) {
+	var cmd = executils.NewCmd("echo", "-n", "hello")
+	err := cmd.Run()
+	t.Log("error:", err)
+	t.Log("stdout:", cmd.Stdout())
+	t.Log("stderr:", cmd.Stderr())
+}
+
+func TestCmd_Process(t *testing.T) {
+	var cmd = executils.NewCmd("echo", "-n", "hello")
+	err := cmd.Run()
+	t.Log("error:", err)
+	t.Log(cmd.Process())
+}
+
+func TestNewTimeoutCmd_String(t *testing.T) {
+	var cmd = executils.NewCmd("echo", "-n", "hello")
+	t.Log("stdout:", cmd.String())
+}