From 952a5d3bd80f8761056a828046d628ad1dc02ca4 Mon Sep 17 00:00:00 2001 From: username Date: Thu, 11 Jan 2024 17:51:49 +0800 Subject: [PATCH] feat: add gen.DB as result type --- internal/generate/interface.go | 5 ++++- internal/parser/parser.go | 5 +++++ internal/template/method.go | 4 ++-- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/internal/generate/interface.go b/internal/generate/interface.go index f43bf56f..f6eec948 100644 --- a/internal/generate/interface.go +++ b/internal/generate/interface.go @@ -127,7 +127,7 @@ func (m *InterfaceMethod) IsRepeatFromSameInterface(newMethod *InterfaceMethod) return m.MethodName == newMethod.MethodName && m.InterfaceName == newMethod.InterfaceName && m.TargetStruct == newMethod.TargetStruct } -//GetParamInTmpl return param list +// GetParamInTmpl return param list func (m *InterfaceMethod) GetParamInTmpl() string { return paramToString(m.Params) } @@ -253,6 +253,9 @@ func (m *InterfaceMethod) checkResult(result []parser.Param) (err error) { param.SetName("rows") m.GormOption = "Raw" param.IsPointer = true + case param.IsGormDB(): + param.SetName("executeSQL") + m.ResultData = param default: if !m.ResultData.IsNull() { return fmt.Errorf("query method cannot return more than 1 data value in [%s.%s]", m.InterfaceName, m.MethodName) diff --git a/internal/parser/parser.go b/internal/parser/parser.go index aec2d6f4..8c1726d0 100644 --- a/internal/parser/parser.go +++ b/internal/parser/parser.go @@ -196,6 +196,11 @@ func (p *Param) IsSQLRows() bool { return (p.Package == "sql" && p.Type == "Rows") || (p.Package == "gen" && p.Type == "SQLRows") } +// IsGormDB ... +func (p *Param) IsGormDB() bool { + return p.Package == "gorm" && p.Type == "DB" +} + // SetName ... func (p *Param) SetName(name string) { p.Name = name diff --git a/internal/template/method.go b/internal/template/method.go index 5d2eda59..34e7be2e 100644 --- a/internal/template/method.go +++ b/internal/template/method.go @@ -16,8 +16,8 @@ func ({{.S}} {{.TargetStruct}}Do){{.FuncSign}}{ result,{{if .ReturnError}}err{{else}}_{{end}} = stmt.ConnPool.ExecContext(stmt.Context,generateSQL.String(){{if .HasSQLData}},params...{{end}}) // ignore_security_alert {{else if .ReturnSQLRow}}row = {{.S}}.UnderlyingDB().Raw(generateSQL.String(){{if .HasSQLData}},params...{{end}}).Row() // ignore_security_alert {{else if .ReturnSQLRows}}rows,{{if .ReturnError}}err{{else}}_{{end}} = {{.S}}.UnderlyingDB().Raw(generateSQL.String(){{if .HasSQLData}},params...{{end}}).Rows() // ignore_security_alert - {{else}}var executeSQL *gorm.DB - executeSQL = {{.S}}.UnderlyingDB().{{.GormOption}}(generateSQL.String(){{if .HasSQLData}},params...{{end}}){{if not .ResultData.IsNull}}.{{.GormRunMethodName}}({{if .HasGotPoint}}&{{end}}{{.ResultData.Name}}){{end}} // ignore_security_alert + {{else}}{{if not .ResultData.IsGormDB}}var executeSQL *gorm.DB{{end}} + executeSQL = {{if and (.ResultData.IsGormDB) (not .ResultData.IsPointer)}}*{{end}}{{.S}}.UnderlyingDB().{{.GormOption}}(generateSQL.String(){{if .HasSQLData}},params...{{end}}){{if and (not .ResultData.IsNull) (not .ResultData.IsGormDB)}}.{{.GormRunMethodName}}({{if .HasGotPoint}}&{{end}}{{.ResultData.Name}}){{end}} // ignore_security_alert {{if .ReturnRowsAffected}}rowsAffected = executeSQL.RowsAffected {{end}}{{if .ReturnError}}err = executeSQL.Error {{end}}{{if .ReturnNothing}}_ = executeSQL