From 58a342057d3e391624cbc42dfb86f7ed8c4420d4 Mon Sep 17 00:00:00 2001
From: Nikita Tokarchuk <tokarchuk+bitbucket@nevisq.com>
Date: Fri, 8 Feb 2019 00:03:39 +0100
Subject: [PATCH] Move protection query to explicit function

---
 mongox/common/deleteone.go |  9 +------
 mongox/common/saveone.go   | 14 +++-------
 mongox/query/compose.go    | 55 ++++++++++++++++++++++++++++++++------
 3 files changed, 52 insertions(+), 26 deletions(-)

diff --git a/mongox/common/deleteone.go b/mongox/common/deleteone.go
index 3b6c068..199ff10 100644
--- a/mongox/common/deleteone.go
+++ b/mongox/common/deleteone.go
@@ -26,14 +26,7 @@ func DeleteOne(db *mongox.Database, target interface{}, filters ...interface{})
 	}
 
 	if protected != nil {
-		if protected.X.IsZero() {
-			composed.And(primitive.M{"_x": primitive.M{"$exists": false}})
-			composed.And(primitive.M{"_v": primitive.M{"$exists": false}})
-		} else {
-			composed.And(primitive.M{"_x": protected.X})
-			composed.And(primitive.M{"_v": protected.V})
-		}
-
+		query.Push(composed, protected)
 		protected.X = primitive.NewObjectID()
 		protected.V = time.Now().Unix()
 	}
diff --git a/mongox/common/saveone.go b/mongox/common/saveone.go
index cd4e5b1..20f87d8 100644
--- a/mongox/common/saveone.go
+++ b/mongox/common/saveone.go
@@ -4,6 +4,7 @@ import (
 	"github.com/mainnika/mongox-go-driver/mongox"
 	"github.com/mainnika/mongox-go-driver/mongox/base"
 	"github.com/mainnika/mongox-go-driver/mongox/errors"
+	"github.com/mainnika/mongox-go-driver/mongox/query"
 	"github.com/mongodb/mongo-go-driver/bson"
 	"github.com/mongodb/mongo-go-driver/bson/primitive"
 	"github.com/mongodb/mongo-go-driver/mongo/options"
@@ -17,25 +18,18 @@ func SaveOne(db *mongox.Database, source interface{}) error {
 	opts := options.FindOneAndReplace()
 	id := base.GetID(source)
 	protected := base.GetProtection(source)
-	filter := bson.M{"_id": id}
+	composed := query.Compose(bson.M{"_id": id})
 
 	opts.SetUpsert(true)
 	opts.SetReturnDocument(options.After)
 
 	if protected != nil {
-		if protected.X.IsZero() {
-			filter["_x"] = primitive.M{"$exists": false}
-			filter["_v"] = primitive.M{"$exists": false}
-		} else {
-			filter["_x"] = protected.X
-			filter["_v"] = protected.V
-		}
-
+		query.Push(composed, protected)
 		protected.X = primitive.NewObjectID()
 		protected.V = time.Now().Unix()
 	}
 
-	result := collection.FindOneAndReplace(db.Context(), filter, source, opts)
+	result := collection.FindOneAndReplace(db.Context(), composed.M(), source, opts)
 	if result.Err() != nil {
 		return errors.NotFoundErrorf("%s", result.Err())
 	}
diff --git a/mongox/query/compose.go b/mongox/query/compose.go
index e529b58..9fd1a76 100644
--- a/mongox/query/compose.go
+++ b/mongox/query/compose.go
@@ -1,8 +1,10 @@
 package query
 
 import (
+	"github.com/mainnika/mongox-go-driver/mongox/base"
 	"github.com/mainnika/mongox-go-driver/mongox/errors"
 	"github.com/mongodb/mongo-go-driver/bson"
+	"github.com/mongodb/mongo-go-driver/bson/primitive"
 )
 
 // ComposeQuery is a function to compose filters into a single query
@@ -11,14 +13,7 @@ func Compose(filters ...interface{}) *Query {
 	q := &Query{}
 
 	for _, f := range filters {
-
-		ok := false
-		ok = ok || applyBson(q, f)
-		ok = ok || applyLimit(q, f)
-		ok = ok || applySort(q, f)
-		ok = ok || applySkip(q, f)
-
-		if !ok {
+		if !Push(q, f) {
 			panic(errors.InternalErrorf("unknown filter %v", f))
 		}
 	}
@@ -26,6 +21,19 @@ func Compose(filters ...interface{}) *Query {
 	return q
 }
 
+// Push applies single filter to a query
+func Push(q *Query, f interface{}) bool {
+
+	ok := false
+	ok = ok || applyBson(q, f)
+	ok = ok || applyLimit(q, f)
+	ok = ok || applySort(q, f)
+	ok = ok || applySkip(q, f)
+	ok = ok || applyProtection(q, f)
+
+	return ok
+}
+
 // applyBson is a fallback for a custom bson.M
 func applyBson(q *Query, f interface{}) bool {
 
@@ -69,3 +77,34 @@ func applySkip(q *Query, f interface{}) bool {
 
 	return false
 }
+
+func applyProtection(q *Query, f interface{}) bool {
+
+	var x *primitive.ObjectID
+	var v *int64
+
+	switch f := f.(type) {
+	case base.Protection:
+		x = &f.X
+		v = &f.V
+	case *base.Protection:
+		if f == nil {
+			return false
+		}
+		x = &f.X
+		v = &f.V
+
+	default:
+		return false
+	}
+
+	if x.IsZero() {
+		q.And(primitive.M{"_x": primitive.M{"$exists": false}})
+		q.And(primitive.M{"_v": primitive.M{"$exists": false}})
+	} else {
+		q.And(primitive.M{"_x": *x})
+		q.And(primitive.M{"_v": *v})
+	}
+
+	return true
+}