Feature/mongo find and save (#316)

* tpl/mongo,orm: add mongo FindAndSave for atomic operation Fix #315 Signed-off-by: scnace <scbizu@gmail.com> * pkg: fix extract from singleResult --------- Signed-off-by: scnace <scbizu@gmail.com>
ezbuy · May 8, 2024 · 31a858b · 31a858b
1 parent 689b4ff
commit 31a858b
Show file tree

Hide file tree

Showing 7 changed files with 137 additions and 6 deletions.
diff --git a/e2e/mongo/nested/gen_User_mongo_orm.go b/e2e/mongo/nested/gen_User_mongo_orm.go
@@ -68,6 +68,27 @@ func (o *User) Id() string {
 	return o.ID.Hex()
 }
 
+// FindAndSave finds by `query`,  and then upsert the result with the current object
+func (o *User) FindAndSave(ctx context.Context, query interface{}) (*mongo.SingleResult, error) {
+	col := UserMgr.GetCol()
+	opts := options.FindOneAndUpdate().SetUpsert(true)
+	opts.SetReturnDocument(options.After)
+	update := bson.M{
+		"$set": bson.M{
+			UserMgoFieldUserId:       o.UserId,
+			UserMgoFieldUsername:     o.Username,
+			UserMgoFieldAge:          o.Age,
+			UserMgoFieldBlogs:        o.Blogs,
+			UserMgoFieldRegisterDate: o.RegisterDate,
+		},
+	}
+	ret := col.FindOneAndUpdate(ctx, query, update, opts)
+	if ret.Err() != nil {
+		return nil, ret.Err()
+	}
+	return ret, nil
+}
+
 func (o *User) Save(ctx context.Context) (*mongo.UpdateResult, error) {
 	isNew := o.isNew
 	if o.ID == primitive.NilObjectID {

diff --git a/e2e/mongo/user/gen_UserBlog_mongo_orm.go b/e2e/mongo/user/gen_UserBlog_mongo_orm.go
@@ -62,6 +62,25 @@ func (o *UserBlog) Id() string {
 	return o.ID.Hex()
 }
 
+// FindAndSave finds by `query`,  and then upsert the result with the current object
+func (o *UserBlog) FindAndSave(ctx context.Context, query interface{}) (*mongo.SingleResult, error) {
+	col := UserBlogMgr.GetCol()
+	opts := options.FindOneAndUpdate().SetUpsert(true)
+	opts.SetReturnDocument(options.After)
+	update := bson.M{
+		"$set": bson.M{
+			UserBlogMgoFieldUserId:  o.UserId,
+			UserBlogMgoFieldBlogId:  o.BlogId,
+			UserBlogMgoFieldContent: o.Content,
+		},
+	}
+	ret := col.FindOneAndUpdate(ctx, query, update, opts)
+	if ret.Err() != nil {
+		return nil, ret.Err()
+	}
+	return ret, nil
+}
+
 func (o *UserBlog) Save(ctx context.Context) (*mongo.UpdateResult, error) {
 	isNew := o.isNew
 	if o.ID == primitive.NilObjectID {

diff --git a/e2e/mongo/user/gen_User_mongo_orm.go b/e2e/mongo/user/gen_User_mongo_orm.go
@@ -81,6 +81,26 @@ func (o *User) Id() string {
 	return o.ID.Hex()
 }
 
+// FindAndSave finds by `query`,  and then upsert the result with the current object
+func (o *User) FindAndSave(ctx context.Context, query interface{}) (*mongo.SingleResult, error) {
+	col := UserMgr.GetCol()
+	opts := options.FindOneAndUpdate().SetUpsert(true)
+	opts.SetReturnDocument(options.After)
+	update := bson.M{
+		"$set": bson.M{
+			UserMgoFieldUserId:       o.UserId,
+			UserMgoFieldUsername:     o.Username,
+			UserMgoFieldAge:          o.Age,
+			UserMgoFieldRegisterDate: o.RegisterDate,
+		},
+	}
+	ret := col.FindOneAndUpdate(ctx, query, update, opts)
+	if ret.Err() != nil {
+		return nil, ret.Err()
+	}
+	return ret, nil
+}
+
 func (o *User) Save(ctx context.Context) (*mongo.UpdateResult, error) {
 	isNew := o.isNew
 	if o.ID == primitive.NilObjectID {

diff --git a/e2e/mongo/user/user_test.go b/e2e/mongo/user/user_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/ezbuy/ezorm/v2/e2e/mongo/user"
 	"github.com/ezbuy/ezorm/v2/pkg/db"
 	"github.com/ezbuy/ezorm/v2/pkg/orm"
+	"github.com/stretchr/testify/assert"
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -84,7 +85,7 @@ func TestSave(t *testing.T) {
 	u1.Username = "username_1"
 
 	t.Log("insert user start")
-	fmt.Println("insert user start")
+
 	if _, err := u1.Save(ctx); err != nil {
 		t.Fatalf("failed to insert user_1 by Save: %s", err)
 	}
@@ -98,7 +99,7 @@ func TestSave(t *testing.T) {
 	}
 
 	t.Log("update user start")
-	fmt.Println("update user start")
+
 	u1.Username = "username_1_new"
 	if _, err := u1.Save(ctx); err != nil {
 		t.Fatalf("failed to update user_1 by Save: %s", err)
@@ -117,6 +118,23 @@ func TestSave(t *testing.T) {
 	}
 }
 
+func TestFindAndSave(t *testing.T) {
+	ctx := context.TODO()
+	u1 := user.Get_UserMgr().NewUser()
+	u1.Username = "username_1"
+	_, err := u1.FindAndSave(ctx, bson.M{
+		user.UserMgoFieldUsername: u1.Username,
+	})
+	assert.NoError(t, err)
+	u1.Username = "username_1_new"
+	res, err := u1.FindAndSave(ctx, bson.M{
+		user.UserMgoFieldUsername: "username_1",
+	})
+	assert.NoError(t, err)
+	_, err = orm.GetIDFromSingleResult(res)
+	assert.NoError(t, err)
+}
+
 func TestInsertUnique(t *testing.T) {
 	const (
 		uniqueUsername = "unique_user_1"

diff --git a/go.work.sum b/go.work.sum
@@ -1,9 +1,22 @@
+cloud.google.com/go/firestore v1.9.0/go.mod h1:HMkjKHNTtRyZNiMzu7YAsLr9K3X2udY2AMwDaMEQiiE=
+cloud.google.com/go/longrunning v0.3.0/go.mod h1:qth9Y41RRSUE69rDcOn6DdK3HfQfsUI0YSmW3iIlLJc=
+github.com/armon/go-metrics v0.4.0/go.mod h1:E6amYzXo6aW1tqzoZGT755KkbgrJsSdpwZ+3JqfkOG4=
 github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/ezbuy/ezorm/v2/pkg v0.0.9/go.mod h1:OJv0Xk6t+tGCClK0OaRMubrdjfDVb/uEgqVLs6btC7E=
+github.com/hashicorp/consul/api v1.18.0/go.mod h1:owRRGJ9M5xReDC5nfT8FTJrNAPbT4NM6p/k+d03q2v4=
+github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
+github.com/hashicorp/go-hclog v1.2.0/go.mod h1:whpDNt7SSdeAju8AWKIWsul05p54N/39EeqMAyrmvFQ=
+github.com/hashicorp/go-immutable-radix v1.3.1/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
 github.com/hashicorp/go-rootcerts v1.0.2/go.mod h1:pqUvnprVnM5bf7AOirdbb01K4ccR319Vf4pU3K5EGc8=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
+github.com/hashicorp/serf v0.10.1/go.mod h1:yL2t6BqATOLGc5HF7qbFkTfXoPIY0WZdWHfEvMqbG+4=
 github.com/pelletier/go-toml v1.3.0 h1:e5+lF2E4Y2WCIxBefVowBuB0iHrUH4HZ8q+6mGF7fJc=
+github.com/sagikazarmark/crypt v0.9.0/go.mod h1:RnH7sEhxfdnPm1z+XMgSLjWTEIjyK4z2dw6+4vHTMuo=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+go.etcd.io/etcd/api/v3 v3.5.6/go.mod h1:KFtNaxGDw4Yx/BA4iPPwevUTAuqcsPxzyX8PHydchN8=
+go.etcd.io/etcd/client/pkg/v3 v3.5.6/go.mod h1:ggrwbk069qxpKPq8/FKkQ3Xq9y39kbFR4LnKszpRXeQ=
+go.etcd.io/etcd/client/v2 v2.305.6/go.mod h1:BHha8XJGe8vCIBfWBpbBLVZ4QjOIlfoouvOwydu63E0=
+go.etcd.io/etcd/client/v3 v3.5.6/go.mod h1:f6GRinRMCsFVv9Ht42EyY7nfsVGwrNO0WEoS2pRKzQk=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
diff --git a/internal/parser/shared/tpl/mongo_orm.gogo b/internal/parser/shared/tpl/mongo_orm.gogo
@@ -76,6 +76,29 @@ func (o *{{.Name}}) Id() string {
     return o.ID.Hex()
 }
 
+// FindAndSave finds by `query`,  and then upsert the result with the current object
+func (o *{{.Name}}) FindAndSave(ctx context.Context, query interface{}) (*mongo.SingleResult, error) {
+    col := {{.Name}}Mgr.GetCol()
+    opts := options.FindOneAndUpdate().SetUpsert(true)
+    opts.SetReturnDocument(options.After)
+    update := bson.M{
+        "$set": bson.M{
+        {{- range $field := .Fields -}}
+        {{- if eq $field.Name "ID" -}}
+            {{continue}}
+        {{- end}}
+            {{$obj.Name}}MgoField{{$field.Name}}: o.{{$field.Name}},
+        {{- end}}
+        },
+    }
+    ret := col.FindOneAndUpdate(ctx, query, update, opts)
+    if ret.Err() != nil {
+        return nil, ret.Err()
+    }
+    return ret, nil
+}
+
+// Save upserts the document , Save itself is concurrent-safe , but maybe it is not atomic together with other operations, such as `Find`
 func (o *{{.Name}}) Save(ctx context.Context) (*mongo.UpdateResult, error) {
     isNew := o.isNew
     {{- range $field := .Fields}}

diff --git a/v2/pkg/orm/orm.go b/v2/pkg/orm/orm.go
@@ -6,6 +6,10 @@ import (
 	"fmt"
 	"net/url"
 	"time"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 var (
@@ -65,7 +69,6 @@ func RemoveEzOrmObj(namespace, classname, id string) (err error) {
 	f, ok := ezOrmObjsRemove[namespace+"."+classname]
 	if !ok {
 		return errors.New(namespace + "." + classname + " remove func not found")
-
 	}
 
 	return f(id)
@@ -82,9 +85,11 @@ type SearchObj interface {
 	GetSearchTip() string
 }
 
-var DateTimeLayout = "2006-01-02 15:04"
-var DateLayout = "2006-01-02"
-var TimeLayout = "15:04"
+var (
+	DateTimeLayout = "2006-01-02 15:04"
+	DateLayout     = "2006-01-02"
+	TimeLayout     = "15:04"
+)
 
 func I64DateTime(c int64) string {
 	if c == 0 {
@@ -346,3 +351,15 @@ func ToJsonString(obj interface{}) string {
 func PrintToJson(obj interface{}) {
 	fmt.Println(ToJsonString(obj))
 }
+
+func GetIDFromSingleResult(single *mongo.SingleResult) (string, error) {
+	var result bson.M
+	err := single.Decode(&result)
+	if err != nil {
+		return "", fmt.Errorf("decode error: %w", err)
+	}
+	if id, ok := result["_id"].(primitive.ObjectID); ok {
+		return id.Hex(), nil
+	}
+	return "", errors.New("id not found")
+}