From 86397885e21c811c1dd21acab1ee93dce2500e22 Mon Sep 17 00:00:00 2001
From: Nikita Tokarchuk <nikita@tokarch.uk>
Date: Wed, 18 Nov 2020 22:33:50 +0100
Subject: [PATCH] Function update one now considers the protection key using
 the update operator

---
 mongox/base/protection/key.go | 17 +++++++++++++++++
 mongox/database/updateone.go  | 15 +++++++++------
 mongox/query/compose.go       | 17 ++++-------------
 3 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/mongox/base/protection/key.go b/mongox/base/protection/key.go
index e385d5a..b64f186 100644
--- a/mongox/base/protection/key.go
+++ b/mongox/base/protection/key.go
@@ -3,6 +3,7 @@ package protection
 import (
 	"time"
 
+	"github.com/modern-go/reflect2"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
@@ -12,6 +13,22 @@ type Key struct {
 	V int64              `bson:"_v" json:"_v"`
 }
 
+// PutToDocument extends the doc with protection key values
+func (k *Key) PutToDocument(doc primitive.M) {
+
+	if reflect2.IsNil(doc) {
+		return
+	}
+
+	if k.X.IsZero() {
+		doc["_x"] = primitive.M{"$exists": false}
+		doc["_v"] = primitive.M{"$exists": false}
+	} else {
+		doc["_x"] = k.X
+		doc["_v"] = k.V
+	}
+}
+
 // Restate creates a new protection key
 func (k *Key) Restate() {
 	k.X = primitive.NewObjectID()
diff --git a/mongox/database/updateone.go b/mongox/database/updateone.go
index b57d275..ee404b8 100644
--- a/mongox/database/updateone.go
+++ b/mongox/database/updateone.go
@@ -1,6 +1,7 @@
 package database
 
 import (
+	"github.com/modern-go/reflect2"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo/options"
 
@@ -32,13 +33,15 @@ 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(),
-		//	},
-		//})
+
 		protected.Restate()
+
+		setCmd, _ := updaterDoc["$set"].(primitive.M)
+		if reflect2.IsNil(setCmd) {
+			setCmd = primitive.M{}
+		}
+		protected.PutToDocument(setCmd)
+		updaterDoc["$set"] = setCmd
 	}
 
 	defer func() {
diff --git a/mongox/query/compose.go b/mongox/query/compose.go
index 3e5cb93..1836ee4 100644
--- a/mongox/query/compose.go
+++ b/mongox/query/compose.go
@@ -107,28 +107,19 @@ func applySkip(query *Query, filter interface{}) (ok bool) {
 
 func applyProtection(query *Query, filter interface{}) (ok bool) {
 
-	var x *primitive.ObjectID
-	var v *int64
+	var keyDoc = primitive.M{}
 
 	switch filter := filter.(type) {
 	case protection.Key:
-		x = &filter.X
-		v = &filter.V
+		filter.PutToDocument(keyDoc)
 	case *protection.Key:
-		x = &filter.X
-		v = &filter.V
+		filter.PutToDocument(keyDoc)
 
 	default:
 		return false
 	}
 
-	if x.IsZero() {
-		query.And(primitive.M{"_x": primitive.M{"$exists": false}})
-		query.And(primitive.M{"_v": primitive.M{"$exists": false}})
-	} else {
-		query.And(primitive.M{"_x": *x})
-		query.And(primitive.M{"_v": *v})
-	}
+	query.And(keyDoc)
 
 	return true
 }