From d35e5de157ee43b3a707854cb0c1c479b308c859 Mon Sep 17 00:00:00 2001
From: Ningzio <ningzi0610@gmail.com>
Date: Sat, 31 Aug 2024 22:45:57 +0800
Subject: [PATCH] feat(goctl): support go work (#4332)

---
 tools/goctl/util/ctx/gomod.go       |  4 +++
 tools/goctl/util/ctx/gowork.go      | 33 ++++++++++++++++++++++++
 tools/goctl/util/ctx/gowork_test.go | 40 +++++++++++++++++++++++++++++
 tools/goctl/util/ctx/modcheck.go    |  6 ++---
 4 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 tools/goctl/util/ctx/gowork.go
 create mode 100644 tools/goctl/util/ctx/gowork_test.go

diff --git a/tools/goctl/util/ctx/gomod.go b/tools/goctl/util/ctx/gomod.go
index 15a301863a29..b7b131460850 100644
--- a/tools/goctl/util/ctx/gomod.go
+++ b/tools/goctl/util/ctx/gomod.go
@@ -51,6 +51,10 @@ func projectFromGoMod(workDir string) (*ProjectContext, error) {
 		return nil, err
 	}
 
+	if err := UpdateGoWorkIfExist(workDir); err != nil {
+		return nil, err
+	}
+
 	m, err := getRealModule(workDir, execx.Run)
 	if err != nil {
 		return nil, err
diff --git a/tools/goctl/util/ctx/gowork.go b/tools/goctl/util/ctx/gowork.go
new file mode 100644
index 000000000000..1caff9452e55
--- /dev/null
+++ b/tools/goctl/util/ctx/gowork.go
@@ -0,0 +1,33 @@
+package ctx
+
+import (
+	"errors"
+	"os"
+
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+)
+
+// UpdateGoWorkIfExist updates go work if workDir is in a go workspace
+func UpdateGoWorkIfExist(workDir string) error {
+	if isGoWork, err := isGoWork(workDir); !isGoWork || err != nil {
+		return err
+	}
+
+	_, err := execx.Run("go work use .", workDir)
+	return err
+}
+
+// isGoWork detect if the workDir is in a go workspace
+func isGoWork(workDir string) (bool, error) {
+	if len(workDir) == 0 {
+		return false, errors.New("the work directory is not found")
+	}
+	if _, err := os.Stat(workDir); err != nil {
+		return false, err
+	}
+	goWorkPath, err := execx.Run("go env GOWORK", workDir)
+	if err != nil {
+		return false, err
+	}
+	return len(goWorkPath) > 0, nil
+}
diff --git a/tools/goctl/util/ctx/gowork_test.go b/tools/goctl/util/ctx/gowork_test.go
new file mode 100644
index 000000000000..c8b13f654731
--- /dev/null
+++ b/tools/goctl/util/ctx/gowork_test.go
@@ -0,0 +1,40 @@
+package ctx
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/stringx"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+func Test_isGoWork(t *testing.T) {
+	dir := filepath.Join("/tmp", stringx.Rand())
+
+	err := pathx.MkdirIfNotExist(dir)
+	assert.Nil(t, err)
+
+	defer os.RemoveAll(dir)
+
+	gowork, err := isGoWork(dir)
+	assert.False(t, gowork)
+	assert.Nil(t, err)
+
+	_, err = execx.Run("go work init", dir)
+	assert.Nil(t, err)
+
+	gowork, err = isGoWork(dir)
+	assert.True(t, gowork)
+	assert.Nil(t, err)
+
+	subDir := filepath.Join(dir, stringx.Rand())
+	err = pathx.MkdirIfNotExist(subDir)
+	assert.Nil(t, err)
+
+	gowork, err = isGoWork(subDir)
+	assert.True(t, gowork)
+	assert.Nil(t, err)
+}
diff --git a/tools/goctl/util/ctx/modcheck.go b/tools/goctl/util/ctx/modcheck.go
index b2f37d2db760..eb5342b66f7a 100644
--- a/tools/goctl/util/ctx/modcheck.go
+++ b/tools/goctl/util/ctx/modcheck.go
@@ -7,7 +7,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
 )
 
-// IsGoMod is used to determine whether workDir is a go module project through command `go list -json -m`
+// IsGoMod is used to determine whether workDir is a go module project through command `go env GOMOD`
 func IsGoMod(workDir string) (bool, error) {
 	if len(workDir) == 0 {
 		return false, errors.New("the work directory is not found")
@@ -16,8 +16,8 @@ func IsGoMod(workDir string) (bool, error) {
 		return false, err
 	}
 
-	data, err := execx.Run("go list -m -f '{{.GoMod}}'", workDir)
-	if err != nil || len(data) == 0 {
+	data, err := execx.Run("go env GOMOD", workDir)
+	if err != nil || data == "/dev/null" {
 		return false, nil
 	}