Gozero sqlgen patch (#119)

* merge upstream

* optimize insert logic

* reactor functions

tools/goctl/model/sql/example/sql/user.sql

@@ -10,6 +10,6 @@ CREATE TABLE `user` (
   `update_time` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
   PRIMARY KEY (`id`),
   UNIQUE KEY `name_index` (`name`),
-  KEY `mobile_index` (`mobile`)
+  UNIQUE KEY `mobile_index` (`mobile`)
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci;
 

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

@@ -20,20 +20,14 @@ func genDelete(table Table, withCache bool) (string, error) {
 		}
 		keyVariableSet.AddStr(key.Variable)
 	}
-	var containsIndexCache = false
-	for _, item := range table.Fields {
-		if item.IsUniqueKey {
-			containsIndexCache = true
-			break
-		}
-	}
+
 	camel := table.Name.ToCamel()
 	output, err := util.With("delete").
 		Parse(template.Delete).
 		Execute(map[string]interface{}{
 			"upperStartCamelObject":     camel,
 			"withCache":                 withCache,
-			"containsIndexCache":        containsIndexCache,
+			"containsIndexCache":        table.ContainsUniqueKey,
 			"lowerStartCamelPrimaryKey": stringx.From(table.PrimaryKey.Name.ToCamel()).UnTitle(),
 			"dataType":                  table.PrimaryKey.DataType,
 			"keys":                      strings.Join(keySet.KeysStr(), "\n"),

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

@@ -9,7 +9,7 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/util/stringx"
 )
 
-func genFindOneByField(table Table, withCache bool) (string, error) {
+func genFindOneByField(table Table, withCache bool) (string, string, error) {
 	t := util.With("findOneByField").Parse(template.FindOneByField)
 	var list []string
 	camelTableName := table.Name.ToCamel()
@@ -30,12 +30,23 @@ func genFindOneByField(table Table, withCache bool) (string, error) {
 			"lowerStartCamelField":      stringx.From(camelFieldName).UnTitle(),
 			"upperStartCamelPrimaryKey": table.PrimaryKey.Name.ToCamel(),
 			"originalField":             field.Name.Source(),
-			"originalPrimaryField":      table.PrimaryKey.Name.Source(),
 		})
 		if err != nil {
-			return "", err
+			return "", "", err
 		}
 		list = append(list, output.String())
 	}
-	return strings.Join(list, "\n"), nil
+	if withCache {
+		out, err := util.With("findOneByFieldExtraMethod").Parse(template.FindOneByFieldExtraMethod).Execute(map[string]interface{}{
+			"upperStartCamelObject": camelTableName,
+			"lowerStartCamelObject": stringx.From(camelTableName).UnTitle(),
+			"originalPrimaryField":  table.PrimaryKey.Name.Source(),
+		})
+		if err != nil {
+			return "", "", err
+		}
+		return strings.Join(list, "\n"), out.String(), nil
+	}
+	return strings.Join(list, "\n"), "", nil
+
 }

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

@@ -113,7 +113,8 @@ func (g *defaultGenerator) genFromDDL(withCache bool) (map[string]string, error)
 type (
 	Table struct {
 		parser.Table
-		CacheKey map[string]Key
+		CacheKey          map[string]Key
+		ContainsUniqueKey bool
 	}
 )
 
@@ -135,6 +136,14 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	var table Table
 	table.Table = in
 	table.CacheKey = m
+	var containsUniqueCache = false
+	for _, item := range table.Fields {
+		if item.IsUniqueKey {
+			containsUniqueCache = true
+			break
+		}
+	}
+	table.ContainsUniqueKey = containsUniqueCache
 
 	varsCode, err := genVars(table, withCache)
 	if err != nil {
@@ -162,7 +171,7 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 		return "", err
 	}
 
-	findOneByFieldCode, err := genFindOneByField(table, withCache)
+	findOneByFieldCode, extraMethod, err := genFindOneByField(table, withCache)
 	if err != nil {
 		return "", err
 	}
@@ -179,14 +188,15 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	}
 
 	output, err := t.Execute(map[string]interface{}{
-		"imports": importsCode,
-		"vars":    varsCode,
-		"types":   typesCode,
-		"new":     newCode,
-		"insert":  insertCode,
-		"find":    strings.Join(findCode, "\n"),
-		"update":  updateCode,
-		"delete":  deleteCode,
+		"imports":     importsCode,
+		"vars":        varsCode,
+		"types":       typesCode,
+		"new":         newCode,
+		"insert":      insertCode,
+		"find":        strings.Join(findCode, "\n"),
+		"update":      updateCode,
+		"delete":      deleteCode,
+		"extraMethod": extraMethod,
 	})
 	if err != nil {
 		return "", err

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

@@ -3,12 +3,23 @@ package gen
 import (
 	"strings"
 
+	"github.com/tal-tech/go-zero/core/collection"
 	"github.com/tal-tech/go-zero/tools/goctl/model/sql/template"
 	"github.com/tal-tech/go-zero/tools/goctl/util"
 	"github.com/tal-tech/go-zero/tools/goctl/util/stringx"
 )
 
 func genInsert(table Table, withCache bool) (string, error) {
+	keySet := collection.NewSet()
+	keyVariableSet := collection.NewSet()
+	for fieldName, key := range table.CacheKey {
+		if fieldName == table.PrimaryKey.Name.Source() {
+			continue
+		}
+		keySet.AddStr(key.DataKeyExpression)
+		keyVariableSet.AddStr(key.Variable)
+	}
+
 	expressions := make([]string, 0)
 	expressionValues := make([]string, 0)
 	for _, filed := range table.Fields {
@@ -27,10 +38,13 @@ func genInsert(table Table, withCache bool) (string, error) {
 		Parse(template.Insert).
 		Execute(map[string]interface{}{
 			"withCache":             withCache,
+			"containsIndexCache":    table.ContainsUniqueKey,
 			"upperStartCamelObject": camel,
 			"lowerStartCamelObject": stringx.From(camel).UnTitle(),
 			"expression":            strings.Join(expressions, ", "),
 			"expressionValues":      strings.Join(expressionValues, ", "),
+			"keys":                  strings.Join(keySet.KeysStr(), "\n"),
+			"keyValues":             strings.Join(keyVariableSet.KeysStr(), ", "),
 		})
 	if err != nil {
 		return "", err

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

@@ -35,18 +35,13 @@ var FindOneByField = `
 func (m *{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}({{.in}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
-	err := m.QueryRowIndex(&resp, {{.cacheKeyVariable}}, func(primary interface{}) string {
-		return fmt.Sprintf("%s%v", {{.primaryKeyLeft}}, primary)
-	}, func(conn sqlx.SqlConn, v interface{}) (i interface{}, e error) {
+	err := m.QueryRowIndex(&resp, {{.cacheKeyVariable}}, m.formatPrimary, func(conn sqlx.SqlConn, v interface{}) (i interface{}, e error) {
 		query := ` + "`" + `select ` + "`" + ` + {{.lowerStartCamelObject}}Rows + ` + "`" + ` from ` + "` + " + `m.table ` + " + `" + ` where {{.originalField}} = ? limit 1` + "`" + `
 		if err := conn.QueryRow(&resp, query, {{.lowerStartCamelField}}); err != nil {
 			return nil, err
 		}
 		return resp.{{.upperStartCamelPrimaryKey}}, nil
-	}, func(conn sqlx.SqlConn, v, primary interface{}) error {
-		query := ` + "`" + `select ` + "`" + ` + {{.lowerStartCamelObject}}Rows + ` + "`" + ` from ` + "` + " + `m.table ` + " + `" + ` where {{.originalPrimaryField}} = ? limit 1` + "`" + `
-		return conn.QueryRow(v, query, primary)
-	})
+	}, m.queryPrimary)
 	switch err {
 	case nil:
 		return &resp, nil
@@ -68,3 +63,13 @@ func (m *{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}({{.in}}) (*{{
 	}
 }{{end}}
 `
+var FindOneByFieldExtraMethod = `
+func (m *{{.upperStartCamelObject}}Model) formatPrimary(primary interface{}) string {
+	return fmt.Sprintf("%s%v", cacheUserIdPrefix, primary)
+}
+
+func (m *{{.upperStartCamelObject}}Model) queryPrimary(conn sqlx.SqlConn, v, primary interface{}) error {
+	query := ` + "`" + `select ` + "`" + ` + {{.lowerStartCamelObject}}Rows + ` + "`" + ` from ` + "` + " + `m.table ` + " + `" + ` where {{.originalPrimaryField}} = ? limit 1` + "`" + `
+	return conn.QueryRow(v, query, primary)
+}
+`

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

@@ -15,7 +15,6 @@ var (
 )
 `
 	ImportsNoCache = `import (
-	"database/sql"
 	"strings"
 	{{if .time}}"time"{{end}}
 

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

@@ -1,8 +1,15 @@
 package template
 
 var Insert = `
-func (m *{{.upperStartCamelObject}}Model) Insert(data {{.upperStartCamelObject}}) (sql.Result, error) {
-	query := ` + "`" + `insert into ` + "`" + ` + m.table + ` + "` (` + " + `{{.lowerStartCamelObject}}RowsExpectAutoSet` + " + `) values ({{.expression}})` " + `
-	return m.{{if .withCache}}ExecNoCache{{else}}conn.Exec{{end}}(query, {{.expressionValues}})
+func (m *{{.upperStartCamelObject}}Model) Insert(data {{.upperStartCamelObject}}) error {
+	{{if .withCache}}{{if .containsIndexCache}}{{.keys}}
+    _, err := m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) {
+		query := ` + "`" + `insert into ` + "`" + ` + m.table + ` + "` (` + " + `{{.lowerStartCamelObject}}RowsExpectAutoSet` + " + `) values ({{.expression}})` " + `
+		return conn.Exec(query, {{.expressionValues}})
+	}, {{.keyValues}}){{else}}query := ` + "`" + `insert into ` + "`" + ` + m.table + ` + "` (` + " + `{{.lowerStartCamelObject}}RowsExpectAutoSet` + " + `) values ({{.expression}})` " + `
+    _,err:=m.ExecNoCache(query, {{.expressionValues}})
+	{{end}}{{else}}query := ` + "`" + `insert into ` + "`" + ` + m.table + ` + "` (` + " + `{{.lowerStartCamelObject}}RowsExpectAutoSet` + " + `) values ({{.expression}})` " + `
+    _,err:=m.conn.Exec(query, {{.expressionValues}}){{end}}
+	return err
 }
 `

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

@@ -9,4 +9,5 @@ var Model = `package model
 {{.find}}
 {{.update}}
 {{.delete}}
+{{.extraMethod}}
 `

tools/goctl/util/project/project.go

@@ -1,6 +1,7 @@
 package project
 
 import (
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -136,5 +137,5 @@ func matchModule(data []byte) (string, error) {
 		return strings.TrimSpace(target[index+6:]), nil
 	}
 
-	return "", nil
+	return "", errors.New("module not matched")
 }