Repository: zeromicro/go-zero
Pull request #1738 (2022-04-01 13:45:43)

feat: let model customizable
Separate the model generation into `xxxmodel_gen.go` and `xxxmodel.go`.

Don't edit `xxxmodel_gen.go`, it will be overwritten on regeneration. And extend the sql operations on `xxxmodel.go`.

`xxxmodel.go` looks like:

<img width="623" alt="image" src="https://user-images.githubusercontent.com/1918356/161276159-bb52cb00-0002-4e54-aa4d-334d8a2f33dd.png">



Patch:
diff --git a/tools/goctl/model/sql/gen/gen.go b/tools/goctl/model/sql/gen/gen.go
index 0c34edb9..6c50f4f8 100644
--- a/tools/goctl/model/sql/gen/gen.go
+++ b/tools/goctl/model/sql/gen/gen.go
@@ -20,16 +20,13 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/stringx"
 )
 
-const (
-	pwd             = "."
-	createTableFlag = `(?m)^(?i)CREATE\s+TABLE` // ignore case
-)
+const pwd = "."
 
 type (
 	defaultGenerator struct {
-		// source string
-		dir string
 		console.Console
+		// source string
+		dir          string
 		pkg          string
 		cfg          *config.Config
 		isPostgreSql bool
@@ -48,6 +45,7 @@ type (
 		updateCode  string
 		deleteCode  string
 		cacheExtra  string
+		tableName   string
 	}
 
 	codeTuple struct {
@@ -154,7 +152,8 @@ func (g *defaultGenerator) createFile(modelList map[string]*codeTuple) error {
 
 	for tableName, codes := range modelList {
 		tn := stringx.From(tableName)
-		modelFilename, err := format.FileNamingFormat(g.cfg.NamingFormat, fmt.Sprintf("%s_model", tn.Source()))
+		modelFilename, err := format.FileNamingFormat(g.cfg.NamingFormat,
+			fmt.Sprintf("%s_model", tn.Source()))
 		if err != nil {
 			return err
 		}
@@ -202,7 +201,8 @@ func (g *defaultGenerator) createFile(modelList map[string]*codeTuple) error {
 }
 
 // ret1: key-table name,value-code
-func (g *defaultGenerator) genFromDDL(filename string, withCache bool, database string) (map[string]*codeTuple, error) {
+func (g *defaultGenerator) genFromDDL(filename string, withCache bool, database string) (
+	map[string]*codeTuple, error) {
 	m := make(map[string]*codeTuple)
 	tables, err := parser.Parse(filename, database)
 	if err != nil {
@@ -249,7 +249,7 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	table.UniqueCacheKey = uniqueKey
 	table.ContainsUniqueCacheKey = len(uniqueKey) > 0
 
-	importsCode, err := genImports(withCache, in.ContainsTime(), table)
+	importsCode, err := genImports(table, withCache, in.ContainsTime())
 	if err != nil {
 		return "", err
 	}
@@ -287,7 +287,8 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	}
 
 	var list []string
-	list = append(list, insertCodeMethod, findOneCodeMethod, ret.findOneInterfaceMethod, updateCodeMethod, deleteCodeMethod)
+	list = append(list, insertCodeMethod, findOneCodeMethod, ret.findOneInterfaceMethod,
+		updateCodeMethod, deleteCodeMethod)
 	typesCode, err := genTypes(table, strings.Join(modelutil.TrimStringSlice(list), pathx.NL), withCache)
 	if err != nil {
 		return "", err
@@ -298,6 +299,11 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 		return "", err
 	}
 
+	tableName, err := genTableName(table)
+	if err != nil {
+		return "", err
+	}
+
 	code := &code{
 		importsCode: importsCode,
 		varsCode:    varsCode,
@@ -308,6 +314,7 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 		updateCode:  updateCode,
 		deleteCode:  deleteCode,
 		cacheExtra:  ret.cacheExtra,
+		tableName:   tableName,
 	}
 
 	output, err := g.executeModel(table, code)
@@ -356,6 +363,7 @@ func (g *defaultGenerator) executeModel(table Table, code *code) (*bytes.Buffer,
 		"update":      code.updateCode,
 		"delete":      code.deleteCode,
 		"extraMethod": code.cacheExtra,
+		"tableName":   code.tableName,
 		"data":        table,
 	})
 	if err != nil {
diff --git a/tools/goctl/model/sql/gen/imports.go b/tools/goctl/model/sql/gen/imports.go
index 7fe98768..a9842c9c 100644
--- a/tools/goctl/model/sql/gen/imports.go
+++ b/tools/goctl/model/sql/gen/imports.go
@@ -6,7 +6,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
 )
 
-func genImports(withCache, timeImport bool, table Table) (string, error) {
+func genImports(table Table, withCache, timeImport bool) (string, error) {
 	if withCache {
 		text, err := pathx.LoadTemplate(category, importsTemplateFile, template.Imports)
 		if err != nil {
diff --git a/tools/goctl/model/sql/gen/tablename.go b/tools/goctl/model/sql/gen/tablename.go
new file mode 100644
index 00000000..a39ebb31
--- /dev/null
+++ b/tools/goctl/model/sql/gen/tablename.go
@@ -0,0 +1,26 @@
+package gen
+
+import (
+	"github.com/zeromicro/go-zero/tools/goctl/model/sql/template"
+	"github.com/zeromicro/go-zero/tools/goctl/util"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+func genTableName(table Table) (string, error) {
+	text, err := pathx.LoadTemplate(category, tableNameTemplateFile, template.TableName)
+	if err != nil {
+		return "", err
+	}
+
+	output, err := util.With("tableName").
+		Parse(text).
+		Execute(map[string]interface{}{
+			"tableName":             table.Name.Source(),
+			"upperStartCamelObject": table.Name.ToCamel(),
+		})
+	if err != nil {
+		return "", nil
+	}
+
+	return output.String(), nil
+}
diff --git a/tools/goctl/model/sql/gen/template.go b/tools/goctl/model/sql/gen/template.go
index cdb0f8cb..e2cecdba 100644
--- a/tools/goctl/model/sql/gen/template.go
+++ b/tools/goctl/model/sql/gen/template.go
@@ -25,6 +25,7 @@ const (
 	modelGenTemplateFile                  = "model-gen.tpl"
 	modelCustomTemplateFile               = "model.tpl"
 	modelNewTemplateFile                  = "model-new.tpl"
+	tableNameTemplateFile                 = "table-name.tpl"
 	tagTemplateFile                       = "tag.tpl"
 	typesTemplateFile                     = "types.tpl"
 	updateTemplateFile                    = "update.tpl"
@@ -49,6 +50,7 @@ var templates = map[string]string{
 	modelGenTemplateFile:                  template.ModelGen,
 	modelCustomTemplateFile:               template.ModelCustom,
 	modelNewTemplateFile:                  template.New,
+	tableNameTemplateFile:                 template.TableName,
 	tagTemplateFile:                       template.Tag,
 	typesTemplateFile:                     template.Types,
 	updateTemplateFile:                    template.Update,
@@ -72,7 +74,7 @@ func GenTemplates(_ *cli.Context) error {
 	return pathx.InitTemplates(category, templates)
 }
 
-// RevertTemplate recovers the delete template files
+// RevertTemplate reverts the deleted template files
 func RevertTemplate(name string) error {
 	content, ok := templates[name]
 	if !ok {
diff --git a/tools/goctl/model/sql/template/delete.go b/tools/goctl/model/sql/template/delete.go
index dea75966..f7701a8a 100644
--- a/tools/goctl/model/sql/template/delete.go
+++ b/tools/goctl/model/sql/template/delete.go
@@ -1,7 +1,8 @@
 package template
 
-// Delete defines a delete template
-var Delete = `
+const (
+	// Delete defines a delete template
+	Delete = `
 func (m *default{{.upperStartCamelObject}}Model) Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error {
 	{{if .withCache}}{{if .containsIndexCache}}data, err:=m.FindOne(ctx, {{.lowerStartCamelPrimaryKey}})
 	if err!=nil{
@@ -18,5 +19,6 @@ func (m *default{{.upperStartCamelObject}}Model) Delete(ctx context.Context, {{.
 }
 `
 
-// DeleteMethod defines a delete template for interface method
-var DeleteMethod = `Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error`
+	// DeleteMethod defines a delete template for interface method
+	DeleteMethod = `Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error`
+)
diff --git a/tools/goctl/model/sql/template/errors.go b/tools/goctl/model/sql/template/errors.go
index 0482453a..e9ccc31b 100644
--- a/tools/goctl/model/sql/template/errors.go
+++ b/tools/goctl/model/sql/template/errors.go
@@ -1,7 +1,7 @@
 package template
 
 // Error defines an error template
-var Error = `package {{.pkg}}
+const Error = `package {{.pkg}}
 
 import "github.com/zeromicro/go-zero/core/stores/sqlx"
 
diff --git a/tools/goctl/model/sql/template/field.go b/tools/goctl/model/sql/template/field.go
index bf927e24..ad695eb2 100644
--- a/tools/goctl/model/sql/template/field.go
+++ b/tools/goctl/model/sql/template/field.go
@@ -1,4 +1,4 @@
 package template
 
 // Field defines a filed template for types
-var Field = `{{.name}} {{.type}} {{.tag}} {{if .hasComment}}// {{.comment}}{{end}}`
+const Field = `{{.name}} {{.type}} {{.tag}} {{if .hasComment}}// {{.comment}}{{end}}`
diff --git a/tools/goctl/model/sql/template/find.go b/tools/goctl/model/sql/template/find.go
index aea44282..8154fad2 100644
--- a/tools/goctl/model/sql/template/find.go
+++ b/tools/goctl/model/sql/template/find.go
@@ -1,7 +1,8 @@
 package template
 
-// FindOne defines find row by id.
-var FindOne = `
+const (
+	// FindOne defines find row by id.
+	FindOne = `
 func (m *default{{.upperStartCamelObject}}Model) FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
@@ -30,8 +31,8 @@ func (m *default{{.upperStartCamelObject}}Model) FindOne(ctx context.Context, {{
 }
 `
 
-// FindOneByField defines find row by field.
-var FindOneByField = `
+	// FindOneByField defines find row by field.
+	FindOneByField = `
 func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
@@ -64,8 +65,8 @@ func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}(ctx co
 }{{end}}
 `
 
-// FindOneByFieldExtraMethod defines find row by field with extras.
-var FindOneByFieldExtraMethod = `
+	// FindOneByFieldExtraMethod defines find row by field with extras.
+	FindOneByFieldExtraMethod = `
 func (m *default{{.upperStartCamelObject}}Model) formatPrimary(primary interface{}) string {
 	return fmt.Sprintf("%s%v", {{.primaryKeyLeft}}, primary)
 }
@@ -76,8 +77,9 @@ func (m *default{{.upperStartCamelObject}}Model) queryPrimary(ctx context.Contex
 }
 `
 
-// FindOneMethod defines find row method.
-var FindOneMethod = `FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error)`
+	// FindOneMethod defines find row method.
+	FindOneMethod = `FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error)`
 
-// FindOneByFieldMethod defines find row by field method.
-var FindOneByFieldMethod = `FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) `
+	// FindOneByFieldMethod defines find row by field method.
+	FindOneByFieldMethod = `FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) `
+)
diff --git a/tools/goctl/model/sql/template/import.go b/tools/goctl/model/sql/template/import.go
index a52e3be1..0d92a818 100644
--- a/tools/goctl/model/sql/template/import.go
+++ b/tools/goctl/model/sql/template/import.go
@@ -1,6 +1,6 @@
 package template
 
-var (
+const (
 	// Imports defines a import template for model in cache case
 	Imports = `import (
 	"context"
diff --git a/tools/goctl/model/sql/template/insert.go b/tools/goctl/model/sql/template/insert.go
index 0339808c..6265b294 100644
--- a/tools/goctl/model/sql/template/insert.go
+++ b/tools/goctl/model/sql/template/insert.go
@@ -1,7 +1,8 @@
 package template
 
-// Insert defines a template for insert code in model
-var Insert = `
+const (
+	// Insert defines a template for insert code in model
+	Insert = `
 func (m *default{{.upperStartCamelObject}}Model) Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error) {
 	{{if .withCache}}{{.keys}}
     ret, err := m.ExecCtx(ctx, func(ctx context.Context, conn sqlx.SqlConn) (result sql.Result, err error) {
@@ -13,5 +14,6 @@ func (m *default{{.upperStartCamelObject}}Model) Insert(ctx context.Context, dat
 }
 `
 
-// InsertMethod defines an interface method template for insert code in model
-var InsertMethod = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error)`
+	// InsertMethod defines an interface method template for insert code in model
+	InsertMethod = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error)`
+)
diff --git a/tools/goctl/model/sql/template/model.go b/tools/goctl/model/sql/template/model.go
index d907d915..71f763c6 100644
--- a/tools/goctl/model/sql/template/model.go
+++ b/tools/goctl/model/sql/template/model.go
@@ -6,6 +6,33 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util"
 )
 
+// ModelCustom defines a template for extension
+const ModelCustom = `package {{.pkg}}
+
+import "github.com/zeromicro/go-zero/core/stores/sqlx"
+
+var _ {{.upperStartCamelObject}}Model = (*custom{{.upperStartCamelObject}}Model)(nil)
+
+type (
+	// {{.upperStartCamelObject}}Model is an interface to be customized, add more methods here,
+	// and implement the added methods in custom{{.upperStartCamelObject}}Model.
+	{{.upperStartCamelObject}}Model interface {
+		{{.lowerStartCamelObject}}Model
+	}
+
+	custom{{.upperStartCamelObject}}Model struct {
+		*default{{.upperStartCamelObject}}Model
+	}
+)
+
+// New{{.upperStartCamelObject}}Model returns a model for the database table.
+func New{{.upperStartCamelObject}}Model(conn sqlx.SqlConn{{if .withCache}}, c cache.CacheConf{{end}}) {{.upperStartCamelObject}}Model {
+	return &custom{{.upperStartCamelObject}}Model{
+		default{{.upperStartCamelObject}}Model: new{{.upperStartCamelObject}}Model(conn{{if .withCache}}, c{{end}}),
+	}
+}
+`
+
 // ModelGen defines a template for model
 var ModelGen = fmt.Sprintf(`%s
 
@@ -19,11 +46,5 @@ package {{.pkg}}
 {{.update}}
 {{.delete}}
 {{.extraMethod}}
+{{.tableName}}
 `, util.DoNotEditHead)
-
-// ModelCustom defines a template for extension
-var ModelCustom = fmt.Sprintf(`package {{.pkg}}
-type {{.upperStartCamelObject}}Model interface {
-	{{.lowerStartCamelObject}}Model
-}
-`)
diff --git a/tools/goctl/model/sql/template/new.go b/tools/goctl/model/sql/template/new.go
index e7c0657c..4bf89ff3 100644
--- a/tools/goctl/model/sql/template/new.go
+++ b/tools/goctl/model/sql/template/new.go
@@ -1,8 +1,8 @@
 package template
 
-// New defines an template for creating model instance
-var New = `
-func New{{.upperStartCamelObject}}Model(conn sqlx.SqlConn{{if .withCache}}, c cache.CacheConf{{end}}) {{.upperStartCamelObject}}Model {
+// New defines the template for creating model instance.
+const New = `
+func new{{.upperStartCamelObject}}Model(conn sqlx.SqlConn{{if .withCache}}, c cache.CacheConf{{end}}) *default{{.upperStartCamelObject}}Model {
 	return &default{{.upperStartCamelObject}}Model{
 		{{if .withCache}}CachedConn: sqlc.NewConn(conn, c){{else}}conn:conn{{end}},
 		table:      {{.table}},
diff --git a/tools/goctl/model/sql/template/tablename.go b/tools/goctl/model/sql/template/tablename.go
new file mode 100644
index 00000000..99210e3a
--- /dev/null
+++ b/tools/goctl/model/sql/template/tablename.go
@@ -0,0 +1,8 @@
+package template
+
+// TableName defines a template that generate the tableName method.
+const TableName = `
+func (m *default{{.upperStartCamelObject}}Model) tableName() string {
+	return m.table
+}
+`
diff --git a/tools/goctl/model/sql/template/tag.go b/tools/goctl/model/sql/template/tag.go
index c04baaaf..5dcaedc5 100644
--- a/tools/goctl/model/sql/template/tag.go
+++ b/tools/goctl/model/sql/template/tag.go
@@ -1,4 +1,4 @@
 package template
 
 // Tag defines a tag template text
-var Tag = "`db:\"{{.field}}\"`"
+const Tag = "`db:\"{{.field}}\"`"
diff --git a/tools/goctl/model/sql/template/types.go b/tools/goctl/model/sql/template/types.go
index 870890e7..0eeabeae 100644
--- a/tools/goctl/model/sql/template/types.go
+++ b/tools/goctl/model/sql/template/types.go
@@ -1,7 +1,7 @@
 package template
 
-// Types defines a template for types in model
-var Types = `
+// Types defines a template for types in model.
+const Types = `
 type (
 	{{.lowerStartCamelObject}}Model interface{
 		{{.method}}
diff --git a/tools/goctl/model/sql/template/update.go b/tools/goctl/model/sql/template/update.go
index 62edaee9..24e172bd 100644
--- a/tools/goctl/model/sql/template/update.go
+++ b/tools/goctl/model/sql/template/update.go
@@ -1,7 +1,8 @@
 package template
 
-// Update defines a template for generating update codes
-var Update = `
+const (
+	// Update defines a template for generating update codes
+	Update = `
 func (m *default{{.upperStartCamelObject}}Model) Update(ctx context.Context, data *{{.upperStartCamelObject}}) error {
 	{{if .withCache}}{{.keys}}
     _, err := m.ExecCtx(ctx, func(ctx context.Context, conn sqlx.SqlConn) (result sql.Result, err error) {
@@ -13,5 +14,6 @@ func (m *default{{.upperStartCamelObject}}Model) Update(ctx context.Context, dat
 }
 `
 
-// UpdateMethod defines an interface method template for generating update codes
-var UpdateMethod = `Update(ctx context.Context, data *{{.upperStartCamelObject}}) error`
+	// UpdateMethod defines an interface method template for generating update codes
+	UpdateMethod = `Update(ctx context.Context, data *{{.upperStartCamelObject}}) error`
+)
diff --git a/tools/goctl/model/sql/template/vars.go b/tools/goctl/model/sql/template/vars.go
index 8ef09f86..f5309cfa 100644
--- a/tools/goctl/model/sql/template/vars.go
+++ b/tools/goctl/model/sql/template/vars.go
@@ -5,8 +5,6 @@ import "fmt"
 // Vars defines a template for var block in model
 var Vars = fmt.Sprintf(`
 var (
-	_ {{.upperStartCamelObject}}Model = (*default{{.upperStartCamelObject}}Model)(nil)
-
 	{{.lowerStartCamelObject}}FieldNames          = builder.RawFieldNames(&{{.upperStartCamelObject}}{}{{if .postgreSql}},true{{end}})
 	{{.lowerStartCamelObject}}Rows                = strings.Join({{.lowerStartCamelObject}}FieldNames, ",")
 	{{.lowerStartCamelObject}}RowsExpectAutoSet   = {{if .postgreSql}}strings.Join(stringx.Remove({{.lowerStartCamelObject}}FieldNames, {{if .autoIncrement}}"{{.originalPrimaryKey}}",{{end}} "%screate_time%s", "%supdate_time%s"), ","){{else}}strings.Join(stringx.Remove({{.lowerStartCamelObject}}FieldNames, {{if .autoIncrement}}"{{.originalPrimaryKey}}",{{end}} "%screate_time%s", "%supdate_time%s"), ","){{end}}
diff --git a/tools/goctl/rpc/cli/zrpc.go b/tools/goctl/rpc/cli/zrpc.go
index fa84d8ea..2550f66c 100644
--- a/tools/goctl/rpc/cli/zrpc.go
+++ b/tools/goctl/rpc/cli/zrpc.go
@@ -118,13 +118,18 @@ func removeGoctlFlag(args []string) []string {
 	for step < len(args) {
 		arg := args[step]
 		switch {
-		case arg == "--style", arg == "--home", arg == "--zrpc_out", arg == "--verbose", arg == "-v":
+		case arg == "--style", arg == "--home",
+			arg == "--zrpc_out", arg == "--verbose",
+			arg == "-v", arg == "--remote",
+			arg == "--branch":
 			step += 2
 			continue
 		case strings.HasPrefix(arg, "--style="),
 			strings.HasPrefix(arg, "--home="),
 			strings.HasPrefix(arg, "--verbose="),
 			strings.HasPrefix(arg, "-v="),
+			strings.HasPrefix(arg, "--remote="),
+			strings.HasPrefix(arg, "--branch="),
 			strings.HasPrefix(arg, "--zrpc_out="):
 			step += 1
 			continue


Test patch:
diff --git a/tools/goctl/rpc/cli/zrpc_test.go b/tools/goctl/rpc/cli/zrpc_test.go
index 051556ad..d841115e 100644
--- a/tools/goctl/rpc/cli/zrpc_test.go
+++ b/tools/goctl/rpc/cli/zrpc_test.go
@@ -83,6 +83,14 @@ func Test_RemoveGoctlFlag(t *testing.T) {
 			source:   strings.Fields(`protoc --go_opt=. --go-grpc_out=. --zrpc_out=. foo.proto`),
 			expected: "protoc --go_opt=. --go-grpc_out=. foo.proto",
 		},
+		{
+			source:   strings.Fields(`protoc --go_opt=. --go-grpc_out=. --zrpc_out=. --remote=foo --branch=bar foo.proto`),
+			expected: "protoc --go_opt=. --go-grpc_out=. foo.proto",
+		},
+		{
+			source:   strings.Fields(`protoc --go_opt=. --go-grpc_out=. --zrpc_out=. --remote foo --branch bar foo.proto`),
+			expected: "protoc --go_opt=. --go-grpc_out=. foo.proto",
+		},
 	}
 	for _, e := range testData {
 		cmd := strings.Join(removeGoctlFlag(e.source), " ")