Repository: zeromicro/go-zero
Pull request #1800 (2022-04-18 01:30:12)

feat(goctl): go work multi-module support
Resolve: #1793


Patch:
diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index 9095d7d3..f3e2a7a2 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -681,7 +681,7 @@ func (s *Redis) HdelCtx(ctx context.Context, key string, fields ...string) (val
 			return err
 		}
 
-		val = v == 1
+		val = v >= 1
 		return nil
 	}, acceptable)
 
@@ -1219,7 +1219,7 @@ func (s *Redis) PfaddCtx(ctx context.Context, key string, values ...interface{})
 			return err
 		}
 
-		val = v == 1
+		val = v >= 1
 		return nil
 	}, acceptable)
 
diff --git a/rest/handler/timeouthandler.go b/rest/handler/timeouthandler.go
index 672f44b6..ca4b4a49 100644
--- a/rest/handler/timeouthandler.go
+++ b/rest/handler/timeouthandler.go
@@ -20,6 +20,8 @@ import (
 const (
 	statusClientClosedRequest = 499
 	reason                    = "Request Timeout"
+	headerUpgrade             = "Upgrade"
+	valueWebsocket            = "websocket"
 )
 
 // TimeoutHandler returns the handler with given timeout.
@@ -52,6 +54,11 @@ func (h *timeoutHandler) errorBody() string {
 }
 
 func (h *timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get(headerUpgrade) == valueWebsocket {
+		h.handler.ServeHTTP(w, r)
+		return
+	}
+
 	ctx, cancelCtx := context.WithTimeout(r.Context(), h.dt)
 	defer cancelCtx()
 
diff --git a/tools/goctl/model/sql/command/command.go b/tools/goctl/model/sql/command/command.go
index 6e484c6d..326be929 100644
--- a/tools/goctl/model/sql/command/command.go
+++ b/tools/goctl/model/sql/command/command.go
@@ -11,6 +11,7 @@ import (
 	"github.com/zeromicro/go-zero/core/stores/postgres"
 	"github.com/zeromicro/go-zero/core/stores/sqlx"
 	"github.com/zeromicro/go-zero/tools/goctl/config"
+	"github.com/zeromicro/go-zero/tools/goctl/model/sql/command/migrationnotes"
 	"github.com/zeromicro/go-zero/tools/goctl/model/sql/gen"
 	"github.com/zeromicro/go-zero/tools/goctl/model/sql/model"
 	"github.com/zeromicro/go-zero/tools/goctl/model/sql/util"
@@ -38,6 +39,7 @@ var errNotMatched = errors.New("sql not matched")
 
 // MysqlDDL generates model code from ddl
 func MysqlDDL(ctx *cli.Context) error {
+	migrationnotes.BeforeCommands(ctx)
 	src := ctx.String(flagSrc)
 	dir := ctx.String(flagDir)
 	cache := ctx.Bool(flagCache)
@@ -66,6 +68,7 @@ func MysqlDDL(ctx *cli.Context) error {
 
 // MySqlDataSource generates model code from datasource
 func MySqlDataSource(ctx *cli.Context) error {
+	migrationnotes.BeforeCommands(ctx)
 	url := strings.TrimSpace(ctx.String(flagURL))
 	dir := strings.TrimSpace(ctx.String(flagDir))
 	cache := ctx.Bool(flagCache)
@@ -95,6 +98,7 @@ func MySqlDataSource(ctx *cli.Context) error {
 
 // PostgreSqlDataSource generates model code from datasource
 func PostgreSqlDataSource(ctx *cli.Context) error {
+	migrationnotes.BeforeCommands(ctx)
 	url := strings.TrimSpace(ctx.String(flagURL))
 	dir := strings.TrimSpace(ctx.String(flagDir))
 	cache := ctx.Bool(flagCache)
diff --git a/tools/goctl/model/sql/command/migrationnotes/migrationnotes.go b/tools/goctl/model/sql/command/migrationnotes/migrationnotes.go
new file mode 100644
index 00000000..d8baa62e
--- /dev/null
+++ b/tools/goctl/model/sql/command/migrationnotes/migrationnotes.go
@@ -0,0 +1,27 @@
+package migrationnotes
+
+import (
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/config"
+	"github.com/zeromicro/go-zero/tools/goctl/util/format"
+)
+
+// BeforeCommands run before comamnd run to show some migration notes
+func BeforeCommands(ctx *cli.Context) error {
+	if err := migrateBefore1_3_4(ctx); err != nil {
+		return err
+	}
+	return nil
+}
+
+func getModelSuffix(style string) (string, error) {
+	cfg, err := config.NewConfig(style)
+	if err != nil {
+		return "", err
+	}
+	baseSuffix, err := format.FileNamingFormat(cfg.NamingFormat, "_model")
+	if err != nil {
+		return "", err
+	}
+	return baseSuffix + ".go", nil
+}
diff --git a/tools/goctl/model/sql/command/migrationnotes/v1.3.4.go b/tools/goctl/model/sql/command/migrationnotes/v1.3.4.go
new file mode 100644
index 00000000..2a99e142
--- /dev/null
+++ b/tools/goctl/model/sql/command/migrationnotes/v1.3.4.go
@@ -0,0 +1,57 @@
+package migrationnotes
+
+import (
+	"fmt"
+	"io/ioutil"
+	"strings"
+
+	"github.com/urfave/cli"
+)
+
+func migrateBefore1_3_4(ctx *cli.Context) error {
+	dir := ctx.String("dir")
+	style := ctx.String("style")
+	ok, err := needShow1_3_4(dir, style)
+	if err != nil {
+		return err
+	}
+	if !ok {
+		return nil
+	}
+
+	fmt.Println(`It seems like that your goctl has just been upgraded to version 1.3.4 or later, which refactored the code of the model module. The original XXXmodel.go has been split into XXXmodel_gen.go (read-only) and XXXmodel.go. You just need to follow these steps to complete the migration:
+1. back up the original XXXmodel.go (make sure the file name is no longer in the current directory)
+2. re-run the generate command (a new XXXmodel.go will be created)
+3. populate XXXmodel.go with the code that is not generated by goctl according to the comments in XXXmodel_gen.go`)
+
+	return nil
+}
+
+func needShow1_3_4(dir, style string) (bool, error) {
+	files, err := ioutil.ReadDir(dir)
+	if err != nil {
+		return false, nil
+	}
+	// Returns false when the directory contains a file with the suffix "_gen.go"
+	// In addition, it returns true if it contains a model file extension.
+	// In other case, false is returned.
+	for _, f := range files {
+		if f.IsDir() {
+			continue
+		}
+		if strings.HasSuffix(f.Name(), "_gen.go") {
+			return false, nil
+		}
+	}
+	modelSuffix, err := getModelSuffix(style)
+	if err != nil {
+		return false, err
+	}
+	for _, f := range files {
+		if !f.IsDir() && strings.HasSuffix(f.Name(), modelSuffix) {
+			return true, nil
+		}
+	}
+	return false, nil
+
+}
diff --git a/tools/goctl/rpc/execx/execx.go b/tools/goctl/rpc/execx/execx.go
index 5b5e10bb..3e87923f 100644
--- a/tools/goctl/rpc/execx/execx.go
+++ b/tools/goctl/rpc/execx/execx.go
@@ -12,6 +12,9 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/vars"
 )
 
+// RunFunc defines a function type of Run function
+type RunFunc func(string, string, ...*bytes.Buffer) (string, error)
+
 // Run provides the execution of shell scripts in golang,
 // which can support macOS, Windows, and Linux operating systems.
 // Other operating systems are currently not supported
diff --git a/tools/goctl/util/ctx/context.go b/tools/goctl/util/ctx/context.go
index 25f8d5a8..1eba01f6 100644
--- a/tools/goctl/util/ctx/context.go
+++ b/tools/goctl/util/ctx/context.go
@@ -2,6 +2,7 @@ package ctx
 
 import (
 	"errors"
+	"fmt"
 	"path/filepath"
 
 	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
@@ -31,6 +32,7 @@ func Prepare(workDir string) (*ProjectContext, error) {
 	if err == nil {
 		return ctx, nil
 	}
+	fmt.Printf("get project context from workdir[%s] failed: %s\n", workDir, err)
 
 	name := filepath.Base(workDir)
 	_, err = execx.Run("go mod init "+name, workDir)
diff --git a/tools/goctl/util/ctx/gomod.go b/tools/goctl/util/ctx/gomod.go
index a81f7347..4c526b09 100644
--- a/tools/goctl/util/ctx/gomod.go
+++ b/tools/goctl/util/ctx/gomod.go
@@ -1,11 +1,14 @@
 package ctx
 
 import (
+	"encoding/json"
 	"errors"
+	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"strings"
 
-	"github.com/zeromicro/go-zero/core/jsonx"
 	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
 	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
 )
@@ -36,16 +39,11 @@ func projectFromGoMod(workDir string) (*ProjectContext, error) {
 		return nil, err
 	}
 
-	data, err := execx.Run("go list -json -m", workDir)
+	m, err := getRealModule(workDir, execx.Run)
 	if err != nil {
 		return nil, err
 	}
 
-	var m Module
-	err = jsonx.Unmarshal([]byte(data), &m)
-	if err != nil {
-		return nil, err
-	}
 	var ret ProjectContext
 	ret.WorkDir = workDir
 	ret.Name = filepath.Base(m.Dir)
@@ -58,3 +56,34 @@ func projectFromGoMod(workDir string) (*ProjectContext, error) {
 	ret.Path = m.Path
 	return &ret, nil
 }
+
+func getRealModule(workDir string, execRun execx.RunFunc) (*Module, error) {
+	data, err := execRun("go list -json -m", workDir)
+	if err != nil {
+		return nil, err
+	}
+	modules, err := decodePackages(strings.NewReader(data))
+	if err != nil {
+		return nil, err
+	}
+	for _, m := range modules {
+		if strings.HasPrefix(workDir, m.Dir) {
+			return &m, nil
+		}
+	}
+	return nil, errors.New("no matched module")
+}
+
+func decodePackages(rc io.Reader) ([]Module, error) {
+	var modules []Module
+	decoder := json.NewDecoder(rc)
+	for decoder.More() {
+		var m Module
+		if err := decoder.Decode(&m); err != nil {
+			return nil, fmt.Errorf("invalid module: %v", err)
+		}
+		modules = append(modules, m)
+	}
+
+	return modules, nil
+}
diff --git a/tools/goctl/util/ctx/modcheck.go b/tools/goctl/util/ctx/modcheck.go
index 9f74c04f..b2f37d2d 100644
--- a/tools/goctl/util/ctx/modcheck.go
+++ b/tools/goctl/util/ctx/modcheck.go
@@ -4,7 +4,6 @@ import (
 	"errors"
 	"os"
 
-	"github.com/zeromicro/go-zero/core/jsonx"
 	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
 )
 
@@ -17,16 +16,10 @@ func IsGoMod(workDir string) (bool, error) {
 		return false, err
 	}
 
-	data, err := execx.Run("go list -json -m", workDir)
-	if err != nil {
+	data, err := execx.Run("go list -m -f '{{.GoMod}}'", workDir)
+	if err != nil || len(data) == 0 {
 		return false, nil
 	}
 
-	var m Module
-	err = jsonx.Unmarshal([]byte(data), &m)
-	if err != nil {
-		return false, err
-	}
-
-	return len(m.GoMod) > 0, nil
+	return true, nil
 }
diff --git a/zrpc/internal/client.go b/zrpc/internal/client.go
index 251f43f4..5e6c75ba 100644
--- a/zrpc/internal/client.go
+++ b/zrpc/internal/client.go
@@ -12,6 +12,7 @@ import (
 	"github.com/zeromicro/go-zero/zrpc/resolver"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 const (
@@ -68,7 +69,7 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 
 	var options []grpc.DialOption
 	if !cliOpts.Secure {
-		options = append([]grpc.DialOption(nil), grpc.WithInsecure())
+		options = append([]grpc.DialOption(nil), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	}
 
 	if !cliOpts.NonBlock {


Test patch:
diff --git a/rest/handler/timeouthandler_test.go b/rest/handler/timeouthandler_test.go
index 74c70687..d4d2ad20 100644
--- a/rest/handler/timeouthandler_test.go
+++ b/rest/handler/timeouthandler_test.go
@@ -79,6 +79,19 @@ func TestTimeoutPanic(t *testing.T) {
 	})
 }
 
+func TestTimeoutWebsocket(t *testing.T) {
+	timeoutHandler := TimeoutHandler(time.Millisecond)
+	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Millisecond * 10)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.Header.Set(headerUpgrade, valueWebsocket)
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusOK, resp.Code)
+}
+
 func TestTimeoutWroteHeaderTwice(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Minute)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
diff --git a/tools/goctl/model/sql/command/migrationnotes/v1.3.4_test.go b/tools/goctl/model/sql/command/migrationnotes/v1.3.4_test.go
new file mode 100644
index 00000000..4d77bbfc
--- /dev/null
+++ b/tools/goctl/model/sql/command/migrationnotes/v1.3.4_test.go
@@ -0,0 +1,75 @@
+package migrationnotes
+
+import (
+	"io/fs"
+	"io/ioutil"
+	"os"
+	"path"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_needShow1_3_4(t *testing.T) {
+	root, err := os.MkdirTemp("", "goctl-model")
+	require.NoError(t, err)
+	defer os.RemoveAll(root)
+
+	dir1 := path.Join(root, "dir1")
+	require.NoError(t, os.Mkdir(dir1, fs.ModePerm))
+	ioutil.WriteFile(filepath.Join(dir1, "foo_gen.go"), nil, fs.ModePerm)
+
+	dir2 := path.Join(root, "dir2")
+	require.NoError(t, os.Mkdir(dir2, fs.ModePerm))
+	ioutil.WriteFile(filepath.Join(dir2, "foomodel.go"), nil, fs.ModePerm)
+
+	dir3 := path.Join(root, "dir3")
+	require.NoError(t, os.Mkdir(dir3, fs.ModePerm))
+	ioutil.WriteFile(filepath.Join(dir3, "irrelevant.go"), nil, fs.ModePerm)
+
+	type args struct {
+		dir   string
+		style string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    bool
+		wantErr bool
+	}{
+		{
+			name: "dir that contains *_gen.go should return false",
+			args: args{
+				dir: dir1,
+			},
+			want: false,
+		},
+		{
+			name: "dir that contains *model.go without *_gen.go should return true",
+			args: args{
+				dir: dir2,
+			},
+			want: true,
+		},
+		{
+			name: "dir that only contains irrelevant files should return false",
+			args: args{
+				dir: dir3,
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := needShow1_3_4(tt.args.dir, tt.args.style)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("needShow1_3_4() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("needShow1_3_4() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/tools/goctl/util/ctx/gomod_test.go b/tools/goctl/util/ctx/gomod_test.go
index 6a433352..11c419ca 100644
--- a/tools/goctl/util/ctx/gomod_test.go
+++ b/tools/goctl/util/ctx/gomod_test.go
@@ -1,9 +1,11 @@
 package ctx
 
 import (
+	"bytes"
 	"go/build"
 	"os"
 	"path/filepath"
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -36,3 +38,75 @@ func TestProjectFromGoMod(t *testing.T) {
 	assert.Equal(t, projectName, ctx.Path)
 	assert.Equal(t, dir, ctx.Dir)
 }
+
+func Test_getRealModule(t *testing.T) {
+	type args struct {
+		workDir string
+		execRun execx.RunFunc
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    *Module
+		wantErr bool
+	}{
+		{
+			name: "single module",
+			args: args{
+				workDir: "/home/foo",
+				execRun: func(arg, dir string, in ...*bytes.Buffer) (string, error) {
+					return `{
+						"Path":"foo",
+						"Dir":"/home/foo",
+						"GoMod":"/home/foo/go.mod",
+						"GoVersion":"go1.16"}`, nil
+				},
+			},
+			want: &Module{
+				Path:      "foo",
+				Dir:       "/home/foo",
+				GoMod:     "/home/foo/go.mod",
+				GoVersion: "go1.16",
+			},
+		},
+		{
+			name: "go work multiple modules",
+			args: args{
+				workDir: "/home/bar",
+				execRun: func(arg, dir string, in ...*bytes.Buffer) (string, error) {
+					return `
+					{
+						"Path":"foo",
+						"Dir":"/home/foo",
+						"GoMod":"/home/foo/go.mod",
+						"GoVersion":"go1.18"
+					}
+					{
+						"Path":"bar",
+						"Dir":"/home/bar",
+						"GoMod":"/home/bar/go.mod",
+						"GoVersion":"go1.18"
+					}`, nil
+				},
+			},
+			want: &Module{
+				Path:      "bar",
+				Dir:       "/home/bar",
+				GoMod:     "/home/bar/go.mod",
+				GoVersion: "go1.18",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := getRealModule(tt.args.workDir, tt.args.execRun)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getRealModule() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("getRealModule() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}