From cc4c4928e006ac1bcd5477067c2aa7b75bba913b Mon Sep 17 00:00:00 2001 From: toutou_o <33993460+kurimi1@users.noreply.github.com> Date: Tue, 1 Mar 2022 19:53:23 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20fix(gen):=20pg=20gen=20of=20inse?= =?UTF-8?q?rt=20(#1591)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: kurimi1 --- tools/goctl/model/sql/gen/gen.go | 2 +- tools/goctl/model/sql/gen/imports.go | 18 ++++++++++---- tools/goctl/model/sql/gen/insert.go | 30 +++++++++++++++++++----- tools/goctl/model/sql/template/import.go | 13 ++++++++++ tools/goctl/model/sql/template/insert.go | 19 +++++++++++++++ 5 files changed, 70 insertions(+), 12 deletions(-) diff --git a/tools/goctl/model/sql/gen/gen.go b/tools/goctl/model/sql/gen/gen.go index 6730142d6e56..49d95a4b5b06 100644 --- a/tools/goctl/model/sql/gen/gen.go +++ b/tools/goctl/model/sql/gen/gen.go @@ -223,7 +223,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(withCache, in.ContainsTime(), table, g.isPostgreSql) if err != nil { return "", err } diff --git a/tools/goctl/model/sql/gen/imports.go b/tools/goctl/model/sql/gen/imports.go index 7fe9876863d5..2fabf48d1704 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(withCache, timeImport bool, table Table, postgreSql bool) (string, error) { if withCache { text, err := pathx.LoadTemplate(category, importsTemplateFile, template.Imports) if err != nil { @@ -23,10 +23,18 @@ func genImports(withCache, timeImport bool, table Table) (string, error) { return buffer.String(), nil } - - text, err := pathx.LoadTemplate(category, importsWithNoCacheTemplateFile, template.ImportsNoCache) - if err != nil { - return "", err + var text string + var err error + if postgreSql { + text, err = pathx.LoadTemplate(category, importsWithNoCacheTemplateFile, template.ImportsNoCachePg) + if err != nil { + return "", err + } + } else { + text, err = pathx.LoadTemplate(category, importsWithNoCacheTemplateFile, template.ImportsNoCache) + if err != nil { + return "", err + } } buffer, err := util.With("import").Parse(text).Execute(map[string]interface{}{ diff --git a/tools/goctl/model/sql/gen/insert.go b/tools/goctl/model/sql/gen/insert.go index 46c287489060..f207c636e11f 100644 --- a/tools/goctl/model/sql/gen/insert.go +++ b/tools/goctl/model/sql/gen/insert.go @@ -46,9 +46,19 @@ func genInsert(table Table, withCache, postgreSql bool) (string, string, error) } camel := table.Name.ToCamel() - text, err := pathx.LoadTemplate(category, insertTemplateFile, template.Insert) - if err != nil { - return "", "", err + var text string + var err error + // if database is postgresql + if postgreSql { + text, err = pathx.LoadTemplate(category, insertTemplateFile, template.InsertPg) + if err != nil { + return "", "", err + } + } else { + text, err = pathx.LoadTemplate(category, insertTemplateFile, template.Insert) + if err != nil { + return "", "", err + } } output, err := util.With("insert"). @@ -69,9 +79,17 @@ func genInsert(table Table, withCache, postgreSql bool) (string, string, error) } // interface method - text, err = pathx.LoadTemplate(category, insertTemplateMethodFile, template.InsertMethod) - if err != nil { - return "", "", err + // if database is postgresql + if postgreSql { + text, err = pathx.LoadTemplate(category, insertTemplateMethodFile, template.InsertMethodPg) + if err != nil { + return "", "", err + } + } else { + text, err = pathx.LoadTemplate(category, insertTemplateMethodFile, template.InsertMethod) + if err != nil { + return "", "", err + } } insertMethodOutput, err := util.With("insertMethod").Parse(text).Execute(map[string]interface{}{ diff --git a/tools/goctl/model/sql/template/import.go b/tools/goctl/model/sql/template/import.go index a52e3be16b6d..944d709405a0 100644 --- a/tools/goctl/model/sql/template/import.go +++ b/tools/goctl/model/sql/template/import.go @@ -30,4 +30,17 @@ var ( "github.com/zeromicro/go-zero/core/stringx" ) ` + // pg + ImportsNoCachePg = `import ( + "context" + "fmt" + "strings" + {{if .time}}"time"{{end}} + + "github.com/zeromicro/go-zero/core/stores/builder" + "github.com/zeromicro/go-zero/core/stores/sqlc" + "github.com/zeromicro/go-zero/core/stores/sqlx" + "github.com/zeromicro/go-zero/core/stringx" + ) + ` ) diff --git a/tools/goctl/model/sql/template/insert.go b/tools/goctl/model/sql/template/insert.go index 115ea7ffe646..67687be9b5c4 100644 --- a/tools/goctl/model/sql/template/insert.go +++ b/tools/goctl/model/sql/template/insert.go @@ -15,5 +15,24 @@ func (m *default{{.upperStartCamelObject}}Model) Insert(ctx context.Context, dat } ` +// pg +var InsertPg = ` +func (m *default{{.upperStartCamelObject}}Model) Insert(data *{{.upperStartCamelObject}}) (int64,error) { + {{if .withCache}}{{if .containsIndexCache}}{{.keys}} + ret, err := m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) { + query := fmt.Sprintf("insert into %s (%s) values ({{.expression}})", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet) + return conn.Exec(query, {{.expressionValues}}) + }, {{.keyValues}}){{else}}query := fmt.Sprintf("insert into %s (%s) values ({{.expression}})", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet) + ret,err:=m.ExecNoCache(query, {{.expressionValues}}) + {{end}}{{else}}query := fmt.Sprintf("insert into %s (%s) values ({{.expression}}) RETURNING id", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet) + var id int64 + err:=m.conn.QueryRow(&id, query, {{.expressionValues}}){{end}} + return id,err +} +` + // InsertMethod defines an interface method template for insert code in model var InsertMethod = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error)` + +// pg +var InsertMethodPg = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (int64,error)`