From fc2a867cbbe2accf8527e43599655dfb2e857c35 Mon Sep 17 00:00:00 2001
From: Nikita Tokarchuk <nikita@tokarch.uk>
Date: Wed, 18 Nov 2020 22:30:16 +0100
Subject: [PATCH] Use update operator not the update pipeline

---
 go.mod                       |  1 +
 go.sum                       |  2 ++
 mongox/database/updateone.go | 22 ++++++++++++----------
 mongox/query/query.go        | 29 +++++++++++++++++++++++++++--
 mongox/query/update.go       |  6 +++---
 5 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/go.mod b/go.mod
index 065075d..f0b1c29 100644
--- a/go.mod
+++ b/go.mod
@@ -6,5 +6,6 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.1
 	github.com/stretchr/testify v1.6.1
+	github.com/valyala/bytebufferpool v1.0.0
 	go.mongodb.org/mongo-driver v1.4.3
 )
diff --git a/go.sum b/go.sum
index e883386..4335323 100644
--- a/go.sum
+++ b/go.sum
@@ -82,6 +82,8 @@ github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/tidwall/pretty v1.0.0 h1:HsD+QiTn7sK6flMKIvNmpqz1qrpP3Ps6jOKIKMooyg4=
 github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c h1:u40Z8hqBAAQyv+vATcGgV0YCnDjqSL7/q/JyPhhJSPk=
 github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c/go.mod h1:lB8K/P019DLNhemzwFU4jHLhdvlE6uDZjXFejJXr49I=
 github.com/xdg/stringprep v0.0.0-20180714160509-73f8eece6fdc h1:n+nNi93yXLkJvKwXNP9d55HC7lGK4H/SRcwB5IaUZLo=
diff --git a/mongox/database/updateone.go b/mongox/database/updateone.go
index 3a50fa7..0ed8aef 100644
--- a/mongox/database/updateone.go
+++ b/mongox/database/updateone.go
@@ -1,8 +1,6 @@
 package database
 
 import (
-	"time"
-
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo/options"
 
@@ -18,10 +16,14 @@ func (d *Database) UpdateOne(target interface{}, filters ...interface{}) (err er
 		return
 	}
 
+	updaterDoc, err := composed.Updater()
+	if err != nil {
+		return
+	}
+
 	collection := d.GetCollectionOf(target)
 	protected := base.GetProtection(target)
 	ctx := query.WithContext(d.Context(), composed)
-	updater := composed.Updater()
 
 	opts := options.FindOneAndUpdate()
 	opts.SetReturnDocument(options.After)
@@ -30,12 +32,12 @@ func (d *Database) UpdateOne(target interface{}, filters ...interface{}) (err er
 		if !protected.X.IsZero() {
 			query.Push(composed, protected)
 		}
-		updater = append(updater, primitive.M{
-			"$set": primitive.M{
-				"_x": primitive.NewObjectID(),
-				"_v": time.Now().Unix(),
-			},
-		})
+		//updater = append(updater, primitive.M{
+		//	"$set": primitive.M{
+		//		"_x": primitive.NewObjectID(),
+		//		"_v": time.Now().Unix(),
+		//	},
+		//})
 	}
 
 	defer func() {
@@ -47,7 +49,7 @@ func (d *Database) UpdateOne(target interface{}, filters ...interface{}) (err er
 		return
 	}()
 
-	result := collection.FindOneAndUpdate(ctx, composed.M(), updater, opts)
+	result := collection.FindOneAndUpdate(ctx, composed.M(), updaterDoc, opts)
 	if result.Err() != nil {
 		return result.Err()
 	}
diff --git a/mongox/query/query.go b/mongox/query/query.go
index 2f44783..3c20053 100644
--- a/mongox/query/query.go
+++ b/mongox/query/query.go
@@ -1,6 +1,9 @@
 package query
 
 import (
+	"github.com/modern-go/reflect2"
+	"github.com/valyala/bytebufferpool"
+	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
@@ -66,13 +69,35 @@ func (q *Query) Skipper() (skip *int64) {
 }
 
 // Updater is an update command for a query
-func (q *Query) Updater() (update primitive.A) {
+func (q *Query) Updater() (update primitive.M, err error) {
 
 	if q.updater == nil {
+		update = primitive.M{}
 		return
 	}
 
-	return q.updater.Update()
+	update = q.updater.Update()
+
+	if reflect2.IsNil(update) {
+		update = primitive.M{}
+		return
+	}
+
+	buffer := bytebufferpool.Get()
+	defer bytebufferpool.Put(buffer)
+
+	// convert update document to bson map values
+	bsonBytes, err := bson.MarshalAppend(buffer.B, update)
+	if err != nil {
+		return
+	}
+	update = primitive.M{}
+	err = bson.Unmarshal(bsonBytes, update)
+	if err != nil {
+		return
+	}
+
+	return
 }
 
 // Preloader is a preloader list for a query
diff --git a/mongox/query/update.go b/mongox/query/update.go
index 6fbc009..64796ba 100644
--- a/mongox/query/update.go
+++ b/mongox/query/update.go
@@ -6,7 +6,7 @@ import (
 
 // Updater is a filter to update the data
 type Updater interface {
-	Update() (update primitive.A)
+	Update() (update primitive.M)
 }
 
 // Update is a simple implementations of the Updater filter
@@ -15,6 +15,6 @@ type Update primitive.M
 var _ Updater = &Update{}
 
 // Update returns an update command
-func (u Update) Update() (update primitive.A) {
-	return primitive.A{primitive.M(u)}
+func (u Update) Update() (update primitive.M) {
+	return primitive.M(u)
 }