From a2df1bb16406d0214af7b48dfbaa9fb2bbb25041 Mon Sep 17 00:00:00 2001
From: foliet <foliet@rainspace.cn>
Date: Sun, 15 Jan 2023 22:11:08 +0800
Subject: [PATCH] fix: modify the generated update function and add return
 values for update and delete functions (#2793)

---
 tools/goctl/model/mongo/readme.md          | 19 ++++++++++---------
 tools/goctl/model/mongo/template/model.tpl | 19 ++++++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/tools/goctl/model/mongo/readme.md b/tools/goctl/model/mongo/readme.md
index 8af5e877..87969f87 100644
--- a/tools/goctl/model/mongo/readme.md
+++ b/tools/goctl/model/mongo/readme.md
@@ -67,6 +67,7 @@ import (
 	"github.com/zeromicro/go-zero/core/stores/monc"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 var prefixUserCacheKey = "cache:user:"
@@ -74,8 +75,8 @@ var prefixUserCacheKey = "cache:user:"
 type userModel interface {
 	Insert(ctx context.Context, data *User) error
 	FindOne(ctx context.Context, id string) (*User, error)
-	Update(ctx context.Context, data *User) error
-	Delete(ctx context.Context, id string) error
+	Update(ctx context.Context, data *User) (*mongo.UpdateResult, error)
+	Delete(ctx context.Context, id string) (int64, error)
 }
 
 type defaultUserModel struct {
@@ -117,21 +118,21 @@ func (m *defaultUserModel) FindOne(ctx context.Context, id string) (*User, error
 	}
 }
 
-func (m *defaultUserModel) Update(ctx context.Context, data *User) error {
+func (m *defaultUserModel) Update(ctx context.Context, data *User) (*mongo.UpdateResult, error) {
 	data.UpdateAt = time.Now()
 	key := prefixUserCacheKey + data.ID.Hex()
-	_, err := m.conn.ReplaceOne(ctx, key, bson.M{"_id": data.ID}, data)
-	return err
+	res, err := m.conn.ReplaceOne(ctx, key, bson.M{"_id": data.ID}, bson.M{"$set": data})
+	return res, err
 }
 
-func (m *defaultUserModel) Delete(ctx context.Context, id string) error {
+func (m *defaultUserModel) Delete(ctx context.Context, id string) (int64, error) {
 	oid, err := primitive.ObjectIDFromHex(id)
 	if err != nil {
-		return ErrInvalidObjectId
+		return 0, ErrInvalidObjectId
 	}
 	key := prefixUserCacheKey + id
-	_, err = m.conn.DeleteOne(ctx, key, bson.M{"_id": oid})
-	return err
+	res, err := m.conn.DeleteOne(ctx, key, bson.M{"_id": oid})
+	return res, err
 }
 
 ```
diff --git a/tools/goctl/model/mongo/template/model.tpl b/tools/goctl/model/mongo/template/model.tpl
index 31fab610..287125d9 100644
--- a/tools/goctl/model/mongo/template/model.tpl
+++ b/tools/goctl/model/mongo/template/model.tpl
@@ -8,6 +8,7 @@ import (
     {{if .Cache}}"github.com/zeromicro/go-zero/core/stores/monc"{{else}}"github.com/zeromicro/go-zero/core/stores/mon"{{end}}
     "go.mongodb.org/mongo-driver/bson"
     "go.mongodb.org/mongo-driver/bson/primitive"
+    "go.mongodb.org/mongo-driver/mongo"
 )
 
 {{if .Cache}}var prefix{{.Type}}CacheKey = "cache:{{.lowerType}}:"{{end}}
@@ -15,8 +16,8 @@ import (
 type {{.lowerType}}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
+    Update(ctx context.Context,data *{{.Type}}) (*mongo.UpdateResult, error)
+    Delete(ctx context.Context,id string) (int64, error)
 }
 
 type default{{.Type}}Model struct {
@@ -59,19 +60,19 @@ func (m *default{{.Type}}Model) FindOne(ctx context.Context, id string) (*{{.Typ
     }
 }
 
-func (m *default{{.Type}}Model) Update(ctx context.Context, data *{{.Type}}) error {
+func (m *default{{.Type}}Model) Update(ctx context.Context, data *{{.Type}}) (*mongo.UpdateResult, error) {
     data.UpdateAt = time.Now()
     {{if .Cache}}key := prefix{{.Type}}CacheKey + data.ID.Hex(){{end}}
-    _, err := m.conn.ReplaceOne(ctx, {{if .Cache}}key, {{end}}bson.M{"_id": data.ID}, data)
-    return err
+    res, err := m.conn.UpdateOne(ctx, {{if .Cache}}key, {{end}}bson.M{"_id": data.ID}, bson.M{"$set": data})
+    return res, err
 }
 
-func (m *default{{.Type}}Model) Delete(ctx context.Context, id string) error {
+func (m *default{{.Type}}Model) Delete(ctx context.Context, id string) (int64, error) {
     oid, err := primitive.ObjectIDFromHex(id)
     if err != nil {
-        return ErrInvalidObjectId
+        return 0, ErrInvalidObjectId
     }
 	{{if .Cache}}key := prefix{{.Type}}CacheKey +id{{end}}
-    _, err = m.conn.DeleteOne(ctx, {{if .Cache}}key, {{end}}bson.M{"_id": oid})
-	return err
+    res, err := m.conn.DeleteOne(ctx, {{if .Cache}}key, {{end}}bson.M{"_id": oid})
+	return res, err
 }