Feature mongo gen (#546)

* add feature: mongo code generation

* upgrade version

* update doc

* format code

* update update.tpl of mysql

tools/goctl/goctl.go

@@ -19,6 +19,7 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/configgen"
 	"github.com/tal-tech/go-zero/tools/goctl/docker"
 	"github.com/tal-tech/go-zero/tools/goctl/kube"
+	"github.com/tal-tech/go-zero/tools/goctl/model/mongo"
 	model "github.com/tal-tech/go-zero/tools/goctl/model/sql/command"
 	"github.com/tal-tech/go-zero/tools/goctl/plugin"
 	rpc "github.com/tal-tech/go-zero/tools/goctl/rpc/cli"
@@ -28,7 +29,7 @@ import (
 )
 
 var (
-	buildVersion = "1.1.5"
+	buildVersion = "1.1.6"
 	commands     = []cli.Command{
 		{
 			Name:   "upgrade",
@@ -447,6 +448,29 @@ var (
 						},
 					},
 				},
+				{
+					Name:  "mongo",
+					Usage: `generate mongo model`,
+					Flags: []cli.Flag{
+						cli.StringSliceFlag{
+							Name:  "type, t",
+							Usage: "specified model type name",
+						},
+						cli.BoolFlag{
+							Name:  "cache, c",
+							Usage: "generate code with cache [optional]",
+						},
+						cli.StringFlag{
+							Name:  "dir, d",
+							Usage: "the target dir",
+						},
+						cli.StringFlag{
+							Name:  "style",
+							Usage: "the file naming format, see [https://github.com/tal-tech/go-zero/tree/master/tools/goctl/config/readme.md]",
+						},
+					},
+					Action: mongo.Action,
+				},
 			},
 		},
 		{

tools/goctl/model/mongo/generate/generate.go

@@ -0,0 +1,69 @@
+package generate
+
+import (
+	"errors"
+	"path/filepath"
+
+	"github.com/tal-tech/go-zero/tools/goctl/config"
+	"github.com/tal-tech/go-zero/tools/goctl/model/mongo/template"
+	"github.com/tal-tech/go-zero/tools/goctl/util"
+	"github.com/tal-tech/go-zero/tools/goctl/util/format"
+)
+
+// Context defines the model generation data what they needs
+type Context struct {
+	Types  []string
+	Cache  bool
+	Output string
+	Cfg    *config.Config
+}
+
+// Do executes model template and output the result into the specified file path
+func Do(ctx *Context) error {
+	if ctx.Cfg == nil {
+		return errors.New("missing config")
+	}
+
+	err := generateModel(ctx)
+	if err != nil {
+		return err
+	}
+
+	return generateError(ctx)
+}
+
+func generateModel(ctx *Context) error {
+	for _, t := range ctx.Types {
+		fn, err := format.FileNamingFormat(ctx.Cfg.NamingFormat, t+"_model")
+		if err != nil {
+			return err
+		}
+
+		text, err := util.LoadTemplate(category, modelTemplateFile, template.Text)
+		if err != nil {
+			return err
+		}
+
+		output := filepath.Join(ctx.Output, fn+".go")
+		err = util.With("model").Parse(text).GoFmt(true).SaveTo(map[string]interface{}{
+			"Type":  t,
+			"Cache": ctx.Cache,
+		}, output, false)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func generateError(ctx *Context) error {
+	text, err := util.LoadTemplate(category, errTemplateFile, template.Error)
+	if err != nil {
+		return err
+	}
+
+	output := filepath.Join(ctx.Output, "error.go")
+
+	return util.With("error").Parse(text).GoFmt(true).SaveTo(ctx, output, false)
+}

tools/goctl/model/mongo/generate/generate_test.go

@@ -0,0 +1,34 @@
+package generate
+
+import (
+	"io/ioutil"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/tools/goctl/config"
+)
+
+var testTypes = `
+	type User struct{}
+	type Class struct{}
+`
+
+func TestDo(t *testing.T) {
+	cfg, err := config.NewConfig(config.DefaultFormat)
+	assert.Nil(t, err)
+
+	tempDir := t.TempDir()
+	typesfile := filepath.Join(tempDir, "types.go")
+	err = ioutil.WriteFile(typesfile, []byte(testTypes), 0666)
+	assert.Nil(t, err)
+
+	err = Do(&Context{
+		Types:  []string{"User", "Class"},
+		Cache:  false,
+		Output: tempDir,
+		Cfg:    cfg,
+	})
+
+	assert.Nil(t, err)
+}

tools/goctl/model/mongo/generate/template.go

@@ -0,0 +1,50 @@
+package generate
+
+import (
+	"fmt"
+
+	"github.com/tal-tech/go-zero/tools/goctl/model/mongo/template"
+	"github.com/tal-tech/go-zero/tools/goctl/util"
+	"github.com/urfave/cli"
+)
+
+const (
+	category          = "mongo"
+	modelTemplateFile = "model.tpl"
+	errTemplateFile   = "err.tpl"
+)
+
+var templates = map[string]string{
+	modelTemplateFile: template.Text,
+	errTemplateFile:   template.Error,
+}
+
+func Category() string {
+	return category
+}
+
+func Clean() error {
+	return util.Clean(category)
+}
+
+func Templates(_ *cli.Context) error {
+	return util.InitTemplates(category, templates)
+}
+
+func RevertTemplate(name string) error {
+	content, ok := templates[name]
+	if !ok {
+		return fmt.Errorf("%s: no such file name", name)
+	}
+
+	return util.CreateTemplate(category, name, content)
+}
+
+func Update() error {
+	err := Clean()
+	if err != nil {
+		return err
+	}
+
+	return util.InitTemplates(category, templates)
+}

tools/goctl/model/mongo/mongo.go

@@ -0,0 +1,39 @@
+package mongo
+
+import (
+	"errors"
+	"path/filepath"
+	"strings"
+
+	"github.com/tal-tech/go-zero/tools/goctl/config"
+	"github.com/tal-tech/go-zero/tools/goctl/model/mongo/generate"
+	"github.com/urfave/cli"
+)
+
+// Command provides the entry for goctl
+func Action(ctx *cli.Context) error {
+	tp := ctx.StringSlice("type")
+	c := ctx.Bool("cache")
+	o := strings.TrimSpace(ctx.String("dir"))
+	s := ctx.String("style")
+	if len(tp) == 0 {
+		return errors.New("missing type")
+	}
+
+	cfg, err := config.NewConfig(s)
+	if err != nil {
+		return err
+	}
+
+	a, err := filepath.Abs(o)
+	if err != nil {
+		return err
+	}
+
+	return generate.Do(&generate.Context{
+		Types:  tp,
+		Cache:  c,
+		Output: a,
+		Cfg:    cfg,
+	})
+}

tools/goctl/model/mongo/readme.md

@@ -0,0 +1,210 @@
+# mongo生成model
+
+## 背景
+
+在业务务开发中，model(dao)数据访问层是一个服务必不可缺的一层，因此数据库访问的CURD也是必须要对外提供的访问方法， 而CURD在go-zero中就仅存在两种情况
+
+* 带缓存model
+* 不带缓存model
+
+从代码结构上来看，C-U-R-D四个方法就是固定的结构，因此我们可以将其交给goctl工具去完成，帮助我们提升开发效率。
+
+## 方案设计
+
+mongo的生成不同于mysql，mysql可以从scheme_information库中读取到一张表的信息（字段名称，数据类型，索引等），
+而mongo是文档型数据库，我们暂时无法从db中读取某一条记录来实现字段信息获取，就算有也不一定是完整信息（某些字段可能是omitempty修饰，可有可无）， 这里采用type自己编写+代码生成方式实现
+
+## 使用示例
+
+假设我们需要生成一个usermodel.go的代码文件，其包含用户信息字段有
+
+|字段名称|字段类型|
+|---|---|
+|_id|bson.ObejctId|
+|name|string|
+
+### 编写types.go
+
+```shell
+$ vim types.go
+```
+
+```golang
+package model
+
+//go:generate goctl model mongo -t User
+import "github.com/globalsign/mgo/bson"
+
+type User struct {
+	ID   bson.ObjectId `bson:"_id"`
+	Name string        `bson:"name"`
+}
+```
+
+### 生成代码
+
+生成代码的方式有两种
+
+* 命令行生成 在types.go所在文件夹执行命令
+    ```shell
+    $ goctl model mongo -t User -style gozero
+    ```
+* 在types.go中添加`//go:generate`，然后点击执行按钮即可生成，内容示例如下：
+  ```golang
+  //go:generate goctl model mongo -t User
+  ```
+
+### 生成示例代码
+
+* usermodel.go
+
+  ```golang
+  package model
+  
+  import (
+      "context"
+  
+      "github.com/globalsign/mgo/bson"
+      cachec "github.com/tal-tech/go-zero/core/stores/cache"
+      "github.com/tal-tech/go-zero/core/stores/mongoc"
+  )
+  
+  type UserModel interface {
+      Insert(data *User, ctx context.Context) error
+      FindOne(id string, ctx context.Context) (*User, error)
+      Update(data *User, ctx context.Context) error
+      Delete(id string, ctx context.Context) error
+  }
+  
+  type defaultUserModel struct {
+      *mongoc.Model
+  }
+  
+  func NewUserModel(url, collection string, c cachec.CacheConf) UserModel {
+      return &defaultUserModel{
+          Model: mongoc.MustNewModel(url, collection, c),
+      }
+  }
+  
+  func (m *defaultUserModel) Insert(data *User, ctx context.Context) error {
+      if !data.ID.Valid() {
+          data.ID = bson.NewObjectId()
+      }
+  
+      session, err := m.TakeSession()
+      if err != nil {
+          return err
+      }
+  
+      defer m.PutSession(session)
+      return m.GetCollection(session).Insert(data)
+  }
+  
+  func (m *defaultUserModel) FindOne(id string, ctx context.Context) (*User, error) {
+      if !bson.IsObjectIdHex(id) {
+          return nil, ErrInvalidObjectId
+      }
+  
+      session, err := m.TakeSession()
+      if err != nil {
+          return nil, err
+      }
+  
+      defer m.PutSession(session)
+      var data User
+  
+      err = m.GetCollection(session).FindOneIdNoCache(&data, bson.ObjectIdHex(id))
+      switch err {
+      case nil:
+          return &data, nil
+      case mongoc.ErrNotFound:
+          return nil, ErrNotFound
+      default:
+          return nil, err
+      }
+  }
+  
+  func (m *defaultUserModel) Update(data *User, ctx context.Context) error {
+      session, err := m.TakeSession()
+      if err != nil {
+          return err
+      }
+  
+      defer m.PutSession(session)
+  
+      return m.GetCollection(session).UpdateIdNoCache(data.ID, data)
+  }
+  
+  func (m *defaultUserModel) Delete(id string, ctx context.Context) error {
+      session, err := m.TakeSession()
+      if err != nil {
+          return err
+      }
+  
+      defer m.PutSession(session)
+  
+      return m.GetCollection(session).RemoveIdNoCache(bson.ObjectIdHex(id))
+  }
+  ```
+
+* error.go
+
+  ```golang
+  package model
+
+  import "errors"
+  
+  var ErrNotFound = errors.New("not found")
+  var ErrInvalidObjectId = errors.New("invalid objectId")
+  ```
+
+### 文件目录预览
+
+```text
+.
+├── error.go
+├── types.go
+└── usermodel.go
+
+```
+
+## 命令预览
+
+```text
+NAME:
+   goctl model - generate model code
+
+USAGE:
+   goctl model command [command options] [arguments...]
+
+COMMANDS:
+   mysql  generate mysql model
+   mongo  generate mongo model
+
+OPTIONS:
+   --help, -h  show help
+```
+
+```text
+NAME:
+   goctl model mongo - generate mongo model
+
+USAGE:
+   goctl model mongo [command options] [arguments...]
+
+OPTIONS:
+   --type value, -t value  specified model type name
+   --cache, -c             generate code with cache [optional]
+   --dir value, -d value   the target dir
+   --style value           the file naming format, see [https://github.com/tal-tech/go-zero/tree/master/tools/goctl/config/readme.md]
+
+```
+
+> 温馨提示
+> 
+> `--type` 支持slice传值，示例 `goctl model mongo -t=User -t=Class`
+## 注意事项
+
+types.go本质上与xxxmodel.go无关，只是将type定义部分交给开发人员自己编写了，在xxxmodel.go中，mongo文档的存储结构必须包含
+`_id`字段，对应到types中的field为`ID`，model中的findOne,update均以data.ID来进行操作的，当然，如果不符合你的命名风格，你也 可以修改模板，只要保证`id`
+在types中的field名称和模板中一致就行。

tools/goctl/model/mongo/template/template.go

@@ -0,0 +1,111 @@
+package template
+
+// Text provides the default template for model to generate
+var Text = `package model
+
+import (
+    "context"
+
+    "github.com/globalsign/mgo/bson"
+     cachec "github.com/tal-tech/go-zero/core/stores/cache"
+	"github.com/tal-tech/go-zero/core/stores/mongoc"
+)
+
+{{if .Cache}}var prefix{{.Type}}CacheKey = "cache#{{.Type}}#"{{end}}
+
+type {{.Type}}Model interface{
+	Insert(ctx context.Context,data *{{.Type}}) error
+	FindOne(ctx context.Context,id string) (*{{.Type}}, error)
+	Update(ctx context.Context,data *{{.Type}}) error
+	Delete(ctx context.Context,id string) error
+}
+
+type default{{.Type}}Model struct {
+    *mongoc.Model
+}
+
+func New{{.Type}}Model(url, collection string, c cachec.CacheConf) {{.Type}}Model {
+	return &default{{.Type}}Model{
+		Model: mongoc.MustNewModel(url, collection, c),
+	}
+}
+
+
+func (m *default{{.Type}}Model) Insert(ctx context.Context, data *{{.Type}}) error {
+    if !data.ID.Valid() {
+        data.ID = bson.NewObjectId()
+    }
+
+    session, err := m.TakeSession()
+    if err != nil {
+        return err
+    }
+
+    defer m.PutSession(session)
+    return m.GetCollection(session).Insert(data)
+}
+
+func (m *default{{.Type}}Model) FindOne(ctx context.Context, id string) (*{{.Type}}, error) {
+    if !bson.IsObjectIdHex(id) {
+        return nil, ErrInvalidObjectId
+    }
+
+    session, err := m.TakeSession()
+    if err != nil {
+        return nil, err
+    }
+
+    defer m.PutSession(session)
+    var data {{.Type}}
+    {{if .Cache}}key := prefix{{.Type}}CacheKey + id
+    err = m.GetCollection(session).FindOneId(&data, key, bson.ObjectIdHex(id))
+	{{- else}}
+	err = m.GetCollection(session).FindOneIdNoCache(&data, bson.ObjectIdHex(id))
+	{{- end}}
+    switch err {
+    case nil:
+        return &data,nil
+    case mongoc.ErrNotFound:
+        return nil,ErrNotFound
+    default:
+        return nil,err
+    }
+}
+
+func (m *default{{.Type}}Model) Update(ctx context.Context, data *{{.Type}}) error {
+    session, err := m.TakeSession()
+    if err != nil {
+        return err
+    }
+
+    defer m.PutSession(session)
+	{{if .Cache}}key := prefix{{.Type}}CacheKey + data.ID.Hex()
+    return m.GetCollection(session).UpdateId(data.ID, data, key)
+	{{- else}}
+	return m.GetCollection(session).UpdateIdNoCache(data.ID, data)
+	{{- end}}
+}
+
+func (m *default{{.Type}}Model) Delete(ctx context.Context, id string) error {
+    session, err := m.TakeSession()
+    if err != nil {
+        return err
+    }
+
+    defer m.PutSession(session)
+    {{if .Cache}}key := prefix{{.Type}}CacheKey + id
+    return m.GetCollection(session).RemoveId(bson.ObjectIdHex(id), key)
+	{{- else}}
+	return m.GetCollection(session).RemoveIdNoCache(bson.ObjectIdHex(id))
+	{{- end}}
+}
+`
+
+var Error = `
+package model
+
+import "errors"
+
+var ErrNotFound = errors.New("not found")
+var ErrInvalidObjectId = errors.New("invalid objectId")
+`

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

@@ -3,6 +3,7 @@ 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"
@@ -23,6 +24,15 @@ func genUpdate(table Table, withCache bool) (string, string, error) {
 		expressionValues = append(expressionValues, "data."+camel)
 	}
 
+	keySet := collection.NewSet()
+	keyVariableSet := collection.NewSet()
+	keySet.AddStr(table.PrimaryCacheKey.DataKeyExpression)
+	keyVariableSet.AddStr(table.PrimaryCacheKey.KeyLeft)
+	for _, key := range table.UniqueCacheKey {
+		keySet.AddStr(key.DataKeyExpression)
+		keyVariableSet.AddStr(key.KeyLeft)
+	}
+
 	expressionValues = append(expressionValues, "data."+table.PrimaryKey.Name.ToCamel())
 	camelTableName := table.Name.ToCamel()
 	text, err := util.LoadTemplate(category, updateTemplateFile, template.Update)
@@ -35,6 +45,8 @@ func genUpdate(table Table, withCache bool) (string, string, error) {
 		Execute(map[string]interface{}{
 			"withCache":             withCache,
 			"upperStartCamelObject": camelTableName,
+			"keys":                  strings.Join(keySet.KeysStr(), "\n"),
+			"keyValues":             strings.Join(keyVariableSet.KeysStr(), ", "),
 			"primaryCacheKey":       table.PrimaryCacheKey.DataKeyExpression,
 			"primaryKeyVariable":    table.PrimaryCacheKey.KeyLeft,
 			"lowerStartCamelObject": stringx.From(camelTableName).Untitle(),

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

@@ -3,11 +3,11 @@ package template
 // Update defines a template for generating update codes
 var Update = `
 func (m *default{{.upperStartCamelObject}}Model) Update(data {{.upperStartCamelObject}}) error {
-	{{if .withCache}}{{.primaryCacheKey}}
+	{{if .withCache}}{{.keys}}
     _, err := m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) {
 		query := fmt.Sprintf("update %s set %s where {{.originalPrimaryKey}} = ?", m.table, {{.lowerStartCamelObject}}RowsWithPlaceHolder)
 		return conn.Exec(query, {{.expressionValues}})
-	}, {{.primaryKeyVariable}}){{else}}query := fmt.Sprintf("update %s set %s where {{.originalPrimaryKey}} = ?", m.table, {{.lowerStartCamelObject}}RowsWithPlaceHolder)
+	}, {{.keyValues}}){{else}}query := fmt.Sprintf("update %s set %s where {{.originalPrimaryKey}} = ?", m.table, {{.lowerStartCamelObject}}RowsWithPlaceHolder)
     _,err:=m.conn.Exec(query, {{.expressionValues}}){{end}}
 	return err
 }

tools/goctl/tpl/templates.go

@@ -8,6 +8,7 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/api/gogen"
 	"github.com/tal-tech/go-zero/tools/goctl/docker"
 	"github.com/tal-tech/go-zero/tools/goctl/kube"
+	mongogen "github.com/tal-tech/go-zero/tools/goctl/model/mongo/generate"
 	modelgen "github.com/tal-tech/go-zero/tools/goctl/model/sql/gen"
 	rpcgen "github.com/tal-tech/go-zero/tools/goctl/rpc/generator"
 	"github.com/tal-tech/go-zero/tools/goctl/util"
@@ -34,6 +35,9 @@ func GenTemplates(ctx *cli.Context) error {
 		func() error {
 			return kube.GenTemplates(ctx)
 		},
+		func() error {
+			return mongogen.Templates(ctx)
+		},
 	); err != nil {
 		return err
 	}
@@ -61,6 +65,15 @@ func CleanTemplates(_ *cli.Context) error {
 		func() error {
 			return rpcgen.Clean()
 		},
+		func() error {
+			return docker.Clean()
+		},
+		func() error {
+			return kube.Clean()
+		},
+		func() error {
+			return mongogen.Clean()
+		},
 	)
 	if err != nil {
 		return err
@@ -90,6 +103,8 @@ func UpdateTemplates(ctx *cli.Context) (err error) {
 		return rpcgen.Update()
 	case modelgen.Category():
 		return modelgen.Update()
+	case mongogen.Category():
+		return mongogen.Update()
 	default:
 		err = fmt.Errorf("unexpected category: %s", category)
 		return
@@ -116,6 +131,8 @@ func RevertTemplates(ctx *cli.Context) (err error) {
 		return rpcgen.RevertTemplate(filename)
 	case modelgen.Category():
 		return modelgen.RevertTemplate(filename)
+	case mongogen.Category():
+		return mongogen.RevertTemplate(filename)
 	default:
 		err = fmt.Errorf("unexpected category: %s", category)
 		return