feat: let model customizable (#1738)

tools/goctl/model/sql/gen/gen.go

@@ -20,16 +20,13 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/stringx"
 )
 
-const (
+const pwd = "."
-	pwd             = "."
-	createTableFlag = `(?m)^(?i)CREATE\s+TABLE` // ignore case
-)
 
 type (
 	defaultGenerator struct {
+		console.Console
 		// source string
 		dir          string
-		console.Console
 		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 {

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 {

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
+}

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 {

tools/goctl/model/sql/template/delete.go

@@ -1,7 +1,8 @@
 package template
 
+const (
 	// Delete defines a delete template
-var Delete = `
+	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{
@@ -19,4 +20,5 @@ 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 = `Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error`
+)

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"
 

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}}`

tools/goctl/model/sql/template/find.go

@@ -1,7 +1,8 @@
 package template
 
+const (
 	// FindOne defines find row by id.
-var FindOne = `
+	FindOne = `
 func (m *default{{.upperStartCamelObject}}Model) FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
@@ -31,7 +32,7 @@ func (m *default{{.upperStartCamelObject}}Model) FindOne(ctx context.Context, {{
 `
 
 	// FindOneByField defines find row by field.
-var FindOneByField = `
+	FindOneByField = `
 func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
@@ -65,7 +66,7 @@ func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}(ctx co
 `
 
 	// FindOneByFieldExtraMethod defines find row by field with extras.
-var FindOneByFieldExtraMethod = `
+	FindOneByFieldExtraMethod = `
 func (m *default{{.upperStartCamelObject}}Model) formatPrimary(primary interface{}) string {
 	return fmt.Sprintf("%s%v", {{.primaryKeyLeft}}, primary)
 }
@@ -77,7 +78,8 @@ 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 = `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 = `FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) `
+)

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"

tools/goctl/model/sql/template/insert.go

@@ -1,7 +1,8 @@
 package template
 
+const (
 	// Insert defines a template for insert code in model
-var Insert = `
+	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) {
@@ -14,4 +15,5 @@ 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 = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error)`
+)

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
-}
-`)

tools/goctl/model/sql/template/new.go

@@ -1,8 +1,8 @@
 package template
 
-// New defines an template for creating model instance
+// New defines the template for creating model instance.
-var New = `
+const New = `
-func New{{.upperStartCamelObject}}Model(conn sqlx.SqlConn{{if .withCache}}, c cache.CacheConf{{end}}) {{.upperStartCamelObject}}Model {
+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}},

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
+}
+`

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}}\"`"

tools/goctl/model/sql/template/types.go

@@ -1,7 +1,7 @@
 package template
 
-// Types defines a template for types in model
+// Types defines a template for types in model.
-var Types = `
+const Types = `
 type (
 	{{.lowerStartCamelObject}}Model interface{
 		{{.method}}

tools/goctl/model/sql/template/update.go

@@ -1,7 +1,8 @@
 package template
 
+const (
 	// Update defines a template for generating update codes
-var Update = `
+	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) {
@@ -14,4 +15,5 @@ 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 = `Update(ctx context.Context, data *{{.upperStartCamelObject}}) error`
+)

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}}