feat：DIY SQL support return sql.Result/sql.Row/sql.Rows (#752)

* feat：DIY SQL support return sql.Result/sql.Row/sql.Rows

* feat：rename SQLRow and add test cases.

* fix: Fix little problems

* fix: fix test case

* fix: fix test case

* fix: fix test case

* fix: fix test case

generator.go

@@ -3,6 +3,7 @@ package gen
 import (
 	"bytes"
 	"context"
+	"database/sql"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -33,6 +34,15 @@ type T interface{}
 // M map[string]interface{}
 type M map[string]interface{}
 
+// SQLResult sql.result
+type SQLResult sql.Result
+
+// SQLRow sql.Row
+type SQLRow sql.Row
+
+// SQLRows sql.Rows
+type SQLRows sql.Rows
+
 // RowsAffected execute affected raws
 type RowsAffected int64
 
@@ -200,12 +210,12 @@ func (g *Generator) genModelObjConfig() *model.Config {
 	}
 }
 
-// ApplyBasic specify models which will implement basic method
+// ApplyBasic specify models which will implement basic .diy_method
 func (g *Generator) ApplyBasic(models ...interface{}) {
 	g.ApplyInterface(func() {}, models...)
 }
 
-// ApplyInterface specifies method interfaces on structures, implment codes will be generated after calling g.Execute()
+// ApplyInterface specifies .diy_method interfaces on structures, implment codes will be generated after calling g.Execute()
 // eg: g.ApplyInterface(func(model.Method){}, model.User{}, model.Company{})
 func (g *Generator) ApplyInterface(fc interface{}, models ...interface{}) {
 	structs, err := generate.ConvertStructs(g.db, models...)

internal/generate/interface.go

@@ -57,6 +57,46 @@ func (m *InterfaceMethod) GormRunMethodName() string {
 	return "Take"
 }
 
+// ReturnSQLResult return sql result
+func (m *InterfaceMethod) ReturnSQLResult() bool {
+	for _, res := range m.Result {
+		if res.IsSQLResult() {
+			return true
+		}
+	}
+	return false
+}
+
+// ReturnSQLRow return sql result
+func (m *InterfaceMethod) ReturnSQLRow() bool {
+	for _, res := range m.Result {
+		if res.IsSQLRow() {
+			return true
+		}
+	}
+	return false
+}
+
+// ReturnSQLRows return sql result
+func (m *InterfaceMethod) ReturnSQLRows() bool {
+	for _, res := range m.Result {
+		if res.IsSQLRows() {
+			return true
+		}
+	}
+	return false
+}
+
+// ReturnNothing not return error and rowAffected
+func (m *InterfaceMethod) ReturnNothing() bool {
+	for _, res := range m.Result {
+		if res.IsError() || res.Name == "rowsAffected" {
+			return false
+		}
+	}
+	return true
+}
+
 // ReturnRowsAffected return rows affected
 func (m *InterfaceMethod) ReturnRowsAffected() bool {
 	for _, res := range m.Result {
@@ -113,7 +153,7 @@ func paramToString(params []parser.Param) string {
 
 // DocComment return comment sql add "//" every line
 func (m *InterfaceMethod) DocComment() string {
-	return strings.Replace(strings.TrimSpace(m.Doc), "\n", "\n//", -1)
+	return strings.Replace(strings.Replace(strings.TrimSpace(m.Doc), "\n", "\n// ", -1), "//  ", "// ", -1)
 }
 
 // checkParams check all parameters
@@ -196,6 +236,23 @@ func (m *InterfaceMethod) checkResult(result []parser.Param) (err error) {
 			param.Package = ""
 			param.SetName("rowsAffected")
 			m.GormOption = "Exec"
+		case param.IsSQLResult():
+			param.Type = "Result"
+			param.Package = "sql"
+			param.SetName("result")
+			m.GormOption = "Statement.ConnPool.ExecContext"
+		case param.IsSQLRow():
+			param.Type = "Row"
+			param.Package = "sql"
+			param.SetName("row")
+			m.GormOption = "Raw"
+			param.IsPointer = true
+		case param.IsSQLRows():
+			param.Type = "Rows"
+			param.Package = "sql"
+			param.SetName("rows")
+			m.GormOption = "Raw"
+			param.IsPointer = true
 		default:
 			if !m.ResultData.IsNull() {
 				return fmt.Errorf("query method cannot return more than 1 data value in [%s.%s]", m.InterfaceName, m.MethodName)

internal/parser/parser.go

@@ -181,6 +181,21 @@ func (p *Param) IsTime() bool {
 	return p.Package == "time" && p.Type == "Time"
 }
 
+// IsSQLResult ...
+func (p *Param) IsSQLResult() bool {
+	return (p.Package == "sql" && p.Type == "Result") || (p.Package == "gen" && p.Type == "SQLResult")
+}
+
+// IsSQLRow ...
+func (p *Param) IsSQLRow() bool {
+	return (p.Package == "sql" && p.Type == "Row") || (p.Package == "gen" && p.Type == "SQLRow")
+}
+
+// IsSQLRows ...
+func (p *Param) IsSQLRows() bool {
+	return (p.Package == "sql" && p.Type == "Rows") || (p.Package == "gen" && p.Type == "SQLRows")
+}
+
 // SetName ...
 func (p *Param) SetName(name string) {
 	p.Name = name

internal/template/method.go

@@ -3,7 +3,7 @@ package template
 // DIYMethod DIY method
 const DIYMethod = `
 
-//{{.DocComment }}
+// {{.DocComment }}
 func ({{.S}} {{.TargetStruct}}Do){{.FuncSign}}{
 	{{if .HasSQLData}}var params []interface{}
 
@@ -12,12 +12,17 @@ func ({{.S}} {{.TargetStruct}}Do){{.FuncSign}}{
 	{{end}}
 
 	{{if .HasNeedNewResult}}result ={{if .ResultData.IsMap}}make{{else}}new{{end}}({{if ne .ResultData.Package ""}}{{.ResultData.Package}}.{{end}}{{.ResultData.Type}}){{end}}
-	{{if or .ReturnRowsAffected .ReturnError}}var executeSQL *gorm.DB
-	{{end}}
-	{{if or .ReturnRowsAffected .ReturnError}}executeSQL{{else}}_{{end}} = {{.S}}.UnderlyingDB().{{.GormOption}}(generateSQL.String(){{if .HasSQLData}},params...{{end}}){{if not .ResultData.IsNull}}.{{.GormRunMethodName}}({{if .HasGotPoint}}&{{end}}{{.ResultData.Name}}){{end}}
+	{{if .ReturnSQLResult}}stmt := {{.S}}.UnderlyingDB().Statement
+	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
 	{{if .ReturnRowsAffected}}rowsAffected = executeSQL.RowsAffected
 	{{end}}{{if .ReturnError}}err = executeSQL.Error
-	{{end}}return
+	{{end}}{{if .ReturnNothing}}_ = executeSQL
+	{{end}}{{end}}
+	return
 }
 
 `