From dfd58c213cdec612b3b19769fdd128122ffe08d4 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 2 Apr 2022 15:36:06 +0800
Subject: [PATCH] fix: model generation bug on with cache (#1743)

* fix: model generation bug on with cache

* chore: refine template

* chore: fix test failure
---
 tools/goctl/model/sql/gen/gen.go        | 9 ++++++---
 tools/goctl/model/sql/gen/gen_test.go   | 7 +++++--
 tools/goctl/model/sql/template/model.go | 9 +++++++--
 3 files changed, 18 insertions(+), 7 deletions(-)

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/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 {"))
 }
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 (