Repository: zeromicro/go-zero
Pull request #1743 (2022-04-02 06:18:50)

fix: model generation bug on with cache


Patch:
diff --git a/tools/goctl/internal/version/version.go b/tools/goctl/internal/version/version.go
index 02b4aae5..ff9e96ba 100644
--- a/tools/goctl/internal/version/version.go
+++ b/tools/goctl/internal/version/version.go
@@ -6,7 +6,7 @@ import (
 )
 
 // BuildVersion is the version of goctl.
-const BuildVersion = "1.3.3"
+const BuildVersion = "1.3.4"
 
 var tag = map[string]int{"pre-alpha": 0, "alpha": 1, "pre-bata": 2, "beta": 3, "released": 4, "": 5}
 
diff --git a/tools/goctl/model/sql/gen/gen.go b/tools/goctl/model/sql/gen/gen.go
index 6c50f4f8..29194225 100644
--- a/tools/goctl/model/sql/gen/gen.go
+++ b/tools/goctl/model/sql/gen/gen.go
@@ -123,7 +123,7 @@ func (g *defaultGenerator) StartFromInformationSchema(tables map[string]*model.T
 		if err != nil {
 			return err
 		}
-		customCode, err := g.genModelCustom(*table)
+		customCode, err := g.genModelCustom(*table, withCache)
 		if err != nil {
 			return err
 		}
@@ -214,7 +214,7 @@ func (g *defaultGenerator) genFromDDL(filename string, withCache bool, database
 		if err != nil {
 			return nil, err
 		}
-		customCode, err := g.genModelCustom(*e)
+		customCode, err := g.genModelCustom(*e, withCache)
 		if err != nil {
 			return nil, err
 		}
@@ -325,22 +325,25 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	return output.String(), nil
 }
 
-func (g *defaultGenerator) genModelCustom(in parser.Table) (string, error) {
+func (g *defaultGenerator) genModelCustom(in parser.Table, withCache bool) (string, error) {
 	text, err := pathx.LoadTemplate(category, modelCustomTemplateFile, template.ModelCustom)
 	if err != nil {
 		return "", err
 	}
+
 	t := util.With("model-custom").
 		Parse(text).
 		GoFmt(true)
 	output, err := t.Execute(map[string]interface{}{
 		"pkg":                   g.pkg,
+		"withCache":             withCache,
 		"upperStartCamelObject": in.Name.ToCamel(),
 		"lowerStartCamelObject": stringx.From(in.Name.ToCamel()).Untitle(),
 	})
 	if err != nil {
 		return "", err
 	}
+
 	return output.String(), nil
 }
 
diff --git a/tools/goctl/model/sql/template/model.go b/tools/goctl/model/sql/template/model.go
index 71f763c6..c134c5b9 100644
--- a/tools/goctl/model/sql/template/model.go
+++ b/tools/goctl/model/sql/template/model.go
@@ -8,9 +8,14 @@ import (
 
 // ModelCustom defines a template for extension
 const ModelCustom = `package {{.pkg}}
-
+{{if .withCache}}
+import (
+	"github.com/zeromicro/go-zero/core/stores/cache"
+	"github.com/zeromicro/go-zero/core/stores/sqlx"
+)
+{{else}}
 import "github.com/zeromicro/go-zero/core/stores/sqlx"
-
+{{end}}
 var _ {{.upperStartCamelObject}}Model = (*custom{{.upperStartCamelObject}}Model)(nil)
 
 type (


Test patch:
diff --git a/tools/goctl/model/sql/gen/gen_test.go b/tools/goctl/model/sql/gen/gen_test.go
index 5faa6f88..8c228bcd 100644
--- a/tools/goctl/model/sql/gen/gen_test.go
+++ b/tools/goctl/model/sql/gen/gen_test.go
@@ -145,7 +145,10 @@ func Test_genPublicModel(t *testing.T) {
 	tables, err := parser.Parse(modelFilename, "")
 	require.Equal(t, 1, len(tables))
 
-	code, err := g.genModelCustom(*tables[0])
+	code, err := g.genModelCustom(*tables[0], false)
 	assert.NoError(t, err)
-	assert.Equal(t, "package model\n\ntype TestUserModel interface {\n\ttestUserModel\n}\n", code)
+	assert.True(t, strings.Contains(code, "package model"))
+	assert.True(t, strings.Contains(code, "TestUserModel interface {\n\t\ttestUserModel\n\t}\n"))
+	assert.True(t, strings.Contains(code, "customTestUserModel struct {\n\t\t*defaultTestUserModel\n\t}\n"))
+	assert.True(t, strings.Contains(code, "func NewTestUserModel(conn sqlx.SqlConn) TestUserModel {"))
 }